mirror of
https://github.com/m-labs/artiq.git
synced 2024-12-25 19:28:26 +08:00
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:
parent
c95a978ab6
commit
38c17622cc
@ -512,8 +512,23 @@ class ARTIQIRGenerator(algorithm.Visitor):
|
||||
# Assuming the value is within bounds.
|
||||
if builtins.is_array(value.type):
|
||||
# Scalar indexing into ndarray.
|
||||
if value.type.find()["num_dims"].value > 1:
|
||||
raise NotImplementedError
|
||||
num_dims = value.type.find()["num_dims"].value
|
||||
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:
|
||||
buffer = self.append(ir.GetAttr(value, "buffer"))
|
||||
return self.append(ir.GetElem(buffer, index))
|
||||
|
@ -738,6 +738,24 @@ class LLVMIRGenerator:
|
||||
else:
|
||||
assert False
|
||||
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])
|
||||
lleltty = self.llty_of_type(builtins.get_iterable_elt(insn.type))
|
||||
llalloc = self.llbuilder.alloca(lleltty, size=llsize)
|
||||
|
@ -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]])
|
||||
assert len(matrix) == 2
|
||||
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]]])
|
||||
assert len(three_tensor) == 1
|
||||
|
Loading…
Reference in New Issue
Block a user