[pypy-commit] pypy default: Merge numpy-single-jitdriver. This branch refactors a bit jitdrivers around so

fijal noreply at buildbot.pypy.org
Fri Feb 3 11:46:57 CET 2012


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r52053:06c10732a5bd
Date: 2012-02-03 12:45 +0200
http://bitbucket.org/pypy/pypy/changeset/06c10732a5bd/

Log:	Merge numpy-single-jitdriver. This branch refactors a bit jitdrivers
	around so we have (mostly) one important jitdriver. This would be
	useful for future optimizations like vectorizing. The rest can be
	incorporated, but "later" if at all, since those are not amenable to
	"easy' vectorization.

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -98,6 +98,10 @@
         ('bitwise_not', 'invert'),
         ('isnan', 'isnan'),
         ('isinf', 'isinf'),
+        ('logical_and', 'logical_and'),
+        ('logical_xor', 'logical_xor'),
+        ('logical_not', 'logical_not'),
+        ('logical_or', 'logical_or'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 
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
@@ -86,8 +86,9 @@
 
     def apply_transformations(self, arr, transformations):
         v = self
-        for transform in transformations:
-            v = v.transform(arr, transform)
+        if transformations is not None:
+            for transform in transformations:
+                v = v.transform(arr, transform)
         return v
 
     def transform(self, arr, t):
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
@@ -3,7 +3,7 @@
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.interpreter.typedef import TypeDef, GetSetProperty
 from pypy.module.micronumpy import (interp_ufuncs, interp_dtype, interp_boxes,
-    signature, support)
+    signature, support, loop)
 from pypy.module.micronumpy.strides import (calculate_slice_strides,
     shape_agreement, find_shape_and_elems, get_shape_from_iterable,
     calc_new_strides, to_coords)
