Author: Brian Kearns <bdkea...@gmail.com>
Branch:
Changeset: r69535:27f7c050271b
Date: 2014-02-27 22:20 -0800
http://bitbucket.org/pypy/pypy/changeset/27f7c050271b/
Log: simplify iter/dot optimizations
diff --git a/pypy/module/micronumpy/concrete.py
b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -15,6 +15,8 @@
class BaseConcreteArray(object):
+ _immutable_fields_ = ['dtype?', 'storage', 'start', 'size', 'shape[*]',
+ 'strides[*]', 'backstrides[*]', 'order']
start = 0
parent = None
@@ -350,6 +352,8 @@
orig_array)
def set_dtype(self, space, dtype):
+ # size/shape/strides shouldn't change
+ assert dtype.elsize == self.dtype.elsize
self.dtype = dtype
def argsort(self, space, w_axis):
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
@@ -79,29 +79,31 @@
class ArrayIter(object):
- _immutable_fields_ = ['array', 'size', 'ndim_m1', 'shape_m1',
- 'strides', 'backstrides']
+ _immutable_fields_ = ['array', 'size', 'indices', 'shape[*]',
+ 'strides[*]', 'backstrides[*]']
def __init__(self, array, size, shape, strides, backstrides):
assert len(shape) == len(strides) == len(backstrides)
self.array = array
self.size = size
- self.ndim_m1 = len(shape) - 1
- self.shape_m1 = [s - 1 for s in shape]
+ self.indices = [0] * len(shape)
+ self.shape = shape
self.strides = strides
self.backstrides = backstrides
self.reset()
+ @jit.unroll_safe
def reset(self):
self.index = 0
- self.indices = [0] * (self.ndim_m1 + 1)
+ for i in xrange(len(self.shape)):
+ self.indices[i] = 0
self.offset = self.array.start
@jit.unroll_safe
def next(self):
self.index += 1
- for i in xrange(self.ndim_m1, -1, -1):
- if self.indices[i] < self.shape_m1[i]:
+ for i in xrange(len(self.shape) - 1, -1, -1):
+ if self.indices[i] < self.shape[i] - 1:
self.indices[i] += 1
self.offset += self.strides[i]
break
@@ -115,14 +117,14 @@
if step == 0:
return
self.index += step
- for i in xrange(self.ndim_m1, -1, -1):
- if self.indices[i] < (self.shape_m1[i] + 1) - step:
+ for i in xrange(len(self.shape) - 1, -1, -1):
+ if self.indices[i] < self.shape[i] - step:
self.indices[i] += step
self.offset += self.strides[i] * step
break
else:
- remaining_step = (self.indices[i] + step) // (self.shape_m1[i]
+ 1)
- this_i_step = step - remaining_step * (self.shape_m1[i] + 1)
+ remaining_step = (self.indices[i] + step) // self.shape[i]
+ this_i_step = step - remaining_step * self.shape[i]
self.indices[i] = self.indices[i] + this_i_step
self.offset += self.strides[i] * this_i_step
step = remaining_step
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
@@ -285,21 +285,18 @@
outi = result.create_iter()
lefti = AllButAxisIter(left.implementation, len(left_shape) - 1)
righti = AllButAxisIter(right.implementation, right_critical_dim)
- n = left.implementation.shape[-1]
- s1 = left.implementation.strides[-1]
- s2 = right.implementation.strides[right_critical_dim]
while not lefti.done():
while not righti.done():
oval = outi.getitem()
i1 = lefti.offset
i2 = righti.offset
- for _ in xrange(n):
+ for _ in xrange(left.implementation.shape[-1]):
dot_driver.jit_merge_point(dtype=dtype)
lval = left.implementation.getitem(i1).convert_to(space, dtype)
rval = right.implementation.getitem(i2).convert_to(space,
dtype)
oval = dtype.itemtype.add(oval, dtype.itemtype.mul(lval, rval))
- i1 += s1
- i2 += s2
+ i1 += left.implementation.strides[-1]
+ i2 += right.implementation.strides[right_critical_dim]
outi.setitem(oval)
outi.next()
righti.next()
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
@@ -59,6 +59,7 @@
if self.graph is None:
interp, graph = self.meta_interp(f, [0],
listops=True,
+ listcomp=True,
backendopt=True,
graph_and_interp_only=True)
self.__class__.interp = interp
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
