[pypy-commit] pypy default: merge ufunc-casting which provides casting of arguments to ufuncs and frompypyfunc

mattip noreply at buildbot.pypy.org
Tue Oct 13 23:57:54 CEST 2015


Author: mattip <matti.picus at gmail.com>
Branch: 
Changeset: r80181:517db24acc3f
Date: 2015-10-14 00:57 +0300
http://bitbucket.org/pypy/pypy/changeset/517db24acc3f/

Log:	merge ufunc-casting which provides casting of arguments to ufuncs
	and frompypyfunc

diff --git a/pypy/doc/whatsnew-head.rst b/pypy/doc/whatsnew-head.rst
--- a/pypy/doc/whatsnew-head.rst
+++ b/pypy/doc/whatsnew-head.rst
@@ -74,3 +74,9 @@
 ffi.new_handle() returns handles that work more like CPython's: they
 remain valid as long as the target exists (unlike the previous
 version, where handles become invalid *before* the __del__ is called).
+
+.. branch: ufunc-casting
+
+allow automatic casting in ufuncs (and frompypyfunc) to cast the
+arguments to the allowed function type declarations, fixes various
+failures in linalg cffi functions
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
@@ -74,10 +74,10 @@
 
 call_many_to_one_driver = jit.JitDriver(
     name='numpy_call_many_to_one',
-    greens=['shapelen', 'nin', 'func', 'res_dtype'],
+    greens=['shapelen', 'nin', 'func', 'in_dtypes', 'res_dtype'],
     reds='auto')
 
-def call_many_to_one(space, shape, func, res_dtype, in_args, out):
+def call_many_to_one(space, shape, func, in_dtypes, res_dtype, in_args, out):
     # out must hav been built. func needs no calc_type, is usually an
     # external ufunc
     nin = len(in_args)
@@ -95,9 +95,9 @@
     vals = [None] * nin
     while not out_iter.done(out_state):
         call_many_to_one_driver.jit_merge_point(shapelen=shapelen, func=func,
-                                     res_dtype=res_dtype, nin=nin)
+                        in_dtypes=in_dtypes, res_dtype=res_dtype, nin=nin)
         for i in range(nin):
-            vals[i] = in_iters[i].getitem(in_states[i])
+            vals[i] = in_dtypes[i].coerce(space, in_iters[i].getitem(in_states[i]))
         w_arglist = space.newlist(vals)
         w_out_val = space.call_args(func, Arguments.frompacked(space, w_arglist))
         out_iter.setitem(out_state, res_dtype.coerce(space, w_out_val))
@@ -108,10 +108,10 @@
 
 call_many_to_many_driver = jit.JitDriver(
     name='numpy_call_many_to_many',
-    greens=['shapelen', 'nin', 'nout', 'func', 'res_dtype'],
+    greens=['shapelen', 'nin', 'nout', 'func', 'in_dtypes', 'out_dtypes'],
     reds='auto')
 
-def call_many_to_many(space, shape, func, res_dtype, in_args, out_args):
+def call_many_to_many(space, shape, func, in_dtypes, out_dtypes, in_args, out_args):
     # out must hav been built. func needs no calc_type, is usually an
     # external ufunc
     nin = len(in_args)
@@ -134,24 +134,29 @@
         out_states[i] = out_state
     shapelen = len(shape)
     vals = [None] * nin
-    while not out_iters[0].done(out_states[0]):
+    test_iter, test_state = in_iters[-1], in_states[-1]
+    if nout > 0:
+        test_iter, test_state = out_iters[0], out_states[0]
+    while not test_iter.done(test_state):
         call_many_to_many_driver.jit_merge_point(shapelen=shapelen, func=func,
-                                     res_dtype=res_dtype, nin=nin, nout=nout)
+                             in_dtypes=in_dtypes, out_dtypes=out_dtypes,
+                             nin=nin, nout=nout)
         for i in range(nin):
-            vals[i] = in_iters[i].getitem(in_states[i])
+            vals[i] = in_dtypes[i].coerce(space, in_iters[i].getitem(in_states[i]))
         w_arglist = space.newlist(vals)
         w_outvals = space.call_args(func, Arguments.frompacked(space, w_arglist))
         # w_outvals should be a tuple, but func can return a single value as well
         if space.isinstance_w(w_outvals, space.w_tuple):
             batch = space.listview(w_outvals)
             for i in range(len(batch)):
