compiler: Implement unary plus/minus for arrays

Implementation is needlessly generic to anticipate
coercion/transcendental functions.

artiq/compiler/transforms/artiq_ir_generator.py

@@ -88,6 +88,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
     necessary. They are kept track of in global dictionaries, with a mangled name
     containing types and operations as key:
 
+    :ivar array_unaryop_funcs: the map from mangled name to implementation of unary
+        operations for arrays
     :ivar array_binop_funcs: the map from mangled name to implementation of binary
         operations between arrays
     """
@@ -118,6 +120,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
         self.function_map = dict()
         self.variable_map = dict()
         self.method_map = defaultdict(lambda: [])
+        self.array_unaryop_funcs = dict()
         self.array_binop_funcs = dict()
 
     def annotate_calls(self, devirtualization):
@@ -1316,6 +1319,68 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 value_tail.append(ir.Branch(tail))
         return phi
 
+    def _make_array_unaryop(self, name, make_op, result_type, arg_type):
+        try:
+            result = ir.Argument(result_type, "result")
+            arg = ir.Argument(arg_type, "arg")
+
+            # TODO: We'd like to use a "C function" here to be able to supply
+            # specialised implementations in a library in the future (and e.g. avoid
+            # passing around the context argument), but the code generator currently
+            # doesn't allow emitting them.
+            args = [result, arg]
+            typ = types.TFunction(args=OrderedDict([(arg.name, arg.type)
+                                                    for arg in args]),
+                                  optargs=OrderedDict(),
+                                  ret=builtins.TNone())
+            env_args = [ir.EnvironmentArgument(self.current_env.type, "ARG.ENV")]
+
+            # TODO: What to use for loc?
+            func = ir.Function(typ, name, env_args + args, loc=None)
+            func.is_internal = True
+            func.is_generated = True
+            self.functions.append(func)
+            old_func, self.current_function = self.current_function, func
+
+            entry = self.add_block("entry")
+            old_block, self.current_block = self.current_block, entry
+
+            old_final_branch, self.final_branch = self.final_branch, None
+            old_unwind, self.unwind_target = self.unwind_target, None
+
+            shape = self.append(ir.GetAttr(arg, "shape"))
+
+            result_buffer = self.append(ir.GetAttr(result, "buffer"))
+            arg_buffer = self.append(ir.GetAttr(arg, "buffer"))
+            num_total_elts = self._get_total_array_len(shape)
+
+            def body_gen(index):
+                a = self.append(ir.GetElem(arg_buffer, index))
+                self.append(
+                    ir.SetElem(result_buffer, index, self.append(make_op(a))))
+                return self.append(
+                    ir.Arith(ast.Add(loc=None), index, ir.Constant(1, self._size_type)))
+
+            self._make_loop(
+                ir.Constant(0, self._size_type), lambda index: self.append(
+                    ir.Compare(ast.Lt(loc=None), index, num_total_elts)), body_gen)
+
+            self.append(ir.Return(ir.Constant(None, builtins.TNone())))
+            return func
+        finally:
+            self.current_function = old_func
+            self.current_block = old_block
+            self.final_branch = old_final_branch
+            self.unwind_target = old_unwind
+
+    def _get_array_unaryop(self, name, make_op, result_type, arg_type):
+        name = "_array_{}_{}".format(
+            name, self._mangle_arrayop_types([result_type, arg_type]))
+        if name not in self.array_unaryop_funcs:
+            self.array_binop_funcs[name] = self._make_array_unaryop(
+                name, make_op, result_type, arg_type)
+        return self.array_binop_funcs[name]
+
     def visit_UnaryOpT(self, node):
         if isinstance(node.op, ast.Not):
             cond = self.coerce_to_bool(self.visit(node.operand))
@@ -1327,9 +1392,18 @@ class ARTIQIRGenerator(algorithm.Visitor):
             return self.append(ir.Arith(ast.BitXor(loc=None),
                                         ir.Constant(-1, operand.type), operand))
         elif isinstance(node.op, ast.USub):
+            def make_sub(val):
+                return ir.Arith(ast.Sub(loc=None),
+                                        ir.Constant(0, val.type), val)
             operand = self.visit(node.operand)
-            return self.append(ir.Arith(ast.Sub(loc=None),
+            if builtins.is_array(operand.type):
-                                        ir.Constant(0, operand.type), operand))
+                shape = self.append(ir.GetAttr(operand, "shape"))
+                result = self._alloate_new_array(node.type.find()["elt"], shape)
+                func = self._get_array_unaryop("USub", make_sub, node.type, operand.type)
+                self._invoke_arrayop(func, [result, operand])
+                return result
+            else:
+                return self.append(make_sub(operand))
         elif isinstance(node.op, ast.UAdd):
             # No-op.
             return self.visit(node.operand)
@@ -1423,10 +1497,7 @@ class ARTIQIRGenerator(algorithm.Visitor):
             self.final_branch = old_final_branch
             self.unwind_target = old_unwind
 
-    def _get_array_binop(self, op, result_type, lhs_type, rhs_type):
+    def _mangle_arrayop_types(self, types):
-        # Currently, we always have any type coercions resolved explicitly in the AST.
-        # In the future, this might no longer be true and the three types might all
-        # differ.
         def name_error(typ):
             assert False, "Internal compiler error: No RPC tag for {}".format(typ)
 
@@ -1438,13 +1509,30 @@ class ARTIQIRGenerator(algorithm.Visitor):
             return (ir.rpc_tag(typ["elt"], name_error).decode() +
                     str(typ["num_dims"].find().value))
 
-        name = "_array_{}_{}_{}_{}".format(
+        return "_".join(mangle_name(t) for t in types)
-            type(op).__name__, *(map(mangle_name, (result_type, lhs_type, rhs_type))))
+
+    def _get_array_binop(self, op, result_type, lhs_type, rhs_type):
+        # Currently, we always have any type coercions resolved explicitly in the AST.
+        # In the future, this might no longer be true and the three types might all
+        # differ.
+        name = "_array_{}_{}".format(
+            type(op).__name__,
+            self._mangle_arrayop_types([result_type, lhs_type, rhs_type]))
         if name not in self.array_binop_funcs:
             self.array_binop_funcs[name] = self._make_array_binop(
                 name, op, result_type, lhs_type, rhs_type)
         return self.array_binop_funcs[name]
 
+    def _invoke_arrayop(self, func, params):
+        closure = self.append(
+            ir.Closure(func, ir.Constant(None, ir.TEnvironment("arrayop", {}))))
+        if self.unwind_target is None:
+            self.append(ir.Call(closure, params, {}))
+        else:
+            after_invoke = self.add_block("arrayop.invoke")
+            self.append(ir.Invoke(func, params, {}, after_invoke, self.unwind_target))
+            self.current_block = after_invoke
+
     def visit_BinOpT(self, node):
         if builtins.is_array(node.type):
             lhs = self.visit(node.left)
@@ -1457,14 +1545,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
             result = self._alloate_new_array(node.type.find()["elt"], shape)
 
             func = self._get_array_binop(node.op, node.type, node.left.type, node.right.type)
-            closure = self.append(ir.Closure(func, ir.Constant(None, ir.TEnvironment("arrayop", {}))))
+            self._invoke_arrayop(func, [result, lhs, rhs])
-            params = [result, lhs, rhs]
+
-            if self.unwind_target is None:
-                insn = self.append(ir.Call(closure, params, {}))
-            else:
-                after_invoke = self.add_block("arrayop.invoke")
-                insn = self.append(ir.Invoke(func, params, {}, after_invoke, self.unwind_target))
-                self.current_block = after_invoke
             return result
         elif builtins.is_numeric(node.type):
             lhs = self.visit(node.left)

artiq/compiler/transforms/inferencer.py

@@ -298,10 +298,21 @@ class Inferencer(algorithm.Visitor):
                     node.operand.loc)
                 self.engine.process(diag)
         else: # UAdd, USub
+            if types.is_var(operand_type):
+                return
+
             if builtins.is_numeric(operand_type):
-                self._unify(node.type, operand_type,
+                self._unify(node.type, operand_type, node.loc, None)
-                            node.loc, None)
+                return
-            elif not types.is_var(operand_type):
+
+            if builtins.is_array(operand_type):
+                elt = operand_type.find()["elt"]
+                if builtins.is_numeric(elt):
+                    self._unify(node.type, operand_type, node.loc, None)
+                    return
+                if types.is_var(elt):
+                    return
+
             diag = diagnostic.Diagnostic("error",
                 "expected unary '{op}' operand to be of numeric type, not {type}",
                 {"op": node.op.loc.source(),
@@ -436,7 +447,8 @@ class Inferencer(algorithm.Visitor):
                     return typ
 
             def map_return(typ):
-                a = builtins.TArray(elt=typ, num_dims=left_dims)
+                elt = builtins.TFloat() if isinstance(op, ast.Div) else typ
+                a = builtins.TArray(elt=elt, num_dims=left_dims)
                 return (a, a, a)
 
             return self._coerce_numeric((left, right),

artiq/compiler/validators/escape.py

@@ -156,7 +156,7 @@ class RegionOf(algorithm.Visitor):
     visit_NameConstantT = visit_immutable
     visit_NumT = visit_immutable
     visit_EllipsisT = visit_immutable
-    visit_UnaryOpT = visit_immutable
+    visit_UnaryOpT = visit_sometimes_allocating   # possibly array op
     visit_CompareT = visit_immutable
 
     # Value lives forever

artiq/test/lit/integration/array_unaryops.py

@@ -0,0 +1,11 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+
+a = array([1, 2])
+
+b = +a
+assert b[0] == 1
+assert b[1] == 2
+
+b = -a
+assert b[0] == -1
+assert b[1] == -2