[pypy-commit] pypy pypy-pyarray: Implement fijal's approach for nonzero()
shmuller
noreply at buildbot.pypy.org
Mon Aug 26 22:03:05 CEST 2013
Author: Stefan H. Muller <shmueller2 at gmail.com>
Branch: pypy-pyarray
Changeset: r66351:81c7341f996b
Date: 2013-08-13 01:06 +0200
http://bitbucket.org/pypy/pypy/changeset/81c7341f996b/
Log: Implement fijal's approach for nonzero()
- Iterators get get_index() method.
- create_iter() get additional keyword argument 'require_index'.
diff --git a/pypy/module/micronumpy/arrayimpl/base.py b/pypy/module/micronumpy/arrayimpl/base.py
--- a/pypy/module/micronumpy/arrayimpl/base.py
+++ b/pypy/module/micronumpy/arrayimpl/base.py
@@ -6,7 +6,7 @@
def base(self):
raise NotImplementedError
- def create_iter(self, shape=None, backward_broadcast=False):
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
raise NotImplementedError
class BaseArrayIterator(object):
diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -282,18 +282,12 @@
def nonzero(self, space, index_type):
s = loop.count_all_true_concrete(self)
box = index_type.itemtype.box
- nd = len(self.shape)
-
- if nd == 1:
- w_res = W_NDimArray.from_shape(space, [s], index_type)
- loop.nonzero_onedim(w_res, self, box)
- return space.newtuple([w_res])
- else:
- w_res = W_NDimArray.from_shape(space, [s, nd], index_type)
- loop.nonzero_multidim(w_res, self, box)
- w_res = w_res.implementation.swapaxes(space, w_res, 0, 1)
- l_w = [w_res.descr_getitem(space, space.wrap(d)) for d in range(nd)]
- return space.newtuple(l_w)
+ nd = len(self.get_shape())
+ w_res = W_NDimArray.from_shape(space, [s, nd], index_type)
+ loop.nonzero(w_res, self, box)
+ w_res = w_res.implementation.swapaxes(space, w_res, 0, 1)
+ l_w = [w_res.descr_getitem(space, space.wrap(d)) for d in range(nd)]
+ return space.newtuple(l_w)
def get_storage_as_int(self, space):
return rffi.cast(lltype.Signed, self.storage) + self.start
@@ -331,13 +325,22 @@
self.backstrides = backstrides
self.storage = storage
- def create_iter(self, shape=None, backward_broadcast=False):
- if shape is None or shape == self.get_shape():
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
+ if shape is not None and shape != self.get_shape():
+ r = calculate_broadcast_strides(self.get_strides(),
+ self.get_backstrides(),
+ self.get_shape(), shape, backward_broadcast)
+ return iter.MultiDimViewIterator(self, self.dtype, self.start, r[0], r[1], shape)
+
+ if not require_index:
return iter.ConcreteArrayIterator(self)
- r = calculate_broadcast_strides(self.get_strides(),
- self.get_backstrides(),
- self.get_shape(), shape, backward_broadcast)
- return iter.MultiDimViewIterator(self, self.dtype, 0, r[0], r[1], shape)
+ else:
+ if len(self.get_shape()) == 1:
+ return iter.OneDimViewIterator(self, self.dtype, self.start,
+ self.get_strides(), self.get_shape())
+ else:
+ return iter.MultiDimViewIterator(self, self.dtype, self.start,
+ self.get_strides(), self.get_backstrides(), self.get_shape())
def fill(self, box):
self.dtype.fill(self.storage, box, 0, self.size)
@@ -400,7 +403,7 @@
def fill(self, box):
loop.fill(self, box.convert_to(self.dtype))
- def create_iter(self, shape=None, backward_broadcast=False):
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
if shape is not None and shape != self.get_shape():
r = calculate_broadcast_strides(self.get_strides(),
self.get_backstrides(),
diff --git a/pypy/module/micronumpy/arrayimpl/scalar.py b/pypy/module/micronumpy/arrayimpl/scalar.py
--- a/pypy/module/micronumpy/arrayimpl/scalar.py
+++ b/pypy/module/micronumpy/arrayimpl/scalar.py
@@ -45,7 +45,7 @@
def get_backstrides(self):
return []
- def create_iter(self, shape=None, backward_broadcast=False):
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
return ScalarIterator(self)
def get_scalar_value(self):
diff --git a/pypy/module/micronumpy/interp_flatiter.py b/pypy/module/micronumpy/interp_flatiter.py
--- a/pypy/module/micronumpy/interp_flatiter.py
+++ b/pypy/module/micronumpy/interp_flatiter.py
@@ -19,7 +19,7 @@
def get_shape(self):
return self.shape
- def create_iter(self, shape=None, backward_broadcast=False):
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
assert isinstance(self.base(), W_NDimArray)
return self.base().create_iter()
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
@@ -230,10 +230,11 @@
s.append('])')
return s.build()
- def create_iter(self, shape=None, backward_broadcast=False):
+ def create_iter(self, shape=None, backward_broadcast=False, require_index=False):
assert isinstance(self.implementation, BaseArrayImplementation)
return self.implementation.create_iter(shape=shape,
- backward_broadcast=backward_broadcast)
+ backward_broadcast=backward_broadcast,
+ require_index=require_index)
def create_axis_iter(self, shape, dim, cum):
return self.implementation.create_axis_iter(shape, dim, cum)
diff --git a/pypy/module/micronumpy/iter.py b/pypy/module/micronumpy/iter.py
--- a/pypy/module/micronumpy/iter.py
+++ b/pypy/module/micronumpy/iter.py
@@ -215,6 +215,9 @@
def reset(self):
self.offset %= self.size
+ def get_index(self, d):
+ return self.index
+
class MultiDimViewIterator(ConcreteArrayIterator):
''' The view iterator dtype can be different from the
array.dtype, this is what makes it a View
@@ -268,6 +271,9 @@
def reset(self):
self.offset %= self.size
+ def get_index(self, d):
+ return self.indexes[d]
+
class AxisIterator(base.BaseArrayIterator):
def __init__(self, array, shape, dim, cumultative):
self.shape = shape
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
@@ -9,8 +9,7 @@
from rpython.rlib import jit
from rpython.rtyper.lltypesystem import lltype, rffi
from pypy.module.micronumpy.base import W_NDimArray
-from pypy.module.micronumpy.iter import PureShapeIterator, OneDimViewIterator, \
- MultiDimViewIterator
+from pypy.module.micronumpy.iter import PureShapeIterator
from pypy.module.micronumpy import constants
from pypy.module.micronumpy.support import int_w
@@ -341,40 +340,21 @@
else:
return count_all_true_concrete(arr.implementation)
-nonzero_driver_onedim = jit.JitDriver(name = 'numpy_nonzero_onedim',
- greens = ['shapelen', 'dtype'],
- reds = 'auto')
+nonzero_driver = jit.JitDriver(name = 'numpy_nonzero',
+ greens = ['shapelen', 'dims', 'dtype'],
+ reds = 'auto')
-def nonzero_onedim(res, arr, box):
+def nonzero(res, arr, box):
res_iter = res.create_iter()
- arr_iter = OneDimViewIterator(arr, arr.dtype, 0,
- arr.strides, arr.shape)
- shapelen = 1
- dtype = arr.dtype
- while not arr_iter.done():
- nonzero_driver_onedim.jit_merge_point(shapelen=shapelen, dtype=dtype)
- if arr_iter.getitem_bool():
- res_iter.setitem(box(arr_iter.index))
- res_iter.next()
- arr_iter.next()
- return res
-
-nonzero_driver_multidim = jit.JitDriver(name = 'numpy_nonzero_onedim',
- greens = ['shapelen', 'dims', 'dtype'],
- reds = 'auto')
-
-def nonzero_multidim(res, arr, box):
- res_iter = res.create_iter()
- arr_iter = MultiDimViewIterator(arr, arr.dtype, 0,
- arr.strides, arr.backstrides, arr.shape)
+ arr_iter = arr.create_iter(require_index=True)
shapelen = len(arr.shape)
dtype = arr.dtype
dims = range(shapelen)
while not arr_iter.done():
- nonzero_driver_multidim.jit_merge_point(shapelen=shapelen, dims=dims, dtype=dtype)
+ nonzero_driver.jit_merge_point(shapelen=shapelen, dims=dims, dtype=dtype)
if arr_iter.getitem_bool():
for d in dims:
- res_iter.setitem(box(arr_iter.indexes[d]))
+ res_iter.setitem(box(arr_iter.get_index(d)))
res_iter.next()
arr_iter.next()
return res
More information about the pypy-commit
mailing list