Author: Maciej Fijalkowski <[email protected]>
Branch:
Changeset: r58097:c62c6904e304
Date: 2012-10-13 12:04 +0200
http://bitbucket.org/pypy/pypy/changeset/c62c6904e304/
Log: an attempt to improve the fancy indexing situation untested
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
@@ -424,39 +424,44 @@
self._done = True
@jit.unroll_safe
- def get_index(self, space):
- return [space.wrap(i) for i in self.indexes]
+ def get_index(self, space, shapelen):
+ return [space.wrap(self.indexes[i]) for i in range(shapelen)]
getitem_int_driver = jit.JitDriver(name = 'numpy_getitem_int',
- greens = ['shapelen', 'indexlen', 'dtype'],
+ greens = ['shapelen', 'indexlen',
+ 'prefixlen', 'dtype'],
reds = ['arr', 'res', 'iter', 'indexes_w',
'prefix_w'])
def getitem_array_int(space, arr, res, iter_shape, indexes_w, prefix_w):
shapelen = len(iter_shape)
+ prefixlen = len(prefix_w)
indexlen = len(indexes_w)
dtype = arr.get_dtype()
iter = PureShapeIterator(iter_shape, indexes_w)
+ indexlen = len(indexes_w)
while not iter.done():
getitem_int_driver.jit_merge_point(shapelen=shapelen,
indexlen=indexlen,
dtype=dtype, arr=arr, res=res,
iter=iter, indexes_w=indexes_w,
- prefix_w=prefix_w)
+ prefix_w=prefix_w,
+ prefixlen=prefixlen)
# prepare the index
- index_w = [None] * len(indexes_w)
- for i in range(len(indexes_w)):
+ index_w = [None] * indexlen
+ for i in range(indexlen):
if iter.idx_w[i] is not None:
index_w[i] = iter.idx_w[i].getitem()
else:
index_w[i] = indexes_w[i]
- res.descr_setitem(space, space.newtuple(prefix_w +
- iter.get_index(space)),
+ res.descr_setitem(space, space.newtuple(prefix_w[:prefixlen] +
+ iter.get_index(space, shapelen)),
arr.descr_getitem(space, space.newtuple(index_w)))
iter.next()
return res
setitem_int_driver = jit.JitDriver(name = 'numpy_setitem_int',
- greens = ['shapelen', 'indexlen', 'dtype'],
+ greens = ['shapelen', 'indexlen',
+ 'prefixlen', 'dtype'],
reds = ['arr', 'iter', 'indexes_w',
'prefix_w', 'val_arr'])
@@ -464,21 +469,24 @@
prefix_w):
shapelen = len(iter_shape)
indexlen = len(indexes_w)
+ prefixlen = len(prefix_w)
dtype = arr.get_dtype()
iter = PureShapeIterator(iter_shape, indexes_w)
while not iter.done():
setitem_int_driver.jit_merge_point(shapelen=shapelen,
indexlen=indexlen,
dtype=dtype, arr=arr,
iter=iter, indexes_w=indexes_w,
- prefix_w=prefix_w, val_arr=val_arr)
+ prefix_w=prefix_w, val_arr=val_arr,
+ prefixlen=prefixlen)
# prepare the index
- index_w = [None] * len(indexes_w)
- for i in range(len(indexes_w)):
+ index_w = [None] * indexlen
+ for i in range(indexlen):
if iter.idx_w[i] is not None:
index_w[i] = iter.idx_w[i].getitem()
else:
index_w[i] = indexes_w[i]
- w_idx = space.newtuple(prefix_w + iter.get_index(space))
+ w_idx = space.newtuple(prefix_w[:prefixlen] + iter.get_index(space,
+ shapelen))
arr.descr_setitem(space, space.newtuple(index_w),
val_arr.descr_getitem(space, w_idx))
iter.next()
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
