Author: Ronan Lamy <ronan.l...@gmail.com> Branch: ufunc-reduce Changeset: r78750:1eb5d76e3f7b Date: 2015-08-02 20:17 +0100 http://bitbucket.org/pypy/pypy/changeset/1eb5d76e3f7b/
Log: Clean up accumulate loops 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 @@ -214,68 +214,68 @@ obj_state = obj_iter.next(obj_state) return cur_value -reduce_cum_driver = jit.JitDriver( - name='numpy_reduce_cum_driver', +accumulate_flat_driver = jit.JitDriver( + name='numpy_accumulate_flat', greens=['shapelen', 'func', 'dtype', 'out_dtype'], reds='auto') -def compute_reduce_cumulative(space, obj, out, calc_dtype, func, identity): - obj_iter, obj_state = obj.create_iter() +def accumulate_flat(space, func, w_arr, calc_dtype, out, identity): + arr_iter, arr_state = w_arr.create_iter() out_iter, out_state = out.create_iter() out_iter.track_index = False if identity is None: - cur_value = obj_iter.getitem(obj_state).convert_to(space, calc_dtype) + cur_value = arr_iter.getitem(arr_state).convert_to(space, calc_dtype) out_iter.setitem(out_state, cur_value) out_state = out_iter.next(out_state) - obj_state = obj_iter.next(obj_state) + arr_state = arr_iter.next(arr_state) else: cur_value = identity.convert_to(space, calc_dtype) - shapelen = len(obj.get_shape()) + shapelen = len(w_arr.get_shape()) out_dtype = out.get_dtype() - while not obj_iter.done(obj_state): - reduce_cum_driver.jit_merge_point( - shapelen=shapelen, func=func, - dtype=calc_dtype, out_dtype=out_dtype) - rval = obj_iter.getitem(obj_state).convert_to(space, calc_dtype) - cur_value = func(calc_dtype, cur_value, rval) + while not arr_iter.done(arr_state): + accumulate_flat_driver.jit_merge_point( + shapelen=shapelen, func=func, dtype=calc_dtype, + out_dtype=out_dtype) + w_item = arr_iter.getitem(arr_state).convert_to(space, calc_dtype) + cur_value = func(calc_dtype, cur_value, w_item) out_iter.setitem(out_state, out_dtype.coerce(space, cur_value)) out_state = out_iter.next(out_state) - obj_state = obj_iter.next(obj_state) + arr_state = arr_iter.next(arr_state) -accumulate_driver = jit.JitDriver(name='numpy_accumulate_flat', - greens=['shapelen', 'func', 'dtype'], - reds='auto') +accumulate_driver = jit.JitDriver( + name='numpy_accumulate', + greens=['shapelen', 'func', 'calc_dtype'], reds='auto') -def do_accumulate(space, func, arr, dtype, axis, out, identity): +def accumulate(space, func, w_arr, axis, calc_dtype, out, identity): out_iter, out_state = out.create_iter() - obj_shape = arr.get_shape() - temp_shape = obj_shape[:axis] + obj_shape[axis + 1:] - temp = W_NDimArray.from_shape(space, temp_shape, dtype, w_instance=arr) - temp_iter = AxisIter(temp.implementation, arr.get_shape(), axis) + arr_shape = w_arr.get_shape() + temp_shape = arr_shape[:axis] + arr_shape[axis + 1:] + temp = W_NDimArray.from_shape(space, temp_shape, calc_dtype, w_instance=w_arr) + temp_iter = AxisIter(temp.implementation, w_arr.get_shape(), axis) temp_state = temp_iter.reset() - arr_iter, arr_state = arr.create_iter() + arr_iter, arr_state = w_arr.create_iter() arr_iter.track_index = False if identity is not None: - identity = identity.convert_to(space, dtype) - shapelen = len(obj_shape) + identity = identity.convert_to(space, calc_dtype) + shapelen = len(arr_shape) while not out_iter.done(out_state): accumulate_driver.jit_merge_point(shapelen=shapelen, func=func, - dtype=dtype) - w_val = arr_iter.getitem(arr_state).convert_to(space, dtype) + calc_dtype=calc_dtype) + w_item = arr_iter.getitem(arr_state).convert_to(space, calc_dtype) arr_state = arr_iter.next(arr_state) out_indices = out_iter.indices(out_state) if out_indices[axis] == 0: if identity is not None: - w_val = func(dtype, identity, w_val) + w_item = func(calc_dtype, identity, w_item) else: - cur = temp_iter.getitem(temp_state) - w_val = func(dtype, cur, w_val) + cur_value = temp_iter.getitem(temp_state) + w_item = func(calc_dtype, cur_value, w_item) - out_iter.setitem(out_state, w_val) + out_iter.setitem(out_state, w_item) out_state = out_iter.next(out_state) - temp_iter.setitem(temp_state, w_val) + temp_iter.setitem(temp_state, w_item) temp_state = temp_iter.next(temp_state) return out 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 @@ -270,11 +270,11 @@ if self.identity is not None: out.fill(space, self.identity.convert_to(space, dtype)) return out - loop.do_accumulate(space, self.func, obj, dtype, axis, - out, self.identity) + loop.accumulate( + space, self.func, obj, axis, dtype, out, self.identity) else: - loop.compute_reduce_cumulative(space, obj, out, dtype, self.func, - self.identity) + loop.accumulate_flat( + space, self.func, obj, dtype, out, self.identity) if call__array_wrap__: out = space.call_method(obj, '__array_wrap__', out) return out _______________________________________________ pypy-commit mailing list pypy-commit@python.org https://mail.python.org/mailman/listinfo/pypy-commit