[pypy-commit] pypy numppy-flatitter: rework next_skip_x, there still is an AnnotationError

[mattip]

Author: mattip
Branch: numppy-flatitter
Changeset: r51806:6d23b2d09136
Date: 2012-01-19 07:29 +0200
http://bitbucket.org/pypy/pypy/changeset/6d23b2d09136/

Log:	rework next_skip_x, there still is an AnnotationError

diff --git a/pypy/module/micronumpy/REVIEW b/pypy/module/micronumpy/REVIEW
deleted file mode 100644
--- a/pypy/module/micronumpy/REVIEW
+++ /dev/null
@@ -1,25 +0,0 @@
-* this file should go away before asking for a review :)
-
-* I think we should wait for indexing-by-arrays-2, since this would clean up
-  the iterator interface
-
-* For getitem, we need to reuse parent getitem, but with some hook
-  that recomputes indexes. That hook would be used for any sort of access,
-  be it slices or be it integers, but in general I would like to avoid code
-  duplication, since the indexing is getting slowly fairly complex.
-
-* iterating over a transposed array still fails.
-
-* next_skip_x is implemented badly. It still iterates over items (calls
-  next_skip n times) and is jit.unroll_safe. This means that the JIT will assume
-  skip is a constant and compile a linear trace unrolling all iterations. This
-  is *BAD*. Instead it would be something like:
-
-  for arrays:
-
-  res.offset = self.offset + skip
-
-  for views something slightly more complex that does appropriate modulo.
-
-  Note that next() has unroll_safe, because the loop only dependas on the
-  length of shape which is a constant for assembler.
\ No newline at end of file
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
@@ -51,9 +51,9 @@
         self.size = size
 
     def next(self, shapelen):
-        return self._next(1)
+        return self.next_skip_x(1)
 
-    def _next(self, ofs):
+    def next_skip_x(self, ofs):
         arr = instantiate(ArrayIterator)
         arr.size = self.size
         arr.offset = self.offset + ofs
@@ -61,7 +61,7 @@
 
     def next_no_increase(self, shapelen):
         # a hack to make JIT believe this is always virtual
-        return self._next(0)
+        return self.next_skip_x(0)
 
     def done(self):
         return self.offset >= self.size
@@ -133,11 +133,34 @@
         res._done = done
         return res
 
-    @jit.unroll_safe
+    #@jit.unroll_safe
     def next_skip_x(self, shapelen, step):
-        res = self.next(shapelen)
-        for x in range(step-1):
-            res = res.next(shapelen)
+        shapelen = jit.promote(len(self.res_shape))
+        offset = self.offset
+        indices = [0] * shapelen
+        for i in range(shapelen):
+            indices[i] = self.indices[i]
+        done = False
+        for i in range(shapelen - 1, -1, -1):
+            if indices[i] < self.res_shape[i] - step:
+                indices[i] += step
+                offset += self.strides[i] * step
+                break
+            else:
+                remaining_step = (indices[i] + step) // self.res_shape[i]
+                this_i_step = step - remaining_step * self.res_shape[i]
+                offset += self.strides[i] * this_i_step
+                indices[i] = indices[i] +  this_i_step
+                step = remaining_step
+        else:
+            done = True
+        res = instantiate(ViewIterator)
+        res.offset = offset
+        res.indices = indices
+        res.strides = self.strides
+        res.backstrides = self.backstrides
+        res.res_shape = self.res_shape
+        res._done = done
         return res
 
     def apply_transformations(self, arr, transformations):
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
@@ -1484,8 +1484,7 @@
         basei = ViewIterator(self.base.start, self.base.strides,
                                self.base.backstrides,self.base.shape)
         shapelen = len(self.base.shape)
-        if start > 0:
-            basei = basei.next_skip_x(shapelen, start)
+        basei = basei.next_skip_x(shapelen, start)
         ri = ArrayIterator(lngth)
         while not ri.done():
             # TODO: add a jit_merge_point
@@ -1502,8 +1501,7 @@
         basei = ViewIterator(self.base.start, self.base.strides,
                                self.base.backstrides,self.base.shape)
         shapelen = len(self.base.shape)
-        if start > 0:
-            basei = basei.next_skip_x(shapelen, start)
+        basei = basei.next_skip_x(shapelen, start)
         for i in range(lngth):
             # TODO: add jit_merge_point
             v = arr.getitem(ai).convert_to(self.base.dtype)