[pypy-commit] pypy numpy-speed: virtualize the array iter in call2

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: numpy-speed
Changeset: r70708:958f25449c81
Date: 2014-04-17 02:48 -0400
http://bitbucket.org/pypy/pypy/changeset/958f25449c81/

Log:	virtualize the array iter in call2

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,8 +79,9 @@
 
 
 class ArrayIter(object):
+    _virtualizable_ = ['index', 'indices[*]', 'offset']
     _immutable_fields_ = ['array', 'size', 'ndim_m1', 'shape_m1[*]',
-                          'strides[*]', 'backstrides[*]']
+                          'strides[*]', 'backstrides[*]', 'indices']
 
     def __init__(self, array, size, shape, strides, backstrides):
         assert len(shape) == len(strides) == len(backstrides)
@@ -90,11 +91,16 @@
         self.shape_m1 = [s - 1 for s in shape]
         self.strides = strides
         self.backstrides = backstrides
-        self.reset()
 
+        self.index = 0
+        self.indices = [0] * len(shape)
+        self.offset = array.start
+
+    @jit.unroll_safe
     def reset(self):
         self.index = 0
-        self.indices = [0] * len(self.shape_m1)
+        for i in xrange(self.ndim_m1, -1, -1):
+            self.indices[i] = 0
         self.offset = self.array.start
 
     @jit.unroll_safe
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
@@ -16,7 +16,8 @@
                              greens = ['shapelen', 'func', 'calc_dtype',
                                        'res_dtype'],
                              reds = ['shape', 'w_lhs', 'w_rhs', 'out',
-                                     'left_iter', 'right_iter', 'out_iter'])
+                                     'left_iter', 'right_iter', 'out_iter'],
+                             virtualizables=['out_iter'])
 
 def call2(space, shape, func, calc_dtype, res_dtype, w_lhs, w_rhs, out):
     # handle array_priority
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
@@ -99,9 +99,11 @@
         assert result == 3 + 3
         self.check_trace_count(1)
         self.check_simple_loop({
+            'arraylen_gc': 2,
+            'cond_call': 2,
             'float_add': 1,
-            'getarrayitem_gc': 3,
-            'getfield_gc': 7,
+            'getarrayitem_gc': 2,
+            'getfield_gc': 4,
             'guard_false': 1,
             'guard_not_invalidated': 1,
             'guard_true': 3,
@@ -111,8 +113,8 @@
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 3,
-            'setfield_gc': 6,
+            'setarrayitem_gc': 2,
+            'setfield_gc': 4,
         })
 
     def define_pow():
@@ -126,13 +128,15 @@
         assert result == 3 ** 2
         self.check_trace_count(1)
         self.check_simple_loop({
+            'arraylen_gc': 2,
             'call': 3,
+            'cond_call': 2,
             'float_add': 1,
             'float_eq': 3,
             'float_mul': 2,
             'float_ne': 1,
-            'getarrayitem_gc': 3,
-            'getfield_gc': 7,
+            'getarrayitem_gc': 2,
+            'getfield_gc': 4,
             'guard_false': 4,
             'guard_not_invalidated': 1,
             'guard_true': 5,
@@ -143,8 +147,8 @@
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 3,
-            'setfield_gc': 6,
+            'setarrayitem_gc': 2,
+            'setfield_gc': 4,
         })
 
     def define_sum():
@@ -526,6 +530,7 @@
         assert result == 1.0
         self.check_trace_count(1)
         self.check_simple_loop({
+            'cond_call': 2,
             'getarrayitem_gc': 2,
             'getfield_gc': 4,
             'guard_not_invalidated': 1,
@@ -564,12 +569,12 @@
             'raw_load': 2,
         })
         self.check_resops({
-            'arraylen_gc': 1,
+            'cond_call': 6,
             'float_add': 2,
             'float_mul': 2,
             'getarrayitem_gc': 7,
             'getarrayitem_gc_pure': 15,
-            'getfield_gc': 35,
+            'getfield_gc': 41,
             'getfield_gc_pure': 39,
             'guard_class': 4,
             'guard_false': 14,
@@ -584,11 +589,11 @@
             'int_lt': 11,
             'int_sub': 4,
             'jump': 3,
-            'new_array': 1,
+            'ptr_ne': 6,
             'raw_load': 6,
             'raw_store': 1,
-            'setarrayitem_gc': 8,
-            'setfield_gc': 15,
+            'setarrayitem_gc': 10,
+            'setfield_gc': 14,
         })
 
     def define_argsort():