[pypy-commit] pypy default: optimize iterators for single dim arrays

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r74800:268850625399
Date: 2014-12-03 18:15 -0500
http://bitbucket.org/pypy/pypy/changeset/268850625399/

Log:	optimize iterators for single dim arrays

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
@@ -159,6 +159,8 @@
         offset = state.offset
         if self.contiguous:
             offset += self.array.dtype.elsize
+        elif self.ndim_m1 == 0:
+            offset += self.strides[0]
         else:
             for i in xrange(self.ndim_m1, -1, -1):
                 idx = indices[i]
@@ -176,6 +178,8 @@
         offset = self.array.start
         if self.contiguous:
             offset += index * self.array.dtype.elsize
+        elif self.ndim_m1 == 0:
+            offset += index * self.strides[0]
         else:
             current = index
             for i in xrange(len(self.shape_m1)):
@@ -188,7 +192,7 @@
         assert state.iterator is self
         assert self.track_index
         indices = state._indices
-        if not self.contiguous:
+        if not (self.contiguous or self.ndim_m1 == 0):
             return indices
         current = state.index
         for i in xrange(len(self.shape_m1)):
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
@@ -102,18 +102,15 @@
         assert result == 3 + 3
         self.check_trace_count(1)
         self.check_simple_loop({
+            'arraylen_gc': 1,
             'float_add': 1,
-            'getarrayitem_gc': 1,
             'guard_false': 1,
             'guard_not_invalidated': 1,
-            'guard_true': 1,
-            'int_add': 5,
+            'int_add': 4,
             'int_ge': 1,
-            'int_lt': 1,
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 1,
         })
 
     def define_pow():
@@ -127,24 +124,22 @@
         assert result == 3 ** 2
         self.check_trace_count(1)
         self.check_simple_loop({
+            'arraylen_gc': 1,
             'call': 1,
             'float_add': 1,
             'float_eq': 3,
             'float_mul': 2,
             'float_ne': 1,
-            'getarrayitem_gc': 1,
             'getarrayitem_raw': 1,     # read the errno
             'guard_false': 4,
             'guard_not_invalidated': 1,
-            'guard_true': 3,
-            'int_add': 5,
+            'guard_true': 2,
+            'int_add': 4,
             'int_ge': 1,
             'int_is_true': 1,
-            'int_lt': 1,
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 1,
             'setarrayitem_raw': 1,     # write the errno
         })
 
@@ -162,18 +157,15 @@
         self.check_trace_count(2)  # extra one for the astype
         del get_stats().loops[0]   # we don't care about it
         self.check_simple_loop({
+            'arraylen_gc': 1,
             'call': 1,
-            'getarrayitem_gc': 1,
             'guard_false': 1,
             'guard_not_invalidated': 1,
-            'guard_true': 1,
-            'int_add': 5,
+            'int_add': 4,
             'int_ge': 1,
-            'int_lt': 1,
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 1,
         })
 
     def define_sum():
@@ -390,18 +382,15 @@
         assert result == 18
         self.check_trace_count(1)
         self.check_simple_loop({
+            'arraylen_gc': 2,
             'float_add': 1,
-            'getarrayitem_gc': 2,
             'guard_false': 1,
             'guard_not_invalidated': 1,
-            'guard_true': 2,
-            'int_add': 6,
+            'int_add': 4,
             'int_ge': 1,
-            'int_lt': 2,
             'jump': 1,
             'raw_load': 2,
             'raw_store': 1,
-            'setarrayitem_gc': 2,
         })
 
     def define_take():
@@ -488,17 +477,14 @@
         assert result == 5.5
         self.check_trace_count(1)
         self.check_simple_loop({
-            'getarrayitem_gc': 1,
+            'arraylen_gc': 1,
             'guard_false': 1,
             'guard_not_invalidated': 1,
-            'guard_true': 1,
-            'int_add': 4,
+            'int_add': 3,
             'int_ge': 1,
-            'int_lt': 1,
             'jump': 1,
             'raw_load': 1,
             'raw_store': 1,
-            'setarrayitem_gc': 1,
         })
 
     def define_virtual_slice():
@@ -618,12 +604,13 @@
             'getarrayitem_gc_pure': 9,
             'getfield_gc_pure': 49,
             'guard_class': 4,
-            'guard_false': 13,
+            'guard_false': 15,
             'guard_not_invalidated': 2,
             'guard_true': 14,
             'int_add': 17,
             'int_ge': 4,
             'int_is_true': 3,
+            'int_is_zero': 2,
             'int_le': 5,
             'int_lt': 8,
             'int_sub': 3,