forked from M-Labs/artiq
1
0
Fork 0

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).
This commit is contained in:
David Nadlinger 2020-07-26 02:46:05 +01:00
parent c95a978ab6
commit 38c17622cc
3 changed files with 44 additions and 2 deletions

View File

@ -512,8 +512,23 @@ class ARTIQIRGenerator(algorithm.Visitor):
# Assuming the value is within bounds. # Assuming the value is within bounds.
if builtins.is_array(value.type): if builtins.is_array(value.type):
# Scalar indexing into ndarray. # Scalar indexing into ndarray.
if value.type.find()["num_dims"].value > 1: num_dims = value.type.find()["num_dims"].value
raise NotImplementedError 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: else:
buffer = self.append(ir.GetAttr(value, "buffer")) buffer = self.append(ir.GetAttr(value, "buffer"))
return self.append(ir.GetElem(buffer, index)) return self.append(ir.GetElem(buffer, index))

View File

@ -738,6 +738,24 @@ class LLVMIRGenerator:
else: else:
assert False assert False
elif builtins.is_listish(insn.type) and not builtins.is_array(insn.type): 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]) llsize = self.map(insn.operands[0])
lleltty = self.llty_of_type(builtins.get_iterable_elt(insn.type)) lleltty = self.llty_of_type(builtins.get_iterable_elt(insn.type))
llalloc = self.llbuilder.alloca(lleltty, size=llsize) llalloc = self.llbuilder.alloca(lleltty, size=llsize)

View File

@ -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]]) matrix = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
assert len(matrix) == 2 assert len(matrix) == 2
assert matrix.shape == (2, 3) 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]]]) three_tensor = array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
assert len(three_tensor) == 1 assert len(three_tensor) == 1