[pypy-commit] pypy numpy-multidim-shards: Make array iterators a once-off immutable things.
fijal
noreply at buildbot.pypy.org
Mon Nov 14 18:50:40 CET 2011
Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim-shards
Changeset: r49408:14f8da6a95bf
Date: 2011-11-14 18:50 +0100
http://bitbucket.org/pypy/pypy/changeset/14f8da6a95bf/
Log: Make array iterators a once-off immutable things.
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
@@ -80,12 +80,12 @@
raise NotImplementedError
class ArrayIterator(BaseIterator):
- def __init__(self, size):
- self.offset = 0
+ def __init__(self, size, offset=0):
+ self.offset = offset
self.size = size
def next(self):
- self.offset += 1
+ return ArrayIterator(self.size, self.offset + 1)
def done(self):
return self.offset >= self.size
@@ -94,24 +94,32 @@
return self.offset
class ViewIterator(BaseIterator):
- def __init__(self, arr):
- self.indices = [0] * len(arr.shape)
- self.offset = arr.start
- self.arr = arr
- self._done = False
+ def __init__(self, arr, offset=0, indices=None, done=False):
+ if indices is None:
+ self.indices = [0] * len(arr.shape)
+ self.offset = arr.start
+ else:
+ self.offset = offset
+ self.indices = indices
+ self.arr = arr
+ self._done = done
@jit.unroll_safe
def next(self):
+ indices = self.indices[:]
+ done = False
+ offset = self.offset
for i in range(len(self.indices)):
- if self.indices[i] < self.arr.shape[i] - 1:
- self.indices[i] += 1
- self.offset += self.arr.shards[i]
+ if indices[i] < self.arr.shape[i] - 1:
+ indices[i] += 1
+ offset += self.arr.shards[i]
break
else:
- self.indices[i] = 0
- self.offset -= self.arr.backshards[i]
+ indices[i] = 0
+ offset -= self.arr.backshards[i]
else:
- self._done = True
+ done = True
+ return ViewIterator(self.arr, offset, indices, done)
def done(self):
return self._done
@@ -125,8 +133,7 @@
self.right = right
def next(self):
- self.left.next()
- self.right.next()
+ return Call2Iterator(self.left.next(), self.right.next())
def done(self):
return self.left.done() or self.right.done()
@@ -141,7 +148,7 @@
self.child = child
def next(self):
- self.child.next()
+ return Call1Iterator(self.child.next())
def done(self):
return self.child.done()
@@ -151,7 +158,7 @@
class ConstantIterator(BaseIterator):
def next(self):
- pass
+ return self
def done(self):
return False
@@ -268,7 +275,7 @@
if dtype.ne(new_best, cur_best):
result = i.get_offset()
cur_best = new_best
- i.next()
+ i = i.next()
return result
def impl(self, space):
size = self.find_size()
@@ -286,7 +293,7 @@
all_driver.jit_merge_point(signature=self.signature, self=self, dtype=dtype, i=i)
if not dtype.bool(self.eval(i)):
return False
- i.next()
+ i = i.next()
return True
def descr_all(self, space):
return space.wrap(self._all())
@@ -299,7 +306,7 @@
dtype=dtype, i=i)
if dtype.bool(self.eval(i)):
return True
- i.next()
+ i = i.next()
return False
def descr_any(self, space):
return space.wrap(self._any())
@@ -633,8 +640,8 @@
result_size=result_size, i=i, ri=ri,
self=self, result=result)
result.dtype.setitem(result.storage, ri.offset, self.eval(i))
- i.next()
- ri.next()
+ i = i.next()
+ ri = ri.next()
return result
def force_if_needed(self):
@@ -811,8 +818,8 @@
source_iter=source_iter)
self.setitem(res_iter.offset, source.eval(source_iter).convert_to(
self.find_dtype()))
- source_iter.next()
- res_iter.next()
+ source_iter = source_iter.next()
+ res_iter = res_iter.next()
def start_iter(self):
return ViewIterator(self)
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
@@ -76,7 +76,7 @@
value=value, obj=obj, i=i,
dtype=dtype)
value = self.func(dtype, value, obj.eval(i).convert_to(dtype))
- i.next()
+ i = i.next()
return value
class W_Ufunc1(W_Ufunc):
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
@@ -29,21 +29,21 @@
def test_create_slice(self):
space = self.space
a = NDimArray(10*5*3, [10, 5, 3], MockDtype())
- s = a._create_slice(space, space.wrap(3))
+ s = a.create_slice(space, space.wrap(3))
assert s.start == 45
assert s.shards == [3, 1]
assert s.backshards == [12, 2]
- s = a._create_slice(space, self.newslice(1, 9, 2))
+ s = a.create_slice(space, self.newslice(1, 9, 2))
assert s.start == 15
assert s.shards == [30, 3, 1]
assert s.backshards == [120, 12, 2]
- s = a._create_slice(space, space.newtuple([
+ s = a.create_slice(space, space.newtuple([
self.newslice(1, 5, 3), self.newslice(1, 2, 1), space.wrap(1)]))
assert s.start == 19
assert s.shape == [2, 1]
assert s.shards == [45, 3]
assert s.backshards == [90, 3]
- s = a._create_slice(space, self.newtuple(
+ s = a.create_slice(space, self.newtuple(
self.newslice(None, None, None), space.wrap(2)))
assert s.start == 6
assert s.shape == [10, 3]
@@ -51,16 +51,16 @@
def test_slice_of_slice(self):
space = self.space
a = NDimArray(10*5*3, [10, 5, 3], MockDtype())
- s = a._create_slice(space, space.wrap(5))
+ s = a.create_slice(space, space.wrap(5))
assert s.start == 15*5
- s2 = s._create_slice(space, space.wrap(3))
+ s2 = s.create_slice(space, space.wrap(3))
assert s2.shape == [3]
assert s2.shards == [1]
assert s2.parent is a
assert s2.backshards == [2]
assert s2.start == 5*15 + 3*3
- s = a._create_slice(space, self.newslice(1, 5, 3))
- s2 = s._create_slice(space, space.newtuple([
+ s = a.create_slice(space, self.newslice(1, 5, 3))
+ s2 = s.create_slice(space, space.newtuple([
self.newslice(None, None, None), space.wrap(2)]))
assert s2.shape == [2, 3]
assert s2.shards == [45, 1]
@@ -70,7 +70,7 @@
def test_negative_step(self):
space = self.space
a = NDimArray(10*5*3, [10, 5, 3], MockDtype())
- s = a._create_slice(space, self.newslice(None, None, -2))
+ s = a.create_slice(space, self.newslice(None, None, -2))
assert s.start == 135
assert s.shards == [-30, 3, 1]
assert s.backshards == [-150, 12, 2]
@@ -79,7 +79,7 @@
a = NDimArray(10*5*3, [10, 5, 3], MockDtype())
r = a._index_of_single_item(self.space, self.newtuple(1, 2, 2))
assert r == 1 * 3 * 5 + 2 * 3 + 2
- s = a._create_slice(self.space, self.newtuple(
+ s = a.create_slice(self.space, self.newtuple(
self.newslice(None, None, None), 2))
r = s._index_of_single_item(self.space, self.newtuple(1, 0))
assert r == a._index_of_single_item(self.space, self.newtuple(1, 2, 0))
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
@@ -254,8 +254,10 @@
result = self.run('multidim')
assert result == 8
self.check_loops({'float_add': 1, 'getarrayitem_raw': 2,
- 'guard_true': 1, 'int_add': 1, 'int_lt': 1,
+ 'guard_false': 1, 'int_add': 3, 'int_ge': 1,
'jump': 1, 'setarrayitem_raw': 1})
+ # int_add might be 1 here if we try slightly harder with
+ # reusing indexes or some optimization
def define_multidim_slice():
return """
More information about the pypy-commit
mailing list