forked from M-Labs/artiq
1
0
Fork 0

Allow indexing tuples in kernel code

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 <jonathan.coates@oxionics.com>
This commit is contained in:
Jonathan Coates 2023-09-12 14:43:38 +01:00 committed by GitHub
parent ccb140a929
commit 5c85cef0c2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 113 additions and 2 deletions

View File

@ -1198,7 +1198,27 @@ class ARTIQIRGenerator(algorithm.Visitor):
finally: finally:
self.current_assign = old_assign 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: try:
old_assign, self.current_assign = self.current_assign, None old_assign, self.current_assign = self.current_assign, None
index = self.visit(node.slice.value) index = self.visit(node.slice.value)

View File

@ -259,7 +259,31 @@ class Inferencer(algorithm.Visitor):
def visit_SubscriptT(self, node): def visit_SubscriptT(self, node):
self.generic_visit(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_tuple(node.slice.value.type):
if types.is_var(node.value.type): if types.is_var(node.value.type):
return return

View File

@ -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))

View File

@ -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))

View File

@ -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]

View File

@ -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