[pypy-commit] pypy numpypy-axisops: zjit improvement

mattip noreply at buildbot.pypy.org
Mon Jan 9 21:38:29 CET 2012


Author: mattip
Branch: numpypy-axisops
Changeset: r51179:0722e568f060
Date: 2012-01-09 22:36 +0200
http://bitbucket.org/pypy/pypy/changeset/0722e568f060/

Log:	zjit improvement

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
@@ -22,6 +22,9 @@
     def done(self):
         raise NotImplementedError
 
+    def axis_done(self):
+        raise NotImplementedError
+
 class ArrayIterator(BaseIterator):
     def __init__(self, size):
         self.offset = 0
@@ -120,7 +123,7 @@
         self.shapelen = len(shape)
         self.indices = [0] * len(shape)
         self._done = False
-        self.axis_done = False
+        self._axis_done = False
         self.offset = arr_start
         self.dim = dim
         self.dim_order = []
@@ -136,13 +139,19 @@
     def done(self):
         return self._done
 
+    def axis_done(self):
+        return self._axis_done
+
+    @jit.unroll_safe
     def next(self, shapelen):
         #shapelen will always be one less than self.shapelen
         offset = self.offset
-        axis_done = False
-        indices = [0] * self.shapelen
-        for i in range(self.shapelen):
-            indices[i] = self.indices[i]
+        _axis_done = False
+        done = False
+        #indices = [0] * self.shapelen
+        #for i in range(self.shapelen):
+        #    indices[i] = self.indices[i]
+        indices = self.indices
         for i in self.dim_order:
             if indices[i] < self.shape[i] - 1:
                 indices[i] += 1
@@ -150,13 +159,13 @@
                 break
             else:
                 if i == self.dim:
-                    axis_done = True
+                    _axis_done = True
                 indices[i] = 0
                 offset -= self.backstrides[i]
         else:
-            self._done = True
+            done = True
         res = instantiate(AxisIterator)
-        res.axis_done = axis_done
+        res._axis_done = _axis_done
         res.offset = offset
         res.indices = indices
         res.strides = self.strides
@@ -165,7 +174,7 @@
         res.shape = self.shape
         res.shapelen = self.shapelen
         res.dim = self.dim
-        res._done = self._done
+        res._done = done
         return res
 
 # ------ other iterators that are not part of the computation frame ----------
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
@@ -39,7 +39,7 @@
 axisreduce_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
     virtualizables=['frame'],
-    reds=['identity', 'self','result', 'ri', 'frame', 'nextval', 'dtype', 'value'],
+    reds=['identity', 'self','result', 'ri', 'frame', 'dtype', 'value'],
     get_printable_location=signature.new_printable_location('axisreduce'),
 )
 
@@ -758,8 +758,6 @@
 
 
 class Reduce(VirtualArray):
-    _immutable_fields_ = ['dim', 'binfunc', 'dtype', 'identity']
-
     def __init__(self, binfunc, name, dim, res_dtype, values, identity=None):
         shape = values.shape[0:dim] + values.shape[dim + 1:len(values.shape)]
         VirtualArray.__init__(self, name, shape, res_dtype)
@@ -803,27 +801,27 @@
         ri = ArrayIterator(result.size)
         frame = sig.create_frame(self.values, dim=self.dim)
         value = self.get_identity(sig, frame, shapelen)
-        nextval = sig.eval(frame, self.values).convert_to(dtype)
+        assert isinstance(sig, signature.ReduceSignature)
         while not frame.done():
             axisreduce_driver.jit_merge_point(frame=frame, self=self,
                                           value=value, sig=sig,
                                           shapelen=shapelen, ri=ri,
-                                          nextval=nextval, dtype=dtype,
+                                          dtype=dtype,
                                           identity=identity,
                                           result=result)
