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

simonb at codespeak.net simonb at codespeak.net
Fri Aug 31 03:15:16 CEST 2007


Author: simonb
Date: Fri Aug 31 03:15:14 2007
New Revision: 46204

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:
multi dimensional get/set item with integer indices

Modified: pypy/dist/pypy/rpython/numpy/aarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/aarray.py	(original)
+++ pypy/dist/pypy/rpython/numpy/aarray.py	Fri Aug 31 03:15:14 2007
@@ -181,9 +181,9 @@
         return v_result
 
 
-class ZeroesCallEntry(ExtRegistryEntry):
-    "Annotation and rtyping of calls to numpy.zeroes"
-    _about_ = numpy.zeros
+class EmptyCallEntry(ExtRegistryEntry):
+    "Annotation and rtyping of calls to numpy.empty"
+    _about_ = numpy.empty
 
     def compute_result_annotation(self, s_tuple, s_dtype=None):
         if isinstance(s_tuple, SomeTuple):

Modified: pypy/dist/pypy/rpython/numpy/rarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/rarray.py	(original)
+++ pypy/dist/pypy/rpython/numpy/rarray.py	Fri Aug 31 03:15:14 2007
@@ -215,20 +215,42 @@
 class __extend__(pairtype(ArrayRepr, IntegerRepr)):
     def rtype_setitem((r_arr, r_int), hop):
         v_array, v_index, v_item = hop.inputargs(r_arr, Signed, r_arr.item_repr)
-        return hop.gendirectcall(ll_setitem, v_array, v_index, v_item)
+        return hop.gendirectcall(ll_setitem1, v_array, v_index, v_item)
 
     def rtype_getitem((r_arr, r_int), hop):
         v_array, v_index = hop.inputargs(r_arr, Signed)
-        return hop.gendirectcall(ll_getitem, v_array, v_index)
-
-#class __extend__(pairtype(ArrayRepr, AbstractTupleRepr)):
-#    def rtype_setitem((r_arr, r_tpl), hop):
-#        v_array, v_tuple, v_item = hop.inputargs(r_arr, r_tpl, r_arr)
-#        return hop.gendirectcall(ll_setitem_from_array, v_array, v_tuple, v_item)
-#
-#    def rtype_getitem((r_arr, r_tpl), hop):
-#        return NotImplemented
+        return hop.gendirectcall(ll_getitem1, v_array, v_index)
 
+def gen_getset_item(ndim):
+    unrolling_dims = unrolling_iterable(range(ndim))
+    def ll_get_item(ARRAY, ao, tpl):
+        array = ll_allocate(ARRAY, ndim)
+        idx = 0
+        for i in unrolling_dims:
+            idx += ao.strides[i] * getattr(tpl, 'item%d'%i)
+        return ao.data[idx]
+
+    def ll_set_item(ARRAY, ao, tpl, value):
+        array = ll_allocate(ARRAY, ndim)
+        idx = 0
+        for i in unrolling_dims:
+            idx += ao.strides[i] * getattr(tpl, 'item%d'%i)
+        ao.data[idx] = value
+
+    return ll_get_item, ll_set_item
+
+class __extend__(pairtype(ArrayRepr, AbstractTupleRepr)):
+    def rtype_getitem((r_arr, r_tpl), hop):
+        v_array, v_tuple = hop.inputargs(r_arr, r_tpl)
+        cARRAY = hop.inputconst(Void, r_arr.ARRAY.TO)
+        get_item, set_item = gen_getset_item(r_arr.ndim)
+        return hop.gendirectcall(get_item, cARRAY, v_array, v_tuple)
+
+    def rtype_setitem((r_arr, r_tpl), hop):
+        v_array, v_tuple, v_item = hop.inputargs(r_arr, r_tpl, hop.args_r[2])
+        cARRAY = hop.inputconst(Void, r_arr.ARRAY.TO)
+        get_item, set_item = gen_getset_item(r_arr.ndim)
+        return hop.gendirectcall(set_item, cARRAY, v_array, v_tuple, v_item)
 
 class __extend__(pairtype(ArrayRepr, AbstractSliceRepr)):
     def rtype_setitem((r_arr, r_slc), hop):
