diff --git a/artiq/compiler/math_fns.py b/artiq/compiler/math_fns.py index 83c1f82b4..4e67dfef0 100644 --- a/artiq/compiler/math_fns.py +++ b/artiq/compiler/math_fns.py @@ -26,6 +26,9 @@ unary_fp_runtime_calls = [ ("arctan", "atan"), ] +#: Array handling builtins (special treatment due to allocations). +numpy_builtins = ["transpose"] + def unary_fp_type(name): return types.TExternalFunction(OrderedDict([("arg", builtins.TFloat())]), @@ -36,6 +39,8 @@ numpy_map = { getattr(numpy, symbol): unary_fp_type(mangle) for symbol, mangle in (unary_fp_intrinsics + unary_fp_runtime_calls) } +for name in numpy_builtins: + numpy_map[getattr(numpy, name)] = types.TBuiltinFunction("numpy." + name) def match(obj): diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index e148ce184..ab162288d 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/artiq/compiler/transforms/artiq_ir_generator.py @@ -2217,6 +2217,51 @@ class ARTIQIRGenerator(algorithm.Visitor): return result else: assert False + elif types.is_builtin(typ, "numpy.transpose"): + if len(node.args) == 1 and len(node.keywords) == 0: + arg, = map(self.visit, node.args) + + num_dims = arg.type.find()["num_dims"].value + if num_dims == 1: + # No-op as per NumPy semantics. + return arg + assert num_dims == 2 + arg_shape = self.append(ir.GetAttr(arg, "shape")) + dim0 = self.append(ir.GetAttr(arg_shape, 0)) + dim1 = self.append(ir.GetAttr(arg_shape, 1)) + shape = self._make_array_shape([dim1, dim0]) + result = self._allocate_new_array(node.type.find()["elt"], shape) + arg_buffer = self.append(ir.GetAttr(arg, "buffer")) + result_buffer = self.append(ir.GetAttr(result, "buffer")) + + def outer_gen(idx1): + arg_base = self.append(ir.Offset(arg_buffer, idx1)) + result_offset = self.append(ir.Arith(ast.Mult(loc=None), idx1, + dim0)) + result_base = self.append(ir.Offset(result_buffer, result_offset)) + + def inner_gen(idx0): + arg_offset = self.append( + ir.Arith(ast.Mult(loc=None), idx0, dim1)) + val = self.append(ir.GetElem(arg_base, arg_offset)) + self.append(ir.SetElem(result_base, idx0, val)) + return self.append( + ir.Arith(ast.Add(loc=None), idx0, ir.Constant(1, + idx0.type))) + + self._make_loop( + ir.Constant(0, self._size_type), lambda idx0: self.append( + ir.Compare(ast.Lt(loc=None), idx0, dim0)), inner_gen) + return self.append( + ir.Arith(ast.Add(loc=None), idx1, ir.Constant(1, idx1.type))) + + self._make_loop( + ir.Constant(0, self._size_type), + lambda idx1: self.append(ir.Compare(ast.Lt(loc=None), idx1, dim1)), + outer_gen) + return result + else: + assert False elif types.is_builtin(typ, "print"): self.polymorphic_print([self.visit(arg) for arg in node.args], separator=" ", suffix="\n") diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py index 445d64bce..d017af963 100644 --- a/artiq/compiler/transforms/inferencer.py +++ b/artiq/compiler/transforms/inferencer.py @@ -1074,6 +1074,45 @@ class Inferencer(algorithm.Visitor): arg1.loc, None) else: diagnose(valid_forms()) + elif types.is_builtin(typ, "numpy.transpose"): + valid_forms = lambda: [ + valid_form("transpose(x: array(elt='a, num_dims=1)) -> array(elt='a, num_dims=1)"), + valid_form("transpose(x: array(elt='a, num_dims=2)) -> array(elt='a, num_dims=2)") + ] + + if len(node.args) == 1 and len(node.keywords) == 0: + arg, = node.args + + if types.is_var(arg.type): + pass # undetermined yet + elif not builtins.is_array(arg.type): + note = diagnostic.Diagnostic( + "note", "this expression has type {type}", + {"type": types.TypePrinter().name(arg.type)}, arg.loc) + diag = diagnostic.Diagnostic( + "error", + "the argument of {builtin}() must be an array", + {"builtin": typ.find().name}, + node.func.loc, + notes=[note]) + self.engine.process(diag) + else: + num_dims = arg.type.find()["num_dims"].value + if num_dims not in (1, 2): + note = diagnostic.Diagnostic( + "note", "argument is {num_dims}-dimensional", + {"num_dims": num_dims}, arg.loc) + diag = diagnostic.Diagnostic( + "error", + "{builtin}() is currently only supported for up to " + "two-dimensional arrays", {"builtin": typ.find().name}, + node.func.loc, + notes=[note]) + self.engine.process(diag) + else: + self._unify(node.type, arg.type, node.loc, None) + else: + diagnose(valid_forms()) elif types.is_builtin(typ, "rtio_log"): valid_forms = lambda: [ valid_form("rtio_log(channel:str, args...) -> None"), diff --git a/artiq/test/lit/embedding/array_transpose.py b/artiq/test/lit/embedding/array_transpose.py new file mode 100644 index 000000000..2ab44bd7d --- /dev/null +++ b/artiq/test/lit/embedding/array_transpose.py @@ -0,0 +1,22 @@ +# RUN: %python -m artiq.compiler.testbench.embedding %s + +from artiq.language.core import * +from artiq.language.types import * +import numpy as np + +@kernel +def entrypoint(): + # FIXME: This needs to be a runtime test (but numpy.* integration is + # currently embedding-only). + a = np.array([1, 2, 3]) + b = np.transpose(a) + assert a.shape == b.shape + for i in range(len(a)): + assert a[i] == b[i] + + c = np.array([[1, 2, 3], [4, 5, 6]]) + d = np.transpose(c) + assert c.shape == d.shape + for i in range(2): + for j in range(3): + assert c[i][j] == d[j][i]