Author: Brian Kearns <[email protected]>
Branch:
Changeset: r73850:bb3689880ac8
Date: 2014-10-08 18:42 -0400
http://bitbucket.org/pypy/pypy/changeset/bb3689880ac8/
Log: add numpy iterator goto method
diff --git a/pypy/module/micronumpy/flatiter.py
b/pypy/module/micronumpy/flatiter.py
--- a/pypy/module/micronumpy/flatiter.py
+++ b/pypy/module/micronumpy/flatiter.py
@@ -51,16 +51,15 @@
if not (space.isinstance_w(w_idx, space.w_int) or
space.isinstance_w(w_idx, space.w_slice)):
raise oefmt(space.w_IndexError, 'unsupported iterator index')
- self.state = self.iter.reset(self.state)
- start, stop, step, length = space.decode_index4(w_idx, self.iter.size)
- self.state = self.iter.next_skip_x(self.state, start)
try:
+ start, stop, step, length = space.decode_index4(w_idx,
self.iter.size)
+ state = self.iter.goto(start)
if length == 1:
- return self.iter.getitem(self.state)
+ return self.iter.getitem(state)
base = self.base
res = W_NDimArray.from_shape(space, [length], base.get_dtype(),
base.get_order(), w_instance=base)
- return loop.flatiter_getitem(res, self.iter, self.state, step)
+ return loop.flatiter_getitem(res, self.iter, state, step)
finally:
self.state = self.iter.reset(self.state)
@@ -69,19 +68,18 @@
space.isinstance_w(w_idx, space.w_slice)):
raise oefmt(space.w_IndexError, 'unsupported iterator index')
start, stop, step, length = space.decode_index4(w_idx, self.iter.size)
- self.state = self.iter.reset(self.state)
- self.state = self.iter.next_skip_x(self.state, start)
try:
+ state = self.iter.goto(start)
dtype = self.base.get_dtype()
if length == 1:
try:
val = dtype.coerce(space, w_value)
except OperationError:
raise oefmt(space.w_ValueError, "Error setting single item
of array.")
- self.iter.setitem(self.state, val)
+ self.iter.setitem(state, val)
return
arr = convert_to_array(space, w_value)
- loop.flatiter_setitem(space, dtype, arr, self.iter, self.state,
step, length)
+ loop.flatiter_setitem(space, dtype, arr, self.iter, state, step,
length)
finally:
self.state = self.iter.reset(self.state)
diff --git a/pypy/module/micronumpy/iterators.py
b/pypy/module/micronumpy/iterators.py
--- a/pypy/module/micronumpy/iterators.py
+++ b/pypy/module/micronumpy/iterators.py
@@ -91,7 +91,7 @@
class ArrayIter(object):
_immutable_fields_ = ['array', 'size', 'ndim_m1', 'shape_m1[*]',
- 'strides[*]', 'backstrides[*]']
+ 'strides[*]', 'backstrides[*]', 'factors[*]']
def __init__(self, array, size, shape, strides, backstrides):
assert len(shape) == len(strides) == len(backstrides)
@@ -102,6 +102,15 @@
self.strides = strides
self.backstrides = backstrides
+ ndim = len(shape)
+ factors = [0] * ndim
+ for i in xrange(ndim):
+ if i == 0:
+ factors[ndim-1] = 1
+ else:
+ factors[ndim-i-1] = factors[ndim-i] * shape[ndim-i]
+ self.factors = factors
+
@jit.unroll_safe
def reset(self, state=None):
if state is None:
@@ -154,6 +163,16 @@
assert step > 0
return IterState(self, index, indices, offset)
+ @jit.unroll_safe
+ def goto(self, index):
+ # XXX simplify if self.contiguous (offset = start + index * elsize)
+ offset = self.array.start
+ current = index
+ for i in xrange(len(self.shape_m1)):
+ offset += (current / self.factors[i]) * self.strides[i]
+ current %= self.factors[i]
+ return IterState(self, index, None, offset)
+
def done(self, state):
assert state.iterator is self
return state.index >= self.size
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
@@ -435,7 +435,7 @@
while not ri.done(rs):
flatiter_getitem_driver.jit_merge_point(dtype=dtype)
ri.setitem(rs, base_iter.getitem(base_state))
- base_state = base_iter.next_skip_x(base_state, step)
+ base_state = base_iter.goto(base_state.index + step)
rs = ri.next(rs)
return res
@@ -454,7 +454,7 @@
val = val.convert_to(space, dtype)
arr_iter.setitem(arr_state, val)
# need to repeat i_nput values until all assignments are done
- arr_state = arr_iter.next_skip_x(arr_state, step)
+ arr_state = arr_iter.goto(arr_state.index + step)
val_state = val_iter.next(val_state)
length -= 1
diff --git a/pypy/module/micronumpy/test/test_iterators.py
b/pypy/module/micronumpy/test/test_iterators.py
--- a/pypy/module/micronumpy/test/test_iterators.py
+++ b/pypy/module/micronumpy/test/test_iterators.py
@@ -98,3 +98,21 @@
assert s.indices == [0,1]
assert s.offset == 3
assert i.done(s)
+
+ def test_iterator_goto(self):
+ shape = [3, 5]
+ strides = [1, 3]
+ backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
+ assert backstrides == [2, 12]
+ a = MockArray()
+ a.start = 42
+ i = ArrayIter(a, support.product(shape), shape,
+ strides, backstrides)
+ s = i.reset()
+ assert s.index == 0
+ assert s.indices == [0, 0]
+ assert s.offset == a.start
+ s = i.goto(11)
+ assert s.index == 11
+ assert s.indices is None
+ assert s.offset == a.start + 5
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
@@ -544,16 +544,23 @@
assert result == 10.0
self.check_trace_count(1)
self.check_simple_loop({
- 'getarrayitem_gc': 2,
+ 'arraylen_gc': 2,
+ 'getarrayitem_gc': 1,
'guard_false': 1,
- 'guard_true': 2,
- 'int_add': 6,
+ 'guard_true': 1,
+ 'int_add': 7,
+ 'int_and': 1,
+ 'int_floordiv': 1,
'int_ge': 1,
- 'int_lt': 2,
+ 'int_lt': 1,
+ 'int_mod': 1,
+ 'int_mul': 2,
+ 'int_rshift': 2,
+ 'int_sub': 1,
'jump': 1,
'raw_load': 1,
'raw_store': 1,
- 'setarrayitem_gc': 2,
+ 'setarrayitem_gc': 1,
})
def define_flat_setitem():
@@ -569,17 +576,23 @@
assert result == 1.0
self.check_trace_count(1)
self.check_simple_loop({
- 'getarrayitem_gc': 2,
+ 'arraylen_gc': 2,
+ 'getarrayitem_gc': 1,
'guard_not_invalidated': 1,
- 'guard_true': 3,
- 'int_add': 6,
+ 'guard_true': 2,
+ 'int_add': 7,
+ 'int_and': 1,
+ 'int_floordiv': 1,
'int_gt': 1,
- 'int_lt': 2,
- 'int_sub': 1,
+ 'int_lt': 1,
+ 'int_mod': 1,
+ 'int_mul': 2,
+ 'int_rshift': 2,
+ 'int_sub': 2,
'jump': 1,
'raw_load': 1,
'raw_store': 1,
- 'setarrayitem_gc': 2,
+ 'setarrayitem_gc': 1,
})
def define_dot():
diff --git a/pypy/module/pypyjit/test_pypy_c/test_micronumpy.py
b/pypy/module/pypyjit/test_pypy_c/test_micronumpy.py
--- a/pypy/module/pypyjit/test_pypy_c/test_micronumpy.py
+++ b/pypy/module/pypyjit/test_pypy_c/test_micronumpy.py
@@ -92,47 +92,50 @@
assert log.result == 42.0
loop, = log.loops_by_filename(self.filepath)
assert loop.match("""
- i123 = int_lt(i112, i44)
- guard_true(i123, descr=...)
- p124 = getfield_gc_pure(p121, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
- setarrayitem_gc(p124, 1, 0, descr=<ArrayS 8>)
- setarrayitem_gc(p124, 0, 0, descr=<ArrayS 8>)
- i126 = int_lt(i112, i65)
+ i125 = int_lt(i117, i44)
+ guard_true(i125, descr=...)
+ i126 = int_lt(i117, i50)
guard_true(i126, descr=...)
- i127 = int_sub(i75, i112)
- i128 = int_lt(0, i127)
- guard_false(i128, descr=...)
- i129 = int_floordiv(i112, i75)
- i130 = int_mul(i129, i75)
- i131 = int_sub(i112, i130)
- i132 = int_rshift(i131, 63)
- i133 = int_add(i129, i132)
- i134 = int_mul(i133, i75)
- i135 = int_sub(i112, i134)
- i136 = int_mul(i91, i135)
- i137 = int_add(i64, i136)
- i138 = int_sub(i98, i133)
- setarrayitem_gc(p124, 1, i135, descr=<ArrayS 8>)
- i139 = int_lt(0, i138)
- guard_true(i139, descr=...)
- i140 = int_mul(i102, i133)
- i141 = int_add(i137, i140)
- f142 = raw_load(i108, i141, descr=<ArrayF 8>)
- i143 = int_add(i112, 1)
- setarrayitem_gc(p124, 1, 0, descr=<ArrayS 8>)
+ i127 = int_floordiv(i117, i61)
+ i128 = int_mul(i127, i61)
+ i129 = int_sub(i117, i128)
+ i130 = int_rshift(i129, 63)
+ i131 = int_add(i127, i130)
+ i132 = int_mul(i131, i71)
+ i133 = int_add(i55, i132)
+ i134 = int_mod(i117, i61)
+ i135 = int_rshift(i134, 63)
+ i136 = int_and(i61, i135)
+ i137 = int_add(i134, i136)
+ i138 = int_floordiv(i137, i80)
+ i139 = int_mul(i138, i80)
+ i140 = int_sub(i137, i139)
+ i141 = int_rshift(i140, 63)
+ i142 = int_add(i138, i141)
+ i143 = int_mul(i142, i88)
+ i144 = int_add(i133, i143)
+ i145 = int_mod(i137, i80)
+ i146 = int_rshift(i145, 63)
+ i147 = int_and(i80, i146)
+ i148 = int_add(i145, i147)
+ f149 = raw_load(i100, i144, descr=<ArrayF 8>)
+ p150 = getfield_gc_pure(p123, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
+ i151 = int_add(i117, 1)
+ setarrayitem_gc(p150, 1, 0, descr=<ArrayS 8>)
+ setarrayitem_gc(p150, 0, 0, descr=<ArrayS 8>)
guard_not_invalidated(descr=...)
- i144 = getfield_raw(ticker_address, descr=<FieldS
pypysig_long_struct.c_value 0>)
- i145 = int_lt(i144, 0)
- guard_false(i145, descr=...)
- p146 = new_with_vtable(...)
- setfield_gc(p146, p49, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_iterator 24>)
- setfield_gc(p146, p124, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
- setfield_gc(p146, i64, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_offset 32>)
- setfield_gc(p146, 0, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_index 8>)
- setfield_gc(p16, p146, descr=<FieldP
pypy.module.micronumpy.flatiter.W_FlatIterator.inst_state 56>)
- i147 = arraylen_gc(p72, descr=<ArrayS 8>)
- i148 = arraylen_gc(p90, descr=<ArrayS 8>)
- jump(p0, p1, p3, p6, p7, p12, p14, p16, i143, f142, p26, i44,
p146, i65, p49, i64, i75, i91, i98, i102, i108, p72, p90, descr=...)
+ i154 = getfield_raw(ticker_address, descr=<FieldS
pypysig_long_struct.c_value 0>)
+ i155 = int_lt(i154, 0)
+ guard_false(i155, descr=...)
+ p156 = new_with_vtable(...)
+ setfield_gc(p156, i55, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_offset 32>)
+ setfield_gc(p156, 0, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_index 8>)
+ setfield_gc(p156, p150, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
+ setfield_gc(p156, p49, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_iterator 24>)
+ setfield_gc(p16, p156, descr=<FieldP
pypy.module.micronumpy.flatiter.W_FlatIterator.inst_state 56>)
+ i157 = arraylen_gc(p60, descr=<ArrayS 8>)
+ i158 = arraylen_gc(p70, descr=<ArrayS 8>)
+ jump(p0, p1, p3, p6, p7, p12, p14, p16, i151, f149, p26, i44, i50,
i61, i71, i55, i80, i88, i100, p156, p49, p60, p70, descr=...)
""")
def test_array_flatiter_setitem_single(self):
@@ -149,47 +152,48 @@
assert log.result == 42.0
loop, = log.loops_by_filename(self.filepath)
assert loop.match("""
- i126 = int_lt(i115, i42)
- guard_true(i126, descr=...)
- i127 = int_lt(i115, i48)
- guard_true(i127, descr=...)
- p128 = getfield_gc_pure(p124, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
- i129 = int_sub(i73, i115)
- setarrayitem_gc(p128, 1, 0, descr=<ArrayS 8>)
- setarrayitem_gc(p128, 0, 0, descr=<ArrayS 8>)
- i131 = int_lt(0, i129)
- guard_false(i131, descr=...)
- i132 = int_floordiv(i115, i73)
- i133 = int_mul(i132, i73)
- i134 = int_sub(i115, i133)
- i135 = int_rshift(i134, 63)
- i136 = int_add(i132, i135)
- i137 = int_mul(i136, i73)
- i138 = int_sub(i115, i137)
- i139 = int_mul(i89, i138)
- i140 = int_add(i66, i139)
- i141 = int_sub(i96, i136)
- setarrayitem_gc(p128, 1, i138, descr=<ArrayS 8>)
- i142 = int_lt(0, i141)
- guard_true(i142, descr=...)
- i143 = int_mul(i100, i136)
- i144 = int_add(i140, i143)
- setarrayitem_gc(p128, 0, i136, descr=<ArrayS 8>)
+ i128 = int_lt(i120, i42)
+ guard_true(i128, descr=...)
+ i129 = int_lt(i120, i48)
+ guard_true(i129, descr=...)
+ i130 = int_floordiv(i120, i59)
+ i131 = int_mul(i130, i59)
+ i132 = int_sub(i120, i131)
+ i133 = int_rshift(i132, 63)
+ i134 = int_add(i130, i133)
+ i135 = int_mul(i134, i69)
+ i136 = int_add(i53, i135)
+ i137 = int_mod(i120, i59)
+ i138 = int_rshift(i137, 63)
+ i139 = int_and(i59, i138)
+ i140 = int_add(i137, i139)
+ i141 = int_floordiv(i140, i78)
+ i142 = int_mul(i141, i78)
+ i143 = int_sub(i140, i142)
+ i144 = int_rshift(i143, 63)
+ i145 = int_add(i141, i144)
+ i146 = int_mul(i145, i86)
+ i147 = int_add(i136, i146)
+ i148 = int_mod(i140, i78)
+ i149 = int_rshift(i148, 63)
+ i150 = int_and(i78, i149)
+ i151 = int_add(i148, i150)
guard_not_invalidated(descr=...)
- raw_store(i111, i144, 42.000000, descr=<ArrayF 8>)
- i146 = int_add(i115, 1)
- i147 = getfield_raw(ticker_address, descr=<FieldS
pypysig_long_struct.c_value 0>)
- setarrayitem_gc(p128, 1, 0, descr=<ArrayS 8>)
- setarrayitem_gc(p128, 0, 0, descr=<ArrayS 8>)
- i149 = int_lt(i147, 0)
- guard_false(i149, descr=...)
- p150 = new_with_vtable(...)
- setfield_gc(p150, p47, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_iterator 24>)
- setfield_gc(p150, p128, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
- setfield_gc(p150, i66, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_offset 32>)
- setfield_gc(p150, 0, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_index 8>)
- setfield_gc(p16, p150, descr=<FieldP
pypy.module.micronumpy.flatiter.W_FlatIterator.inst_state 56>)
- i151 = arraylen_gc(p70, descr=<ArrayS 8>)
- i152 = arraylen_gc(p88, descr=<ArrayS 8>)
- jump(p0, p1, p3, p6, p7, p12, p14, p16, i146, i42, i48, p150, i73,
i66, p47, i89, i96, i100, i111, p70, p88, descr=...)
+ raw_store(i103, i147, 42.000000, descr=<ArrayF 8>)
+ p152 = getfield_gc_pure(p126, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
+ i153 = int_add(i120, 1)
+ i154 = getfield_raw(ticker_address, descr=<FieldS
pypysig_long_struct.c_value 0>)
+ setarrayitem_gc(p152, 1, 0, descr=<ArrayS 8>)
+ setarrayitem_gc(p152, 0, 0, descr=<ArrayS 8>)
+ i157 = int_lt(i154, 0)
+ guard_false(i157, descr=...)
+ p158 = new_with_vtable(...)
+ setfield_gc(p158, i53, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_offset 32>)
+ setfield_gc(p158, 0, descr=<FieldS
pypy.module.micronumpy.iterators.IterState.inst_index 8>)
+ setfield_gc(p158, p152, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_indices 16>)
+ setfield_gc(p158, p47, descr=<FieldP
pypy.module.micronumpy.iterators.IterState.inst_iterator 24>)
+ setfield_gc(p16, p158, descr=<FieldP
pypy.module.micronumpy.flatiter.W_FlatIterator.inst_state 56>)
+ i159 = arraylen_gc(p58, descr=<ArrayS 8>)
+ i160 = arraylen_gc(p68, descr=<ArrayS 8>)
+ jump(p0, p1, p3, p6, p7, p12, p14, p16, i153, i42, i48, i59, i69,
i53, i78, i86, p47, i103, p158, p58, p68, descr=...)
""")
_______________________________________________
pypy-commit mailing list
[email protected]
https://mail.python.org/mailman/listinfo/pypy-commit