-                out_iters[i].setitem(out_states[i], res_dtype.coerce(space, batch[i]))
+                out_iters[i].setitem(out_states[i], out_dtypes[i].coerce(space, batch[i]))
                 out_states[i] = out_iters[i].next(out_states[i])
-        else:
-            out_iters[0].setitem(out_states[0], res_dtype.coerce(space, w_outvals))
+        elif nout > 0:
+            out_iters[0].setitem(out_states[0], out_dtypes[0].coerce(space, w_outvals))
             out_states[0] = out_iters[0].next(out_states[0])
         for i in range(nin):
             in_states[i] = in_iters[i].next(in_states[i])
+        test_state = test_iter.next(test_state)
     return space.newtuple([convert_to_array(space, o) for o in out_args])
 
 setslice_driver = jit.JitDriver(name='numpy_setslice',
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
@@ -159,8 +159,7 @@
         af2 = ufunc(af)
         assert all(af2 == af * 2)
         ac = arange(10, dtype=complex)
-        skip('casting not implemented yet')
-        ac1 = ufunc(ac)
+        raises(TypeError, ufunc, ac)
 
     def test_frompyfunc_2d_sig(self):
         import sys
@@ -199,6 +198,10 @@
         ai2 = ufunc(aiV)
         assert (ai2 == aiV * 2).all()
 
+        ai = arange(0).reshape(0, 1, 1)
+        ao = ufunc(ai)
+        assert ao.shape == (0, 1, 1)
+
     def test_frompyfunc_needs_nditer(self):
         import sys
         from numpy import frompyfunc, dtype, arange
@@ -268,6 +271,54 @@
         assert out0.shape == in0.shape
         assert (out0 == in0 * 2).all()
 
+    def test_frompyfunc_casting(self):
+        import sys
+        import numpy as np
+        if '__pypy__' not in sys.builtin_module_names:
+            skip('PyPy only frompyfunc extension')
+
+        def times2_int(in0, out0):
+            assert in0.dtype == int
+            assert out0.dtype == int
+            # hack to assing to a 0-dim array
+            out0.real = in0 * 2
+
+        def times2_complex(in0, out0):
+            assert in0.dtype == complex
+            assert out0.dtype == complex
+            out0.real = in0.real * 2
+            out0.imag = in0.imag
+
+        def times2_complex0(in0):
+            assert in0.dtype == complex
+            return in0 * 2
+
+        def times2_int0(in0):
+            assert in0.dtype == int
+            return in0 * 2
+
+        times2stacked = np.frompyfunc([times2_int, times2_complex], 1, 1,
+                            dtypes=[np.dtype(int), np.dtype(int),
+                                np.dtype(complex), np.dtype(complex)],
+                            stack_inputs=True, signature='()->()',
+                          )
+        times2 = np.frompyfunc([times2_int0, times2_complex0], 1, 1,
+                            dtypes=[np.dtype(int), np.dtype(int),
+                                np.dtype(complex), np.dtype(complex)],
+                            stack_inputs=False,
+                          )
+        for d in [np.dtype(float), np.dtype('uint8'), np.dtype('complex64')]:
+            in0 = np.arange(4, dtype=d)
+            out0 = times2stacked(in0)
+            assert out0.shape == in0.shape
+            assert out0.dtype in (int, complex) 
+            assert (out0 == in0 * 2).all()
+
+            out0 = times2(in0)
+            assert out0.shape == in0.shape
+            assert out0.dtype in (int, complex) 
+            assert (out0 == in0 * 2).all()
+
     def test_ufunc_kwargs(self):
         from numpy import ufunc, frompyfunc, arange, dtype
         def adder(a, b):
@@ -1393,7 +1444,7 @@
     def test_add_doc(self):
         import sys
         if '__pypy__' not in sys.builtin_module_names:
-            skip('')
+            skip('cpython sets docstrings differently')
         try:
             from numpy import set_docstring
         except ImportError:
diff --git a/pypy/module/micronumpy/ufuncs.py b/pypy/module/micronumpy/ufuncs.py
--- a/pypy/module/micronumpy/ufuncs.py
+++ b/pypy/module/micronumpy/ufuncs.py
@@ -709,6 +709,32 @@
             raise oefmt(space.w_TypeError,
                 "ufunc '%s' not supported for the input types", self.name)
 
+def _match_dtypes(space, indtypes, targetdtypes, i_target, casting):
+    allok = True
+    for i in range(len(indtypes)):
+        origin = indtypes[i]
+        target = targetdtypes[i + i_target]
+        if origin is None:
+            continue
+        if target is None:
+            continue
+        if not can_cast_type(space, origin, target, casting):
+            allok = False
+            break
+    return allok
+
+def _raise_err_msg(self, space, dtypes0, dtypes1):
+    dtypesstr = ''
+    for d in dtypes0:
+        if d is None:
+            dtypesstr += 'None,'
+        else:
+            dtypesstr += '%s%s%s,' % (d.byteorder, d.kind, d.elsize)
+    _dtypesstr = ','.join(['%s%s%s' % (d.byteorder, d.kind, d.elsize) \
+                    for d in dtypes1])
+    raise oefmt(space.w_TypeError,
+         "input dtype [%s] did not match any known dtypes [%s] ",
+         dtypesstr,_dtypesstr)
 
 
 class W_UfuncGeneric(W_Ufunc):
@@ -799,29 +825,36 @@
             outargs0 = outargs[0]
             assert isinstance(inargs0, W_NDimArray)
             assert isinstance(outargs0, W_NDimArray)
+            nin = self.nin
+            assert nin >= 0
             res_dtype = outargs0.get_dtype()
             new_shape = inargs0.get_shape()
             # XXX use _find_array_wrap and wrap outargs using __array_wrap__
+            if self.stack_inputs:
+                loop.call_many_to_many(space, new_shape, func,
+                                         dtypes, [], inargs + outargs, [])
+                if len(outargs) < 2:
+                    return outargs[0]
+                return space.newtuple(outargs)
             if len(outargs) < 2:
                 return loop.call_many_to_one(space, new_shape, func,
-                                             res_dtype, inargs, outargs[0])
+                         dtypes[:nin], dtypes[-1], inargs, outargs[0])
             return loop.call_many_to_many(space, new_shape, func,
-                                             res_dtype, inargs, outargs)
+                         dtypes[:nin], dtypes[nin:], inargs, outargs)
+        w_casting = space.w_None
+        w_op_dtypes = space.w_None
         for tf in need_to_cast:
             if tf:
-                raise oefmt(space.w_NotImplementedError, "casting not supported yet")
+                w_casting = space.wrap('safe')
+                w_op_dtypes = space.newtuple([space.wrap(d) for d in dtypes])
+                
         w_flags = space.w_None # NOT 'external_loop', we do coalescing by core_num_dims
-        w_op_flags = space.newtuple([space.wrap(r) for r in ['readonly'] * len(inargs)] + \
-                                    [space.wrap(r) for r in ['readwrite'] * len(outargs)])
-        w_op_dtypes = space.w_None
-        w_casting = space.w_None
+        w_ro = space.newtuple([space.wrap('readonly'), space.wrap('copy')])
+        w_rw = space.newtuple([space.wrap('readwrite'), space.wrap('updateifcopy')])
+        
+        w_op_flags = space.newtuple([w_ro] * len(inargs) + [w_rw] * len(outargs))
         w_op_axes = space.w_None
 
-        #print '\nsignature', sig
-        #print [(d, getattr(self,d)) for d in dir(self) if 'core' in d or 'broad' in d]
-        #print [(d, locals()[d]) for d in locals() if 'core' in d or 'broad' in d]
-        #print 'shapes',[d.get_shape() for d in inargs + outargs]
-        #print 'steps',[d.implementation.strides for d in inargs + outargs]
         if isinstance(func, W_GenericUFuncCaller):
             # Use GeneralizeUfunc interface with signature
             # Unlike numpy, we will not broadcast dims before
@@ -934,19 +967,32 @@
         # linear_search_type_resolver in numpy ufunc_type_resolutions.c
         # type_tup can be '', a tuple of dtypes, or a string
         # of the form d,t -> D where the letters are dtype specs
-        nop = len(inargs) + len(outargs)
+
+        # XXX why does the next line not pass translation?
+        # dtypes = [i.get_dtype() for i in inargs]
         dtypes = []
+        for i in inargs:
+            if isinstance(i, W_NDimArray):
+                dtypes.append(i.get_dtype())
+            else:
+                dtypes.append(None)
+        for i in outargs:
+            if isinstance(i, W_NDimArray):
+                dtypes.append(i.get_dtype())
+            else:
+                dtypes.append(None)
         if isinstance(type_tup, str) and len(type_tup) > 0:
             try:
                 if len(type_tup) == 1:
