[pypy-svn] r46124 - in pypy/dist/pypy/rpython/numpy: . test

simonb at codespeak.net simonb at codespeak.net
Wed Aug 29 02:33:27 CEST 2007


Author: simonb
Date: Wed Aug 29 02:33:26 2007
New Revision: 46124

Modified:
   pypy/dist/pypy/rpython/numpy/aarray.py
   pypy/dist/pypy/rpython/numpy/rarray.py
   pypy/dist/pypy/rpython/numpy/test/test_array.py
Log:
build multi dimensional arrays, annotating of multi-indexes 

Modified: pypy/dist/pypy/rpython/numpy/aarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/aarray.py	(original)
+++ pypy/dist/pypy/rpython/numpy/aarray.py	Wed Aug 29 02:33:26 2007
@@ -1,11 +1,11 @@
 from pypy.rpython.extregistry import ExtRegistryEntry
-from pypy.annotation.pairtype import pairtype
+from pypy.annotation.pairtype import pair, pairtype
 from pypy.annotation.model import SomeExternalObject, SomeList, SomeImpossibleValue
-from pypy.annotation.model import SomeObject, SomeInteger, SomeFloat, SomeString, SomeChar,\
-    SomeTuple
+from pypy.annotation.model import SomeObject, SomeInteger, SomeFloat, SomeString, SomeChar, SomeTuple, SomeSlice
 from pypy.tool.error import AnnotatorError
 from pypy.rpython.lltypesystem import rffi
 from pypy.rlib import rarithmetic
+from pypy.annotation import listdef
 
 import numpy
 
@@ -25,25 +25,15 @@
         #'f' : SomeFloat(), # XX single precision float XX
         'd' : SomeFloat(),
     }
-    def __init__(self, knowntype, typecode, ndim=1):
-        self.knowntype = knowntype # == numpy.ndarray (do we need this for anything?)
+    def __init__(self, typecode, ndim=1):
+        if not typecode in self.typecode_to_item:
+            raise AnnotatorError("bad typecode: %r"%typecode)
         self.typecode = typecode
         self.ndim = ndim
 
     def can_be_none(self):
         return True
 
-    def return_annotation(self):
-        """Returns either 'self' or the annotation of the unwrapped version
-        of this ctype, following the logic used when ctypes operations
-        return a value.
-        """
-        from pypy.rpython import extregistry
-        assert extregistry.is_registered_type(self.knowntype)
-        entry = extregistry.lookup_type(self.knowntype)
-        # special case for returning primitives or c_char_p
-        return getattr(entry, 's_return_trick', self)
-
     def get_item_type(self):
         return self.typecode_to_item[self.typecode]
 
@@ -60,7 +50,7 @@
         return s
 
     def method_transpose(self):
-        return SomeArray(self.knowntype, self.typecode, self.ndim)
+        return SomeArray(self.typecode, self.ndim)
 
 class __extend__(pairtype(SomeArray, SomeArray)):
 
@@ -77,18 +67,53 @@
                     break
         if typecode is None:
             raise AnnotatorError()
-        return SomeArray(s_arr1.knowntype, typecode)
+        return SomeArray(typecode)
 
     add = sub = mul = div = truediv = union
 
 
 class __extend__(pairtype(SomeArray, SomeInteger)):
-    def setitem((s_cto, s_index), s_value):
-        pass
+    def setitem((s_array, s_index), s_value):
+        if s_array.ndim == 0:
+            raise AnnotatorError()
+        if isinstance(s_value, SomeArray):
+            assert s_array.ndim == s_value.ndim + 1
 
-    def getitem((s_cto, s_index)):
-        # TODO: higher ndimed arrays have getitem returns SomeArray
-        return s_cto.get_item_type()
+    def getitem((s_array, s_index)):
+        if s_array.ndim == 0:
+            raise AnnotatorError()
+        if s_array.ndim > 1:
+            return SomeArray(s_array.typecode, s_array.ndim-1)
+        return s_array.get_item_type()
+
+class __extend__(pairtype(SomeArray, SomeTuple)):
+    def get_leftover_dim((s_array, s_index)):
+        ndim = s_array.ndim
+        for s_item in s_index.items:
+            if isinstance(s_item, SomeInteger):
+                ndim -= 1
+            elif isinstance(s_item, SomeSlice):
+                pass
+            else:
+                raise AnnotatorError("cannot index with %s"%s_item)
+        return ndim
+
+    def setitem((s_array, s_index), s_value):
+        ndim = pair(s_array, s_index).get_leftover_dim()
+        if isinstance(s_value, SomeArray):
+            if s_value.ndim + ndim != s_array.ndim:
+                # XX allow broadcasting..
+                raise AnnotatorError("shape mismatch")
+        elif ndim > 0:
+            raise AnnotatorError("need to set from array")
+
+    def getitem((s_array, s_index)):
+        ndim = pair(s_array, s_index).get_leftover_dim()
+        if s_array.ndim == 0 and len(s_index.items):
+            raise AnnotatorError("indexing rank zero array with nonempty tuple")
+        if ndim > 0:
+            return SomeArray(s_array.typecode, ndim)
+        return s_array.get_item_type()
 
 numpy_typedict = {
     (SomeInteger, rffi.r_signedchar) : 'b', 
@@ -107,48 +132,78 @@
 }
 valid_typecodes='bhilqBHILQfd'
 
-class CallEntry(ExtRegistryEntry):
-    "Annotation and rtyping of calls to numpy.array."
+class ArrayCallEntry(ExtRegistryEntry):
+    "Annotation and rtyping of calls to numpy.array"
     _about_ = numpy.array
 
-    def compute_result_annotation(self, arg_list, *args_s, **kwds_s):
-        if not isinstance(arg_list, SomeList):
-            raise AnnotatorError("numpy.array expects SomeList")
-
-        # First guess type from input list
-        listitem = arg_list.listdef.listitem
-        key = type(listitem.s_value), listitem.s_value.knowntype
-        typecode = numpy_typedict.get(key, None)
+    def compute_result_annotation(self, s_list, s_dtype=None):
+        if isinstance(s_list, SomeList):
+            # First guess type from input list
+            listitem = s_list.listdef.listitem
+            key = type(listitem.s_value), listitem.s_value.knowntype
+            typecode = numpy_typedict.get(key, None)
+            ndim = 1
+        elif isinstance(s_list, SomeArray):
+            typecode = s_list.typecode
+            ndim = s_list.ndim
+        else:
+            raise AnnotatorError("cannot build array from %s"%s_list)
 
         # now see if the dtype arg over-rides the typecode
-        dtype = None
-        if len(args_s)>0:
-            dtype = args_s[0]
-        if "dtype" in kwds_s:
-            dtype = kwds_s["dtype"]
-        if isinstance(dtype,SomeChar) and dtype.is_constant():
-            typecode = dtype.const
-            dtype = None
-        if dtype is not None:
+        if isinstance(s_dtype, SomeChar) and s_dtype.is_constant():
+            typecode = s_dtype.const
+            s_dtype = None
+        if s_dtype is not None:
             raise AnnotatorError("dtype is not a valid type specification")
         if typecode is None or typecode not in valid_typecodes:
             raise AnnotatorError("List item type not supported")
-        knowntype = numpy.ndarray
-        return SomeArray(knowntype, typecode)
+        return SomeArray(typecode, ndim)
 
     def specialize_call(self, hop):
         r_array = hop.r_result
-        [v_lst] = hop.inputargs(r_array)
-        v_result = r_array.allocate_instance(hop.llops, v_lst)
+        [v_lst] = hop.inputargs(r_array) # coerce list arg to array arg
+        v_result = r_array.build_from_array(hop.llops, v_lst)
         return v_result
 
-class NumpyObjEntry(ExtRegistryEntry):
-    "Annotation and rtyping of numpy array instances."
-    _type_ = numpy.ndarray
-
-    def get_repr(self, rtyper, s_array):
-        from pypy.rpython.numpy.rarray import ArrayRepr
-        return ArrayRepr(rtyper, s_array)
+
+class ZeroesCallEntry(ExtRegistryEntry):
+    "Annotation and rtyping of calls to numpy.zeroes"
+    _about_ = numpy.zeros
+
+    def compute_result_annotation(self, s_tuple, s_dtype=None):
+        if isinstance(s_tuple, SomeTuple):
+            for s_item in s_tuple.items:
+                if not isinstance(s_item, SomeInteger):
+                    raise AnnotatorError("shape must be tuple of integers")
+            ndim = len(s_tuple.items)
+        else:
+            # XX also build from single int arg
+            raise AnnotatorError("could not build array shape from %s"%s_list)
+
+        typecode = 'd'
+        if isinstance(s_dtype, SomeChar) and s_dtype.is_constant():
+            typecode = s_dtype.const
+            s_dtype = None
+        return SomeArray(typecode, ndim)
+
+#    def specialize_call(self, hop):
+#        ldef = listdef.ListDef(None, SomeInteger())
+#        r_lst = hop.rtyper.getrepr(SomeList(ldef))
+#        # XX TyperError: don't know how to convert from 
+#        # <TupleRepr * GcStruct tuple2 { item0, item1 }> to 
+#        # <FixedSizeListRepr * GcForwardReference>
+#        [v_lst] = hop.inputargs(r_lst)
+#        r_array = hop.r_result
+#        v_result = r_array.build_from_shape(hop.llops, r_lst, v_lst)
+#        return v_result
+
+    def specialize_call(self, hop):
+        r_tpl = hop.args_r[0]
+        # XX also call with single int arg
+        [v_tpl] = hop.inputargs(r_tpl)
+        r_array = hop.r_result
+        v_result = r_array.build_from_shape(hop.llops, r_tpl, v_tpl)
+        return v_result
 
 
 

Modified: pypy/dist/pypy/rpython/numpy/rarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/rarray.py	(original)
+++ pypy/dist/pypy/rpython/numpy/rarray.py	Wed Aug 29 02:33:26 2007
@@ -4,58 +4,129 @@
 from pypy.rpython.rlist import AbstractBaseListRepr
 from pypy.rpython.error import TyperError
 from pypy.rpython.lltypesystem import lltype, llmemory
+from pypy.rpython.lltypesystem.rtupletype import TUPLE_TYPE
 from pypy.rpython.lltypesystem.lltype import \
      GcArray, GcStruct, Signed, Ptr, Unsigned, malloc, Void
 from pypy.annotation.model import SomeObject, SomeInteger
 from pypy.rpython.numpy.aarray import SomeArray
 from pypy.annotation.pairtype import pairtype
+from pypy.rlib.unroll import unrolling_iterable
+from pypy.annotation import listdef
+from pypy.rpython.memory.lltypelayout import sizeof
+
+def gen_build_from_shape(ndim):
+    unrolling_dims = unrolling_iterable(reversed(range(ndim)))
+    def ll_build_from_shape(ARRAY, shape):
+        array = ll_allocate(ARRAY, ndim)
+        itemsize = 1
+        for i in unrolling_dims:
+            attr = 'item%d'%i
+            size = getattr(shape, attr)
+            array.shape[i] = size
+            array.strides[i] = itemsize
+            itemsize *= size
+        array.data = malloc(ARRAY.data.TO, itemsize)
+        return array
+    return ll_build_from_shape
+
+def gen_get_shape(ndim):
+    unrolling_dims = unrolling_iterable(range(ndim))
+    def ll_get_shape(ARRAY, TUPLE, array):
+        shape = malloc(TUPLE)
+        for i in unrolling_dims:
+            size = array.shape[i]
+            attr = 'item%d'%i
+            setattr(shape, attr, size)
+        return shape
+    return ll_get_shape
 
 
 class ArrayRepr(Repr):
     def __init__(self, rtyper, s_array):
+        self.s_array = s_array
         self.s_value = s_array.get_item_type()
         self.item_repr = rtyper.getrepr(self.s_value)
         ITEM = self.item_repr.lowleveltype
-        ITEMARRAY = GcArray(ITEM)
-        SIZEARRAY = GcArray(Signed)
+        ITEMARRAY = GcArray(ITEM, hints={'nolength':True})
+        SIZEARRAY = GcArray(Signed, hints={'nolength':True})
+        self.PTR_SIZEARRAY = Ptr(SIZEARRAY)
+        self.itemsize = sizeof(ITEM)
+        self.ndim = s_array.ndim
         self.ARRAY = Ptr(
             GcStruct("array",
                 ("data", Ptr(ITEMARRAY)), # pointer to raw data buffer 
-                ("nd", Signed), # number of dimensions, also called ndim
-                ("dimensions", Ptr(SIZEARRAY)), # size in each dimension
-                ("strides", Ptr(SIZEARRAY)), # bytes to jump to get to the
-                                             # next element in each dimension 
+                ("ndim", Signed), # number of dimensions
+                ("shape", self.PTR_SIZEARRAY), # size in each dimension
+                ("strides", self.PTR_SIZEARRAY), # bytes (?) to jump to get to the
+                                                 # next element in each dimension 
             ))
         self.lowleveltype = self.ARRAY
 
