forked from M-Labs/artiq
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
parent
56010c49fb
commit
575be2aeca
@ -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))
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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]])
|
||||
|
||||
Loading…
Reference in New Issue
Block a user