[pypy-commit] pypy default: Merge numpy-faster-setslice, uses memcpy for setslice operations if possible

fijal noreply at buildbot.pypy.org
Wed Dec 21 20:52:10 CET 2011


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r50810:441db5bbe505
Date: 2011-12-21 21:51 +0200
http://bitbucket.org/pypy/pypy/changeset/441db5bbe505/

Log:	Merge numpy-faster-setslice, uses memcpy for setslice operations if
	possible

diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -102,3 +102,27 @@
 class ConstantIterator(BaseIterator):
     def next(self, shapelen):
         return self
+
+# ------ other iterators that are not part of the computation frame ----------
+
+class AxisIterator(object):
+    """ This object will return offsets of each start of the last stride
+    """
+    def __init__(self, arr):
+        self.arr = arr
+        self.indices = [0] * (len(arr.shape) - 1)
+        self.done = False
+        self.offset = arr.start
+
+    def next(self):
+        for i in range(len(self.arr.shape) - 2, -1, -1):
+            if self.indices[i] < self.arr.shape[i] - 1:
+                self.indices[i] += 1
+                self.offset += self.arr.strides[i]
+                break
+            else:
+                self.indices[i] = 0
+                self.offset -= self.arr.backstrides[i]
+        else:
+            self.done = True
+        
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
@@ -9,7 +9,7 @@
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
 from pypy.module.micronumpy.interp_iter import ArrayIterator,\
-     view_iter_from_arr, OneDimIterator
+     view_iter_from_arr, OneDimIterator, AxisIterator
 
 numpy_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
@@ -390,10 +390,10 @@
         return space.wrap(self.size)
 
     def descr_copy(self, space):
-        return self.copy()
+        return self.copy(space)
 
-    def copy(self):
-        return self.get_concrete().copy()
+    def copy(self, space):
+        return self.get_concrete().copy(space)
 
     def descr_len(self, space):
         if len(self.shape):
@@ -536,7 +536,7 @@
                               new_shape, self)
         else:
             # Create copy with contiguous data
-            arr = concrete.copy()
+            arr = concrete.copy(space)
             arr.setshape(space, new_shape)
         return arr
 
@@ -606,6 +606,9 @@
                                                        space.w_False]))
         return w_d
 
+    def supports_fast_slicing(self):
+        return False
+
 def convert_to_array(space, w_obj):
     if isinstance(w_obj, BaseArray):
         return w_obj
@@ -639,7 +642,7 @@
     def to_str(self, space, comma, builder, indent=' ', use_ellipsis=False):
         builder.append(self.dtype.itemtype.str_format(self.value))
 
-    def copy(self):
+    def copy(self, space):
         return Scalar(self.dtype, self.value)
 
     def create_sig(self, res_shape):
@@ -790,6 +793,9 @@
     def get_concrete(self):
         return self
 
+    def supports_fast_slicing(self):
+        return self.order == 'C' and self.strides[-1] == 1
+
     def find_dtype(self):
         return self.dtype
 
@@ -929,39 +935,35 @@
             item += v * self.strides[i]
         return item
 
-
-class ViewArray(ConcreteArray):
-    def copy(self):
-        array = W_NDimArray(self.size, self.shape[:], self.find_dtype())
-        iter = view_iter_from_arr(self)
-        a_iter = ArrayIterator(array.size)
-        while not iter.done():
-            array.setitem(a_iter.offset, self.getitem(iter.offset))
-            iter = iter.next(len(self.shape))
-            a_iter = a_iter.next(len(array.shape))
-        return array
-
-    def create_sig(self, res_shape):
-        return signature.ViewSignature(self.dtype)
-
-
-class W_NDimSlice(ViewArray):
-    def __init__(self, start, strides, backstrides, shape, parent):
-        assert isinstance(parent, ConcreteArray)
-        if isinstance(parent, W_NDimSlice):
-            parent = parent.parent
-        size = 1
-        for sh in shape:
-            size *= sh
-        self.strides = strides
-        self.backstrides = backstrides
-        ViewArray.__init__(self, size, shape, parent.dtype, parent.order,
-                               parent)
-        self.start = start
-
     def setslice(self, space, w_value):
         res_shape = shape_agreement(space, self.shape, w_value.shape)