-            if frame.iterators[0].axis_done:
+            if frame.axis_done():
+                result.dtype.setitem(result.storage, ri.offset, value)
                 if identity is None:
                     value = sig.eval(frame, self.values).convert_to(dtype)
                     frame.next(shapelen)
                 else:
                     value = identity.convert_to(dtype)
                 ri = ri.next(shapelen)
-            assert isinstance(sig, signature.ReduceSignature)
-            nextval = sig.eval(frame, self.values).convert_to(dtype)
-            value = self.binfunc(dtype, value, nextval)
-            result.dtype.setitem(result.storage, ri.offset, value)
+            value = self.binfunc(dtype, value, 
+                               sig.eval(frame, self.values).convert_to(dtype))
             frame.next(shapelen)
         assert ri.done
+        result.dtype.setitem(result.storage, ri.offset, value)
         return result
 
 
diff --git a/pypy/module/micronumpy/signature.py b/pypy/module/micronumpy/signature.py
--- a/pypy/module/micronumpy/signature.py
+++ b/pypy/module/micronumpy/signature.py
@@ -59,6 +59,12 @@
         for i in range(len(self.iterators)):
             self.iterators[i] = self.iterators[i].next(shapelen)
 
+    def axis_done(self):
+        final_iter = promote(self.final_iter)
+        if final_iter < 0:
+            return False
+        return self.iterators[final_iter].axis_done()
+
 def _add_ptr_to_cache(ptr, cache):
     i = 0
     for p in cache:
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
@@ -127,16 +127,17 @@
     def test_axissum(self):
         result = self.run("axissum")
         assert result == 30
-        self.check_simple_loop({'arraylen_gc': 1,
-                                'call': 1,
-                                'getfield_gc': 3,
-                                "getinteriorfield_raw": 1, 
-                                "guard_class": 1,
-                                "guard_false": 2,
-                                'guard_no_exception': 1,
-                                "float_add": 1,
-                                "jump": 1, 
-                                'setinteriorfield_raw': 1,
+        self.check_simple_loop({\
+                                'setarrayitem_gc': 1,
+                                'getarrayitem_gc': 5,
+                                'getinteriorfield_raw': 1,
+                                'arraylen_gc': 2,
+                                'guard_true': 1,
+                                'int_sub': 1,
+                                'int_lt': 1,
+                                'jump': 1,
+                                'float_add': 1,
+                                'int_add': 2,
                                 })
 
     def define_prod():
@@ -236,7 +237,8 @@
     def test_ufunc(self):
         result = self.run("ufunc")
         assert result == -6
-        self.check_simple_loop({"getinteriorfield_raw": 2, "float_add": 1, "float_neg": 1,
+        self.check_simple_loop({"getinteriorfield_raw": 2, "float_add": 1,
+                                "float_neg": 1,
                                 "setinteriorfield_raw": 1, "int_add": 2,
                                 "int_ge": 1, "guard_false": 1, "jump": 1,
                                 'arraylen_gc': 1})
@@ -346,7 +348,7 @@
         result = self.run("setslice")
         assert result == 11.0
         self.check_trace_count(1)
-        self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add' : 1,
+        self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
                                 'setinteriorfield_raw': 1, 'int_add': 3,
                                 'int_lt': 1, 'guard_true': 1, 'jump': 1,
                                 'arraylen_gc': 3})
@@ -363,11 +365,12 @@
         result = self.run("virtual_slice")
         assert result == 4
         self.check_trace_count(1)
-        self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add' : 1,
+        self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
                                 'setinteriorfield_raw': 1, 'int_add': 2,
                                 'int_ge': 1, 'guard_false': 1, 'jump': 1,
                                 'arraylen_gc': 1})
 
+
 class TestNumpyOld(LLJitMixin):
     def setup_class(cls):
         py.test.skip("old")
@@ -401,4 +404,3 @@
 
         result = self.meta_interp(f, [5], listops=True, backendopt=True)
         assert result == f(5)
-


More information about the pypy-commit mailing list