@@ -12,39 +12,11 @@
 from pypy.rpython.lltypesystem import lltype, rffi
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
-from pypy.module.micronumpy.interp_iter import (ArrayIterator, OneDimIterator,
+from pypy.module.micronumpy.interp_iter import (ArrayIterator,
     SkipLastAxisIterator, Chunk, ViewIterator)
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
 
 
-numpy_driver = jit.JitDriver(
-    greens=['shapelen', 'sig'],
-    virtualizables=['frame'],
-    reds=['result_size', 'frame', 'ri', 'self', 'result'],
-    get_printable_location=signature.new_printable_location('numpy'),
-    name='numpy',
-)
-all_driver = jit.JitDriver(
-    greens=['shapelen', 'sig'],
-    virtualizables=['frame'],
-    reds=['frame', 'self', 'dtype'],
-    get_printable_location=signature.new_printable_location('all'),
-    name='numpy_all',
-)
-any_driver = jit.JitDriver(
-    greens=['shapelen', 'sig'],
-    virtualizables=['frame'],
-    reds=['frame', 'self', 'dtype'],
-    get_printable_location=signature.new_printable_location('any'),
-    name='numpy_any',
-)
-slice_driver = jit.JitDriver(
-    greens=['shapelen', 'sig'],
-    virtualizables=['frame'],
-    reds=['self', 'frame', 'arr'],
-    get_printable_location=signature.new_printable_location('slice'),
-    name='numpy_slice',
-)
 count_driver = jit.JitDriver(
     greens=['shapelen'],
     virtualizables=['frame'],
@@ -173,6 +145,8 @@
     descr_prod = _reduce_ufunc_impl("multiply", True)
     descr_max = _reduce_ufunc_impl("maximum")
     descr_min = _reduce_ufunc_impl("minimum")
+    descr_all = _reduce_ufunc_impl('logical_and')
+    descr_any = _reduce_ufunc_impl('logical_or')
 
     def _reduce_argmax_argmin_impl(op_name):
         reduce_driver = jit.JitDriver(
@@ -212,40 +186,6 @@
             return space.wrap(loop(self))
         return func_with_new_name(impl, "reduce_arg%s_impl" % op_name)
 
-    def _all(self):
-        dtype = self.find_dtype()
-        sig = self.find_sig()
-        frame = sig.create_frame(self)
-        shapelen = len(self.shape)
-        while not frame.done():
-            all_driver.jit_merge_point(sig=sig,
-                                       shapelen=shapelen, self=self,
-                                       dtype=dtype, frame=frame)
-            if not dtype.itemtype.bool(sig.eval(frame, self)):
-                return False
-            frame.next(shapelen)
-        return True
-
-    def descr_all(self, space):
-        return space.wrap(self._all())
-
-    def _any(self):
-        dtype = self.find_dtype()
-        sig = self.find_sig()
-        frame = sig.create_frame(self)
-        shapelen = len(self.shape)
-        while not frame.done():
-            any_driver.jit_merge_point(sig=sig, frame=frame,
-                                       shapelen=shapelen, self=self,
-                                       dtype=dtype)
-            if dtype.itemtype.bool(sig.eval(frame, self)):
-                return True
-            frame.next(shapelen)
-        return False
-
-    def descr_any(self, space):
-        return space.wrap(self._any())
-
     descr_argmax = _reduce_argmax_argmin_impl("max")
     descr_argmin = _reduce_argmax_argmin_impl("min")
 
@@ -685,6 +625,9 @@
         raise OperationError(space.w_NotImplementedError, space.wrap(
             "non-int arg not supported"))
 
+    def compute_first_step(self, sig, frame):
+        pass
+
 def convert_to_array(space, w_obj):
     if isinstance(w_obj, BaseArray):
         return w_obj
@@ -750,22 +693,9 @@
         raise NotImplementedError
 
     def compute(self):
-        result = W_NDimArray(self.size, self.shape, self.find_dtype())
-        shapelen = len(self.shape)
-        sig = self.find_sig()
-        frame = sig.create_frame(self)
-        ri = ArrayIterator(self.size)
-        while not ri.done():
-            numpy_driver.jit_merge_point(sig=sig,
-                                         shapelen=shapelen,
-                                         result_size=self.size,
-                                         frame=frame,
-                                         ri=ri,
-                                         self=self, result=result)
-            result.setitem(ri.offset, sig.eval(frame, self))
-            frame.next(shapelen)
-            ri = ri.next(shapelen)
-        return result
+        ra = ResultArray(self, self.size, self.shape, self.res_dtype)
+        loop.compute(ra)
+        return ra.left
 
     def force_if_needed(self):
         if self.forced_result is None:
@@ -823,7 +753,8 @@
     def create_sig(self):
         if self.forced_result is not None:
             return self.forced_result.create_sig()
-        return signature.Call1(self.ufunc, self.name, self.values.create_sig())
+        return signature.Call1(self.ufunc, self.name, self.calc_dtype,
+                               self.values.create_sig())
 
 class Call2(VirtualArray):
     """
@@ -864,6 +795,66 @@
         return signature.Call2(self.ufunc, self.name, self.calc_dtype,
                                self.left.create_sig(), self.right.create_sig())
 
+class ResultArray(Call2):
+    def __init__(self, child, size, shape, dtype, res=None, order='C'): 
+        if res is None:
+            res = W_NDimArray(size, shape, dtype, order)
+        Call2.__init__(self, None, 'assign', shape, dtype, dtype, res, child)
+
+    def create_sig(self):
+        return signature.ResultSignature(self.res_dtype, self.left.create_sig(),
+                                         self.right.create_sig())
+
+def done_if_true(dtype, val):
+    return dtype.itemtype.bool(val)
+
+def done_if_false(dtype, val):
+    return not dtype.itemtype.bool(val)
+
+class ReduceArray(Call2):
+    def __init__(self, func, name, identity, child, dtype):
+        self.identity = identity
+        Call2.__init__(self, func, name, [1], dtype, dtype, None, child)
+
+    def compute_first_step(self, sig, frame):
+        assert isinstance(sig, signature.ReduceSignature)
+        if self.identity is None:
+            frame.cur_value = sig.right.eval(frame, self.right).convert_to(
+                self.calc_dtype)
+            frame.next(len(self.right.shape))
+        else:
+            frame.cur_value = self.identity.convert_to(self.calc_dtype)
+    
+    def create_sig(self):
+        if self.name == 'logical_and':
+            done_func = done_if_false
+        elif self.name == 'logical_or':
+            done_func = done_if_true
+        else:
+            done_func = None
+        return signature.ReduceSignature(self.ufunc, self.name, self.res_dtype,
+                                 signature.ScalarSignature(self.res_dtype),
+                                         self.right.create_sig(), done_func)
+
+class AxisReduce(Call2):
+    _immutable_fields_ = ['left', 'right']
+
+    def __init__(self, ufunc, name, identity, shape, dtype, left, right, dim):
+        Call2.__init__(self, ufunc, name, shape, dtype, dtype,
+                       left, right)
+        self.dim = dim
+        self.identity = identity
+
+    def compute_first_step(self, sig, frame):
+        if self.identity is not None:
+            frame.identity = self.identity.convert_to(self.calc_dtype)
+
+    def create_sig(self):
+        return signature.AxisReduceSignature(self.ufunc, self.name,
+                                             self.res_dtype,
+                                 signature.ScalarSignature(self.res_dtype),
+                                             self.right.create_sig())
+
 class SliceArray(Call2):
     def __init__(self, shape, dtype, left, right, no_broadcast=False):
         self.no_broadcast = no_broadcast
@@ -882,18 +873,6 @@
                                             self.calc_dtype,
                                             lsig, rsig)
 
-class AxisReduce(Call2):
-    """ NOTE: this is only used as a container, you should never
-    encounter such things in the wild. Remove this comment
-    when we'll make AxisReduce lazy
-    """
-    _immutable_fields_ = ['left', 'right']
-
-    def __init__(self, ufunc, name, shape, dtype, left, right, dim):
-        Call2.__init__(self, ufunc, name, shape, dtype, dtype,
-                       left, right)
-        self.dim = dim
-
 class ConcreteArray(BaseArray):
     """ An array that have actual storage, whether owned or not
     """
@@ -979,7 +958,7 @@
             self._fast_setslice(space, w_value)
         else:
             arr = SliceArray(self.shape, self.dtype, self, w_value)
-            self._sliceloop(arr)
+            loop.compute(arr)
 
     def _fast_setslice(self, space, w_value):
         assert isinstance(w_value, ConcreteArray)
@@ -1003,17 +982,6 @@
                 source.next()
                 dest.next()
 
-    def _sliceloop(self, arr):
-        sig = arr.find_sig()
-        frame = sig.create_frame(arr)
-        shapelen = len(self.shape)
-        while not frame.done():
-            slice_driver.jit_merge_point(sig=sig, frame=frame, self=self,
-                                         arr=arr,
-                                         shapelen=shapelen)
-            sig.eval(frame, arr)
-            frame.next(shapelen)
-
     def copy(self, space):
         array = W_NDimArray(self.size, self.shape[:], self.dtype, self.order)
         array.setslice(space, self)
@@ -1039,9 +1007,9 @@
                            parent.order, parent)
         self.start = start
 
-    def create_iter(self):
+    def create_iter(self, transforms=None):
         return ViewIterator(self.start, self.strides, self.backstrides,
-                            self.shape)
+                            self.shape).apply_transformations(self, transforms)
 
     def setshape(self, space, new_shape):
         if len(self.shape) < 1:
@@ -1090,8 +1058,8 @@
         self.shape = new_shape
         self.calc_strides(new_shape)
 
-    def create_iter(self):
-        return ArrayIterator(self.size)
+    def create_iter(self, transforms=None):
+        return ArrayIterator(self.size).apply_transformations(self, transforms)
 
     def create_sig(self):
         return signature.ArraySignature(self.dtype)
@@ -1427,6 +1395,12 @@
     def create_sig(self):
         return signature.FlatSignature(self.base.dtype)
 
+    def create_iter(self, transforms=None):
+        return ViewIterator(self.base.start, self.base.strides, 
+                    self.base.backstrides,
+                    self.base.shape).apply_transformations(self.base,
+                                                           transforms)
+
     def descr_base(self, space):
         return space.wrap(self.base)
 
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -2,31 +2,10 @@
 from pypy.interpreter.error import OperationError, operationerrfmt
 from pypy.interpreter.gateway import interp2app, unwrap_spec, NoneNotWrapped
 from pypy.interpreter.typedef import TypeDef, GetSetProperty, interp_attrproperty
-from pypy.module.micronumpy import interp_boxes, interp_dtype, support
-from pypy.module.micronumpy.signature import (ReduceSignature, find_sig,
-    new_printable_location, AxisReduceSignature, ScalarSignature)
-from pypy.rlib import jit
+from pypy.module.micronumpy import interp_boxes, interp_dtype, support, loop
 from pypy.rlib.rarithmetic import LONG_BIT
 from pypy.tool.sourcetools import func_with_new_name
 
-
-reduce_driver = jit.JitDriver(
-    greens=['shapelen', "sig"],
-    virtualizables=["frame"],
-    reds=["frame", "self", "dtype", "value", "obj"],
-    get_printable_location=new_printable_location('reduce'),
-    name='numpy_reduce',
-)
-
-axisreduce_driver = jit.JitDriver(
-    greens=['shapelen', 'sig'],
-    virtualizables=['frame'],
-    reds=['self','arr', 'identity', 'frame'],
-    name='numpy_axisreduce',
-    get_printable_location=new_printable_location('axisreduce'),
-)
-
-
 class W_Ufunc(Wrappable):
     _attrs_ = ["name", "promote_to_float", "promote_bools", "identity"]
     _immutable_fields_ = ["promote_to_float", "promote_bools", "name"]
@@ -140,7 +119,7 @@
     def reduce(self, space, w_obj, multidim, promote_to_largest, dim,
                keepdims=False):
         from pypy.module.micronumpy.interp_numarray import convert_to_array, \
-                                                           Scalar
+                                                           Scalar, ReduceArray
         if self.argcount != 2:
             raise OperationError(space.w_ValueError, space.wrap("reduce only "
                 "supported for binary functions"))
@@ -151,96 +130,37 @@
         if isinstance(obj, Scalar):
             raise OperationError(space.w_TypeError, space.wrap("cannot reduce "
                 "on a scalar"))
-
         size = obj.size
-        dtype = find_unaryop_result_dtype(
-            space, obj.find_dtype(),
-            promote_to_float=self.promote_to_float,
-            promote_to_largest=promote_to_largest,
-            promote_bools=True
-        )
+        if self.comparison_func:
+            dtype = interp_dtype.get_dtype_cache(space).w_booldtype
+        else:
+            dtype = find_unaryop_result_dtype(
+                space, obj.find_dtype(),
+                promote_to_float=self.promote_to_float,
+                promote_to_largest=promote_to_largest,
+                promote_bools=True
+            )
         shapelen = len(obj.shape)
         if self.identity is None and size == 0:
             raise operationerrfmt(space.w_ValueError, "zero-size array to "
                     "%s.reduce without identity", self.name)
         if shapelen > 1 and dim >= 0:
-            res = self.do_axis_reduce(obj, dtype, dim, keepdims)
-            return space.wrap(res)
-        scalarsig = ScalarSignature(dtype)
-        sig = find_sig(ReduceSignature(self.func, self.name, dtype,
-                                       scalarsig,
-                                       obj.create_sig()), obj)
-        frame = sig.create_frame(obj)
-        if self.identity is None:
-            value = sig.eval(frame, obj).convert_to(dtype)
-            frame.next(shapelen)
-        else:
-            value = self.identity.convert_to(dtype)
-        return self.reduce_loop(shapelen, sig, frame, value, obj, dtype)
+            return self.do_axis_reduce(obj, dtype, dim, keepdims)
+        arr = ReduceArray(self.func, self.name, self.identity, obj, dtype)
+        return loop.compute(arr)
 
     def do_axis_reduce(self, obj, dtype, dim, keepdims):
         from pypy.module.micronumpy.interp_numarray import AxisReduce,\
              W_NDimArray
-
         if keepdims:
             shape = obj.shape[:dim] + [1] + obj.shape[dim + 1:]
         else:
             shape = obj.shape[:dim] + obj.shape[dim + 1:]
         result = W_NDimArray(support.product(shape), shape, dtype)
-        rightsig = obj.create_sig()
-        # note - this is just a wrapper so signature can fetch
-        #        both left and right, nothing more, especially
-        #        this is not a true virtual array, because shapes
-        #        don't quite match
-        arr = AxisReduce(self.func, self.name, obj.shape, dtype,
+        arr = AxisReduce(self.func, self.name, self.identity, obj.shape, dtype,
                          result, obj, dim)
-        scalarsig = ScalarSignature(dtype)
-        sig = find_sig(AxisReduceSignature(self.func, self.name, dtype,
-                                           scalarsig, rightsig), arr)
-        assert isinstance(sig, AxisReduceSignature)
-        frame = sig.create_frame(arr)
-        shapelen = len(obj.shape)
-        if self.identity is not None:
-            identity = self.identity.convert_to(dtype)
-        else:
-            identity = None
-        self.reduce_axis_loop(frame, sig, shapelen, arr, identity)
-        return result
-
-    def reduce_axis_loop(self, frame, sig, shapelen, arr, identity):
-        # note - we can be advanterous here, depending on the exact field
-        # layout. For now let's say we iterate the original way and
-        # simply follow the original iteration order
-        while not frame.done():
-            axisreduce_driver.jit_merge_point(frame=frame, self=self,
-                                              sig=sig,
-                                              identity=identity,
-                                              shapelen=shapelen, arr=arr)
-            iterator = frame.get_final_iter()
-            v = sig.eval(frame, arr).convert_to(sig.calc_dtype)
-            if iterator.first_line:
-                if identity is not None:
-                    value = self.func(sig.calc_dtype, identity, v)
-                else:
-                    value = v
-            else:
-                cur = arr.left.getitem(iterator.offset)
-                value = self.func(sig.calc_dtype, cur, v)
-            arr.left.setitem(iterator.offset, value)
-            frame.next(shapelen)
-
-    def reduce_loop(self, shapelen, sig, frame, value, obj, dtype):
-        while not frame.done():
-            reduce_driver.jit_merge_point(sig=sig,
-                                          shapelen=shapelen, self=self,
-                                          value=value, obj=obj, frame=frame,
-                                          dtype=dtype)
-            assert isinstance(sig, ReduceSignature)
-            value = sig.binfunc(dtype, value,
-                                sig.eval(frame, obj).convert_to(dtype))
-            frame.next(shapelen)
-        return value
-
+        loop.compute(arr)
+        return arr.left
 
 class W_Ufunc1(W_Ufunc):
     argcount = 1
@@ -312,7 +232,6 @@
                 w_lhs.value.convert_to(calc_dtype),
                 w_rhs.value.convert_to(calc_dtype)
             ))
-
         new_shape = shape_agreement(space, w_lhs.shape, w_rhs.shape)
         w_res = Call2(self.func, self.name,
                       new_shape, calc_dtype,
@@ -482,6 +401,13 @@
             ("isnan", "isnan", 1, {"bool_result": True}),
             ("isinf", "isinf", 1, {"bool_result": True}),
 
+            ('logical_and', 'logical_and', 2, {'comparison_func': True,
+                                               'identity': 1}),
+            ('logical_or', 'logical_or', 2, {'comparison_func': True,
+                                             'identity': 0}),
+            ('logical_xor', 'logical_xor', 2, {'comparison_func': True}),
+            ('logical_not', 'logical_not', 1, {'bool_result': True}),
+
             ("maximum", "max", 2),
             ("minimum", "min", 2),
 
diff --git a/pypy/module/micronumpy/loop.py b/pypy/module/micronumpy/loop.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/micronumpy/loop.py
@@ -0,0 +1,83 @@
+
+""" This file is the main run loop as well as evaluation loops for various
+signatures
+"""
+
+from pypy.rlib.jit import JitDriver, hint, unroll_safe, promote
+from pypy.module.micronumpy.interp_iter import ConstantIterator
+
+class NumpyEvalFrame(object):
+    _virtualizable2_ = ['iterators[*]', 'final_iter', 'arraylist[*]',
+                        'value', 'identity', 'cur_value']
+
+    @unroll_safe
+    def __init__(self, iterators, arrays):
+        self = hint(self, access_directly=True, fresh_virtualizable=True)
+        self.iterators = iterators[:]
+        self.arrays = arrays[:]
+        for i in range(len(self.iterators)):
+            iter = self.iterators[i]
+            if not isinstance(iter, ConstantIterator):
+                self.final_iter = i
+                break
+        else:
+            self.final_iter = -1
+        self.cur_value = None
+        self.identity = None
+
+    def done(self):
+        final_iter = promote(self.final_iter)
+        if final_iter < 0:
+            assert False
+        return self.iterators[final_iter].done()
+
+    @unroll_safe
+    def next(self, shapelen):
+        for i in range(len(self.iterators)):
+            self.iterators[i] = self.iterators[i].next(shapelen)
+
+    @unroll_safe
+    def next_from_second(self, shapelen):
+        """ Don't increase the first iterator
+        """
+        for i in range(1, len(self.iterators)):
+            self.iterators[i] = self.iterators[i].next(shapelen)
+
+    def next_first(self, shapelen):
+        self.iterators[0] = self.iterators[0].next(shapelen)
+
+    def get_final_iter(self):
+        final_iter = promote(self.final_iter)
+        if final_iter < 0:
+            assert False
+        return self.iterators[final_iter]
+
+def get_printable_location(shapelen, sig):
+    return 'numpy ' + sig.debug_repr() + ' [%d dims]' % (shapelen,)
+
+numpy_driver = JitDriver(
+    greens=['shapelen', 'sig'],
+    virtualizables=['frame'],
+    reds=['frame', 'arr'],
+    get_printable_location=get_printable_location,
+    name='numpy',
+)
+
+class ComputationDone(Exception):
+    def __init__(self, value):
+        self.value = value
+
+def compute(arr):
+    sig = arr.find_sig()
+    shapelen = len(arr.shape)
+    frame = sig.create_frame(arr)
+    try:
+        while not frame.done():
+            numpy_driver.jit_merge_point(sig=sig,
+                                         shapelen=shapelen,
+                                         frame=frame, arr=arr)
+            sig.eval(frame, arr)
+            frame.next(shapelen)
+        return frame.cur_value
+    except ComputationDone, e:
+        return e.value
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
@@ -1,9 +1,9 @@
 from pypy.rlib.objectmodel import r_dict, compute_identity_hash, compute_hash
 from pypy.rlib.rarithmetic import intmask
-from pypy.module.micronumpy.interp_iter import ViewIterator, ArrayIterator, \
-     ConstantIterator, AxisIterator, ViewTransform,\
-     BroadcastTransform
-from pypy.rlib.jit import hint, unroll_safe, promote
+from pypy.module.micronumpy.interp_iter import ConstantIterator, AxisIterator,\
+     ViewTransform, BroadcastTransform
+from pypy.tool.pairtype import extendabletype
+from pypy.module.micronumpy.loop import ComputationDone
 
 """ Signature specifies both the numpy expression that has been constructed
 and the assembler to be compiled. This is a very important observation -
@@ -54,50 +54,6 @@
         known_sigs[sig] = sig
         return sig
 
-class NumpyEvalFrame(object):
-    _virtualizable2_ = ['iterators[*]', 'final_iter', 'arraylist[*]',
-                        'value', 'identity']
-
-    @unroll_safe
-    def __init__(self, iterators, arrays):
-        self = hint(self, access_directly=True, fresh_virtualizable=True)
-        self.iterators = iterators[:]
-        self.arrays = arrays[:]
-        for i in range(len(self.iterators)):
-            iter = self.iterators[i]
-            if not isinstance(iter, ConstantIterator):
-                self.final_iter = i
-                break
-        else:
-            self.final_iter = -1
-
-    def done(self):
-        final_iter = promote(self.final_iter)
-        if final_iter < 0:
-            assert False
-        return self.iterators[final_iter].done()
-
-    @unroll_safe
-    def next(self, shapelen):
-        for i in range(len(self.iterators)):
-            self.iterators[i] = self.iterators[i].next(shapelen)
-
-    @unroll_safe
-    def next_from_second(self, shapelen):
-        """ Don't increase the first iterator
-        """
-        for i in range(1, len(self.iterators)):
-            self.iterators[i] = self.iterators[i].next(shapelen)
-
-    def next_first(self, shapelen):
-        self.iterators[0] = self.iterators[0].next(shapelen)
-
-    def get_final_iter(self):
-        final_iter = promote(self.final_iter)
-        if final_iter < 0:
-            assert False
-        return self.iterators[final_iter]
-
 def _add_ptr_to_cache(ptr, cache):
     i = 0
     for p in cache:
@@ -113,6 +69,8 @@
     return r_dict(sigeq_no_numbering, sighash)
 
 class Signature(object):
+    __metaclass_ = extendabletype
+    
     _attrs_ = ['iter_no', 'array_no']
     _immutable_fields_ = ['iter_no', 'array_no']
 
@@ -138,11 +96,15 @@
         self.iter_no = no
 
     def create_frame(self, arr):
+        from pypy.module.micronumpy.loop import NumpyEvalFrame
+        
         iterlist = []
         arraylist = []
         self._create_iter(iterlist, arraylist, arr, [])
-        return NumpyEvalFrame(iterlist, arraylist)
-
+        f = NumpyEvalFrame(iterlist, arraylist)
+        # hook for cur_value being used by reduce
+        arr.compute_first_step(self, f)
+        return f
 
 class ConcreteSignature(Signature):
     _immutable_fields_ = ['dtype']
@@ -182,13 +144,10 @@
         assert isinstance(concr, ConcreteArray)
         storage = concr.storage
         if self.iter_no >= len(iterlist):
-            iterlist.append(self.allocate_iter(concr, transforms))
+            iterlist.append(concr.create_iter(transforms))
         if self.array_no >= len(arraylist):
             arraylist.append(storage)
 
-    def allocate_iter(self, arr, transforms):
-        return ArrayIterator(arr.size).apply_transformations(arr, transforms)
-
     def eval(self, frame, arr):
         iter = frame.iterators[self.iter_no]
         return self.dtype.getitem(frame.arrays[self.array_no], iter.offset)
@@ -220,22 +179,10 @@
         allnumbers.append(no)
         self.iter_no = no
 
-    def allocate_iter(self, arr, transforms):
-        return ViewIterator(arr.start, arr.strides, arr.backstrides,
-                            arr.shape).apply_transformations(arr, transforms)
-
 class FlatSignature(ViewSignature):
     def debug_repr(self):
         return 'Flat'
 
-    def allocate_iter(self, arr, transforms):
-        from pypy.module.micronumpy.interp_numarray import W_FlatIterator
-        assert isinstance(arr, W_FlatIterator)
-        return ViewIterator(arr.base.start, arr.base.strides, 
-                    arr.base.backstrides,
-                    arr.base.shape).apply_transformations(arr.base,
-                                                         transforms)
-
 class VirtualSliceSignature(Signature):
     def __init__(self, child):
         self.child = child
@@ -269,12 +216,13 @@
         return self.child.eval(frame, arr.child)
 
 class Call1(Signature):
-    _immutable_fields_ = ['unfunc', 'name', 'child']
+    _immutable_fields_ = ['unfunc', 'name', 'child', 'dtype']
 
-    def __init__(self, func, name, child):
+    def __init__(self, func, name, dtype, child):
         self.unfunc = func
         self.child = child
         self.name = name
+        self.dtype = dtype
 
     def hash(self):
         return compute_hash(self.name) ^ intmask(self.child.hash() << 1)
@@ -359,6 +307,17 @@
         return 'Call2(%s, %s, %s)' % (self.name, self.left.debug_repr(),
                                       self.right.debug_repr())
 
+class ResultSignature(Call2):
+    def __init__(self, dtype, left, right):
+        Call2.__init__(self, None, 'assign', dtype, left, right)
+
+    def eval(self, frame, arr):
+        from pypy.module.micronumpy.interp_numarray import ResultArray
+
+        assert isinstance(arr, ResultArray)
+        offset = frame.get_final_iter().offset
+        arr.left.setitem(offset, self.right.eval(frame, arr.right))
+
 class BroadcastLeft(Call2):
     def _invent_numbering(self, cache, allnumbers):
         self.left._invent_numbering(new_cache(), allnumbers)
@@ -400,20 +359,24 @@
         self.right._create_iter(iterlist, arraylist, arr.right, rtransforms)
 
 class ReduceSignature(Call2):
-    def _create_iter(self, iterlist, arraylist, arr, transforms):
-        self.right._create_iter(iterlist, arraylist, arr, transforms)
-
-    def _invent_numbering(self, cache, allnumbers):
-        self.right._invent_numbering(cache, allnumbers)
-
-    def _invent_array_numbering(self, arr, cache):
-        self.right._invent_array_numbering(arr, cache)
-
+    _immutable_fields_ = ['binfunc', 'name', 'calc_dtype',
+                          'left', 'right', 'done_func']
+    
+    def __init__(self, func, name, calc_dtype, left, right,
+                 done_func):
+        Call2.__init__(self, func, name, calc_dtype, left, right)
+        self.done_func = done_func
+        
     def eval(self, frame, arr):
-        return self.right.eval(frame, arr)
+        from pypy.module.micronumpy.interp_numarray import ReduceArray
+        assert isinstance(arr, ReduceArray)
+        rval = self.right.eval(frame, arr.right).convert_to(self.calc_dtype)
+        if self.done_func is not None and self.done_func(self.calc_dtype, rval):
+            raise ComputationDone(rval)
+        frame.cur_value = self.binfunc(self.calc_dtype, frame.cur_value, rval)
 
     def debug_repr(self):
-        return 'ReduceSig(%s, %s)' % (self.name, self.right.debug_repr())
+        return 'ReduceSig(%s)' % (self.name, self.right.debug_repr())
 
 class SliceloopSignature(Call2):
     def eval(self, frame, arr):
@@ -467,7 +430,17 @@
         from pypy.module.micronumpy.interp_numarray import AxisReduce
 
         assert isinstance(arr, AxisReduce)
-        return self.right.eval(frame, arr.right).convert_to(self.calc_dtype)
+        iterator = frame.get_final_iter()
+        v = self.right.eval(frame, arr.right).convert_to(self.calc_dtype)
+        if iterator.first_line:
+            if frame.identity is not None:
+                value = self.binfunc(self.calc_dtype, frame.identity, v)
+            else:
+                value = v
+        else:
+            cur = arr.left.getitem(iterator.offset)
+            value = self.binfunc(self.calc_dtype, cur, v)
+        arr.left.setitem(iterator.offset, value)
     
     def debug_repr(self):
         return 'AxisReduceSig(%s, %s)' % (self.name, self.right.debug_repr())
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -347,8 +347,9 @@
         raises((ValueError, TypeError), add.reduce, 1)
 
     def test_reduce_1d(self):
-        from _numpypy import add, maximum
+        from _numpypy import add, maximum, less
 
+        assert less.reduce([5, 4, 3, 2, 1])
         assert add.reduce([1, 2, 3]) == 6
         assert maximum.reduce([1]) == 1
         assert maximum.reduce([1, 2, 3]) == 3
@@ -433,3 +434,14 @@
         assert (isnan(array([0.2, float('inf'), float('nan')])) == [False, False, True]).all()
         assert (isinf(array([0.2, float('inf'), float('nan')])) == [False, True, False]).all()
         assert isinf(array([0.2])).dtype.kind == 'b'
+
+    def test_logical_ops(self):
+        from _numpypy import logical_and, logical_or, logical_xor, logical_not
+
+        assert (logical_and([True, False , True, True], [1, 1, 3, 0])
+                == [True, False, True, False]).all()
+        assert (logical_or([True, False, True, False], [1, 2, 0, 0])
+                == [True, True, True, False]).all()
+        assert (logical_xor([True, False, True, False], [1, 2, 0, 0])
+                == [False, True, True, False]).all()
+        assert (logical_not([True, False]) == [False, True]).all()
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
@@ -84,7 +84,7 @@
     def test_add(self):
         result = self.run("add")
         self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
-                                'setinteriorfield_raw': 1, 'int_add': 2,
+                                'setinteriorfield_raw': 1, 'int_add': 1,
                                 'int_ge': 1, 'guard_false': 1, 'jump': 1,
                                 'arraylen_gc': 1})
         assert result == 3 + 3
@@ -99,7 +99,7 @@
         result = self.run("float_add")
         assert result == 3 + 3
         self.check_simple_loop({"getinteriorfield_raw": 1, "float_add": 1,
-                                "setinteriorfield_raw": 1, "int_add": 2,
+                                "setinteriorfield_raw": 1, "int_add": 1,
                                 "int_ge": 1, "guard_false": 1, "jump": 1,
                                 'arraylen_gc': 1})
 
@@ -198,7 +198,8 @@
         result = self.run("any")
         assert result == 1
         self.check_simple_loop({"getinteriorfield_raw": 2, "float_add": 1,
-                                "float_ne": 1, "int_add": 1,
+                                "int_and": 1, "int_add": 1,
+                                'convert_float_to_int': 1,
                                 "int_ge": 1, "jump": 1,
                                 "guard_false": 2, 'arraylen_gc': 1})
 
@@ -239,7 +240,7 @@
         assert result == -6
         self.check_simple_loop({"getinteriorfield_raw": 2, "float_add": 1,
                                 "float_neg": 1,
-                                "setinteriorfield_raw": 1, "int_add": 2,
+                                "setinteriorfield_raw": 1, "int_add": 1,
                                 "int_ge": 1, "guard_false": 1, "jump": 1,
                                 'arraylen_gc': 1})
 
@@ -321,7 +322,7 @@
         # int_add might be 1 here if we try slightly harder with
         # reusing indexes or some optimization
         self.check_simple_loop({'float_add': 1, 'getinteriorfield_raw': 2,
-                                'guard_false': 1, 'int_add': 2, 'int_ge': 1,
+                                'guard_false': 1, 'int_add': 1, 'int_ge': 1,
                                 'jump': 1, 'setinteriorfield_raw': 1,
                                 'arraylen_gc': 1})
 
@@ -387,7 +388,7 @@
         assert result == 4
         self.check_trace_count(1)
         self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
-                                'setinteriorfield_raw': 1, 'int_add': 2,
+                                'setinteriorfield_raw': 1, 'int_add': 1,
                                 'int_ge': 1, 'guard_false': 1, 'jump': 1,
                                 'arraylen_gc': 1})
     def define_flat_iter():
@@ -403,7 +404,7 @@
         assert result == 6
         self.check_trace_count(1)
         self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
-                                'setinteriorfield_raw': 1, 'int_add': 3,
+                                'setinteriorfield_raw': 1, 'int_add': 2,
                                 'int_ge': 1, 'guard_false': 1,
                                 'arraylen_gc': 1, 'jump': 1})
 
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -181,6 +181,22 @@
     def ge(self, v1, v2):
         return v1 >= v2
 
+    @raw_binary_op
+    def logical_and(self, v1, v2):
+        return bool(v1) and bool(v2)
+
+    @raw_binary_op
+    def logical_or(self, v1, v2):
+        return bool(v1) or bool(v2)
+
+    @raw_unary_op
+    def logical_not(self, v):
+        return not bool(v)
+
+    @raw_binary_op
+    def logical_xor(self, v1, v2):
+        return bool(v1) ^ bool(v2)
+
     def bool(self, v):
         return bool(self.for_computation(self.unbox(v)))
 


More information about the pypy-commit mailing list