[pypy-commit] pypy default: simplify iter/dot optimizations

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r69535:27f7c050271b
Date: 2014-02-27 22:20 -0800
http://bitbucket.org/pypy/pypy/changeset/27f7c050271b/

Log:	simplify iter/dot optimizations

diff --git a/pypy/module/micronumpy/concrete.py b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -15,6 +15,8 @@
 
 
 class BaseConcreteArray(object):
+    _immutable_fields_ = ['dtype?', 'storage', 'start', 'size', 'shape[*]',
+                          'strides[*]', 'backstrides[*]', 'order']
     start = 0
     parent = None
 
@@ -350,6 +352,8 @@
                           orig_array)
 
     def set_dtype(self, space, dtype):
+        # size/shape/strides shouldn't change
+        assert dtype.elsize == self.dtype.elsize
         self.dtype = dtype
 
     def argsort(self, space, w_axis):
diff --git a/pypy/module/micronumpy/iterators.py b/pypy/module/micronumpy/iterators.py
--- a/pypy/module/micronumpy/iterators.py
+++ b/pypy/module/micronumpy/iterators.py
@@ -79,29 +79,31 @@
 
 
 class ArrayIter(object):
-    _immutable_fields_ = ['array', 'size', 'ndim_m1', 'shape_m1',
-                          'strides', 'backstrides']
+    _immutable_fields_ = ['array', 'size', 'indices', 'shape[*]',
+                          'strides[*]', 'backstrides[*]']
 
     def __init__(self, array, size, shape, strides, backstrides):
         assert len(shape) == len(strides) == len(backstrides)
         self.array = array
         self.size = size
-        self.ndim_m1 = len(shape) - 1
-        self.shape_m1 = [s - 1 for s in shape]
+        self.indices = [0] * len(shape)
+        self.shape = shape
         self.strides = strides
         self.backstrides = backstrides
         self.reset()
 
+    @jit.unroll_safe
     def reset(self):
         self.index = 0
-        self.indices = [0] * (self.ndim_m1 + 1)
+        for i in xrange(len(self.shape)):
+            self.indices[i] = 0
         self.offset = self.array.start
 
     @jit.unroll_safe
     def next(self):
         self.index += 1
-        for i in xrange(self.ndim_m1, -1, -1):
-            if self.indices[i] < self.shape_m1[i]:
+        for i in xrange(len(self.shape) - 1, -1, -1):
+            if self.indices[i] < self.shape[i] - 1:
                 self.indices[i] += 1
                 self.offset += self.strides[i]
                 break
@@ -115,14 +117,14 @@
         if step == 0:
             return
         self.index += step
-        for i in xrange(self.ndim_m1, -1, -1):
-            if self.indices[i] < (self.shape_m1[i] + 1) - step:
+        for i in xrange(len(self.shape) - 1, -1, -1):
+            if self.indices[i] < self.shape[i] - step:
                 self.indices[i] += step
                 self.offset += self.strides[i] * step
                 break
             else:
-                remaining_step = (self.indices[i] + step) // (self.shape_m1[i] + 1)
-                this_i_step = step - remaining_step * (self.shape_m1[i] + 1)
+                remaining_step = (self.indices[i] + step) // self.shape[i]
+                this_i_step = step - remaining_step * self.shape[i]
                 self.indices[i] = self.indices[i] + this_i_step
                 self.offset += self.strides[i] * this_i_step
                 step = remaining_step
diff --git a/pypy/module/micronumpy/loop.py b/pypy/module/micronumpy/loop.py
--- a/pypy/module/micronumpy/loop.py
+++ b/pypy/module/micronumpy/loop.py
@@ -285,21 +285,18 @@
     outi = result.create_iter()
     lefti = AllButAxisIter(left.implementation, len(left_shape) - 1)
     righti = AllButAxisIter(right.implementation, right_critical_dim)
-    n = left.implementation.shape[-1]
-    s1 = left.implementation.strides[-1]
-    s2 = right.implementation.strides[right_critical_dim]
     while not lefti.done():
         while not righti.done():
             oval = outi.getitem()
             i1 = lefti.offset
             i2 = righti.offset
-            for _ in xrange(n):
+            for _ in xrange(left.implementation.shape[-1]):
                 dot_driver.jit_merge_point(dtype=dtype)
                 lval = left.implementation.getitem(i1).convert_to(space, dtype)
                 rval = right.implementation.getitem(i2).convert_to(space, dtype)
                 oval = dtype.itemtype.add(oval, dtype.itemtype.mul(lval, rval))
-                i1 += s1
-                i2 += s2
+                i1 += left.implementation.strides[-1]
+                i2 += right.implementation.strides[right_critical_dim]
             outi.setitem(oval)
             outi.next()
             righti.next()
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
@@ -59,6 +59,7 @@
         if self.graph is None:
             interp, graph = self.meta_interp(f, [0],
                                              listops=True,
+                                             listcomp=True,
                                              backendopt=True,
                                              graph_and_interp_only=True)
             self.__class__.interp = interp