compiler: Basic support for creation of multidimensional arrays

Breaks all uses of array(), as indexing is not yet implemented.

artiq/compiler/builtins.py

@@ -86,6 +86,10 @@ class TArray(types.TMono):
         if elt is None:
             elt = types.TVar()
         super().__init__("array", {"elt": elt})
+        self.attributes = OrderedDict([
+            ("shape", TList(TInt32())),
+            ("buffer", TList(elt)),
+        ])
 
 def _array_printer(typ, printer, depth, max_depth):
     return "numpy.array(elt={})".format(printer.name(typ["elt"], depth, max_depth))

artiq/compiler/transforms/artiq_ir_generator.py

@@ -1637,7 +1637,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 return self.append(ir.Coerce(arg, node.type))
             else:
                 assert False
-        elif (types.is_builtin(typ, "list") or types.is_builtin(typ, "array") or
+        elif (types.is_builtin(typ, "list") or
               types.is_builtin(typ, "bytearray") or types.is_builtin(typ, "bytes")):
             if len(node.args) == 0 and len(node.keywords) == 0:
                 length = ir.Constant(0, builtins.TInt32())
@@ -1660,6 +1660,77 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 return result
             else:
                 assert False
+        elif types.is_builtin(typ, "array"):
+            if len(node.args) == 1 and len(node.keywords) == 0:
+                result_type = node.type.find()
+                arg = self.visit(node.args[0])
+
+                num_dims = 0
+                result_elt = result_type["elt"].find()
+                inner_type = arg.type.find()
+                while True:
+                    if inner_type == result_elt:
+                        # TODO: What about types needing coercion (e.g. int32 to int64)?
+                        break
+                    assert builtins.is_iterable(inner_type)
+                    num_dims += 1
+                    inner_type = builtins.get_iterable_elt(inner_type)
+
+                # Derive shape from first element on each level (currently, type
+                # inference make sure arrays are always rectangular; in the future, we
+                # might want to insert a runtime check here).
+                #
+                # While we are at it, also total up overall number of elements
+                shape = self.append(
+                    ir.Alloc([ir.Constant(num_dims, self._size_type)],
+                             result_type.attributes["shape"]))
+                first_elt = arg
+                dim_idx = 0
+                num_total_elts = None
+                while True:
+                    length = self.iterable_len(first_elt)
+                    self.append(
+                        ir.SetElem(shape, ir.Constant(dim_idx, length.type), length))
+                    if num_total_elts is None:
+                        num_total_elts = length
+                    else:
+                        num_total_elts = self.append(
+                            ir.Arith(ast.Mult(loc=None), num_total_elts, length))
+
+                    dim_idx += 1
+                    if dim_idx == num_dims:
+                        break
+                    first_elt = self.iterable_get(first_elt,
+                                                  ir.Constant(0, length.type))
+
+                # Assign buffer from nested iterables.
+                buffer = self.append(
+                    ir.Alloc([num_total_elts], result_type.attributes["buffer"]))
+                def body_gen(index):
+                    # TODO: This is hilariously inefficient; we really want to emit a
+                    # nested loop for the source and keep one running index for the
+                    # target buffer.
+                    indices = []
+                    mod_idx = index
+                    for dim_idx in reversed(range(1, num_dims)):
+                        dim_len = self.append(ir.GetElem(shape, ir.Constant(dim_idx, self._size_type)))
+                        indices.append(self.append(ir.Arith(ast.Mod(loc=None), mod_idx, dim_len)))
+                        mod_idx = self.append(ir.Arith(ast.FloorDiv(loc=None), mod_idx, dim_len))
+                    indices.append(mod_idx)
+
+                    elt = arg
+                    for idx in reversed(indices):
+                        elt = self.iterable_get(elt, idx)
+                    self.append(ir.SetElem(buffer, index, elt))
+                    return self.append(
+                        ir.Arith(ast.Add(loc=None), index, ir.Constant(1, length.type)))
+                self._make_loop(
+                    ir.Constant(0, length.type), lambda index: self.append(
+                        ir.Compare(ast.Lt(loc=None), index, num_total_elts)), body_gen)
+
+                return self.append(ir.Alloc([shape, buffer], node.type))
+            else:
+                assert False
         elif types.is_builtin(typ, "range"):
             elt_typ = builtins.get_iterable_elt(node.type)
             if len(node.args) == 1 and len(node.keywords) == 0:

artiq/compiler/transforms/llvm_ir_generator.py

@@ -246,6 +246,10 @@ class LLVMIRGenerator:
             return ll.IntType(builtins.get_int_width(typ))
         elif builtins.is_float(typ):
             return lldouble
+        elif builtins.is_array(typ):
+            llshapety = self.llty_of_type(typ.attributes["shape"])
+            llbufferty = self.llty_of_type(typ.attributes["buffer"])
+            return ll.LiteralStructType([llshapety, llbufferty])
         elif builtins.is_listish(typ):
             lleltty = self.llty_of_type(builtins.get_iterable_elt(typ))
             return ll.LiteralStructType([lleltty.as_pointer(), lli32])
@@ -733,7 +737,7 @@ class LLVMIRGenerator:
                                                    name=insn.name)
             else:
                 assert False
-        elif builtins.is_listish(insn.type):
+        elif builtins.is_listish(insn.type) and not builtins.is_array(insn.type):
             llsize = self.map(insn.operands[0])
             lleltty = self.llty_of_type(builtins.get_iterable_elt(insn.type))
             llalloc = self.llbuilder.alloca(lleltty, size=llsize)
@@ -741,7 +745,8 @@ class LLVMIRGenerator:
             llvalue = self.llbuilder.insert_value(llvalue, llalloc, 0, name=insn.name)
             llvalue = self.llbuilder.insert_value(llvalue, llsize, 1)
             return llvalue
-        elif not builtins.is_allocated(insn.type) or ir.is_keyword(insn.type):
+        elif (not builtins.is_allocated(insn.type) or ir.is_keyword(insn.type)
+              or builtins.is_array(insn.type)):
             llvalue = ll.Constant(self.llty_of_type(insn.type), ll.Undefined)
             for index, elt in enumerate(insn.operands):
                 llvalue = self.llbuilder.insert_value(llvalue, self.map(elt), index)

artiq/test/lit/integration/array.py

@@ -2,4 +2,7 @@
 # REQUIRES: exceptions
 
 ary = array([1, 2, 3])
-assert [x*x for x in ary] == [1, 4, 9]
+# FIXME: Implement ndarray indexing
+# assert [x*x for x in ary] == [1, 4, 9]
+
+matrix = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])