Author: Maciej Fijalkowski <[email protected]>
Branch:
Changeset: r51858:caa2d47782ba
Date: 2012-01-27 21:00 +0200
http://bitbucket.org/pypy/pypy/changeset/caa2d47782ba/
Log: (mattip, fijal reviewing) Merge numppy-flatitter branch (yes, with
all the typos). This improves flat iterator and adds __getitem__ and
__setitem__ to it.
diff --git a/pypy/module/micronumpy/compile.py
b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -32,7 +32,8 @@
class BadToken(Exception):
pass
-SINGLE_ARG_FUNCTIONS = ["sum", "prod", "max", "min", "all", "any", "unegative"]
+SINGLE_ARG_FUNCTIONS = ["sum", "prod", "max", "min", "all", "any",
+ "unegative", "flat"]
class FakeSpace(object):
w_ValueError = None
@@ -396,6 +397,8 @@
elif self.name == "unegative":
neg = interp_ufuncs.get(interp.space).negative
w_res = neg.call(interp.space, [arr])
+ elif self.name == "flat":
+ w_res = arr.descr_get_flatiter(interp.space)
else:
assert False # unreachable code
if isinstance(w_res, BaseArray):
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
@@ -4,6 +4,49 @@
from pypy.module.micronumpy.strides import calculate_broadcast_strides,\
calculate_slice_strides
+""" This is a mini-tutorial on iterators, strides, and
+memory layout. It assumes you are familiar with the terms, see
+http://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html
+for a more gentle introduction.
+
+Given an array x: x.shape == [5,6],
+
+At which byte in x.data does the item x[3,4] begin?
+if x.strides==[1,5]:
+ pData = x.pData + (x.start + 3*1 + 4*5)*sizeof(x.pData[0])
+ pData = x.pData + (x.start + 24) * sizeof(x.pData[0])
+so the offset of the element is 24 elements after the first
+
+What is the next element in x after coordinates [3,4]?
+if x.order =='C':
+ next == [3,5] => offset is 28
+if x.order =='F':
+ next == [4,4] => offset is 24
+so for the strides [1,5] x is 'F' contiguous
+likewise, for the strides [6,1] x would be 'C' contiguous.
+
+Iterators have an internal representation of the current coordinates
+(indices), the array, strides, and backstrides. A short digression to
+explain backstrides: what is the coordinate and offset after [3,5] in
+the example above?
+if x.order == 'C':
+ next == [4,0] => offset is 4
+if x.order == 'F':
+ next == [4,5] => offset is 25
+Note that in 'C' order we stepped BACKWARDS 24 while 'overflowing' a
+shape dimension
+ which is back 25 and forward 1,
+ which is x.strides[1] * (x.shape[1] - 1) + x.strides[0]
+so if we precalculate the overflow backstride as
+[x.strides[i] * (x.shape[i] - 1) for i in range(len(x.shape))]
+we can go faster.
+All the calculations happen in next()
+
+next_step_x() tries to do the iteration for a number of steps at once,
+but then we cannot gaurentee that we only overflow one single shape
+dimension, perhaps we could overflow times in one big step.
+"""
+
# structures to describe slicing
class Chunk(object):
@@ -55,9 +98,9 @@
self.size = size
def next(self, shapelen):
- return self._next(1)
+ return self.next_skip_x(1)
- def _next(self, ofs):
+ def next_skip_x(self, ofs):
arr = instantiate(ArrayIterator)
arr.size = self.size
arr.offset = self.offset + ofs
@@ -65,7 +108,7 @@
def next_no_increase(self, shapelen):
# a hack to make JIT believe this is always virtual
- return self._next(0)
+ return self.next_skip_x(0)
def done(self):
return self.offset >= self.size
@@ -137,6 +180,36 @@
res._done = done
return res
+ @jit.unroll_safe
+ def next_skip_x(self, shapelen, step):
+ shapelen = jit.promote(len(self.res_shape))
+ offset = self.offset
+ indices = [0] * shapelen
+ for i in range(shapelen):
+ indices[i] = self.indices[i]
+ done = False
+ for i in range(shapelen - 1, -1, -1):
+ if indices[i] < self.res_shape[i] - step:
+ indices[i] += step
+ offset += self.strides[i] * step
+ break
+ else:
+ remaining_step = (indices[i] + step) // self.res_shape[i]
+ this_i_step = step - remaining_step * self.res_shape[i]
+ offset += self.strides[i] * this_i_step
+ indices[i] = indices[i] + this_i_step
+ step = remaining_step
+ else:
+ done = True
+ res = instantiate(ViewIterator)
+ res.offset = offset
+ res.indices = indices
+ res.strides = self.strides
+ res.backstrides = self.backstrides
+ res.res_shape = self.res_shape
+ res._done = done
+ return res
+
def apply_transformations(self, arr, transformations):
v = BaseIterator.apply_transformations(self, arr, transformations)
if len(arr.shape) == 1:
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
@@ -4,7 +4,7 @@
from pypy.interpreter.typedef import TypeDef, GetSetProperty
from pypy.module.micronumpy import interp_ufuncs, interp_dtype, signature,\
interp_boxes
-from pypy.module.micronumpy.strides import calculate_slice_strides
+from pypy.module.micronumpy.strides import calculate_slice_strides, to_coords
from pypy.rlib import jit
from pypy.rpython.lltypesystem import lltype, rffi
from pypy.tool.sourcetools import func_with_new_name
@@ -59,6 +59,18 @@
name='numpy_filterset',
)
+flat_get_driver = jit.JitDriver(
+ greens=['shapelen', 'base'],
+ reds=['step', 'ri', 'basei', 'res'],
+ name='numpy_flatget',
+)
+
+flat_set_driver = jit.JitDriver(
+ greens=['shapelen', 'base'],
+ reds=['step', 'ai', 'lngth', 'arr', 'basei'],
+ name='numpy_flatset',
+)
+
def _find_shape_and_elems(space, w_iterable):
shape = [space.len_w(w_iterable)]
batch = space.listview(w_iterable)
@@ -1455,31 +1467,119 @@
@jit.unroll_safe
def __init__(self, arr):
arr = arr.get_concrete()
- size = 1
- for sh in arr.shape:
- size *= sh
- self.strides = [arr.strides[-1]]
- self.backstrides = [arr.backstrides[-1]]
- ViewArray.__init__(self, size, [size], arr.dtype, arr.order,
- arr)
self.shapelen = len(arr.shape)
- self.iter = OneDimIterator(arr.start, self.strides[0],
- self.shape[0])
+ sig = arr.find_sig()
+ self.iter = sig.create_frame(arr).get_final_iter()
+ self.base = arr
+ self.index = 0
+ ViewArray.__init__(self, arr.size, [arr.size], arr.dtype, arr.order,
arr)
def descr_next(self, space):
if self.iter.done():
raise OperationError(space.w_StopIteration, space.w_None)
- result = self.getitem(self.iter.offset)
+ result = self.base.getitem(self.iter.offset)
self.iter = self.iter.next(self.shapelen)
+ self.index += 1
return result
def descr_iter(self):
return self
+ def descr_index(self, space):
+ return space.wrap(self.index)
+
+ def descr_coords(self, space):
+ coords, step, lngth = to_coords(space, self.base.shape,
+ self.base.size, self.base.order,
+ space.wrap(self.index))
+ return space.newtuple([space.wrap(c) for c in coords])
+
+ @jit.unroll_safe
+ def descr_getitem(self, space, w_idx):
+ if not (space.isinstance_w(w_idx, space.w_int) or
+ space.isinstance_w(w_idx, space.w_slice)):
+ raise OperationError(space.w_IndexError,
+ space.wrap('unsupported iterator index'))
+ base = self.base
+ start, stop, step, lngth = space.decode_index4(w_idx, base.size)
+ # setslice would have been better, but flat[u:v] for arbitrary
+ # shapes of array a cannot be represented as a[x1:x2, y1:y2]
+ basei = ViewIterator(base.start, base.strides,
+ base.backstrides,base.shape)
+ shapelen = len(base.shape)
+ basei = basei.next_skip_x(shapelen, start)
+ if lngth <2:
+ return base.getitem(basei.offset)
+ ri = ArrayIterator(lngth)
+ res = W_NDimArray(lngth, [lngth], base.dtype,
+ base.order)
+ while not ri.done():
+ flat_get_driver.jit_merge_point(shapelen=shapelen,
+ base=base,
+ basei=basei,
+ step=step,
+ res=res,
+ ri=ri,
+ )
+ w_val = base.getitem(basei.offset)
+ res.setitem(ri.offset,w_val)
+ basei = basei.next_skip_x(shapelen, step)
+ ri = ri.next(shapelen)
+ return res
+
+ def descr_setitem(self, space, w_idx, w_value):
+ if not (space.isinstance_w(w_idx, space.w_int) or
+ space.isinstance_w(w_idx, space.w_slice)):
+ raise OperationError(space.w_IndexError,
+ space.wrap('unsupported iterator index'))
+ base = self.base
+ start, stop, step, lngth = space.decode_index4(w_idx, base.size)
+ arr = convert_to_array(space, w_value)
+ ai = 0
+ basei = ViewIterator(base.start, base.strides,
+ base.backstrides,base.shape)
+ shapelen = len(base.shape)
+ basei = basei.next_skip_x(shapelen, start)
+ while lngth > 0:
+ flat_set_driver.jit_merge_point(shapelen=shapelen,
+ basei=basei,
+ base=base,
+ step=step,
+ arr=arr,
+ ai=ai,
+ lngth=lngth,
+ )
+ v = arr.getitem(ai).convert_to(base.dtype)
+ base.setitem(basei.offset, v)
+ # need to repeat input values until all assignments are done
+ ai = (ai + 1) % arr.size
+ basei = basei.next_skip_x(shapelen, step)
+ lngth -= 1
+
+ def create_sig(self):
+ return signature.FlatSignature(self.base.dtype)
+
+ def descr_base(self, space):
+ return space.wrap(self.base)
+
W_FlatIterator.typedef = TypeDef(
'flatiter',
+ #__array__ = #MISSING
+ __iter__ = interp2app(W_FlatIterator.descr_iter),
+ __getitem__ = interp2app(W_FlatIterator.descr_getitem),
+ __setitem__ = interp2app(W_FlatIterator.descr_setitem),
+ __eq__ = interp2app(BaseArray.descr_eq),
+ __ne__ = interp2app(BaseArray.descr_ne),
+ __lt__ = interp2app(BaseArray.descr_lt),
+ __le__ = interp2app(BaseArray.descr_le),
+ __gt__ = interp2app(BaseArray.descr_gt),
+ __ge__ = interp2app(BaseArray.descr_ge),
+ #__sizeof__ #MISSING
+ base = GetSetProperty(W_FlatIterator.descr_base),
+ index = GetSetProperty(W_FlatIterator.descr_index),
+ coords = GetSetProperty(W_FlatIterator.descr_coords),
next = interp2app(W_FlatIterator.descr_next),
- __iter__ = interp2app(W_FlatIterator.descr_iter),
+
)
W_FlatIterator.acceptable_as_base_class = False
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
@@ -224,6 +224,18 @@
return ViewIterator(arr.start, arr.strides, arr.backstrides,
arr.shape).apply_transformations(arr, transforms)
+class FlatSignature(ViewSignature):
+ def debug_repr(self):
+ return 'Flat'
+
+ def allocate_iter(self, arr, transforms):
+ from pypy.module.micronumpy.interp_numarray import W_FlatIterator
+ assert isinstance(arr, W_FlatIterator)
+ return ViewIterator(arr.base.start, arr.base.strides,
+ arr.base.backstrides,
+ arr.base.shape).apply_transformations(arr.base,
+ transforms)
+
class VirtualSliceSignature(Signature):
def __init__(self, child):
self.child = child
diff --git a/pypy/module/micronumpy/test/test_compile.py
b/pypy/module/micronumpy/test/test_compile.py
--- a/pypy/module/micronumpy/test/test_compile.py
+++ b/pypy/module/micronumpy/test/test_compile.py
@@ -245,3 +245,11 @@
a -> 3
""")
assert interp.results[0].value == 11
+
+ def test_flat_iter(self):
+ interp = self.run('''
+ a = |30|
+ b = flat(a)
+ b -> 3
+ ''')
+ assert interp.results[0].value == 3
diff --git a/pypy/module/micronumpy/test/test_iter.py
b/pypy/module/micronumpy/test/test_iter.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/micronumpy/test/test_iter.py
@@ -0,0 +1,88 @@
+from pypy.module.micronumpy.interp_iter import ViewIterator
+
+class TestIterDirect(object):
+ def test_C_viewiterator(self):
+ #Let's get started, simple iteration in C order with
+ #contiguous layout => strides[-1] is 1
+ start = 0
+ shape = [3, 5]
+ strides = [5, 1]
+ backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
+ assert backstrides == [10, 4]
+ i = ViewIterator(start, strides, backstrides, shape)
+ i = i.next(2)
+ i = i.next(2)
+ i = i.next(2)
+ assert i.offset == 3
+ assert not i.done()
+ assert i.indices == [0,3]
+ #cause a dimension overflow
+ i = i.next(2)
+ i = i.next(2)
+ assert i.offset == 5
+ assert i.indices == [1,0]
+
+ #Now what happens if the array is transposed? strides[-1] != 1
+ # therefore layout is non-contiguous
+ strides = [1, 3]
+ backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
+ assert backstrides == [2, 12]
+ i = ViewIterator(start, strides, backstrides, shape)
+ i = i.next(2)
+ i = i.next(2)
+ i = i.next(2)
+ assert i.offset == 9
+ assert not i.done()
+ assert i.indices == [0,3]
+ #cause a dimension overflow
+ i = i.next(2)
+ i = i.next(2)
+ assert i.offset == 1
+ assert i.indices == [1,0]
+
+ def test_C_viewiterator_step(self):
+ #iteration in C order with #contiguous layout => strides[-1] is 1
+ #skip less than the shape
+ start = 0
+ shape = [3, 5]
+ strides = [5, 1]
+ backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
+ assert backstrides == [10, 4]
+ i = ViewIterator(start, strides, backstrides, shape)
+ i = i.next_skip_x(2,2)
+ i = i.next_skip_x(2,2)
+ i = i.next_skip_x(2,2)
+ assert i.offset == 6
+ assert not i.done()
+ assert i.indices == [1,1]
+ #And for some big skips
+ i = i.next_skip_x(2,5)
+ assert i.offset == 11
+ assert i.indices == [2,1]
+ i = i.next_skip_x(2,5)
+ # Note: the offset does not overflow but recycles,
+ # this is good for broadcast
+ assert i.offset == 1
+ assert i.indices == [0,1]
+ assert i.done()
+
+ #Now what happens if the array is transposed? strides[-1] != 1
+ # therefore layout is non-contiguous
+ strides = [1, 3]
+ backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
+ assert backstrides == [2, 12]
+ i = ViewIterator(start, strides, backstrides, shape)
+ i = i.next_skip_x(2,2)
+ i = i.next_skip_x(2,2)
+ i = i.next_skip_x(2,2)
+ assert i.offset == 4
+ assert i.indices == [1,1]
+ assert not i.done()
+ i = i.next_skip_x(2,5)
+ assert i.offset == 5
+ assert i.indices == [2,1]
+ assert not i.done()
+ i = i.next_skip_x(2,5)
+ assert i.indices == [0,1]
+ assert i.offset == 3
+ assert i.done()
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
@@ -301,6 +301,12 @@
for i in xrange(5):
assert a[i] == b[i]
+ def test_getitem_nd(self):
+ from _numpypy import arange
+ a = arange(15).reshape(3, 5)
+ assert a[1, 3] == 8
+ assert a.T[1, 2] == 11
+
def test_setitem(self):
from _numpypy import array
a = array(range(5))
@@ -974,7 +980,7 @@
assert debug_repr(a + a) == 'Call2(add, Array, Array)'
assert debug_repr(a[::2]) == 'Slice'
assert debug_repr(a + 2) == 'Call2(add, Array, Scalar)'
- assert debug_repr(a + a.flat) == 'Call2(add, Array, Slice)'
+ assert debug_repr(a + a.flat) == 'Call2(add, Array, Flat)'
assert debug_repr(sin(a)) == 'Call1(sin, Array)'
b = a + a
@@ -1343,7 +1349,7 @@
assert(b[:, 0] == a[0, :]).all()
def test_flatiter(self):
- from _numpypy import array, flatiter
+ from _numpypy import array, flatiter, arange
a = array([[10, 30], [40, 60]])
f_iter = a.flat
assert f_iter.next() == 10
@@ -1356,6 +1362,9 @@
for k in a.flat:
s += k
assert s == 140
+ a = arange(10).reshape(5, 2)
+ raises(IndexError, 'a.flat[(1, 2)]')
+ assert a.flat.base is a
def test_flatiter_array_conv(self):
from _numpypy import array, dot
@@ -1367,6 +1376,75 @@
a = ones((2, 2))
assert list(((a + a).flat)) == [2, 2, 2, 2]
+ def test_flatiter_getitem(self):
+ from _numpypy import arange
+ a = arange(10)
+ assert a.flat[3] == 3
+ assert a[2:].flat[3] == 5
+ assert (a + a).flat[3] == 6
+ assert a[::2].flat[3] == 6
+ assert a.reshape(2,5).flat[3] == 3
+ b = a.reshape(2,5).flat
+ b.next()
+ b.next()
+ b.next()
+ assert b[3] == 3
+ assert (b[::3] == [0, 3, 6, 9]).all()
+ assert (b[2::5] == [2, 7]).all()
+ assert b[-2] == 8
+ raises(IndexError, "b[11]")
+ raises(IndexError, "b[-11]")
+ raises(IndexError, 'b[0, 1]')
+ assert b.index == 3
+ assert b.coords == (0,3)
+
+ def test_flatiter_setitem(self):
+ from _numpypy import arange, array
+ a = arange(12).reshape(3,4)
+ b = a.T.flat
+ b[6::2] = [-1, -2]
+ assert (a == [[0, 1, -1, 3], [4, 5, 6, -1], [8, 9, -2, 11]]).all()
+ b[0:2] = [[[100]]]
+ assert(a[0,0] == 100)
+ assert(a[1,0] == 100)
+ raises(IndexError, 'b[array([10, 11])] == [-20, -40]')
+
+ def test_flatiter_ops(self):
+ from _numpypy import arange, array
+ a = arange(12).reshape(3,4)
+ b = a.T.flat
+ assert (b == [0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11]).all()
+ assert not (b != [0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11]).any()
+ assert ((b >= range(12)) == [True, True, True,False, True, True,
+ False, False, True, False, False, True]).all()
+ assert ((b < range(12)) != [True, True, True,False, True, True,
+ False, False, True, False, False, True]).all()
+ assert ((b <= range(12)) != [False, True, True,False, True, True,
+ False, False, True, False, False, False]).all()
+ assert ((b > range(12)) == [False, True, True,False, True, True,
+ False, False, True, False, False, False]).all()
+ def test_flatiter_view(self):
+ from _numpypy import arange
+ a = arange(10).reshape(5, 2)
+ #no == yet.
+ # a[::2].flat == [0, 1, 4, 5, 8, 9]
+ isequal = True
+ for y,z in zip(a[::2].flat, [0, 1, 4, 5, 8, 9]):
+ if y != z:
+ isequal = False
+ assert isequal == True
+
+ def test_flatiter_transpose(self):
+ from _numpypy import arange
+ a = arange(10).reshape(2,5).T
+ b = a.flat
+ assert (b[:5] == [0, 5, 1, 6, 2]).all()
+ b.next()
+ b.next()
+ b.next()
+ assert b.index == 3
+ assert b.coords == (1,1)
+
def test_slice_copy(self):
from _numpypy import zeros
a = zeros((10, 10))
diff --git a/pypy/module/micronumpy/test/test_zjit.py
b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -12,7 +12,7 @@
from pypy.module.micronumpy.compile import (FakeSpace,
IntObject, Parser, InterpreterState)
from pypy.module.micronumpy.interp_numarray import (W_NDimArray,
- BaseArray)
+ BaseArray, W_FlatIterator)
from pypy.rlib.nonconst import NonConstant
@@ -369,7 +369,71 @@
'setinteriorfield_raw': 1, 'int_add': 2,
'int_ge': 1, 'guard_false': 1, 'jump': 1,
'arraylen_gc': 1})
+ def define_flat_iter():
+ return '''
+ a = |30|
+ b = flat(a)
+ c = b + a
+ c -> 3
+ '''
+ def test_flat_iter(self):
+ result = self.run("flat_iter")
+ assert result == 6
+ self.check_trace_count(1)
+ self.check_simple_loop({'getinteriorfield_raw': 2, 'float_add': 1,
+ 'setinteriorfield_raw': 1, 'int_add': 3,
+ 'int_ge': 1, 'guard_false': 1,
+ 'arraylen_gc': 1, 'jump': 1})
+
+ def define_flat_getitem():
+ return '''
+ a = |30|
+ b = flat(a)
+ b -> 4: -> 6
+ '''
+
+ def test_flat_getitem(self):
+ result = self.run("flat_getitem")
+ assert result == 10.0
+ self.check_trace_count(1)
+ self.check_simple_loop({'getinteriorfield_raw': 1,
+ 'setinteriorfield_raw': 1,
+ 'int_lt': 1,
+ 'int_ge': 1,
+ 'int_add': 3,
+ 'guard_true': 1,
+ 'guard_false': 1,
+ 'arraylen_gc': 2,
+ 'jump': 1})
+
+ def define_flat_setitem():
+ return '''
+ a = |30|
+ b = flat(a)
+ b[4:] = a->:26
+ a -> 5
+ '''
+
+ def test_flat_setitem(self):
+ result = self.run("flat_setitem")
+ assert result == 1.0
+ self.check_trace_count(1)
+ # XXX not ideal, but hey, let's ignore it for now
+ self.check_simple_loop({'getinteriorfield_raw': 1,
+ 'setinteriorfield_raw': 1,
+ 'int_lt': 1,
+ 'int_gt': 1,
+ 'int_add': 4,
+ 'guard_true': 2,
+ 'arraylen_gc': 2,
+ 'jump': 1,
+ 'int_sub': 1,
+ # XXX bad part
+ 'int_and': 1,
+ 'int_mod': 1,
+ 'int_rshift': 1,
+ })
class TestNumpyOld(LLJitMixin):
def setup_class(cls):
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
