Optimize indexing when slicing memviews & move to utility

Cython/Compiler/ExprNodes.py

@@ -1811,7 +1811,6 @@ class NameNode(AtomicExprNode):
             lhs_pos=self.pos,
             rhs=rhs,
             code=code,
-            incref_rhs=rhs.is_name,
             have_gil=not self.in_nogil_context)
 
     def generate_acquire_buffer(self, rhs, code):
@@ -2571,7 +2570,6 @@ class IndexNode(ExprNode):
             # All indices with all start/stop/step for slices.
             # We need to keep this around
             self.indices = new_indices
-
             self.env = env
 
         elif self.base.type.is_buffer:
@@ -4525,8 +4523,7 @@ class AttributeNode(ExprNode):
             elif self.type.is_memoryviewslice:
                 import MemoryView
                 MemoryView.put_assign_to_memviewslice(
-                        select_code, rhs.result(), self.type, code,
-                        incref_rhs=rhs.is_name)
+                        select_code, rhs, rhs.result(), self.type, code)
 
             if not self.type.is_memoryviewslice:
                 code.putln(

Cython/Compiler/MemoryView.py

@@ -79,7 +79,7 @@ def mangle_dtype_name(dtype):
 #    return "".join([access[0].upper()+packing[0] for (access, packing) in axes])
 
 def put_acquire_memoryviewslice(lhs_cname, lhs_type, lhs_pos, rhs, code,
-                                incref_rhs=False, have_gil=False):
+                                have_gil=False):
     assert rhs.type.is_memoryviewslice
 
     pretty_rhs = isinstance(rhs, NameNode) or rhs.result_in_temp()
