Author: Maciej Fijalkowski <fij...@gmail.com>
Branch: numpypy-axisops
Changeset: r51260:d15dcd1e48b1
Date: 2012-01-12 15:19 +0200
http://bitbucket.org/pypy/pypy/changeset/d15dcd1e48b1/
Log: in-progress. Now I understand what's exactly missing from making it
nicely work with lazy evaluation. breaks tests
diff --git a/pypy/module/micronumpy/interp_iter.py
b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -22,9 +22,6 @@
def done(self):
raise NotImplementedError
- def axis_done(self):
- raise NotImplementedError
-
class ArrayIterator(BaseIterator):
def __init__(self, size):
self.offset = 0
@@ -106,74 +103,56 @@
def next(self, shapelen):
return self
-def axis_iter_from_arr(arr, dim=-1):
- # The assert is needed for zjit tests
- from pypy.module.micronumpy.interp_numarray import ConcreteArray
- assert isinstance(arr, ConcreteArray)
- return AxisIterator(arr.start, arr.strides, arr.backstrides, arr.shape,
- dim)
-
class AxisIterator(BaseIterator):
""" Accept an addition argument dim
Redorder the dimensions to iterate over dim most often.
Set a flag at the end of each run over dim.
"""
- def __init__(self, arr_start, strides, backstrides, shape, dim):
+ def __init__(self, dim, shape, strides, backstrides):
self.shape = shape
- self.shapelen = len(shape)
self.indices = [0] * len(shape)
self._done = False
- self._axis_done = False
- self.offset = arr_start
+ self.axis_done = False
+ self.offset = -1
self.dim = dim
- self.dim_order = []
- if self.dim >= 0:
- self.dim_order.append(self.dim)
- for i in range(self.shapelen - 1, -1, -1):
- if i == self.dim:
- continue
- self.dim_order.append(i)
- self.strides = strides
- self.backstrides = backstrides
+ self.strides = strides[:dim] + [0] + strides[dim:]
+ self.backstrides = backstrides[:dim] + [0] + backstrides[dim:]
+ self.dim_order = [dim]
+ for i in range(len(shape) - 1, -1, -1):
+ if i != self.dim:
+ self.dim_order.append(i)
def done(self):
return self._done
- def axis_done(self):
- return self._axis_done
-
@jit.unroll_safe
def next(self, shapelen):
- #shapelen will always be one less than self.shapelen
offset = self.offset
- _axis_done = False
done = False
- #indices = [0] * self.shapelen
- #for i in range(self.shapelen):
- # indices[i] = self.indices[i]
- indices = self.indices
+ indices = [0] * shapelen
+ for i in range(shapelen):
+ indices[i] = self.indices[i]
+ axis_done = False
for i in self.dim_order:
if indices[i] < self.shape[i] - 1:
indices[i] += 1
- offset += self.strides[i]
break
else:
if i == self.dim:
- _axis_done = True
+ axis_done = True
+ offset += 1
indices[i] = 0
- offset -= self.backstrides[i]
else:
done = True
res = instantiate(AxisIterator)
- res._axis_done = _axis_done
+ res.axis_done = axis_done
+ res.strides = self.strides
+ res.backstrides = self.backstrides
res.offset = offset
res.indices = indices
- res.strides = self.strides
+ res.shape = self.shape
+ res.dim = self.dim
res.dim_order = self.dim_order
- res.backstrides = self.backstrides
- res.shape = self.shape
- res.shapelen = self.shapelen
- res.dim = self.dim
res._done = done
return res
diff --git a/pypy/module/micronumpy/interp_numarray.py
b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -9,7 +9,7 @@
from pypy.tool.sourcetools import func_with_new_name
from pypy.rlib.rstring import StringBuilder
from pypy.module.micronumpy.interp_iter import ArrayIterator, OneDimIterator,\
- view_iter_from_arr, SkipLastAxisIterator
+ view_iter_from_arr, SkipLastAxisIterator, AxisIterator
numpy_driver = jit.JitDriver(
greens=['shapelen', 'sig'],
@@ -36,13 +36,6 @@
get_printable_location=signature.new_printable_location('slice'),
)
-axisreduce_driver = jit.JitDriver(
- greens=['shapelen', 'sig'],
- virtualizables=['frame'],
- reds=['identity', 'self','result', 'ri', 'frame', 'dtype', 'value'],
- get_printable_location=signature.new_printable_location('axisreduce'),
-)
-
def _find_shape_and_elems(space, w_iterable):
shape = [space.len_w(w_iterable)]
@@ -695,7 +688,6 @@
# to allow garbage-collecting them
raise NotImplementedError
- @jit.unroll_safe
def compute(self):
result = W_NDimArray(self.size, self.shape, self.find_dtype())
shapelen = len(self.shape)
@@ -760,74 +752,6 @@
self.child = None
-class Reduce(VirtualArray):
- def __init__(self, binfunc, name, dim, res_dtype, values, identity=None):
- shape = values.shape[0:dim] + values.shape[dim + 1:len(values.shape)]
- VirtualArray.__init__(self, name, shape, res_dtype)
- self.values = values
- self.size = 1
- for s in shape:
- self.size *= s
- self.binfunc = binfunc
- self.dtype = res_dtype
- self.dim = dim
- self.identity = identity
-
- def _del_sources(self):
- self.values = None
-
- def create_sig(self, res_shape):
- if self.forced_result is not None:
- return self.forced_result.create_sig(res_shape)
- return signature.AxisReduceSignature(self.binfunc, self.name,
self.dtype,
- signature.ViewSignature(self.dtype),
- self.values.create_sig(res_shape))
-
- def get_identity(self, sig, frame, shapelen):
- #XXX does this allocate? Yes :(
- #XXX is this inlinable? Yes :)
- if self.identity is None:
- value = sig.eval(frame, self.values).convert_to(self.dtype)
- frame.next(shapelen)
- else:
- value = self.identity.convert_to(self.dtype)
- return value
-
- @jit.unroll_safe
- def compute(self):
- dtype = self.dtype
- result = W_NDimArray(self.size, self.shape, dtype)
- self.values = self.values.get_concrete()
- shapelen = len(result.shape)
- identity = self.identity
- sig = self.find_sig(res_shape=result.shape, arr=self.values)
- ri = ArrayIterator(result.size)
- frame = sig.create_frame(self.values, dim=self.dim)
- value = self.get_identity(sig, frame, shapelen)
- assert isinstance(sig, signature.AxisReduceSignature)
- while not frame.done():
- axisreduce_driver.jit_merge_point(frame=frame, self=self,
- value=value, sig=sig,
- shapelen=shapelen, ri=ri,
- dtype=dtype,
- identity=identity,
- result=result)
- if frame.axis_done():
- result.dtype.setitem(result.storage, ri.offset, value)
- if identity is None:
- value = sig.eval(frame, self.values).convert_to(dtype)
- frame.next(shapelen)
- else:
- value = identity.convert_to(dtype)
- ri = ri.next(shapelen)
- value = self.binfunc(dtype, value,
- sig.eval(frame, self.values).convert_to(dtype))
- frame.next(shapelen)
- assert ri.done
- result.dtype.setitem(result.storage, ri.offset, value)
- return result
-
-
class Call1(VirtualArray):
def __init__(self, ufunc, name, shape, res_dtype, values):
VirtualArray.__init__(self, name, shape, res_dtype)
@@ -869,6 +793,16 @@
self.left.create_sig(res_shape),
self.right.create_sig(res_shape))
+class AxisReduce(Call2):
+ """ NOTE: this is only used as a container, you should never
+ encounter such things in the wild. Remove this comment
+ when we'll make AxisReduce lazy
+ """
+ def __init__(self, ufunc, name, shape, dtype, left, right, dim):
+ Call2.__init__(self, ufunc, name, shape, dtype, dtype,
+ left, right)
+ self.dim = dim
+
class ConcreteArray(BaseArray):
""" An array that have actual storage, whether owned or not
"""
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
@@ -3,8 +3,8 @@
from pypy.interpreter.gateway import interp2app
from pypy.interpreter.typedef import TypeDef, GetSetProperty,
interp_attrproperty
from pypy.module.micronumpy import interp_boxes, interp_dtype
-from pypy.module.micronumpy.signature import ReduceSignature, ScalarSignature,\
- find_sig, new_printable_location
+from pypy.module.micronumpy.signature import ReduceSignature,\
+ find_sig, new_printable_location, AxisReduceSignature, ScalarSignature
from pypy.rlib import jit
from pypy.rlib.rarithmetic import LONG_BIT
from pypy.tool.sourcetools import func_with_new_name
@@ -16,6 +16,14 @@
get_printable_location=new_printable_location('reduce'),
)
+axisreduce_driver = jit.JitDriver(
+ greens=['shapelen', 'sig'],
+ virtualizables=['frame'],
+ reds=['self','arr', 'frame', 'shapelen'],
+# name='axisreduce',
+ get_printable_location=new_printable_location('axisreduce'),
+)
+
class W_Ufunc(Wrappable):
_attrs_ = ["name", "promote_to_float", "promote_bools", "identity"]
@@ -129,14 +137,13 @@
raise operationerrfmt(space.w_ValueError, "zero-size array to "
"%s.reduce without identity", self.name)
if shapelen > 1 and dim >= 0:
- from pypy.module.micronumpy.interp_numarray import Reduce
- res = Reduce(self.func, self.name, dim, dtype, obj, self.identity)
- obj.add_invalidates(res)
+ res = self.do_axis_reduce(obj, dtype, dim)
return space.wrap(res)
+ scalarsig = ScalarSignature(dtype)
sig = find_sig(ReduceSignature(self.func, self.name, dtype,
- ScalarSignature(dtype),
+ scalarsig,
obj.create_sig(obj.shape)), obj)
- frame = sig.create_frame(obj, dim=-1)
+ frame = sig.create_frame(obj)
if self.identity is None:
value = sig.eval(frame, obj).convert_to(dtype)
frame.next(shapelen)
@@ -144,6 +151,46 @@
value = self.identity.convert_to(dtype)
return self.reduce_loop(shapelen, sig, frame, value, obj, dtype)
+ def do_axis_reduce(self, obj, dtype, dim):
+ from pypy.module.micronumpy.interp_numarray import AxisReduce,\
+ W_NDimArray
+
+ shape = obj.shape[0:dim] + obj.shape[dim + 1:len(obj.shape)]
+ size = 1
+ for s in shape:
+ size *= s
+ result = W_NDimArray(size, shape, dtype)
+ rightsig = obj.create_sig(obj.shape)
+ # note - this is just a wrapper so signature can fetch
+ # both left and right, nothing more, especially
+ # this is not a true virtual array, because shapes
+ # don't quite match
+ arr = AxisReduce(self.func, self.name, shape, dtype,
+ result, obj, dim)
+ scalarsig = ScalarSignature(dtype)
+ sig = find_sig(AxisReduceSignature(self.func, self.name, dtype,
+ scalarsig, rightsig), arr)
+ frame = sig.create_frame(arr)
+ shapelen = len(obj.shape)
+ if self.identity is None:
+ frame.identity = sig.eval(frame, arr).convert_to(dtype)
+ frame.next(shapelen)
+ else:
+ frame.identity = self.identity.convert_to(dtype)
+ frame.value = frame.identity
+ self.reduce_axis_loop(frame, sig, shapelen, arr)
+ return result
+
+ def reduce_axis_loop(self, frame, sig, shapelen, arr):
+ while not frame.done():
+ axisreduce_driver.jit_merge_point(frame=frame, self=self,
+ sig=sig,
+ shapelen=shapelen, arr=arr)
+ sig.eval(frame, arr)
+ frame.next(shapelen)
+ # store the last value, when everything is done
+ arr.left.setitem(frame.iterators[0].offset, frame.value)
+
def reduce_loop(self, shapelen, sig, frame, value, obj, dtype):
while not frame.done():
reduce_driver.jit_merge_point(sig=sig,
diff --git a/pypy/module/micronumpy/signature.py
b/pypy/module/micronumpy/signature.py
--- a/pypy/module/micronumpy/signature.py
+++ b/pypy/module/micronumpy/signature.py
@@ -1,7 +1,7 @@
from pypy.rlib.objectmodel import r_dict, compute_identity_hash, compute_hash
from pypy.rlib.rarithmetic import intmask
from pypy.module.micronumpy.interp_iter import ViewIterator, ArrayIterator, \
- OneDimIterator, ConstantIterator, axis_iter_from_arr
+ OneDimIterator, ConstantIterator, AxisIterator
from pypy.module.micronumpy.strides import calculate_slice_strides
from pypy.rlib.jit import hint, unroll_safe, promote
@@ -33,7 +33,8 @@
return sig
class NumpyEvalFrame(object):
- _virtualizable2_ = ['iterators[*]', 'final_iter', 'arraylist[*]']
+ _virtualizable2_ = ['iterators[*]', 'final_iter', 'arraylist[*]',
+ 'value', 'identity']
@unroll_safe
def __init__(self, iterators, arrays):
@@ -47,6 +48,8 @@
break
else:
self.final_iter = -1
+ self.value = None
+ self.identity = None
def done(self):
final_iter = promote(self.final_iter)
@@ -57,13 +60,7 @@
@unroll_safe
def next(self, shapelen):
for i in range(len(self.iterators)):
- self.iterators[i] = self.iterators[i].next(shapelen)
-
- def axis_done(self):
- final_iter = promote(self.final_iter)
- if final_iter < 0:
- return False
- return self.iterators[final_iter].axis_done()
+ self.iterators[i] = self.iterators[i].next(shapelen)
def _add_ptr_to_cache(ptr, cache):
i = 0
@@ -101,13 +98,13 @@
allnumbers.append(no)
self.iter_no = no
- def create_frame(self, arr, res_shape=None, chunks=None, dim=-1):
+ def create_frame(self, arr, res_shape=None, chunks=None):
if chunks is None:
chunks = []
res_shape = res_shape or arr.shape
iterlist = []
arraylist = []
- self._create_iter(iterlist, arraylist, arr, res_shape, chunks, dim)
+ self._create_iter(iterlist, arraylist, arr, res_shape, chunks)
return NumpyEvalFrame(iterlist, arraylist)
@@ -150,7 +147,7 @@
assert concr.dtype is self.dtype
self.array_no = _add_ptr_to_cache(concr.storage, cache)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
from pypy.module.micronumpy.interp_numarray import ConcreteArray
concr = arr.get_concrete()
assert isinstance(concr, ConcreteArray)
@@ -178,7 +175,7 @@
def _invent_array_numbering(self, arr, cache):
pass
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
if self.iter_no >= len(iterlist):
iter = ConstantIterator()
iterlist.append(iter)
@@ -219,12 +216,12 @@
assert isinstance(other, VirtualSliceSignature)
return self.child.eq(other.child, compare_array_no)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
from pypy.module.micronumpy.interp_numarray import VirtualSlice
assert isinstance(arr, VirtualSlice)
chunklist.append(arr.chunks)
self.child._create_iter(iterlist, arraylist, arr.child, res_shape,
- chunklist, dim)
+ chunklist)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import VirtualSlice
@@ -260,11 +257,11 @@
assert isinstance(arr, Call1)
self.child._invent_array_numbering(arr.values, cache)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
from pypy.module.micronumpy.interp_numarray import Call1
assert isinstance(arr, Call1)
self.child._create_iter(iterlist, arraylist, arr.values, res_shape,
- chunklist, dim)
+ chunklist)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import Call1
@@ -305,14 +302,14 @@
self.left._invent_numbering(cache, allnumbers)
self.right._invent_numbering(cache, allnumbers)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
from pypy.module.micronumpy.interp_numarray import Call2
assert isinstance(arr, Call2)
self.left._create_iter(iterlist, arraylist, arr.left, res_shape,
- chunklist, dim)
+ chunklist)
self.right._create_iter(iterlist, arraylist, arr.right, res_shape,
- chunklist, dim)
+ chunklist)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import Call2
@@ -327,9 +324,9 @@
class ReduceSignature(Call2):
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
self.right._create_iter(iterlist, arraylist, arr, res_shape,
- chunklist, dim)
+ chunklist)
def _invent_numbering(self, cache, allnumbers):
self.right._invent_numbering(cache, allnumbers)
@@ -341,33 +338,39 @@
return self.right.eval(frame, arr)
def debug_repr(self):
- return 'ReduceSig(%s, %s, %s)' % (self.name, self.left.debug_repr(),
- self.right.debug_repr())
+ return 'ReduceSig(%s, %s)' % (self.name, self.right.debug_repr())
class AxisReduceSignature(Call2):
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist,
dim):
- from pypy.module.micronumpy.interp_numarray import ConcreteArray
- concr = arr.get_concrete()
- assert isinstance(concr, ConcreteArray)
- storage = concr.storage
- if self.iter_no >= len(iterlist):
- _iter = axis_iter_from_arr(concr, dim)
- from interp_iter import AxisIterator
- assert isinstance(_iter, AxisIterator)
- iterlist.append(_iter)
- if self.array_no >= len(arraylist):
- arraylist.append(storage)
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ from pypy.module.micronumpy.interp_numarray import AxisReduce
+
+ assert isinstance(arr, AxisReduce)
+ assert not iterlist # we assume that later in eval
+ iterlist.append(AxisIterator(arr.dim, arr.right.shape,
+ arr.left.strides,
+ arr.left.backstrides))
+ self.right._create_iter(iterlist, arraylist, arr.right,
arr.right.shape,
+ chunklist)
def _invent_numbering(self, cache, allnumbers):
+ no = len(allnumbers)
+ allnumbers.append(no)
self.right._invent_numbering(cache, allnumbers)
def _invent_array_numbering(self, arr, cache):
- self.right._invent_array_numbering(arr, cache)
+ from pypy.module.micronumpy.interp_numarray import AxisReduce
+
+ assert isinstance(arr, AxisReduce)
+ self.right._invent_array_numbering(arr.right, cache)
def eval(self, frame, arr):
- return self.right.eval(frame, arr)
+ if frame.iterators[0].axis_done:
+ arr.left.setitem(frame.iterators[0].offset, frame.value)
+ frame.value = frame.identity
+ v = self.right.eval(frame, arr.right).convert_to(self.calc_dtype)
+ print v.value, frame.value.value
+ frame.value = self.binfunc(self.calc_dtype, frame.value, v)
+ return frame.value
def debug_repr(self):
- return 'AxisReduceSig(%s, %s, %s)' % (self.name,
self.left.debug_repr(),
- self.right.debug_repr())
-
+ return 'AxisReduceSig(%s, %s)' % (self.name, self.right.debug_repr())
diff --git a/pypy/module/micronumpy/test/test_numarray.py
b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -745,7 +745,7 @@
raises(TypeError, 'a.sum(2, 3)')
- def test_reduceND(self):
+ def test_reduce_nd(self):
from numpypy import arange
a = arange(15).reshape(5, 3)
assert a.sum() == 105
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
http://mail.python.org/mailman/listinfo/pypy-commit