-    def allocate_instance(self, llops, v_array):
-        c1 = inputconst(lltype.Void, self.lowleveltype.TO) 
-        return llops.gendirectcall(ll_allocate, c1, v_array)
+    def build_from_array(self, llops, v_array):
+        cARRAY = inputconst(lltype.Void, self.lowleveltype.TO) 
+        return llops.gendirectcall(ll_build_alias, cARRAY, v_array)
+
+#    def build_from_shape(self, llops, r_tuple, v_tuple):
+#        cARRAY = inputconst(lltype.Void, self.lowleveltype.TO) 
+#        cTUPLE = inputconst(lltype.Void, r_tuple.lowleveltype.TO)
+#        ndim = self.ndim
+#        c_ndim = inputconst(lltype.Signed, ndim)
+#        assert ndim == len(r_tuple.items_r)
+#        v_array = llops.gendirectcall(ll_allocate, cARRAY, c_ndim)
+#        c_attr = inputconst(lltype.Void, 'shape')
+#        v_shape = llops.genop('getfield', [v_array, c_attr], self.PTR_SIZEARRAY)
+#        for i in range(ndim):
+#            v_size = r_tuple.getitem_internal(llops, v_tuple, i)
+#            v_i = inputconst(lltype.Signed, i)
+#            llops.genop('setarrayitem', [v_shape, v_i, v_size])
+#        return v_array
+        
+    def build_from_shape(self, llops, r_tuple, v_tuple):
+        cARRAY = inputconst(lltype.Void, self.lowleveltype.TO) 
+        cTUPLE = inputconst(lltype.Void, r_tuple.lowleveltype.TO)
+        ndim = self.s_array.ndim
+        ll_build_from_shape = gen_build_from_shape(ndim)
+        c_ndim = inputconst(lltype.Signed, ndim)
+        assert ndim == len(r_tuple.items_r)
+        rval = llops.gendirectcall(ll_build_from_shape, cARRAY, v_tuple)
+        return rval
 
     def rtype_method_transpose(self, hop):
         [v_self] = hop.inputargs(self)
         cARRAY = hop.inputconst(Void, hop.r_result.ARRAY.TO)
         return hop.gendirectcall(ll_transpose, cARRAY, v_self)
 
+    def get_ndim(self, hop, v_array):
+        cname = inputconst(Void, 'ndim')
+        return hop.llops.genop('getfield', [v_array, cname], resulttype=Signed)
+
+    def get_shape(self, hop, v_array):
+        cname = inputconst(Void, 'shape')
+        TUPLE = TUPLE_TYPE([Signed]*self.ndim)
+        cARRAY = inputconst(lltype.Void, self.lowleveltype.TO) 
+        cTUPLE = inputconst(lltype.Void, TUPLE.TO)
+        ll_get_shape = gen_get_shape(self.ndim)
+        return hop.llops.gendirectcall(ll_get_shape, cARRAY, cTUPLE, v_array)
+        return llops.genop('getfield', [v_array, cname], resulttype=TUPLE)
+
     def rtype_getattr(self, hop):
         s_attr = hop.args_s[1]
         if s_attr.is_constant() and isinstance(s_attr.const, str):
             attr = s_attr.const
             s_obj = hop.args_s[0]
-            vobj, vattr = hop.inputargs(self, Void)
-            llops = hop.llops
-            if attr == 'ndim':
-                cname = inputconst(Void, 'nd')
-                return llops.genop('getfield', [vobj, cname], resulttype=Signed)
+            v_array, vattr = hop.inputargs(self, Void)
+            getter = getattr(self, 'get_'+attr, None)
+            if getter:
+                return getter(hop, v_array)
         return Repr.rtype_getattr(self, hop)
 
 
 class __extend__(SomeArray):
     def rtyper_makerepr(self, rtyper):
