M-Labs
/
artiq
mirror of https://github.com/m-labs/artiq.git
2
0
Fork 2
Code Issues Projects Releases Wiki Activity

compiler: Basic support for creation of multidimensional arrays 

Breaks all uses of array(), as indexing is not yet implemented.
This commit is contained in:
David Nadlinger 2020-07-25 22:19:12 +01:00
parent 56010c49fb
commit 575be2aeca
4 changed files with 87 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
artiq/compiler/builtins.py
View File

@ -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))

73
artiq/compiler/transforms/artiq_ir_generator.py
View File

@ -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:

9
artiq/compiler/transforms/llvm_ir_generator.py
View File

@ -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)

5
artiq/test/lit/integration/array.py
View File

@ -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]])