-        self._sliceloop(w_value, res_shape)
+        if (res_shape == w_value.shape and self.supports_fast_slicing() and
+            w_value.supports_fast_slicing() and
+            self.dtype is w_value.find_dtype()):
+            self._fast_setslice(space, w_value)
+        else:
+            self._sliceloop(w_value, res_shape)
+
+    def _fast_setslice(self, space, w_value):
+        assert isinstance(w_value, ConcreteArray)
+        itemsize = self.dtype.itemtype.get_element_size()
+        if len(self.shape) == 1:
+            rffi.c_memcpy(
+                rffi.ptradd(self.storage, self.start * itemsize),
+                rffi.ptradd(w_value.storage, w_value.start * itemsize),
+                self.size * itemsize
+            )
+        else:
+            dest = AxisIterator(self)
+            source = AxisIterator(w_value)
+            while not dest.done:
+                rffi.c_memcpy(
+                    rffi.ptradd(self.storage, dest.offset * itemsize),
+                    rffi.ptradd(w_value.storage, source.offset * itemsize),
+                    self.shape[-1] * itemsize
+                )
+                source.next()
+                dest.next()
 
     def _sliceloop(self, source, res_shape):
         sig = source.find_sig(res_shape)
@@ -979,6 +981,31 @@
             frame.next(shapelen)
             res_iter = res_iter.next(shapelen)
 
+    def copy(self, space):
+        array = W_NDimArray(self.size, self.shape[:], self.dtype, self.order)
+        array.setslice(space, self)
+        return array
+
+
+class ViewArray(ConcreteArray):
+    def create_sig(self, res_shape):
+        return signature.ViewSignature(self.dtype)
+
+
+class W_NDimSlice(ViewArray):
+    def __init__(self, start, strides, backstrides, shape, parent):
+        assert isinstance(parent, ConcreteArray)
+        if isinstance(parent, W_NDimSlice):
+            parent = parent.parent
+        size = 1
+        for sh in shape:
+            size *= sh
+        self.strides = strides
+        self.backstrides = backstrides
+        ViewArray.__init__(self, size, shape, parent.dtype, parent.order,
+                               parent)
+        self.start = start
+
     def setshape(self, space, new_shape):
         if len(self.shape) < 1:
             return
@@ -1017,15 +1044,6 @@
     """ A class representing contiguous array. We know that each iteration
     by say ufunc will increase the data index by one
     """
-    def copy(self):
-        array = W_NDimArray(self.size, self.shape[:], self.dtype, self.order)
-        rffi.c_memcpy(
-            array.storage,
-            self.storage,
-            self.size * self.dtype.itemtype.get_element_size()
-        )
-        return array
-
     def setitem(self, item, value):
         self.invalidated()
         self.dtype.setitem(self.storage, item, value)
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
@@ -1077,6 +1077,17 @@
         a = ones((1, 2, 3))
         assert a[0, 1, 2] == 1.0
 
+    def test_multidim_setslice(self):
+        from numpypy import zeros, ones
+        a = zeros((3, 3))
+        b = ones((3, 3))
+        a[:,1:3] = b[:,1:3]
+        assert (a == [[0, 1, 1], [0, 1, 1], [0, 1, 1]]).all()
+        a = zeros((3, 3))
+        b = ones((3, 3))
+        a[:,::2] = b[:,::2]
+        assert (a == [[1, 0, 1], [1, 0, 1], [1, 0, 1]]).all()
+
     def test_broadcast_ufunc(self):
         from numpypy import array
         a = array([[1, 2], [3, 4], [5, 6]])


More information about the pypy-commit mailing list