compiler: Properly implement NumPy array slicing

Strided slicing of one-dimensional arrays (i.e. with non-trivial
steps) might have previously been working, but would have had
different semantics, as all slices were copies rather than a view
into the original data.

Fixing this in the future will require adding support for an index
stride field/tuple to our array representation (and all the
associated indexing logic).

GitHub: Fixes #1627.

artiq/compiler/transforms/artiq_ir_generator.py

@@ -1116,7 +1116,11 @@ class ARTIQIRGenerator(algorithm.Visitor):
                                                                   _readable_name(index))))
             if self.current_assign is None:
                 return indexed
-        else: # Slice
+        else:
+            # This is a slice. The endpoint checking logic is the same for both lists
+            # and NumPy arrays, but the actual implementations differ – while slices of
+            # built-in lists are always copies in Python, they are views sharing the
+            # same backing storage in NumPy.
             length = self.iterable_len(value, node.slice.type)
 
             if node.slice.lower is not None:
@@ -1141,91 +1145,127 @@ class ARTIQIRGenerator(algorithm.Visitor):
             mapped_stop_index = self._map_index(length, stop_index, one_past_the_end=True,
                                                 loc=node.begin_loc)
 
-            if node.slice.step is not None:
+            if builtins.is_array(node.type):
-                try:
+                # To implement strided slicing with the proper NumPy reference
-                    old_assign, self.current_assign = self.current_assign, None
+                # semantics, the pointer/length array representation will need to be
-                    step = self.visit(node.slice.step)
+                # extended by another field to hold a variable stride.
-                finally:
+                assert node.slice.step is None, (
-                    self.current_assign = old_assign
+                    "array slices with non-trivial step "
+                    "should have been disallowed during type inference")
 
+                # One-dimensionally slicing an array only affects the outermost
+                # dimension.
+                shape = self.append(ir.GetAttr(value, "shape"))
+                lengths = [
+                    self.append(ir.GetAttr(shape, i))
+                    for i in range(len(shape.type.elts))
+                ]
+
+                # Compute outermost length – zero for "backwards" indices.
+                raw_len = self.append(
+                    ir.Arith(ast.Sub(loc=None), mapped_stop_index, mapped_start_index))
+                is_neg_len = self.append(
+                    ir.Compare(ast.Lt(loc=None), raw_len, ir.Constant(0, raw_len.type)))
+                outer_len = self.append(
+                    ir.Select(is_neg_len, ir.Constant(0, raw_len.type), raw_len))
+                new_shape = self._make_array_shape([outer_len] + lengths[1:])
+
+                # Offset buffer pointer by start index (times stride for inner dims).
+                stride = reduce(
+                    lambda l, r: self.append(ir.Arith(ast.Mult(loc=None), l, r)),
+                    lengths[1:], ir.Constant(1, lengths[0].type))
+                offset = self.append(
+                    ir.Arith(ast.Mult(loc=None), stride, mapped_start_index))
+                buffer = self.append(ir.GetAttr(value, "buffer"))
+                new_buffer = self.append(ir.Offset(buffer, offset))
+
+                return self.append(ir.Alloc([new_buffer, new_shape], node.type))
+            else:
+                if node.slice.step is not None:
+                    try:
+                        old_assign, self.current_assign = self.current_assign, None
+                        step = self.visit(node.slice.step)
+                    finally:
+                        self.current_assign = old_assign
+
+                    self._make_check(
+                        self.append(ir.Compare(ast.NotEq(loc=None), step, ir.Constant(0, step.type))),
+                        lambda: self.alloc_exn(builtins.TException("ValueError"),
+                            ir.Constant("step cannot be zero", builtins.TStr())),
+                        loc=node.slice.step.loc)
+                else:
+                    step = ir.Constant(1, node.slice.type)
+                counting_up = self.append(ir.Compare(ast.Gt(loc=None), step,
+                                                    ir.Constant(0, step.type)))
+
+                unstepped_size = self.append(ir.Arith(ast.Sub(loc=None),
+                                                    mapped_stop_index, mapped_start_index))
+                slice_size_a = self.append(ir.Arith(ast.FloorDiv(loc=None), unstepped_size, step))
+                slice_size_b = self.append(ir.Arith(ast.Mod(loc=None), unstepped_size, step))
+                rem_not_empty = self.append(ir.Compare(ast.NotEq(loc=None), slice_size_b,
+                                                    ir.Constant(0, slice_size_b.type)))
+                slice_size_c = self.append(ir.Arith(ast.Add(loc=None), slice_size_a,
+                                                    ir.Constant(1, slice_size_a.type)))
+                slice_size = self.append(ir.Select(rem_not_empty,
+                                                slice_size_c, slice_size_a,
+                                                name="slice.size"))
                 self._make_check(
-                    self.append(ir.Compare(ast.NotEq(loc=None), step, ir.Constant(0, step.type))),
+                    self.append(ir.Compare(ast.LtE(loc=None), slice_size, length)),
-                    lambda: self.alloc_exn(builtins.TException("ValueError"),
+                    lambda slice_size, length: self.alloc_exn(builtins.TException("ValueError"),
-                        ir.Constant("step cannot be zero", builtins.TStr())),
+                        ir.Constant("slice size {0} is larger than iterable length {1}",
-                    loc=node.slice.step.loc)
+                                    builtins.TStr()),
-            else:
+                        slice_size, length),
-                step = ir.Constant(1, node.slice.type)
+                    params=[slice_size, length],
-            counting_up = self.append(ir.Compare(ast.Gt(loc=None), step,
+                    loc=node.slice.loc)
-                                                 ir.Constant(0, step.type)))
 
