From 575be2aeca88af00f69fe447b7a2fb1d1bcee8d4 Mon Sep 17 00:00:00 2001 From: David Nadlinger Date: Sat, 25 Jul 2020 22:19:12 +0100 Subject: [PATCH] compiler: Basic support for creation of multidimensional arrays Breaks all uses of array(), as indexing is not yet implemented. --- artiq/compiler/builtins.py | 4 + .../compiler/transforms/artiq_ir_generator.py | 73 ++++++++++++++++++- .../compiler/transforms/llvm_ir_generator.py | 9 ++- artiq/test/lit/integration/array.py | 5 +- 4 files changed, 87 insertions(+), 4 deletions(-) diff --git a/artiq/compiler/builtins.py b/artiq/compiler/builtins.py index ce2e23e27..d881396a2 100644 --- a/artiq/compiler/builtins.py +++ b/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)) diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index 4e3c3394b..da593315f 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/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: diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py index 84dcac2dd..3da6037a1 100644 --- a/artiq/compiler/transforms/llvm_ir_generator.py +++ b/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) diff --git a/artiq/test/lit/integration/array.py b/artiq/test/lit/integration/array.py index aec874497..3cb7fc58d 100644 --- a/artiq/test/lit/integration/array.py +++ b/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]])