From 38c17622cce77a49c2387f1e8d5d2107e8d3be01 Mon Sep 17 00:00:00 2001
From: David Nadlinger <code@klickverbot.at>
Date: Sun, 26 Jul 2020 02:46:05 +0100
Subject: [PATCH] compiler: Axis-wise iteration of ndarrays

Matches NumPy. Slicing a TList reallocates, this doesn't; offsetting
couldn't be handled in the IR without introducing new semantics
(the Alloc kludge; could/should be made its own IR type).
---
 .../compiler/transforms/artiq_ir_generator.py | 19 +++++++++++++++++--
 .../compiler/transforms/llvm_ir_generator.py  | 18 ++++++++++++++++++
 artiq/test/lit/integration/array.py           |  9 +++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
index afd0da1a9..8a1981279 100644
--- a/artiq/compiler/transforms/artiq_ir_generator.py
+++ b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -512,8 +512,23 @@ class ARTIQIRGenerator(algorithm.Visitor):
         # Assuming the value is within bounds.
         if builtins.is_array(value.type):
             # Scalar indexing into ndarray.
-            if value.type.find()["num_dims"].value > 1:
-                raise NotImplementedError
+            num_dims = value.type.find()["num_dims"].value
+            if num_dims > 1:
+                old_shape = self.append(ir.GetAttr(value, "shape"))
+                lengths = [self.append(ir.GetAttr(old_shape, i)) for i in range(1, num_dims)]
+                new_shape = self.append(ir.Alloc(lengths, types.TTuple(old_shape.type.elts[1:])))
+
+                stride = reduce(
+                    lambda l, r: self.append(ir.Arith(ast.Mult(loc=None), l, r)),
+                    lengths[1:], lengths[0])
+                offset = self.append(ir.Arith(ast.Mult(loc=None), stride, index))
+                old_buffer = self.append(ir.GetAttr(value, "buffer"))
+                # KLUDGE: Represent offsetting by Alloc with two arguments.
+                new_buffer = self.append(ir.Alloc([old_buffer, offset], old_buffer.type))
+
+                result_type = builtins.TArray(value.type.find()["elt"],
+                    types.TValue(num_dims - 1))
+                return self.append(ir.Alloc([new_shape, new_buffer], result_type))
             else:
                 buffer = self.append(ir.GetAttr(value, "buffer"))
                 return self.append(ir.GetElem(buffer, index))
diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
index 58615a43c..b8e7e8457 100644
--- a/artiq/compiler/transforms/llvm_ir_generator.py
+++ b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -738,6 +738,24 @@ class LLVMIRGenerator:
             else:
                 assert False
         elif builtins.is_listish(insn.type) and not builtins.is_array(insn.type):
+            if builtins.is_listish(insn.operands[0].type):
+                # KLUDGE: Offsetting is represented as Alloc with base list in the first
+                # argument and offset in the second. Should probably move this to a
+                # seprate node type (or make it possible to construct lists from
+                # pointer/length).
+                llbase = self.map(insn.operands[0])
+                lloldbase = self.llbuilder.extract_value(llbase, 0)
+                lloldsize = self.llbuilder.extract_value(llbase, 1)
+
+                lloffset = self.map(insn.operands[1])
+                llbase = self.llbuilder.gep(lloldbase, [lloffset], inbounds=True)
+                llsize = self.llbuilder.sub(lloldsize, lloffset)
+
+                llvalue = ll.Constant(self.llty_of_type(insn.type), ll.Undefined)
+                llvalue = self.llbuilder.insert_value(llvalue, llbase, 0)
+                llvalue = self.llbuilder.insert_value(llvalue, llsize, 1)
+                return llvalue
+
             llsize = self.map(insn.operands[0])
             lleltty = self.llty_of_type(builtins.get_iterable_elt(insn.type))
             llalloc = self.llbuilder.alloca(lleltty, size=llsize)
diff --git a/artiq/test/lit/integration/array.py b/artiq/test/lit/integration/array.py
index 81763c29e..1bca6304c 100644
--- a/artiq/test/lit/integration/array.py
+++ b/artiq/test/lit/integration/array.py
@@ -16,6 +16,15 @@ assert [x * x for x in empty_array] == []
 matrix = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
 assert len(matrix) == 2
 assert matrix.shape == (2, 3)
+assert matrix[0][0] == 1.0
+assert matrix[0][1] == 2.0
+assert matrix[0][2] == 3.0
+assert matrix[1][0] == 4.0
+assert matrix[1][1] == 5.0
+assert matrix[1][2] == 6.0
+# FIXME: Need to decide on a solution for array comparisons —
+# NumPy returns an array of bools!
+# assert [x for x in matrix] == [array([1.0, 2.0, 3.0]), array([4.0, 5.0, 6.0])]
 
 three_tensor = array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
 assert len(three_tensor) == 1