-            unstepped_size = self.append(ir.Arith(ast.Sub(loc=None),
+                if self.current_assign is None:
-                                                  mapped_stop_index, mapped_start_index))
+                    is_neg_size = self.append(ir.Compare(ast.Lt(loc=None),
-            slice_size_a = self.append(ir.Arith(ast.FloorDiv(loc=None), unstepped_size, step))
+                                                        slice_size, ir.Constant(0, slice_size.type)))
-            slice_size_b = self.append(ir.Arith(ast.Mod(loc=None), unstepped_size, step))
+                    abs_slice_size = self.append(ir.Select(is_neg_size,
-            rem_not_empty = self.append(ir.Compare(ast.NotEq(loc=None), slice_size_b,
+                                                        ir.Constant(0, slice_size.type), slice_size))
-                                                   ir.Constant(0, slice_size_b.type)))
+                    other_value = self.append(ir.Alloc([abs_slice_size], value.type,
-            slice_size_c = self.append(ir.Arith(ast.Add(loc=None), slice_size_a,
+                                                    name="slice.result"))
-                                                ir.Constant(1, slice_size_a.type)))
+                else:
-            slice_size = self.append(ir.Select(rem_not_empty,
+                    other_value = self.current_assign
-                                               slice_size_c, slice_size_a,
-                                               name="slice.size"))
-            self._make_check(
-                self.append(ir.Compare(ast.LtE(loc=None), slice_size, length)),
-                lambda slice_size, length: self.alloc_exn(builtins.TException("ValueError"),
-                    ir.Constant("slice size {0} is larger than iterable length {1}",
-                                builtins.TStr()),
-                    slice_size, length),
-                params=[slice_size, length],
-                loc=node.slice.loc)
 
-            if self.current_assign is None:
+                prehead = self.current_block
-                is_neg_size = self.append(ir.Compare(ast.Lt(loc=None),
-                                                     slice_size, ir.Constant(0, slice_size.type)))
-                abs_slice_size = self.append(ir.Select(is_neg_size,
-                                                       ir.Constant(0, slice_size.type), slice_size))
-                other_value = self.append(ir.Alloc([abs_slice_size], value.type,
-                                                   name="slice.result"))
-            else:
-                other_value = self.current_assign
 
-            prehead = self.current_block
+                head = self.current_block = self.add_block("slice.head")
+                prehead.append(ir.Branch(head))
 
-            head = self.current_block = self.add_block("slice.head")
+                index = self.append(ir.Phi(node.slice.type,
-            prehead.append(ir.Branch(head))
+                                        name="slice.index"))
+                index.add_incoming(mapped_start_index, prehead)
+                other_index = self.append(ir.Phi(node.slice.type,
+                                                name="slice.resindex"))
+                other_index.add_incoming(ir.Constant(0, node.slice.type), prehead)
 
-            index = self.append(ir.Phi(node.slice.type,
+                # Still within bounds?
-                                       name="slice.index"))
+                bounded_up = self.append(ir.Compare(ast.Lt(loc=None), index, mapped_stop_index))
-            index.add_incoming(mapped_start_index, prehead)
+                bounded_down = self.append(ir.Compare(ast.Gt(loc=None), index, mapped_stop_index))
-            other_index = self.append(ir.Phi(node.slice.type,
+                within_bounds = self.append(ir.Select(counting_up, bounded_up, bounded_down))
-                                             name="slice.resindex"))
-            other_index.add_incoming(ir.Constant(0, node.slice.type), prehead)
 
