[pypy-commit] pypy numpypy-nditer: progress toward extern_loop

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: numpypy-nditer
Changeset: r64819:730edc51d980
Date: 2013-06-07 17:19 +0300
http://bitbucket.org/pypy/pypy/changeset/730edc51d980/

Log:	progress toward extern_loop

diff --git a/pypy/module/micronumpy/interp_nditer.py b/pypy/module/micronumpy/interp_nditer.py
--- a/pypy/module/micronumpy/interp_nditer.py
+++ b/pypy/module/micronumpy/interp_nditer.py
@@ -4,7 +4,7 @@
 from pypy.interpreter.error import OperationError
 from pypy.module.micronumpy.base import W_NDimArray, convert_to_array
 from pypy.module.micronumpy.strides import (calculate_broadcast_strides,
-                                             shape_agreement_multiple)
+                                          shape_agreement_multiple, calc_steps)
 from pypy.module.micronumpy.iter import MultiDimViewIterator, SliceIterator
 from pypy.module.micronumpy import support
 from pypy.module.micronumpy.arrayimpl.concrete import SliceArray
@@ -229,11 +229,17 @@
         self.iters=[]
         self.shape = iter_shape = shape_agreement_multiple(space, self.seq)
         if self.external_loop:
-            #XXX find longest contiguous shape
-            iter_shape = iter_shape[1:]
+            steps = []
+            for seq in self.seq:
+                impl = seq.implementation
+                steps.append(calc_steps(impl.shape, impl.strides, self.order))
+            #XXX #find longest contiguous shape
+            print 'steps',steps,'tier_shape',iter_shape
+            iter_shape = [1]
         for i in range(len(self.seq)):
             self.iters.append(BoxIterator(get_iter(space, self.order,
-                            self.seq[i].implementation, iter_shape), self.op_flags[i]))
+                            self.seq[i].implementation, iter_shape),
+                            self.op_flags[i]))
 
     def descr_iter(self, space):
         return space.wrap(self)
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -51,8 +51,8 @@
             rstrides.append(strides[i])
             rbackstrides.append(backstrides[i])
     if backwards:
-        rstrides = rstrides + [0] * (len(res_shape) - len(orig_shape))  
-        rbackstrides = rbackstrides + [0] * (len(res_shape) - len(orig_shape)) 
+        rstrides = rstrides + [0] * (len(res_shape) - len(orig_shape))
+        rbackstrides = rbackstrides + [0] * (len(res_shape) - len(orig_shape))
     else:
         rstrides = [0] * (len(res_shape) - len(orig_shape)) + rstrides
         rbackstrides = [0] * (len(res_shape) - len(orig_shape)) + rbackstrides
@@ -62,7 +62,7 @@
     if (is_rec_type and space.isinstance_w(w_elem, space.w_tuple)):
         return True
     if (space.isinstance_w(w_elem, space.w_tuple) or
-        isinstance(w_elem, W_NDimArray) or    
+        isinstance(w_elem, W_NDimArray) or
         space.isinstance_w(w_elem, space.w_list)):
         return False
     return True
@@ -293,3 +293,27 @@
             rbackstrides[i] = backstrides[j]
             j += 1
     return rstrides, rbackstrides
+
+ at jit.unroll_safe
+def calc_steps(shape, strides, order='C'):
+    steps = []
+    if order == 'K':
+        if strides[0] < strides[-1]:
+            order = 'F'
+        else:
+            order = 'C'
+    if order == 'F' or order == 'A':
+        last_step = strides[0]
+        for i in range(len(shape)):
+            steps.append(strides[i] / last_step)
+            last_step *= shape[i]
+        if order == 'A':
+            pass
+            #XXX test for all(steps==steps[0])
+    elif order == 'C':
+        last_step = strides[-1]
+        for i in range(len(shape) - 1, -1, -1):
+            steps.insert(0, strides[i] / last_step)
+            last_step *= shape[i]
+    return steps
+