-        return ArrayRepr( rtyper, self )
+        return ArrayRepr(rtyper, self)
 
     def rtyper_makekey(self):
-        return self.__class__, self.knowntype
+        return self.__class__, self.typecode, self.ndim
 
 
 class __extend__(pairtype(ArrayRepr, ArrayRepr)):
@@ -79,30 +150,42 @@
         if r_lst.listitem is None:
             return NotImplemented
         if r_lst.item_repr != r_arr.item_repr:
+            assert 0, (r_lst, r_arr.item_repr)
             return NotImplemented
-        c1 = inputconst(lltype.Void, r_arr.lowleveltype.TO) 
-        return llops.gendirectcall(ll_build_array, c1, v)
+        cARRAY = inputconst(lltype.Void, r_arr.lowleveltype.TO) 
+        return llops.gendirectcall(ll_build_from_list, cARRAY, v)
 
 class __extend__(pairtype(AbstractRangeRepr, ArrayRepr)):
     def convert_from_to((r_rng, r_arr), v, llops):
-        c1 = inputconst(lltype.Void, r_arr.lowleveltype.TO) 
-        return llops.gendirectcall(ll_build_array, c1, v)
+        cARRAY = inputconst(lltype.Void, r_arr.lowleveltype.TO) 
+        return llops.gendirectcall(ll_build_from_list, cARRAY, v)
 
-def ll_build_array(ARRAY, lst):
-    size = lst.ll_length()
+def ll_allocate(ARRAY, ndim):
     array = malloc(ARRAY)
+    array.ndim = ndim
+    array.shape = malloc(ARRAY.shape.TO, array.ndim)
+    array.strides = malloc(ARRAY.strides.TO, array.ndim)
+    return array
+
+def ll_build_from_list(ARRAY, lst):
+    size = lst.ll_length()
+    array = ll_allocate(ARRAY, 1)
+    for i in range(array.ndim):
+        array.shape[i] = size
+        array.strides[i] = 1
     data = array.data = malloc(ARRAY.data.TO, size)
     i = 0
     while i < size:
         data[i] = lst.ll_getitem_fast(i)
         i += 1
-    array.nd = 1
     return array
 
-def ll_allocate(ARRAY, array):
-    new_array = malloc(ARRAY)
+def ll_build_alias(ARRAY, array):
+    new_array = ll_allocate(ARRAY, array.ndim)
     new_array.data = array.data # alias data
-    new_array.nd = array.nd
+    for i in range(array.ndim):
+        new_array.shape[i] = array.shape[i]
+        new_array.strides[i] = array.strides[i]
     return new_array
 
 def ll_setitem(l, index, item):
@@ -112,8 +195,8 @@
     return l.data[index]
 
 def ll_add(ARRAY, a1, a2):
-    size = len(a1.data)
-    if size != len(a2.data):
+    size = a1.shape[0]
+    if size != a2.shape[0]:
         raise ValueError
     array = malloc(ARRAY)
     array.data = malloc(ARRAY.data.TO, size)
@@ -124,7 +207,7 @@
     return array
 
 def ll_transpose(ARRAY, a1):
-    a2 = ll_allocate(ARRAY, a1)
+    a2 = ll_build_alias(ARRAY, a1)
     # XX do something to a2
     return a2
     

Modified: pypy/dist/pypy/rpython/numpy/test/test_array.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/test/test_array.py	(original)
+++ pypy/dist/pypy/rpython/numpy/test/test_array.py	Wed Aug 29 02:33:26 2007
@@ -5,8 +5,9 @@
 import py
 import pypy.rpython.numpy.implementation
 from pypy.annotation import model as annmodel
-from pypy.annotation.model import SomeTuple
+from pypy.annotation.model import SomeObject, SomeTuple
 from pypy.annotation.annrpython import RPythonAnnotator
+from pypy.tool.error import AnnotatorError
 from pypy.translator.translator import TranslationContext
 from pypy import conftest
 import sys
@@ -69,6 +70,30 @@
         s = a.build_types(access_with_variable, [])
         assert s.knowntype == rffi.r_int
 