@@ -91,16 +91,16 @@ def put_acquire_memoryviewslice(lhs_cname, lhs_type, lhs_pos, rhs, code,
 
     # Allow uninitialized assignment
     #code.putln(code.put_error_if_unbound(lhs_pos, rhs.entry))
-    put_assign_to_memviewslice(lhs_cname, rhstmp, lhs_type, code, incref_rhs,
+    put_assign_to_memviewslice(lhs_cname, rhs, rhstmp, lhs_type, code,
                                have_gil=have_gil)
 
     if not pretty_rhs:
         code.funcstate.release_temp(rhstmp)
 
-def put_assign_to_memviewslice(lhs_cname, rhs_cname, memviewslicetype, code,
-                               incref_rhs=False, have_gil=False):
+def put_assign_to_memviewslice(lhs_cname, rhs, rhs_cname, memviewslicetype, code,
+                               have_gil=False):
     code.put_xdecref_memoryviewslice(lhs_cname, have_gil=have_gil)
-    if incref_rhs:
+    if rhs.is_name:
         code.put_incref_memoryviewslice(rhs_cname, have_gil=have_gil)
 
     code.putln("%s = %s;" % (lhs_cname, rhs_cname))
@@ -269,59 +269,52 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
         Simply call __pyx_memoryview_slice_memviewslice with the right
         arguments.
         """
-        slicefunc = "__pyx_memoryview_slice_memviewslice"
         new_ndim = 0
-        cname = self.cname
-
-        suboffset_dim = code.funcstate.allocate_temp(PyrexTypes.c_int_type,
-                                                     False)
-
-        index_code = ("%(slicefunc)s(&%(cname)s, &%(dst)s, %(have_gil)d, "
-                                    "%(dim)d, %(new_ndim)d, &%(suboffset_dim)s, "
-                                    "%(idx)s, 0, 0, 0, 0, 0, 0)")
-
-        slice_code = ("%(slicefunc)s(&%(cname)s, &%(dst)s, %(have_gil)d, "
-                                    "/* dim */ %(dim)d, "
-                                    "/* new_ndim */ %(new_ndim)d, "
-                                    "/* suboffset_dim */ &%(suboffset_dim)s, "
-                                    "/* start */ %(start)s, "
-                                    "/* stop */ %(stop)s, "
-                                    "/* step */ %(step)s, "
-                                    "/* have_start */ %(have_start)d, "
-                                    "/* have_stop */ %(have_stop)d, "
-                                    "/* have_step */ %(have_step)d, "
-                                    "/* is_slice */ 1)")
-
-        def generate_slice_call(expr):
-            pos = index.pos
-
-            if have_gil:
-                code.putln(code.error_goto_if(expr, pos))
-            else:
-                code.putln("{")
-                code.putln(    "const char *__pyx_t_result = %s;" % expr)
-
-                code.putln(    "if (unlikely(__pyx_t_result)) {")
-                code.put_ensure_gil()
-                code.putln(        "PyErr_Format(PyExc_IndexError, "
-                                                "__pyx_t_result, %d);" % dim)
-                code.put_release_ensured_gil()
-                code.putln(code.error_goto(pos))
-                code.putln(    "}")
-
-                code.putln("}")
-
-        code.putln("%s = -1;" % suboffset_dim)
-        code.putln("%(dst)s.data = %(cname)s.data;" % locals())
-        code.putln("%(dst)s.memview = %(cname)s.memview;" % locals())
+        src = self.cname
+
+        def load_slice_util(name, dict):
+            proto, impl = TempitaUtilityCode.load_as_string(
+                        name, "MemoryView_C.c", context=dict)
+            return impl
+
+        all_dimensions_direct = True
+        for access, packing in self.type.axes:
+            if access != 'direct':
+                all_dimensions_direct = False
+                break
+
+        have_slices = False
+        for index in indices:
+            if isinstance(index, ExprNodes.SliceNode):
+                have_slices = True
+                break
+
+        no_suboffset_dim = all_dimensions_direct and not have_slices
+        if not no_suboffset_dim:
+            suboffset_dim = code.funcstate.allocate_temp(
+                             PyrexTypes.c_int_type, False)
+            code.putln("%s = -1;" % suboffset_dim)
+
+        code.putln("%(dst)s.data = %(src)s.data;" % locals())
+        code.putln("%(dst)s.memview = %(src)s.memview;" % locals())
         code.put_incref_memoryviewslice(dst)
 
         for dim, index in enumerate(indices):
+            error_goto = code.error_goto(index.pos)
+
             if not isinstance(index, ExprNodes.SliceNode):
+                # normal index
                 idx = index.result()
-                generate_slice_call(index_code % locals())
+                d = locals()
+                access, packing = self.type.axes[dim]
+                if access != 'direct':
+                    return error(index.pos,
+                                 "Dimension cannot be indexed away, "
+                                 "must be direct")
+                code.put(load_slice_util("SliceIndex", d))
             else:
-                d = {}
+                # slice, unspecified dimension, or part of ellipsis
+                d = locals()
                 for s in "start stop step".split():
                     idx = getattr(index, s)
                     have_idx = d['have_' + s] = not idx.is_none
@@ -330,11 +323,21 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
                     else:
                         d[s] = "0"
 
-                d.update(locals())
-                generate_slice_call(slice_code % d)
+                if (not d['have_start'] and
+                    not d['have_stop'] and
+                    not d['have_step']):
+                    # full slice (:), simply copy over the extent, stride
+                    # and suboffset. Also update suboffset_dim if needed
+                    access, packing = self.type.axes[dim]
+                    d['access'] = access
+                    code.put(load_slice_util("SimpleSlice", d))
+                else:
+                    code.put(load_slice_util("ToughSlice", d))
+
                 new_ndim += 1
 
-        code.funcstate.release_temp(suboffset_dim)
+        if not no_suboffset_dim:
+            code.funcstate.release_temp(suboffset_dim)
 
 
 def empty_slice(pos):

Cython/Compiler/Nodes.py

@@ -5168,7 +5168,6 @@ class ReturnStatNode(StatNode):
                         lhs_pos=self.value.pos,
                         rhs=self.value,
                         code=code,
-                        incref_rhs=self.value.is_name,
                         have_gil=self.in_nogil_context)
             else:
                 self.value.make_owned_reference(code)

Cython/Utility/MemoryView.pyx

@@ -583,7 +583,7 @@ cdef memoryview memview_slice(memoryview memview, object indices):
 
     for dim, index in enumerate(indices):
         if PyIndex_Check(index):
-            slice_memviewslice(p_src, p_dst, True, dim, new_ndim, p_suboffset_dim,
+            slice_memviewslice(p_src, p_dst, dim, new_ndim, p_suboffset_dim,
                                index, 0, 0, 0, 0, 0, False)
         else:
             start = index.start or 0
@@ -594,7 +594,7 @@ cdef memoryview memview_slice(memoryview memview, object indices):
             have_stop = index.stop is not None
             have_step = index.step is not None
 
-            slice_memviewslice(p_src, p_dst, True, dim, new_ndim, p_suboffset_dim,
+            slice_memviewslice(p_src, p_dst, dim, new_ndim, p_suboffset_dim,
                                start, stop, step, have_start, have_stop, have_step,
                                True)
             new_ndim += 1
@@ -628,9 +628,8 @@ cdef extern from "pystate.h":
     PyObject *PyErr_Format(PyObject *exc, char *msg, ...) nogil
 
 @cname('__pyx_memoryview_slice_memviewslice')
-cdef char *slice_memviewslice({{memviewslice_name}} *src,
+cdef int slice_memviewslice({{memviewslice_name}} *src,
                               {{memviewslice_name}} *dst,
-                              bint have_gil,
                               int dim,
                               int new_ndim,
                               int *suboffset_dim,
@@ -640,11 +639,10 @@ cdef char *slice_memviewslice({{memviewslice_name}} *src,
                               int have_start,
                               int have_stop,
                               int have_step,
-                              bint is_slice) nogil except *:
+                              bint is_slice) nogil except -1:
     """
     Create a new slice dst given slice src.
 
-    have_gil        - if true, the GIL must be held and exceptions may be raised
     dim             - the current src dimension (indexing will make dimensions
                                                  disappear)
     new_dim         - the new dst dimension
@@ -657,19 +655,6 @@ cdef char *slice_memviewslice({{memviewslice_name}} *src,
         Py_ssize_t new_shape
         bint negative_step
 
-        # Somehow these pointers are NULL when set as globals... this needs investigation
-        char *ERR_OOB = "Index out of bounds (axis %d)"
-        char *ERR_STEP = "Step must not be zero (axis %d)"
-        char *ERR_INDIRECT_GIL = ("Cannot make indirect dimension %d disappear "
-                              "through indexing, consider slicing with %d:%d")
-        char *ERR_INDIRECT_NOGIL = ("Cannot make indirect dimension %d disappear "
-                                "through indexing")
-        PyObject *exc = <PyObject *> IndexError
-
-    if have_gil:
-        # Assert the GIL
-        PyThreadState_Get()
-
     shape = src.shape[dim]
     stride = src.strides[dim]
     suboffset = src.suboffsets[dim]
@@ -679,19 +664,15 @@ cdef char *slice_memviewslice({{memviewslice_name}} *src,
         if start < 0:
             start += shape
         if not 0 <= start < shape:
-            if have_gil:
-                PyErr_Format(exc, ERR_OOB, dim)
-            return ERR_OOB
+            with gil:
+                raise IndexError("Index out of bounds (axis %d)")
     else:
         # index is a slice
         negative_step = have_step != 0 and step < 0
 
         if have_step and step == 0:
-            if have_gil:
-                # ValueError might be more appropriate, but this will make it consistent
-                # with nogil slicing
-                PyErr_SetString(exc, ERR_STEP)
-            return ERR_STEP
+            with gil:
+                raise ValueError("Step may not be zero (axis %d)" % dim)
 
         # check our bounds and set defaults
         if have_start:
@@ -749,13 +730,16 @@ cdef char *slice_memviewslice({{memviewslice_name}} *src,
 
     if suboffset >= 0:
         if not is_slice:
-            if have_gil:
-                PyErr_Format(exc, ERR_INDIRECT_GIL, dim, start, start + 1)
-            return ERR_INDIRECT_NOGIL
+            with gil:
+                raise IndexError(
+                    "Cannot make indirect dimension %d disappear "
+                    "through indexing, consider slicing with %d:%d" %
+                                                (dim, start, start + 1))
+
         else:
             suboffset_dim[0] = new_ndim
 
-    return NULL
+    return 0
 
 #
 ### Index a memoryview

Cython/Utility/MemoryView_C.c

@@ -688,3 +688,70 @@ int {{set_function}}(const char *itemp, PyObject *obj) {
     *(PyObject **) itemp = obj;
     return 1;
 }
+
+/////////// ToughSlice //////////
+/* Dimension is indexed with 'start:stop:step' */
+
+if (unlikely(__pyx_memoryview_slice_memviewslice(
+    &{{src}},
+    &{{dst}},
+    {{dim}},
+    {{new_ndim}},
+    &{{suboffset_dim}},
+    {{start}},
+    {{stop}},
+    {{step}},
+    {{int(have_start)}},
+    {{int(have_stop)}},
+    {{int(have_step)}},
+    1) < 0))
+{
+    {{error_goto}}
+}
+
+////////// SimpleSlice //////////
+/* Dimension is indexed with ':' only */
+
+{{dst}}.shape[{{new_ndim}}] = {{src}}.shape[{{dim}}];
+{{dst}}.strides[{{new_ndim}}] = {{src}}.strides[{{dim}}];
+
+{{if access == 'direct'}}
+    {{dst}}.suboffsets[{{new_ndim}}] = -1;
+{{else}}
+    {{dst}}.suboffsets[{{new_ndim}}] = {{src}}.suboffsets[{{dim}}];
+    if ({{src}}.suboffsets[{{dim}}] >= 0)
+        {{suboffset_dim}} = {{new_ndim}};
+{{endif}}
+
+////////// SliceIndex //////////
+/* Dimension is indexed with an integer, we could use the ToughSlice */
+/* approach, but this is faster */
+
+{
+    Py_ssize_t __pyx_tmp_idx = {{idx}};
+
+    if (__pyx_tmp_idx < 0)
+        __pyx_tmp_idx += {{src}}.shape[{{dim}}];
+
+    if (__pyx_tmp_idx < 0 || __pyx_tmp_idx >= {{src}}.shape[{{dim}}]) {
+    {{if not have_gil}}
+        #ifdef WITH_THREAD
+        PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure();
+        #endif
+    {{endif}}
+        PyErr_SetString(PyExc_IndexError, "Index out of bounds (axis {{dim}})");
+    {{if not have_gil}}
+        #ifdef WITH_THREAD
+        PyGILState_Release(__pyx_gilstate_save);
+        #endif
+    {{endif}}
+        {{error_goto}}
+    }
+
+    {{if all_dimensions_direct}}
+        {{dst}}.data += __pyx_tmp_idx * {{src}}.strides[{{dim}}];
+    {{else}}
+        if ({{suboffset_dim}} < 0) {{dst}}.data += __pyx_tmp_idx * {{src}}.strides[{{dim}}];
+        else {{dst}}.suboffsets[{{suboffset_dim}}] = __pyx_tmp_idx * {{src}}.strides[{{dim}}];
+    {{endif}}
+}