[pypy-commit] pypy ufunc-casting: percolate casting into ufuncs, call_many_to_one, call_many_to_many
mattip
noreply at buildbot.pypy.org
Tue Oct 13 17:12:06 CEST 2015
Author: mattip <matti.picus at gmail.com>
Branch: ufunc-casting
Changeset: r80162:43673ac286b4
Date: 2015-10-13 18:08 +0300
http://bitbucket.org/pypy/pypy/changeset/43673ac286b4/
Log: percolate casting into ufuncs, call_many_to_one, call_many_to_many
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]))
+ print 'vals', vals
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_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,7 +159,7 @@
af2 = ufunc(af)
assert all(af2 == af * 2)
ac = arange(10, dtype=complex)
- ac1 = ufunc(ac)
+ raises(TypeError, ufunc, ac)
def test_frompyfunc_2d_sig(self):
import sys
@@ -267,6 +267,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):
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):
@@ -794,6 +820,7 @@
iter_shape, arg_shapes, matched_dims = self.verify_args(space, inargs, outargs)
inargs, outargs, need_to_cast = self.alloc_args(space, inargs, outargs, dtypes,
arg_shapes)
+ print 'call', self.external_loop, need_to_cast, dtypes, [a.get_dtype() for a in inargs], [a.get_dtype() for a in outargs]
if not self.external_loop:
inargs0 = inargs[0]
outargs0 = outargs[0]
@@ -802,26 +829,31 @@
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 outargs
if len(outargs) < 2:
return loop.call_many_to_one(space, new_shape, func,
- res_dtype, inargs, outargs[0])
+ dtypes[:self.nin], dtypes[-1], inargs, outargs[0])
return loop.call_many_to_many(space, new_shape, func,
- res_dtype, inargs, outargs)
+ dtypes[:self.nin], dtypes[self.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 +966,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 +1000,26 @@
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:
+ dtypes = _dtypes[i:i+self.nargs]
+ 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:
More information about the pypy-commit
mailing list