Author: Maciej Fijalkowski <fij...@gmail.com>
Branch: numpy-reintroduce-jit-drivers
Changeset: r57662:c0b1fcc2562e
Date: 2012-09-29 18:08 +0200
http://bitbucket.org/pypy/pypy/changeset/c0b1fcc2562e/
Log: start adding jitdrivers
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
@@ -258,7 +258,7 @@
return out
shape = shape_agreement(space, w_obj.get_shape(), out,
broadcast_down=False)
- return loop.call1(shape, self.func, self.name, calc_dtype, res_dtype,
+ return loop.call1(shape, self.func, calc_dtype, res_dtype,
w_obj, out)
@@ -322,7 +322,7 @@
return out
new_shape = shape_agreement(space, w_lhs.get_shape(), w_rhs)
new_shape = shape_agreement(space, new_shape, out,
broadcast_down=False)
- return loop.call2(new_shape, self.func, self.name, calc_dtype,
+ return loop.call2(new_shape, self.func, calc_dtype,
res_dtype, w_lhs, w_rhs, out)
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
@@ -1,6 +1,7 @@
""" This file is the main run loop as well as evaluation loops for various
-signatures
+operations. This is the place to look for all the computations that iterate
+over all the array elements.
"""
from pypy.rlib.objectmodel import specialize
@@ -9,13 +10,26 @@
from pypy.rpython.lltypesystem import lltype, rffi
from pypy.module.micronumpy.base import W_NDimArray
-def call2(shape, func, name, calc_dtype, res_dtype, w_lhs, w_rhs, out):
+call2_driver = jit.JitDriver(name='numpy_call2',
+ greens = ['shapelen', 'func', 'calc_dtype',
+ 'res_dtype'],
+ reds = ['shape', 'w_lhs', 'w_rhs', 'out',
+ 'left_iter', 'right_iter', 'out_iter'])
+
+def call2(shape, func, calc_dtype, res_dtype, w_lhs, w_rhs, out):
if out is None:
out = W_NDimArray.from_shape(shape, res_dtype)
left_iter = w_lhs.create_iter(shape)
right_iter = w_rhs.create_iter(shape)
out_iter = out.create_iter(shape)
+ shapelen = len(shape)
while not out_iter.done():
+ call2_driver.jit_merge_point(shapelen=shapelen, func=func,
+ calc_dtype=calc_dtype,
res_dtype=res_dtype,
+ shape=shape, w_lhs=w_lhs, w_rhs=w_rhs,
+ out=out,
+ left_iter=left_iter,
right_iter=right_iter,
+ out_iter=out_iter)
w_left = left_iter.getitem().convert_to(calc_dtype)
w_right = right_iter.getitem().convert_to(calc_dtype)
out_iter.setitem(func(calc_dtype, w_left, w_right).convert_to(
@@ -25,30 +39,56 @@
out_iter.next()
return out
-def call1(shape, func, name, calc_dtype, res_dtype, w_obj, out):
+call1_driver = jit.JitDriver(name='numpy_call1',
+ greens = ['shapelen', 'func', 'calc_dtype',
+ 'res_dtype'],
+ reds = ['shape', 'w_obj', 'out', 'obj_iter',
+ 'out_iter'])
+
+def call1(shape, func, calc_dtype, res_dtype, w_obj, out):
if out is None:
out = W_NDimArray.from_shape(shape, res_dtype)
obj_iter = w_obj.create_iter(shape)
out_iter = out.create_iter(shape)
+ shapelen = len(shape)
while not out_iter.done():
+ call1_driver.jit_merge_point(shapelen=shapelen, func=func,
+ calc_dtype=calc_dtype,
res_dtype=res_dtype,
+ shape=shape, w_obj=w_obj, out=out,
+ obj_iter=obj_iter, out_iter=out_iter)
elem = obj_iter.getitem().convert_to(calc_dtype)
out_iter.setitem(func(calc_dtype, elem).convert_to(res_dtype))
out_iter.next()
obj_iter.next()
return out
+setslice_driver = jit.JitDriver(name='numpy_setslice',
+ greens = ['shapelen', 'dtype'],
+ reds = ['target', 'source', 'target_iter',
+ 'source_iter'])
+
def setslice(shape, target, source):
# note that unlike everything else, target and source here are
# array implementations, not arrays
target_iter = target.create_iter(shape)
source_iter = source.create_iter(shape)
dtype = target.dtype
+ shapelen = len(shape)
while not target_iter.done():
+ setslice_driver.jit_merge_point(shapelen=shapelen, dtype=dtype,
+ target=target, source=source,
+ target_iter=target_iter,
+ source_iter=source_iter)
target_iter.setitem(source_iter.getitem().convert_to(dtype))
target_iter.next()
source_iter.next()
return target
+reduce_driver = jit.JitDriver(name='numpy_reduce',
+ greens = ['shapelen', 'func', 'calc_dtype',
+ 'identity', 'done_func'],
+ reds = ['obj', 'obj_iter', 'cur_value'])
+
def compute_reduce(obj, calc_dtype, func, done_func, identity):
obj_iter = obj.create_iter(obj.get_shape())
if identity is None:
@@ -56,7 +96,12 @@
obj_iter.next()
else:
cur_value = identity.convert_to(calc_dtype)
+ shapelen = len(obj.get_shape())
while not obj_iter.done():
+ reduce_driver.jit_merge_point(shapelen=shapelen, func=func,
+ calc_dtype=calc_dtype, identity=identity,
+ done_func=done_func, obj=obj,
+ obj_iter=obj_iter, cur_value=cur_value)
rval = obj_iter.getitem().convert_to(calc_dtype)
if done_func is not None and done_func(calc_dtype, rval):
return rval
@@ -70,6 +115,8 @@
arr_iter.setitem(box)
arr_iter.next()
+#where_driver = jit.JitDriver()
+
def where(out, shape, arr, x, y, dtype):
out_iter = out.create_iter(shape)
arr_iter = arr.create_iter(shape)
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit
