[pypy-commit] pypy ufuncapi: wip

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r72037:baa6d5815247
Date: 2014-06-12 23:45 +0300
http://bitbucket.org/pypy/pypy/changeset/baa6d5815247/

Log:	wip

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
@@ -85,6 +85,29 @@
                                 greens = ['shapelen', 'dtype'],
                                 reds = 'auto')
 
+call_many_to_one_driver = jit.JitDriver(
+    name='numpy_call_many_to_one',
+    greens=['shapelen', 'func', 'res_dtype'],
+    reds='auto')
+
+def call_many_to_one(space, shape, func, res_dtype, w_in, out):
+    # out must hav been built. func needs no calc_type, is usually an
+    # external ufunc
+    iters_and_states = [i.create_iter(shape) for i in w_in]
+    shapelen = len(shape)
+    while not out_iter.done(out_state):
+        call_many_to_one_driver.jit_merge_point(shapelen=shapelen, func=func,
+                                     res_dtype=res_dtype)
+        vals = [None] + [i_s[0].getitem(i_s[1]) for i_s in iters_and_states]
+        arglist = space.wrap(vals)
+        out_val = space.call_args(func, Arguments.frompacked(space, arglist))
+        out_iter.setitem(out_state, out_val.convert_to(space, res_dtype))
+        for i in range(len(iters_and_states)):
+            iters_and_states[i][1] = iters_and_states[i][0].next(iters_and_states[i][1])
+        out_state = out_iter.next(out_state)
+    return out
+
+
 def setslice(space, shape, target, source):
     # note that unlike everything else, target and source here are
     # array implementations, not arrays
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
@@ -208,9 +208,9 @@
         except NotImplementedError as e:
             assert 'object' in str(e)
             # Use pypy specific extension for out_dtype
-            myufunc = frompyfunc(adder, 2, 1, out_dtype='match')
-            int_func22 = frompyfunc(int, 2, 2, out_dtype='match')
-            int_func12 = frompyfunc(int, 1, 2, out_dtype='match')
+            myufunc = frompyfunc(adder, 2, 1, dtypes=['match'])
+            int_func22 = frompyfunc(int, 2, 2, dtypes=['match'])
+            int_func12 = frompyfunc(int, 1, 2, dtypes=['match'])
             retype = dtype(int)
         assert isinstance(myufunc, ufunc)
         res = myufunc(arange(10), arange(10))
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
@@ -527,8 +527,11 @@
         index = self.type_resolver(space, inargs, outargs)
         self.alloc_outargs(space, index, inargs, outargs)
         # XXX handle inner-loop indexing
-        # XXX JIT_me
-        raise oefmt(space.w_NotImplementedError, 'not implemented yet')
+        if len(outargs) < 2:
+            return loop.call_many_to_one(space, new_shape, self.func,
+                                         res_dtype, inargs, outargs[0])
+        return loop.call_many_to_many(space, new_shape, self.func,
+                                     res_dtype, inargs, out)
 
     def type_resolver(self, space, index, outargs):
         # Find a match for the inargs.dtype in self.dtypes, like
@@ -951,8 +954,7 @@
     if space.is_none(w_dtypes) and not signature:
         raise oefmt(space.w_NotImplementedError,
              'object dtype requested but not implemented')
-        #dtypes=[descriptor.get_dtype_cache(space).w_objectdtype]
-    elif space.isinstance_w(w_dtypes, space.w_str):
+    if space.isinstance_w(w_dtypes, space.w_str):
         if not space.str_w(w_dtypes) == 'match':
             raise oefmt(space.w_ValueError,
                 'unknown out_dtype value "%s"', space.str_w(w_dtypes))
@@ -974,14 +976,17 @@
         raise oefmt(space.w_ValueError,
             'identity must be 0, 1, or None')
     if nin==1 and nout==1 and dtypes == 'match':
-        w_ret = W_Ufunc1(func[0], name)
+        w_ret = W_Ufunc1(wrap_ext_func(func[0], name)
     elif nin==2 and nout==1 and dtypes == 'match':
-        def _func(calc_dtype, w_left, w_right):
-            arglist = space.wrap([w_left, w_right])
-            return space.call_args(func[0], Arguments.frompacked(space, arglist))
-        w_ret = W_Ufunc2(_func, name)
+        w_ret = W_Ufunc2(wrap_ext_func(func[0]), name)
     else:
         w_ret = W_UfuncGeneric(space, func, name, identity, nin, nout, dtypes, signature)
     if doc:
         w_ret.w_doc = space.wrap(doc)
     return w_ret
+
+def wrap_ext_func(func):
+    def _func(calc_dtype, w_left, w_right):
+        arglist = space.wrap([w_left, w_right])
+        return space.call_args(func, Arguments.frompacked(space, arglist))
+    return _func