diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index afa75b61b..afccb5aba 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/artiq/compiler/transforms/artiq_ir_generator.py @@ -390,6 +390,14 @@ class ARTIQIRGenerator(algorithm.Visitor): def visit_AugAssign(self, node): lhs = self.visit(node.target) rhs = self.visit(node.value) + + if builtins.is_array(lhs.type): + name = type(node.op).__name__ + def make_op(l, r): + return self.append(ir.Arith(node.op, l, r)) + self._broadcast_binop(name, make_op, lhs.type, lhs, rhs, assign_to_lhs=True) + return + value = self.append(ir.Arith(node.op, lhs, rhs)) try: self.current_assign = value @@ -1715,7 +1723,7 @@ class ARTIQIRGenerator(algorithm.Visitor): return self.append(ir.Alloc([result_buffer, shape], node.type)) return self.append(ir.GetElem(result_buffer, ir.Constant(0, self._size_type))) - def _broadcast_binop(self, name, make_op, result_type, lhs, rhs): + def _broadcast_binop(self, name, make_op, result_type, lhs, rhs, assign_to_lhs): # Broadcast scalars (broadcasting higher dimensions is not yet allowed in the # language). broadcast = False @@ -1736,9 +1744,11 @@ class ARTIQIRGenerator(algorithm.Visitor): builtins.TException("ValueError"), ir.Constant("operands could not be broadcast together", builtins.TStr()))) - - elt = result_type.find()["elt"] - result, _ = self._allocate_new_array(elt, shape) + if assign_to_lhs: + result = lhs + else: + elt = result_type.find()["elt"] + result, _ = self._allocate_new_array(elt, shape) func = self._get_array_elementwise_binop(name, make_op, result_type, lhs.type, rhs.type) self._invoke_arrayop(func, [result, lhs, rhs]) @@ -1753,7 +1763,8 @@ class ARTIQIRGenerator(algorithm.Visitor): name = type(node.op).__name__ def make_op(l, r): return self.append(ir.Arith(node.op, l, r)) - return self._broadcast_binop(name, make_op, node.type, lhs, rhs) + return self._broadcast_binop(name, make_op, node.type, lhs, rhs, + assign_to_lhs=False) elif builtins.is_numeric(node.type): lhs = self.visit(node.left) rhs = self.visit(node.right) @@ -2446,7 +2457,8 @@ class ARTIQIRGenerator(algorithm.Visitor): self._invoke_arrayop(func, [result, args[0]]) insn = result elif len(args) == 2: - insn = self._broadcast_binop(name, make_call, node.type, *args) + insn = self._broadcast_binop(name, make_call, node.type, *args, + assign_to_lhs=False) else: assert False, "Broadcasting for {} arguments not implemented".format(len) else: diff --git a/artiq/test/lit/integration/array_binops.py b/artiq/test/lit/integration/array_binops.py index e60052277..cfca32d92 100644 --- a/artiq/test/lit/integration/array_binops.py +++ b/artiq/test/lit/integration/array_binops.py @@ -8,31 +8,61 @@ assert c[0] == 5 assert c[1] == 7 assert c[2] == 9 +c += a +assert c[0] == 6 +assert c[1] == 9 +assert c[2] == 12 + c = b - a assert c[0] == 3 assert c[1] == 3 assert c[2] == 3 +c -= a +assert c[0] == 2 +assert c[1] == 1 +assert c[2] == 0 + c = a * b assert c[0] == 4 assert c[1] == 10 assert c[2] == 18 +c *= a +assert c[0] == 4 +assert c[1] == 20 +assert c[2] == 54 + c = b // a assert c[0] == 4 assert c[1] == 2 assert c[2] == 2 +c //= a +assert c[0] == 4 +assert c[1] == 1 +assert c[2] == 0 + c = a ** b assert c[0] == 1 assert c[1] == 32 assert c[2] == 729 +c **= a +assert c[0] == 1 +assert c[1] == 1024 +assert c[2] == 387420489 + c = b % a assert c[0] == 0 assert c[1] == 1 assert c[2] == 0 +c %= a +assert c[0] == 0 +assert c[1] == 1 +assert c[2] == 0 + cf = b / a assert cf[0] == 4.0 assert cf[1] == 2.5 @@ -43,6 +73,16 @@ assert cf2[0] == 5.0 assert cf2[1] == 4.5 assert cf2[2] == 5.0 +cf2 += a +assert cf2[0] == 6.0 +assert cf2[1] == 6.5 +assert cf2[2] == 8.0 + +cf /= a +assert cf[0] == 4.0 +assert cf[1] == 1.25 +assert cf[2] == 2.0 / 3.0 + d = array([[1, 2], [3, 4]]) e = array([[5, 6], [7, 8]]) f = d + e @@ -50,3 +90,9 @@ assert f[0][0] == 6 assert f[0][1] == 8 assert f[1][0] == 10 assert f[1][1] == 12 + +f += d +assert f[0][0] == 7 +assert f[0][1] == 10 +assert f[1][0] == 13 +assert f[1][1] == 16