Commits

Alex Gaynor  committed cf91e94

Unroll some functions in numpy correctly.

  • Participants
  • Parent commits aa43421

Comments (0)

Files changed (3)

File pypy/module/micronumpy/interp_ufuncs.py

 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, loop
+from pypy.rlib import jit
 from pypy.rlib.rarithmetic import LONG_BIT
 from pypy.tool.sourcetools import func_with_new_name
 
+
 class W_Ufunc(Wrappable):
     _attrs_ = ["name", "promote_to_float", "promote_bools", "identity"]
     _immutable_fields_ = ["promote_to_float", "promote_bools", "name"]
                 elif out.shape != shape:
                     raise operationerrfmt(space.w_ValueError,
                         'output parameter shape mismatch, expecting [%s]' +
-                        ' , got [%s]', 
+                        ' , got [%s]',
                         ",".join([str(x) for x in shape]),
                         ",".join([str(x) for x in out.shape]),
                         )
         else:
             arr = ReduceArray(self.func, self.name, self.identity, obj, dtype)
             val = loop.compute(arr)
-        return val 
+        return val
 
     def do_axis_reduce(self, obj, dtype, axis, result):
         from pypy.module.micronumpy.interp_numarray import AxisReduce
         if isinstance(w_obj, Scalar):
             arr = self.func(calc_dtype, w_obj.value.convert_to(calc_dtype))
             if isinstance(out,Scalar):
-                out.value=arr
+                out.value = arr
             elif isinstance(out, BaseArray):
                 out.fill(space, arr)
             else:
             if not broadcast_shape or broadcast_shape != out.shape:
                 raise operationerrfmt(space.w_ValueError,
                     'output parameter shape mismatch, could not broadcast [%s]' +
-                    ' to [%s]', 
+                    ' to [%s]',
                     ",".join([str(x) for x in w_obj.shape]),
                     ",".join([str(x) for x in out.shape]),
                     )
         self.func = func
         self.comparison_func = comparison_func
 
+    @jit.unroll_safe
     def call(self, space, args_w):
         from pypy.module.micronumpy.interp_numarray import (Call2,
             convert_to_array, Scalar, shape_agreement, BaseArray)
-        if len(args_w)>2:
+        if len(args_w) > 2:
             [w_lhs, w_rhs, w_out] = args_w
         else:
             [w_lhs, w_rhs] = args_w
                 w_rhs.value.convert_to(calc_dtype)
             )
             if isinstance(out,Scalar):
-                out.value=arr
+                out.value = arr
             elif isinstance(out, BaseArray):
                 out.fill(space, arr)
             else:
         if out and out.shape != shape_agreement(space, new_shape, out.shape):
             raise operationerrfmt(space.w_ValueError,
                 'output parameter shape mismatch, could not broadcast [%s]' +
-                ' to [%s]', 
+                ' to [%s]',
                 ",".join([str(x) for x in new_shape]),
                 ",".join([str(x) for x in out.shape]),
                 )
         w_lhs.add_invalidates(w_res)
         w_rhs.add_invalidates(w_res)
         if out:
-            #out.add_invalidates(w_res) #causes a recursion loop
             w_res.get_concrete()
         return w_res
 

File pypy/module/micronumpy/strides.py

 from pypy.rlib import jit
 from pypy.interpreter.error import OperationError
 
+@jit.look_inside_iff(lambda chunks: jit.isconstant(len(chunks)))
 def enumerate_chunks(chunks):
     result = []
     i = -1
         space.isinstance_w(w_item_or_slice, space.w_slice)):
         raise OperationError(space.w_IndexError,
                              space.wrap('unsupported iterator index'))
-            
+
     start, stop, step, lngth = space.decode_index4(w_item_or_slice, size)
-        
+
     coords = [0] * len(shape)
     i = start
     if order == 'C':

File pypy/module/micronumpy/support.py

+from pypy.rlib import jit
+
+
+@jit.look_inside_iff(lambda s: jit.isconstant(len(s)))
 def product(s):
     i = 1
     for x in s: