From 5c85cef0c208588189061c36d46477008b4af26b Mon Sep 17 00:00:00 2001 From: Jonathan Coates Date: Tue, 12 Sep 2023 14:43:38 +0100 Subject: [PATCH] Allow indexing tuples in kernel code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This only allows for indexing with a constant value (e.g. x[0]). While slices would be possible to implement, it's not clear how to preserve type inference here. The current typing rule is: Γ ⊢ x : τ Γ ⊢ a : Int Γ ⊢ b : Int ------------------------------------ Γ ⊢ x[a:b] : τ However, tuples would require a different typing rule, and so we'd need to defer type inference if τ is a tyvar. I'm not confident that this won't change behaviour, so we leave as-is for now. Signed-off-by: Jonathan Coates --- .../compiler/transforms/artiq_ir_generator.py | 22 +++++++++++++++- artiq/compiler/transforms/inferencer.py | 26 ++++++++++++++++++- .../lit/embedding/error_tuple_index_assign.py | 15 +++++++++++ artiq/test/lit/embedding/index_tuple.py | 16 ++++++++++++ .../test/lit/inferencer/error_tuple_index.py | 14 ++++++++++ artiq/test/lit/integration/tuple_index.py | 22 ++++++++++++++++ 6 files changed, 113 insertions(+), 2 deletions(-) create mode 100644 artiq/test/lit/embedding/error_tuple_index_assign.py create mode 100644 artiq/test/lit/embedding/index_tuple.py create mode 100644 artiq/test/lit/inferencer/error_tuple_index.py create mode 100644 artiq/test/lit/integration/tuple_index.py diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index 4a3d9fbbc..7ede45531 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/artiq/compiler/transforms/artiq_ir_generator.py @@ -1198,7 +1198,27 @@ class ARTIQIRGenerator(algorithm.Visitor): finally: self.current_assign = old_assign - if isinstance(node.slice, ast.Index): + if types.is_tuple(node.value.type): + assert isinstance(node.slice, ast.Index), \ + "Internal compiler error: tuple index should be an Index" + assert isinstance(node.slice.value, ast.Num), \ + "Internal compiler error: tuple index should be a constant" + + if self.current_assign is not None: + diag = diagnostic.Diagnostic("error", + "cannot assign to a tuple element", + {}, node.loc) + self.engine.process(diag) + + index = node.slice.value.n + indexed = self.append( + ir.GetAttr(value, index, name="{}.e{}".format(value.name, index)), + loc=node.loc + ) + + return indexed + + elif isinstance(node.slice, ast.Index): try: old_assign, self.current_assign = self.current_assign, None index = self.visit(node.slice.value) diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py index fdc268d4f..57bbedf82 100644 --- a/artiq/compiler/transforms/inferencer.py +++ b/artiq/compiler/transforms/inferencer.py @@ -259,7 +259,31 @@ class Inferencer(algorithm.Visitor): def visit_SubscriptT(self, node): self.generic_visit(node) - if isinstance(node.slice, ast.Index): + + if types.is_tuple(node.value.type): + if (not isinstance(node.slice, ast.Index) or + not isinstance(node.slice.value, ast.Num)): + diag = diagnostic.Diagnostic( + "error", "tuples can only be indexed by a constant", {}, + node.slice.loc, [] + ) + self.engine.process(diag) + return + + tuple_type = node.value.type.find() + index = node.slice.value.n + if index < 0 or index >= len(tuple_type.elts): + diag = diagnostic.Diagnostic( + "error", + "index {index} is out of range for tuple of size {size}", + {"index": index, "size": len(tuple_type.elts)}, + node.slice.loc, [] + ) + self.engine.process(diag) + return + + self._unify(node.type, tuple_type.elts[index], node.loc, node.value.loc) + elif isinstance(node.slice, ast.Index): if types.is_tuple(node.slice.value.type): if types.is_var(node.value.type): return diff --git a/artiq/test/lit/embedding/error_tuple_index_assign.py b/artiq/test/lit/embedding/error_tuple_index_assign.py new file mode 100644 index 000000000..ffa68a6b6 --- /dev/null +++ b/artiq/test/lit/embedding/error_tuple_index_assign.py @@ -0,0 +1,15 @@ +# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t +# RUN: OutputCheck %s --file-to-check=%t + +from artiq.language.core import * +from artiq.language.types import * + +@kernel +def modify(x): + # CHECK-L: ${LINE:+1}: error: cannot assign to a tuple element + x[0] = 2 + +@kernel +def entrypoint(): + modify((1, "foo", True)) + modify((2, "bar", False)) diff --git a/artiq/test/lit/embedding/index_tuple.py b/artiq/test/lit/embedding/index_tuple.py new file mode 100644 index 000000000..fced90201 --- /dev/null +++ b/artiq/test/lit/embedding/index_tuple.py @@ -0,0 +1,16 @@ +# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s +# RUN: OutputCheck %s --file-to-check=%t.ll + +from artiq.language.core import * +from artiq.language.types import * + +# CHECK-L: void @_Z16testbench.unpackzz({ i32, { i8*, i32 }, i1 } %ARG.x) + +@kernel +def unpack(x): + print(x[0]) + +@kernel +def entrypoint(): + unpack((1, "foo", True)) + unpack((2, "bar", False)) diff --git a/artiq/test/lit/inferencer/error_tuple_index.py b/artiq/test/lit/inferencer/error_tuple_index.py new file mode 100644 index 000000000..dd6f2dca9 --- /dev/null +++ b/artiq/test/lit/inferencer/error_tuple_index.py @@ -0,0 +1,14 @@ +# RUN: %python -m artiq.compiler.testbench.inferencer +diag %s >%t +# RUN: OutputCheck %s --file-to-check=%t + +i = 0 +x = (1, "foo", True) + +# CHECK-L: ${LINE:+1}: error: tuples can only be indexed by a constant +x[i] + +# CHECK-L: ${LINE:+1}: error: tuples can only be indexed by a constant +x[0:2] + +# CHECK-L: ${LINE:+1}: error: index 3 is out of range for tuple of size 3 +x[3] diff --git a/artiq/test/lit/integration/tuple_index.py b/artiq/test/lit/integration/tuple_index.py new file mode 100644 index 000000000..46ebb1428 --- /dev/null +++ b/artiq/test/lit/integration/tuple_index.py @@ -0,0 +1,22 @@ +# RUN: %python -m artiq.compiler.testbench.jit %s +# RUN: %python %s + +# Basic indexing +a = (1, "xyz", True) + +assert a[0] == 1 +assert a[1] == "xyz" +assert a[2] == True + +# Nested indexing +b = (a, 2, (3, "abc", a)) + +assert b[0][0] == 1 +assert b[1] == 2 +assert b[2][2][0] == 1 + +# Usage on the LHS of an assignment +c = (1, 2, [1, 2, 3]) + +c[2][0] = 456 +assert c[2][0] == 456