From 33d931a5b7c457a627117a5e1bbf3c0603732c65 Mon Sep 17 00:00:00 2001 From: David Nadlinger Date: Sun, 9 Aug 2020 03:28:36 +0100 Subject: [PATCH] compiler: Implement multi-dimensional indexing of arrays This generates rather more code than necessary, but has the advantage of automatically handling incomplete multi-dimensional subscripts which still leave arrays behind. --- .../compiler/transforms/artiq_ir_generator.py | 33 ++++++++++---- artiq/compiler/transforms/inferencer.py | 44 +++++++++++++++---- artiq/test/lit/inferencer/error_array.py | 4 ++ artiq/test/lit/integration/array.py | 28 ++++++------ 4 files changed, 77 insertions(+), 32 deletions(-) diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index 9d5201800..ac916f984 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/artiq/compiler/transforms/artiq_ir_generator.py @@ -1092,16 +1092,31 @@ class ARTIQIRGenerator(algorithm.Visitor): finally: self.current_assign = old_assign - length = self.iterable_len(value, index.type) - mapped_index = self._map_index(length, index, - loc=node.begin_loc) - if self.current_assign is None: - result = self.iterable_get(value, mapped_index) - result.set_name("{}.at.{}".format(value.name, _readable_name(index))) - return result + # For multi-dimensional indexes, just apply them sequentially. This + # works, as they are only supported for types where we do not + # immediately need to distinguish between the Get and Set cases + # (i.e. arrays, which are reference types). + if types.is_tuple(index.type): + num_idxs = len(index.type.find().elts) + indices = [ + self.append(ir.GetAttr(index, i)) for i in range(num_idxs) + ] else: - self.append(ir.SetElem(value, mapped_index, self.current_assign, - name="{}.at.{}".format(value.name, _readable_name(index)))) + indices = [index] + indexed = value + for i, idx in enumerate(indices): + length = self.iterable_len(indexed, idx.type) + mapped_index = self._map_index(length, idx, loc=node.begin_loc) + if self.current_assign is None or i < len(indices) - 1: + indexed = self.iterable_get(indexed, mapped_index) + indexed.set_name("{}.at.{}".format(indexed.name, + _readable_name(idx))) + else: + self.append(ir.SetElem(indexed, mapped_index, self.current_assign, + name="{}.at.{}".format(value.name, + _readable_name(index)))) + if self.current_assign is None: + return indexed else: # Slice length = self.iterable_len(value, node.slice.type) diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py index 8ac5fab15..172dfb1fe 100644 --- a/artiq/compiler/transforms/inferencer.py +++ b/artiq/compiler/transforms/inferencer.py @@ -208,10 +208,9 @@ class Inferencer(algorithm.Visitor): self.generic_visit(node) value = node.value if types.is_tuple(value.type): - diag = diagnostic.Diagnostic("error", - "multi-dimensional slices are not supported", {}, - node.loc, []) - self.engine.process(diag) + for elt in value.type.find().elts: + self._unify(elt, builtins.TInt(), + value.loc, None) else: self._unify(value.type, builtins.TInt(), value.loc, None) @@ -237,12 +236,39 @@ class Inferencer(algorithm.Visitor): def visit_SubscriptT(self, node): self.generic_visit(node) if isinstance(node.slice, ast.Index): - self._unify_iterable(element=node, collection=node.value) + if types.is_tuple(node.slice.value.type): + if not builtins.is_array(node.value.type): + diag = diagnostic.Diagnostic( + "error", + "multi-dimensional slices only supported for arrays, not {type}", + {"type": types.TypePrinter().name(node.value.type)}, + node.loc, []) + self.engine.process(diag) + return + num_idxs = len(node.slice.value.type.find().elts) + array_type = node.value.type.find() + num_dims = array_type["num_dims"].value + remaining_dims = num_dims - num_idxs + if remaining_dims < 0: + diag = diagnostic.Diagnostic( + "error", + "too many indices for array of dimension {num_dims}", + {"num_dims": num_dims}, node.slice.loc, []) + self.engine.process(diag) + return + if remaining_dims == 0: + self._unify(node.type, array_type["elt"], node.loc, + node.value.loc) + else: + self._unify( + node.type, + builtins.TArray(array_type["elt"], remaining_dims)) + else: + self._unify_iterable(element=node, collection=node.value) elif isinstance(node.slice, ast.Slice): - self._unify(node.type, node.value.type, - node.loc, node.value.loc) - else: # ExtSlice - pass # error emitted above + self._unify(node.type, node.value.type, node.loc, node.value.loc) + else: # ExtSlice + pass # error emitted above def visit_IfExpT(self, node): self.generic_visit(node) diff --git a/artiq/test/lit/inferencer/error_array.py b/artiq/test/lit/inferencer/error_array.py index 787ae9294..5e79f99e6 100644 --- a/artiq/test/lit/inferencer/error_array.py +++ b/artiq/test/lit/inferencer/error_array.py @@ -5,5 +5,9 @@ a = array() b = array([1, 2, 3]) + +# CHECK-L: ${LINE:+1}: error: too many indices for array of dimension 1 +b[1, 2] + # CHECK-L: ${LINE:+1}: error: array attributes cannot be assigned to b.shape = (5, ) diff --git a/artiq/test/lit/integration/array.py b/artiq/test/lit/integration/array.py index 8f97c4002..039e76372 100644 --- a/artiq/test/lit/integration/array.py +++ b/artiq/test/lit/integration/array.py @@ -26,21 +26,21 @@ assert matrix.shape == (2, 3) # 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])] -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 +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 -matrix[0][0] = 7.0 -matrix[1][1] = 8.0 -assert matrix[0][0] == 7.0 -assert matrix[0][1] == 2.0 -assert matrix[0][2] == 3.0 -assert matrix[1][0] == 4.0 -assert matrix[1][1] == 8.0 -assert matrix[1][2] == 6.0 +matrix[0, 0] = 7.0 +matrix[1, 1] = 8.0 +assert matrix[0, 0] == 7.0 +assert matrix[0, 1] == 2.0 +assert matrix[0, 2] == 3.0 +assert matrix[1, 0] == 4.0 +assert matrix[1, 1] == 8.0 +assert matrix[1, 2] == 6.0 three_tensor = array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]) assert len(three_tensor) == 1