[pypy-commit] pypy numpypy-out: add yet another signature, fix ViewIterator.apply_transformations() optimization

[mattip]

Author: mattip
Branch: numpypy-out
Changeset: r52626:f6e33237468c
Date: 2012-02-19 01:34 +0200
http://bitbucket.org/pypy/pypy/changeset/f6e33237468c/

Log:	add yet another signature, fix ViewIterator.apply_transformations()
	optimization

diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -214,7 +214,7 @@
 
     def apply_transformations(self, arr, transformations):
         v = BaseIterator.apply_transformations(self, arr, transformations)
-        if len(arr.shape) == 1:
+        if len(arr.shape) == 1 and len(v.res_shape) == 1:
             return OneDimIterator(self.offset, self.strides[0],
                                   self.res_shape[0])
         return v
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -726,6 +726,11 @@
         ra = ResultArray(self, self.size, self.shape, self.res_dtype,
                                                                 self.res)
         loop.compute(ra)
+        if self.res:
+            broadcast_dims = len(self.res.shape) - len(self.shape)
+            chunks = [Chunk(0,0,0,0)] * broadcast_dims + \
+                     [Chunk(0, i, 1, i) for i in self.shape]
+            return ra.left.create_slice(chunks)
         return ra.left
 
     def force_if_needed(self):
@@ -840,13 +845,19 @@
     def __init__(self, child, size, shape, dtype, res=None, order='C'):
         if res is None:
             res = W_NDimArray(size, shape, dtype, order)
-        assert isinstance(res, BaseArray)
-        concr = res.get_concrete()
-        Call2.__init__(self, None, 'assign', shape, dtype, dtype, concr, child)
+        else:
+            assert isinstance(res, BaseArray)
+            #Make sure it is not a virtual array i.e. out=a+a
+            res = res.get_concrete()
+        Call2.__init__(self, None, 'assign', shape, dtype, dtype, res, child)
 
     def create_sig(self):
-        sig = signature.ResultSignature(self.res_dtype, self.left.create_sig(),
-                                         self.right.create_sig())
+        if self.left.shape != self.right.shape:
+            sig = signature.BroadcastResultSignature(self.res_dtype,
+                        self.left.create_sig(), self.right.create_sig())
+        else:
+            sig = signature.ResultSignature(self.res_dtype, 
+                        self.left.create_sig(), self.right.create_sig())
         return sig
 
 def done_if_true(dtype, val):
diff --git a/pypy/module/micronumpy/signature.py b/pypy/module/micronumpy/signature.py
--- a/pypy/module/micronumpy/signature.py
+++ b/pypy/module/micronumpy/signature.py
@@ -340,7 +340,17 @@
 
         assert isinstance(arr, ResultArray)
         offset = frame.get_final_iter().offset
-        arr.left.setitem(offset, self.right.eval(frame, arr.right))
+        val = self.right.eval(frame, arr.right)
+        arr.left.setitem(offset, val)
+
+class BroadcastResultSignature(ResultSignature):
+    def _create_iter(self, iterlist, arraylist, arr, transforms):
+        from pypy.module.micronumpy.interp_numarray import ResultArray
+
+        assert isinstance(arr, ResultArray)
+        rtransforms = transforms + [BroadcastTransform(arr.left.shape)]
+        self.left._create_iter(iterlist, arraylist, arr.left, transforms)
+        self.right._create_iter(iterlist, arraylist, arr.right, rtransforms)
 
 class BroadcastLeft(Call2):
     def _invent_numbering(self, cache, allnumbers):
diff --git a/pypy/module/micronumpy/test/test_outarg.py b/pypy/module/micronumpy/test/test_outarg.py
--- a/pypy/module/micronumpy/test/test_outarg.py
+++ b/pypy/module/micronumpy/test/test_outarg.py
@@ -68,7 +68,16 @@
         assert (c == out).all()
         assert c.shape == a.shape
         assert c.dtype is a.dtype
+        c[0,0] = 100
+        assert out[0, 0] == 100
         raises(ValueError, 'c = add(a, a, out=out[1])')
+        c = add(a[0], a[1], out=out[1])
+        assert (c == out[1]).all()
+        assert (c == [4, 6]).all()
+        c = add(a[0], a[1], out=out)
+        assert (c == out[1]).all()
+        assert (c == out[0]).all()
+        
 
     def test_ufunc_cast(self):
         from _numpypy import array, negative