@@ -242,7 +264,6 @@
         cnext0 = hop.inputconst(Void, iter_next0)
         cnew1 = hop.inputconst(Void, iter_new1)
         cnext1 = hop.inputconst(Void, iter_next1)
-#        return hop.gendirectcall(ll_setitem_from_array, cnext, cITER0, v_array, v_slc, cITER1, v_item)
         return hop.gendirectcall(ll_array_unary_op,
             cnew0, cnext0, cITER0, v_array, cnew1, cnext1, cITER1, v_item)
         
@@ -302,10 +323,10 @@
         new_array.strides[i] = array.strides[i]
     return new_array
 
-def ll_setitem(l, index, item):
+def ll_setitem1(l, index, item):
     l.data[index] = item
 
-def ll_getitem(l, index):
+def ll_getitem1(l, index):
     return l.data[index]
 
 def ll_add(ARRAY, a1, 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	Fri Aug 31 03:15:14 2007
@@ -70,9 +70,9 @@
         s = a.build_types(access_with_variable, [])
         assert s.knowntype == rffi.r_int
 
-    def test_annotate_zeros(self):
+    def test_annotate_empty(self):
         def f():
-            a = numpy.zeros((3,4,5))
+            a = numpy.empty((3,4,5))
             return a
 
         t = TranslationContext()
@@ -83,7 +83,7 @@
 
     def test_annotate_indexing(self):
         def f():
-            a = numpy.zeros((3,4,5))
+            a = numpy.empty((3,4,5))
             b = a[0]
             a[0,1,2] = 1.
             b[0,1] = 2.
@@ -185,9 +185,9 @@
         assert res.data[0] == 1
         assert res.data[1] == 20
 
-    def test_specialize_array_zeros(self):
+    def test_specialize_array_empty(self):
         def f(n, m):
-            a = numpy.zeros((n, m))
+            a = numpy.empty((n, m))
             return a
 
         res = interpret(f, [3, 4])
@@ -235,7 +235,7 @@
 
     def test_specialize_array_attr_shape(self):
         def f():
-            a = numpy.zeros((2,3))
+            a = numpy.empty((2,3))
             return list(a.shape)
 
         res = interpret(f, [])
@@ -244,7 +244,7 @@
 
     def test_specialize_array_strides(self):
         def f():
-            a = numpy.zeros((3,4,5))
+            a = numpy.empty((3,4,5))
             return a
 
         res = interpret(f, [])
@@ -264,7 +264,7 @@
 
     def test_specialize_indexing(self):
         def f():
-            a = numpy.zeros((3,), dtype='i')
+            a = numpy.empty((3,), dtype='i')
             b = numpy.array([5,55,555])
             a[:] = b
             return a
@@ -273,9 +273,19 @@
         assert res.data[1] == 55
         assert res.data[2] == 555
 
+    def test_specialize_multi(self):
+        def f(ii, jj):
+            a = numpy.empty((4, 5), dtype='i')
+            for i in range(4):
+                for j in range(5):
+                    a[i, j] = i*j
+            return a[ii, jj]
+        assert interpret(f, [0, 0]) == 0
+        assert interpret(f, [3, 4]) == 12
+
     def test_malloc_remove(self):
         def f():
-            a = numpy.zeros((3,), dtype='i')
+            a = numpy.empty((3,), dtype='i')
             b = numpy.array([5,55,555])
             a[:] = b
             return a
@@ -304,7 +314,7 @@
 
     def test_compile_array_access(self):
         def access_array(index):
-            a = numpy.zeros((3,), dtype='i')
+            a = numpy.empty((3,), dtype='i')
             b = numpy.array([5,55,555])
             a[:] = b
             a[0] = 1
@@ -316,12 +326,14 @@
         assert fn(2) == 555
         
     def test_compile_2d(self):
-        py.test.skip()
-        def access_array(index):
-            a = numpy.zeros((5,6))
-            a[0,0] = 2
-            return 0
+        def f(ii, jj):
+            a = numpy.empty((4, 5), dtype='i')
+            for i in range(4):
+                for j in range(5):
+                    a[i, j] = i*j
+            return a[ii, jj]
 
-        fn = self.compile(access_array, [int])
+        fn = self.compile(f, [int, int])
+        assert fn(2,3) == 2*3
         
 



More information about the Pypy-commit mailing list