-            # Still within bounds?
+                body = self.current_block = self.add_block("slice.body")
-            bounded_up = self.append(ir.Compare(ast.Lt(loc=None), index, mapped_stop_index))
-            bounded_down = self.append(ir.Compare(ast.Gt(loc=None), index, mapped_stop_index))
-            within_bounds = self.append(ir.Select(counting_up, bounded_up, bounded_down))
 
-            body = self.current_block = self.add_block("slice.body")
+                if self.current_assign is None:
+                    elem = self.iterable_get(value, index)
+                    self.append(ir.SetElem(other_value, other_index, elem))
+                else:
+                    elem = self.append(ir.GetElem(self.current_assign, other_index))
+                    self.append(ir.SetElem(value, index, elem))
 
-            if self.current_assign is None:
+                next_index = self.append(ir.Arith(ast.Add(loc=None), index, step))
-                elem = self.iterable_get(value, index)
+                index.add_incoming(next_index, body)
-                self.append(ir.SetElem(other_value, other_index, elem))
+                next_other_index = self.append(ir.Arith(ast.Add(loc=None), other_index,
-            else:
+                                                        ir.Constant(1, node.slice.type)))
-                elem = self.append(ir.GetElem(self.current_assign, other_index))
+                other_index.add_incoming(next_other_index, body)
-                self.append(ir.SetElem(value, index, elem))
+                self.append(ir.Branch(head))
 
-            next_index = self.append(ir.Arith(ast.Add(loc=None), index, step))
+                tail = self.current_block = self.add_block("slice.tail")
-            index.add_incoming(next_index, body)
+                head.append(ir.BranchIf(within_bounds, body, tail))
-            next_other_index = self.append(ir.Arith(ast.Add(loc=None), other_index,
-                                                    ir.Constant(1, node.slice.type)))
-            other_index.add_incoming(next_other_index, body)
-            self.append(ir.Branch(head))
 
-            tail = self.current_block = self.add_block("slice.tail")
+                if self.current_assign is None:
-            head.append(ir.BranchIf(within_bounds, body, tail))
+                    return other_value
-
-            if self.current_assign is None:
-                return other_value
 
     def visit_TupleT(self, node):
         if self.current_assign is None:

artiq/compiler/transforms/inferencer.py

@@ -269,6 +269,14 @@ class Inferencer(algorithm.Visitor):
             else:
                 self._unify_iterable(element=node, collection=node.value)
         elif isinstance(node.slice, ast.Slice):
+            if builtins.is_array(node.value.type):
+                if node.slice.step is not None:
+                    diag = diagnostic.Diagnostic(
+                        "error",
+                        "strided slicing not yet supported for NumPy arrays", {},
+                        node.slice.step.loc, [])
+                    self.engine.process(diag)
+                    return
             self._unify(node.type, node.value.type, node.loc, node.value.loc)
         else:  # ExtSlice
             pass  # error emitted above

artiq/test/lit/inferencer/error_array.py

@@ -9,5 +9,8 @@ b = array([1, 2, 3])
 # CHECK-L: ${LINE:+1}: error: too many indices for array of dimension 1
 b[1, 2]
 
+# CHECK-L: ${LINE:+1}: error: strided slicing not yet supported for NumPy arrays
+b[::-1]
+
 # CHECK-L: ${LINE:+1}: error: array attributes cannot be assigned to
 b.shape = (5, )

artiq/test/lit/integration/array_slice.py

@@ -0,0 +1,24 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+
+a = array([0, 1, 2, 3])
+
+b = a[2:3]
+assert b.shape == (1,)
+assert b[0] == 2
+b[0] = 5
+assert a[2] == 5
+
+b = a[3:2]
+assert b.shape == (0,)
+
+c = array([[0, 1], [2, 3]])
+
+d = c[:1]
+assert d.shape == (1, 2)
+assert d[0, 0] == 0
+assert d[0, 1] == 1
+d[0, 0] = 5
+assert c[0, 0] == 5
+
+d = c[1:0]
+assert d.shape == (0, 2)

artiq/test/lit/regression/issue_1627.py

@@ -0,0 +1,13 @@
+# RUN: %python -m artiq.compiler.testbench.embedding %s
+
+from artiq.language.core import *
+from artiq.language.types import *
+import numpy as np
+
+n = 2
+data = np.zeros((n, n))
+
+
+@kernel
+def entrypoint():
+    print(data[:n])