-                    dtypes = [get_dtype_cache(space).dtypes_by_name[type_tup]] * self.nargs
+                    s_dtypes = [get_dtype_cache(space).dtypes_by_name[type_tup]] * self.nargs
                 elif len(type_tup) == self.nargs + 2:
+                    s_dtypes = []
                     for i in range(self.nin):
-                        dtypes.append(get_dtype_cache(space).dtypes_by_name[type_tup[i]])
+                        s_dtypes.append(get_dtype_cache(space).dtypes_by_name[type_tup[i]])
                     #skip the '->' in the signature
                     for i in range(self.nout):
                         j = i + self.nin + 2
-                        dtypes.append(get_dtype_cache(space).dtypes_by_name[type_tup[j]])
+                        s_dtypes.append(get_dtype_cache(space).dtypes_by_name[type_tup[j]])
                 else:
                     raise oefmt(space.w_TypeError, "a type-string for %s " \
                         "requires 1 typecode or %d typecode(s) before and %d" \
@@ -955,42 +1001,29 @@
             except KeyError:
                 raise oefmt(space.w_ValueError, "unknown typecode in" \
                         " call to %s with type-string '%s'", self.name, type_tup)
-        else:
-            # XXX why does the next line not pass translation?
-            # dtypes = [i.get_dtype() for i in inargs]
-            for i in inargs:
-                if isinstance(i, W_NDimArray):
-                    dtypes.append(i.get_dtype())
-                else:
-                    dtypes.append(None)
-            for i in outargs:
-                if isinstance(i, W_NDimArray):
-                    dtypes.append(i.get_dtype())
-                else:
-                    dtypes.append(None)
+            # Make sure args can be cast to dtypes
+            if not _match_dtypes(space, dtypes, s_dtypes, 0, "safe"):
+                _raise_err_msg(self, space, dtypes, s_dtypes)
+            dtypes = s_dtypes    
         #Find the first matchup of dtypes with _dtypes
         for i in range(0, len(_dtypes), self.nargs):
-            allok = True
-            for j in range(self.nargs):
-                if dtypes[j] is not None and dtypes[j] != _dtypes[i+j]:
-                    allok = False
+            allok = _match_dtypes(space, dtypes, _dtypes, i, "no")
             if allok:
                 break
         else:
-            if len(self.funcs) > 1:
-
-                dtypesstr = ''
-                for d in dtypes:
-                    if d is None:
-                        dtypesstr += 'None,'
-                    else:
-                        dtypesstr += '%s%s%s,' % (d.byteorder, d.kind, d.elsize)
-                _dtypesstr = ','.join(['%s%s%s' % (d.byteorder, d.kind, d.elsize) \
-                                for d in _dtypes])
-                raise oefmt(space.w_TypeError,
-                     "input dtype [%s] did not match any known dtypes [%s] ",
-                     dtypesstr,_dtypesstr)
-            i = 0
+            # No exact matches, can we cast?
+            for i in range(0, len(_dtypes), self.nargs):
+                allok = _match_dtypes(space, dtypes, _dtypes, i, "safe")
+                if allok:
+                    end = i + self.nargs
+                    assert i >= 0
+                    assert end >=0
+                    dtypes = _dtypes[i:end]
+                    break
+            else:
+                if len(self.funcs) > 1:
+                    _raise_err_msg(self, space, dtypes, _dtypes)
+                i = 0
         # Fill in empty dtypes
         for j in range(self.nargs):
             if dtypes[j] is None:
@@ -1086,7 +1119,7 @@
             for j in range(offset, len(iter_shape)):
                 x = iter_shape[j + offset]
                 y = dims_to_broadcast[j]
-                if (x > y and x % y) or y %x:
+                if y != 0 and x != 0 and ((x > y and x % y) or y %x):
                     raise oefmt(space.w_ValueError, "%s: %s operand %d has a "
                         "mismatch in its broadcast dimension %d "
                         "(size %d is different from %d)",
@@ -1123,7 +1156,7 @@
         # the current op (signalling it can handle ndarray's).
 
         # TODO parse and handle subok
-        # TODO handle flags, op_flags
+        # TODO handle more flags, op_flags
         #print 'iter_shape',iter_shape,'arg_shapes',arg_shapes,'matched_dims',matched_dims
         return iter_shape, arg_shapes, matched_dims
 


More information about the pypy-commit mailing list