[pypy-commit] pypy numpypy-axisops: turn AxisIterator inside out
mattip
noreply at buildbot.pypy.org
Fri Jan 6 12:23:58 CET 2012
Author: mattip
Branch: numpypy-axisops
Changeset: r51065:44fd891a3a1c
Date: 2012-01-06 13:01 +0200
http://bitbucket.org/pypy/pypy/changeset/44fd891a3a1c/
Log: turn AxisIterator inside out
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
@@ -103,60 +103,87 @@
def next(self, shapelen):
return self
-def axis_iter_from_arr(arr, dim=-1, start=None):
- if start is None:
- start = []
+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, start)
+ dim)
class AxisIterator(object):
- """ This object will return offsets of each start of a stride on the
- desired dimension, starting at "start" which is an index along
- each axis
+ """ 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, start):
+ def __init__(self, arr_start, strides, backstrides, shape, dim):
self.shape = shape
self.shapelen = len(shape)
self.indices = [0] * len(shape)
- self.done = False
+ self._done = False
+ self.axis_done = False
self.offset = arr_start
- self.dim = len(shape) - 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
- if dim >= 0:
- self.dim = dim
- if len(start) == len(shape):
- for i in range(len(start)):
- self.offset += strides[i] * start[i]
+
+ def done(self):
+ return self._done
def next(self, shapelen):
#shapelen will always be one less than self.shapelen
offset = self.offset
+ axis_done = False
indices = [0] * self.shapelen
for i in range(self.shapelen):
indices[i] = self.indices[i]
- for i in range(self.shapelen - 1, -1, -1):
- if i == self.dim:
- continue
+ 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
indices[i] = 0
offset -= self.backstrides[i]
else:
- self.done = True
+ self._done = True
res = instantiate(AxisIterator)
+ res.axis_done = axis_done
res.offset = offset
res.indices = indices
res.strides = self.strides
+ res.dim_order = self.dim_order
res.backstrides = self.backstrides
res.shape = self.shape
res.shapelen = self.shapelen
res.dim = self.dim
- res.done = self.done
+ res._done = self._done
return res
+
+# ------ other iterators that are not part of the computation frame ----------
+class SkipLastAxisIterator(object):
+ def __init__(self, arr):
+ self.arr = arr
+ self.indices = [0] * (len(arr.shape) - 1)
+ self.done = False
+ self.offset = arr.start
+ def next(self):
+ for i in range(len(self.arr.shape) - 2, -1, -1):
+ if self.indices[i] < self.arr.shape[i] - 1:
+ self.indices[i] += 1
+ self.offset += self.arr.strides[i]
+ break
+ else:
+ self.indices[i] = 0
+ self.offset -= self.arr.backstrides[i]
+ else:
+ self.done = True
+
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
@@ -8,8 +8,8 @@
from pypy.rpython.lltypesystem import lltype, rffi
from pypy.tool.sourcetools import func_with_new_name
from pypy.rlib.rstring import StringBuilder
-from pypy.module.micronumpy.interp_iter import ArrayIterator,\
- view_iter_from_arr, OneDimIterator, axis_iter_from_arr
+from pypy.module.micronumpy.interp_iter import ArrayIterator, OneDimIterator,\
+ view_iter_from_arr, SkipLastAxisIterator
numpy_driver = jit.JitDriver(
greens=['shapelen', 'sig'],
@@ -603,11 +603,12 @@
def getitem(self, item):
raise NotImplementedError
- def find_sig(self, res_shape=None):
+ def find_sig(self, res_shape=None, arr=None):
""" find a correct signature for the array
"""
res_shape = res_shape or self.shape
- return signature.find_sig(self.create_sig(res_shape), self)
+ arr = arr or self
+ return signature.find_sig(self.create_sig(res_shape), arr)
def descr_array_iface(self, space):
if not self.shape:
@@ -756,9 +757,10 @@
for s in shape:
self.size *= s
self.binfunc = binfunc
- self.res_dtype = res_dtype
+ self.dtype = res_dtype
self.dim = dim
self.identity = identity
+ self.computing = False
def _del_sources(self):
self.values = None
@@ -767,46 +769,42 @@
def create_sig(self, res_shape):
if self.forced_result is not None:
return self.forced_result.create_sig(res_shape)
- return signature.ReduceSignature(self.binfunc, self.name, self.res_dtype,
- signature.ViewSignature(self.res_dtype),
- self.values.create_sig(res_shape))
+ return signature.ReduceSignature(self.binfunc, self.name, self.dtype,
+ signature.ScalarSignature(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
def compute(self):
- dtype = self.res_dtype
+ self.computing = True
+ dtype = self.dtype
result = W_NDimArray(self.size, self.shape, dtype)
self.values = self.values.get_concrete()
shapelen = len(result.shape)
objlen = len(self.values.shape)
- target_len = self.values.shape[self.dim]
- sig = self.values.find_sig(result.shape)
- #sig = self.create_sig(result.shape)
+ sig = self.find_sig(res_shape=result.shape,arr=self.values)
ri = ArrayIterator(result.size)
- si = axis_iter_from_arr(self.values, self.dim)
- while not ri.done():
- # explanation: we want to start the frame at the beginning of
- # an axis: use si.indices to create a chunk (slice)
- # in self.values
- chunks = []
- for i in range(objlen):
- if i == self.dim:
- chunks.append((0, target_len, 1, target_len))
- else:
- chunks.append((si.indices[i], 0, 0, 1))
- frame = sig.create_frame(self.values,
- res_shape=[target_len], chunks = [chunks, ])
- if self.identity is None:
- value = sig.eval(frame, self.values).convert_to(dtype)
- frame.next(shapelen)
- else:
- value = self.identity.convert_to(dtype)
- while not frame.done():
- assert isinstance(sig, signature.ViewSignature)
- nextval = sig.eval(frame, self.values).convert_to(dtype)
- value = self.binfunc(dtype, value, nextval)
- frame.next(shapelen)
+ frame = sig.create_frame(self.values, dim=self.dim)
+ value = self.get_identity(sig, frame, shapelen)
+ while not frame.done():
+ #XXX add jit_merge_point ?
+ if frame.iterators[0].axis_done:
+ value = self.get_identity(sig, frame, shapelen)
+ ri = ri.next(shapelen)
+ assert isinstance(sig, signature.ReduceSignature)
+ nextval = sig.eval(frame, self.values).convert_to(dtype)
+ value = self.binfunc(dtype, value, nextval)
result.dtype.setitem(result.storage, ri.offset, value)
- ri = ri.next(shapelen)
- si = si.next(shapelen)
+ frame.next(shapelen)
+ assert ri.done
return result
@@ -1036,19 +1034,19 @@
self.size * itemsize
)
else:
- dest = axis_iter_from_arr(self)
- source = axis_iter_from_arr(w_value)
+ dest = SkipLastAxisIterator(self)
+ source = SkipLastAxisIterator(w_value)
while not dest.done:
rffi.c_memcpy(
rffi.ptradd(self.storage, dest.offset * itemsize),
rffi.ptradd(w_value.storage, source.offset * itemsize),
self.shape[-1] * itemsize
)
- source = source.next(shapelen)
- dest = dest.next(shapelen)
+ source.next()
+ dest.next()
def _sliceloop(self, source, res_shape):
- sig = source.find_sig(res_shape)
+ sig = source.find_sig(res_shape=res_shape)
frame = sig.create_frame(source, res_shape)
res_iter = view_iter_from_arr(self)
shapelen = len(res_shape)
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
@@ -135,7 +135,7 @@
sig = find_sig(ReduceSignature(self.func, self.name, dtype,
ScalarSignature(dtype),
obj.create_sig(obj.shape)), obj)
- frame = sig.create_frame(obj)
+ frame = sig.create_frame(obj,dim=-1)
if self.identity is None:
value = sig.eval(frame, obj).convert_to(dtype)
frame.next(shapelen)
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
+ OneDimIterator, ConstantIterator, axis_iter_from_arr
from pypy.module.micronumpy.strides import calculate_slice_strides
from pypy.rlib.jit import hint, unroll_safe, promote
@@ -95,13 +95,13 @@
allnumbers.append(no)
self.iter_no = no
- def create_frame(self, arr, res_shape=None, chunks=None):
+ def create_frame(self, arr, res_shape=None, chunks=None, dim=-1):
if chunks is None:
chunks = []
res_shape = res_shape or arr.shape
iterlist = []
arraylist = []
- self._create_iter(iterlist, arraylist, arr, res_shape, chunks)
+ self._create_iter(iterlist, arraylist, arr, res_shape, chunks, dim)
return NumpyEvalFrame(iterlist, arraylist)
@@ -143,7 +143,7 @@
assert isinstance(concr, ConcreteArray)
self.array_no = _add_ptr_to_cache(concr.storage, cache)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ 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)
@@ -171,7 +171,7 @@
def _invent_array_numbering(self, arr, cache):
pass
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist, dim):
if self.iter_no >= len(iterlist):
iter = ConstantIterator()
iterlist.append(iter)
@@ -212,12 +212,12 @@
assert isinstance(other, VirtualSliceSignature)
return self.child.eq(other.child, compare_array_no)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist, dim):
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)
+ chunklist, dim)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import VirtualSlice
@@ -253,11 +253,11 @@
assert isinstance(arr, Call1)
self.child._invent_array_numbering(arr.values, cache)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist, dim):
from pypy.module.micronumpy.interp_numarray import Call1
assert isinstance(arr, Call1)
self.child._create_iter(iterlist, arraylist, arr.values, res_shape,
- chunklist)
+ chunklist, dim)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import Call1
@@ -298,14 +298,14 @@
self.left._invent_numbering(cache, allnumbers)
self.right._invent_numbering(cache, allnumbers)
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist, dim):
from pypy.module.micronumpy.interp_numarray import Call2
assert isinstance(arr, Call2)
self.left._create_iter(iterlist, arraylist, arr.left, res_shape,
- chunklist)
+ chunklist, dim)
self.right._create_iter(iterlist, arraylist, arr.right, res_shape,
- chunklist)
+ chunklist, dim)
def eval(self, frame, arr):
from pypy.module.micronumpy.interp_numarray import Call2
@@ -319,8 +319,21 @@
self.right.debug_repr())
class ReduceSignature(Call2):
- def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist):
- self.right._create_iter(iterlist, arraylist, arr, res_shape, chunklist)
+ def _create_iter(self, iterlist, arraylist, arr, res_shape, chunklist, dim):
+ if dim<0:
+ self.right._create_iter(iterlist, arraylist, arr, res_shape, chunklist, dim)
+ else:
+ 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 _invent_numbering(self, cache, allnumbers):
self.right._invent_numbering(cache, allnumbers)
diff --git a/pypy/module/micronumpy/test/test_iterators.py b/pypy/module/micronumpy/test/test_iterators.py
deleted file mode 100644
--- a/pypy/module/micronumpy/test/test_iterators.py
+++ /dev/null
@@ -1,61 +0,0 @@
-
-from pypy.module.micronumpy.interp_iter import axis_iter_from_arr
-from pypy.module.micronumpy.interp_numarray import W_NDimArray
-
-
-class MockDtype(object):
- def malloc(self, size):
- return None
-
-
-class TestAxisIteratorDirect(object):
- def test_axis_iterator(self):
- a = W_NDimArray(5 * 3, [5, 3], MockDtype(), 'C')
- i = axis_iter_from_arr(a)
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(1)
- assert ret == [0, 3, 6, 9, 12]
- a = W_NDimArray(7 * 5 * 3, [7, 5, 3], MockDtype(), 'C')
- i = axis_iter_from_arr(a)
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(1)
- assert ret == [3 * v for v in range(7 * 5)]
- i = axis_iter_from_arr(a, 2)
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(1)
- assert ret == [3 * v for v in range(7 * 5)]
- i = axis_iter_from_arr(a, 1)
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(1)
- assert ret == [ 0, 1, 2, 15, 16, 17, 30, 31, 32, 45, 46, 47,
- 60, 61, 62, 75, 76, 77, 90, 91, 92]
-
- def test_axis_iterator_with_start(self):
- a = W_NDimArray(7 * 5 * 3, [7, 5, 3], MockDtype(), 'C')
- i = axis_iter_from_arr(a, start=[0, 0, 0])
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(2)
- assert ret == [3 * v for v in range(7 * 5)]
- i = axis_iter_from_arr(a, start=[1, 1, 0])
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(2)
- assert ret == [3 * v + 18 for v in range(7 * 5)]
- i = axis_iter_from_arr(a, 1, [2, 0, 2])
- ret = []
- while not i.done:
- ret.append(i.offset)
- i = i.next(2)
- assert ret == [v + 32 for v in [ 0, 1, 2, 15, 16, 17, 30, 31, 32,
- 45, 46, 47, 60, 61, 62, 75, 76, 77, 90, 91, 92]]
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
@@ -246,6 +246,10 @@
c = b.copy()
assert (c == b).all()
+ a = arange(15).reshape(5,3)
+ b = a.copy()
+ assert (b == a).all()
+
def test_iterator_init(self):
from numpypy import array
a = array(range(5))
More information about the pypy-commit
mailing list