+    def test_annotate_zeros(self):
+        def f():
+            a = numpy.zeros((3,4,5))
+            return a
+
+        t = TranslationContext()
+        a = t.buildannotator()
+        s = a.build_types(f, [])
+        assert s.typecode == 'd'
+        assert s.ndim == 3
+
+    def test_annotate_indexing(self):
+        def f():
+            a = numpy.zeros((3,4,5))
+            b = a[0]
+            a[0,1,2] = 1.
+            b[0,1] = a[2]
+            return b
+
+        t = TranslationContext()
+        a = t.buildannotator()
+        s = a.build_types(f, [])
+        assert s.ndim == 2
+
     def test_annotate_array_add(self):
         def f():
             a1 = numpy.array([1,2])
@@ -91,6 +116,27 @@
         s = a.build_types(f, [])
         assert s.typecode == 'd'
 
+    def test_annotate_array_dtype(self):
+        def f():
+            a1 = numpy.array([1,2], dtype='d')
+            return a1
+
+        t = TranslationContext()
+        a = t.buildannotator()
+        s = a.build_types(f, [])
+        assert s.typecode == 'd'
+
+    def test_annotate_array_array(self):
+        def f():
+            a1 = numpy.array([1,2], dtype='d')
+            a2 = numpy.array(a1)
+            return a2
+
+        t = TranslationContext()
+        a = t.buildannotator()
+        s = a.build_types(f, [])
+        assert s.typecode == 'd'
+
     def test_annotate_array_attr(self):
         def f():
             a1 = numpy.array([1,2])
@@ -116,11 +162,21 @@
 class Test_specialization:
     def test_specialize_array_create(self):
         def create_array():
-            return numpy.array([1,2])
+            a = numpy.array([1,20])
+            b = numpy.array(a)
+            return b
 
         res = interpret(create_array, [])
         assert res.data[0] == 1
-        assert res.data[1] == 2
+        assert res.data[1] == 20
+
+    def test_specialize_array_zeros(self):
+        def create_array(n, m):
+            a = numpy.zeros((n, m))
+            return a
+
+        res = interpret(create_array, [3, 4])
+        assert res.ndim == 2
 
     def test_specialize_array_access(self):
         def access_with_variable():
@@ -137,7 +193,7 @@
 
     def test_specialize_array_add(self):
         def create_array():
-            a1 = numpy.array([1,2])
+            a1 = numpy.array([1.,2.])
             a2 = numpy.array([6,9])
             return a1 + a2
 
@@ -153,6 +209,26 @@
         res = interpret(create_array, [])
         assert res == 1
 
+    def test_specialize_array_attr_shape(self):
+        def create_array():
+            a = numpy.zeros((2,3))
+            return list(a.shape)
+
+        res = interpret(create_array, [])
+        assert res[0] == 2
+        assert res[1] == 3
+
+    def test_specialize_array_strides(self):
+        def create_array():
+            a = numpy.zeros((3,4,5))
+            return a
+
+        res = interpret(create_array, [])
+        assert res.strides[0] == 20
+        assert res.strides[1] == 5
+        assert res.strides[2] == 1
+        #assert len(res.data) == 3*4*5 # GcArray has nolength
+
     def test_specialize_array_method(self):
         def create_array():
             a = numpy.array([1,2])
@@ -162,6 +238,14 @@
         assert res.data[0] == 1
         assert res.data[1] == 2
 
+    def X_test_specialize_view(self):
+        t = TranslationContext()
+        a = t.buildannotator()
+        a = a.build_types(f, [])
+        r = t.buildrtyper()
+        r.specialize()
+        t.view()
+
 class Test_compile:
     def setup_class(self):
         if not test_c_compile:
@@ -172,12 +256,20 @@
 
     def test_compile_array_access(self):
         def access_array(index):
-            my_array = numpy.array([3,99,2])
-            my_array[0] = 1
-            return my_array[index]
+            a = numpy.array([3,99,2])
+            a[0] = 1
+            return a[index]
 
         fn = self.compile(access_array, [int])
         assert fn(0) == 1
         assert fn(1) == 99
         
+    def test_compile_2d(self):
+        def access_array(index):
+            a = numpy.zeros((5,6))
+            a[0,0] = 2
+            return 0
+
+        fn = self.compile(access_array, [int])
+        
 



More information about the Pypy-commit mailing list