[pypy-commit] pypy numpy-multidim: some work on broadcast iterator. test_zjit requires some support from

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim
Changeset: r49727:3d9f6f9f6c19
Date: 2011-11-24 15:27 +0200
http://bitbucket.org/pypy/pypy/changeset/3d9f6f9f6c19/

Log:	some work on broadcast iterator. test_zjit requires some support
	from jit targets

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
@@ -211,26 +211,36 @@
             else:
                 self.strides.append(arr.strides[i])
                 self.backstrides.append(arr.backstrides[i])
-        self.shape_len = len(res_shape)
         self.res_shape = res_shape
         self.strides = [0] * (len(res_shape) - len(arr.shape)) + self.strides
         self.backstrides = [0] * (len(res_shape) - len(arr.shape)) + self.backstrides
         self._done = False
-        self.arr = arr
 
     @jit.unroll_safe
     def next(self, shapelen):
+        offset = self.offset
+        indices = [0] * shapelen
+        _done = False
+        for i in range(shapelen):
+            indices[i] = self.indices[i]
         for i in range(shapelen - 1, -1, -1):
-            if self.indices[i] < self.res_shape[i] - 1:
-                self.indices[i] += 1
-                self.offset += self.strides[i]
+            if indices[i] < self.res_shape[i] - 1:
+                indices[i] += 1
+                offset += self.strides[i]
                 break
             else:
-                self.indices[i] = 0
-                self.offset -= self.backstrides[i]
+                indices[i] = 0
+                offset -= self.backstrides[i]
         else:
-            self._done = True
-        return self
+            _done = True
+        res = instantiate(BroadcastIterator)
+        res.indices = indices
+        res.offset = offset
+        res._done = _done
+        res.strides = self.strides
+        res.backstrides = self.backstrides
+        res.res_shape = self.res_shape
+        return res
 
     def done(self):
         return self._done
diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -275,7 +275,20 @@
         # XXX the bridge here is scary. Hopefully jit-targets will fix that,
         #     otherwise it looks kind of good
         self.check_loops({})
-    
+
+    def define_broadcast():
+        return """
+        a = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]
+        b = [1, 2, 3, 4]
+        c = a + b
+        c -> 1 -> 2
+        """
+
+    def test_broadcast(self):
+        result = self.run("broadcast")
+        assert result == 10
+        py.test.skip("improve")
+        self.check_loops({})
 
 class TestNumpyOld(LLJitMixin):
     def setup_class(cls):