[spirv] Translate arithmetic assign operators (#493)

* Covers +=, -=, *=, /=, %= for scalars and vectors

tools/clang/lib/SPIRV/EmitSPIRVAction.cpp

@@ -618,6 +618,12 @@ public:
       return translateAPFloat(floatLiteral->getValue(), expr->getType());
     }
 
+    // CompoundAssignOperator is a subclass of BinaryOperator. It should be
+    // checked before BinaryOperator.
+    if (const auto *compoundAssignOp = dyn_cast<CompoundAssignOperator>(expr)) {
+      return doCompoundAssignOperator(compoundAssignOp);
+    }
+
     if (const auto *binOp = dyn_cast<BinaryOperator>(expr)) {
       return doBinaryOperator(binOp);
     }
@@ -687,6 +693,46 @@ public:
     return 0;
   }
 
+  uint32_t doCompoundAssignOperator(const CompoundAssignOperator *expr) {
+    const auto opcode = expr->getOpcode();
+
+    // Try to optimize floatN *= float case
+    if (opcode == BO_MulAssign) {
+      if (const uint32_t result = tryToGenFloatVectorScale(expr))
+        return result;
+    }
+
+    const auto *rhs = expr->getRHS();
+    const auto *lhs = expr->getLHS();
+
+    switch (opcode) {
+    case BO_AddAssign:
+    case BO_SubAssign:
+    case BO_MulAssign:
+    case BO_DivAssign:
+    case BO_RemAssign: {
+      const uint32_t resultType = typeTranslator.translateType(expr->getType());
+
+      // Evalute rhs before lhs
+      const uint32_t rhsVal = doExpr(rhs);
+      const uint32_t lhsPtr = doExpr(lhs);
+      const uint32_t lhsVal = theBuilder.createLoad(resultType, lhsPtr);
+
+      const spv::Op spvOp = translateOp(opcode, expr->getType());
+      const uint32_t result =
+          theBuilder.createBinaryOp(spvOp, resultType, lhsVal, rhsVal);
+      theBuilder.createStore(lhsPtr, result);
+
+      // Compound assign operators return lvalues.
+      return lhsPtr;
+    }
+    default:
+      emitError("CompoundAssignOperator '%0' unimplemented")
+          << expr->getStmtClassName();
+      return 0;
+    }
+  }
+
   uint32_t doUnaryOperator(const UnaryOperator *expr) {
     const auto opcode = expr->getOpcode();
     const auto *subExpr = expr->getSubExpr();
@@ -848,10 +894,22 @@ public:
         if (cast->getCastKind() == CK_HLSLVectorSplat) {
           const uint32_t vecType =
               typeTranslator.translateType(expr->getType());
-          const uint32_t lhsId = doExpr(lhs);
-          const uint32_t rhsId = doExpr(cast->getSubExpr());
-          return theBuilder.createBinaryOp(spv::Op::OpVectorTimesScalar,
-                                           vecType, lhsId, rhsId);
+          if (isa<CompoundAssignOperator>(expr)) {
+            // For floatN * float cases. We'll need to do the load/store and
+            // return the lhs.
+            const uint32_t rhsVal = doExpr(cast->getSubExpr());
+            const uint32_t lhsPtr = doExpr(lhs);
+            const uint32_t lhsVal = theBuilder.createLoad(vecType, lhsPtr);
+            const uint32_t result = theBuilder.createBinaryOp(
+                spv::Op::OpVectorTimesScalar, vecType, lhsVal, rhsVal);
+            theBuilder.createStore(lhsPtr, result);
+            return lhsPtr;
+          } else {
+            const uint32_t lhsId = doExpr(lhs);
+            const uint32_t rhsId = doExpr(cast->getSubExpr());
+            return theBuilder.createBinaryOp(spv::Op::OpVectorTimesScalar,
+                                             vecType, lhsId, rhsId);
+          }
         }
       }
     }
@@ -910,9 +968,13 @@ case BO_##kind : {                                                             \
 
     switch (op) {
       BIN_OP_CASE_INT_FLOAT(Add, IAdd, FAdd);
+      BIN_OP_CASE_INT_FLOAT(AddAssign, IAdd, FAdd);
       BIN_OP_CASE_INT_FLOAT(Sub, ISub, FSub);
+      BIN_OP_CASE_INT_FLOAT(SubAssign, ISub, FSub);
       BIN_OP_CASE_INT_FLOAT(Mul, IMul, FMul);
+      BIN_OP_CASE_INT_FLOAT(MulAssign, IMul, FMul);
       BIN_OP_CASE_SINT_UINT_FLOAT(Div, SDiv, UDiv, FDiv);
+      BIN_OP_CASE_SINT_UINT_FLOAT(DivAssign, SDiv, UDiv, FDiv);
       // According to HLSL spec, "the modulus operator returns the remainder of
       // a division." "The % operator is defined only in cases where either both
       // sides are positive or both sides are negative."
@@ -928,6 +990,7 @@ case BO_##kind : {                                                             \
       //
       // Note there is no OpURem in SPIR-V.
       BIN_OP_CASE_SINT_UINT_FLOAT(Rem, SRem, UMod, FRem);
+      BIN_OP_CASE_SINT_UINT_FLOAT(RemAssign, SRem, UMod, FRem);
       BIN_OP_CASE_SINT_UINT_FLOAT(LT, SLessThan, ULessThan, FOrdLessThan);
       BIN_OP_CASE_SINT_UINT_FLOAT(LE, SLessThanEqual, ULessThanEqual,
                                   FOrdLessThanEqual);

tools/clang/test/CodeGenSPIRV/binary-op.arith-assign.mixed.hlsl

@@ -0,0 +1,54 @@
+// Run: %dxc -T vs_6_0 -E main
+
+void main() {
+// CHECK-LABEL: %bb_entry = OpLabel
+
+    float4 a;
+    float s;
+
+    int3 c;
+    int t;
+
+// CHECK:      [[s0:%\d+]] = OpLoad %float %s
+// CHECK-NEXT: [[cc0:%\d+]] = OpCompositeConstruct %v4float [[s0]] [[s0]] [[s0]] [[s0]]
+// CHECK-NEXT: [[a0:%\d+]] = OpLoad %v4float %a
+// CHECK-NEXT: [[add0:%\d+]] = OpFAdd %v4float [[a0]] [[cc0]]
+// CHECK-NEXT: OpStore %a [[add0]]
+    a += s;
+
+// CHECK-NEXT: [[s2:%\d+]] = OpLoad %float %s
+// CHECK-NEXT: [[cc2:%\d+]] = OpCompositeConstruct %v4float [[s2]] [[s2]] [[s2]] [[s2]]
+// CHECK-NEXT: [[a2:%\d+]] = OpLoad %v4float %a
+// CHECK-NEXT: [[sub0:%\d+]] = OpFSub %v4float [[a2]] [[cc2]]
+// CHECK-NEXT: OpStore %a [[sub0]]
+    a -= s;
+
+    // Use OpVectorTimesScalar for floatN * float
+// CHECK-NEXT: [[s4:%\d+]] = OpLoad %float %s
+// CHECK-NEXT: [[a4:%\d+]] = OpLoad %v4float %a
+// CHECK-NEXT: [[mul0:%\d+]] = OpVectorTimesScalar %v4float [[a4]] [[s4]]
+// CHECK-NEXT: OpStore %a [[mul0]]
+    a *= s;
+
+// CHECK-NEXT: [[s6:%\d+]] = OpLoad %float %s
+// CHECK-NEXT: [[cc6:%\d+]] = OpCompositeConstruct %v4float [[s6]] [[s6]] [[s6]] [[s6]]
+// CHECK-NEXT: [[a6:%\d+]] = OpLoad %v4float %a
+// CHECK-NEXT: [[div0:%\d+]] = OpFDiv %v4float [[a6]] [[cc6]]
+// CHECK-NEXT: OpStore %a [[div0]]
+    a /= s;
+
+// CHECK-NEXT: [[s8:%\d+]] = OpLoad %float %s
+// CHECK-NEXT: [[cc8:%\d+]] = OpCompositeConstruct %v4float [[s8]] [[s8]] [[s8]] [[s8]]
+// CHECK-NEXT: [[a8:%\d+]] = OpLoad %v4float %a
+// CHECK-NEXT: [[mod0:%\d+]] = OpFRem %v4float [[a8]] [[cc8]]
+// CHECK-NEXT: OpStore %a [[mod0]]
+    a %= s;
+
+    // Use normal OpCompositeConstruct and OpIMul for intN * int
+// CHECK-NEXT: [[t0:%\d+]] = OpLoad %int %t
+// CHECK-NEXT: [[cc10:%\d+]] = OpCompositeConstruct %v3int [[t0]] [[t0]] [[t0]]
+// CHECK-NEXT: [[c0:%\d+]] = OpLoad %v3int %c
+// CHECK-NEXT: [[mul2:%\d+]] = OpIMul %v3int [[c0]] [[cc10]]
+// CHECK-NEXT: OpStore %c [[mul2]]
+    c *= t;
+}

tools/clang/test/CodeGenSPIRV/binary-op.arith-assign.scalar.hlsl

@@ -0,0 +1,102 @@
+// Run: %dxc -T ps_6_0 -E main
+
+void main() {
+// CHECK-LABEL: %bb_entry = OpLabel
+    int a, b, c, d;
+    uint i, j;
+    float o, p;
+
+// CHECK:      [[a0:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b0:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[add0:%\d+]] = OpIAdd %int [[b0]] [[a0]]
+// CHECK-NEXT: OpStore %b [[add0]]
+    b += a;
+// CHECK-NEXT: [[i0:%\d+]] = OpLoad %uint %i
+// CHECK-NEXT: [[j0:%\d+]] = OpLoad %uint %j
+// CHECK-NEXT: [[add1:%\d+]] = OpIAdd %uint [[j0]] [[i0]]
+// CHECK-NEXT: OpStore %j [[add1]]
+    j += i;
+// CHECK-NEXT: [[o0:%\d+]] = OpLoad %float %o
+// CHECK-NEXT: [[p0:%\d+]] = OpLoad %float %p
+// CHECK-NEXT: [[add2:%\d+]] = OpFAdd %float [[p0]] [[o0]]
+// CHECK-NEXT: OpStore %p [[add2]]
+    p += o;
+
+// CHECK-NEXT: [[a1:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b1:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[sub0:%\d+]] = OpISub %int [[b1]] [[a1]]
+// CHECK-NEXT: OpStore %b [[sub0]]
+    b -= a;
+// CHECK-NEXT: [[i1:%\d+]] = OpLoad %uint %i
+// CHECK-NEXT: [[j1:%\d+]] = OpLoad %uint %j
+// CHECK-NEXT: [[sub1:%\d+]] = OpISub %uint [[j1]] [[i1]]
+// CHECK-NEXT: OpStore %j [[sub1]]
+    j -= i;
+// CHECK-NEXT: [[o1:%\d+]] = OpLoad %float %o
+// CHECK-NEXT: [[p1:%\d+]] = OpLoad %float %p
+// CHECK-NEXT: [[sub2:%\d+]] = OpFSub %float [[p1]] [[o1]]
+// CHECK-NEXT: OpStore %p [[sub2]]
+    p -= o;
+
+// CHECK-NEXT: [[a2:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b2:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[mul0:%\d+]] = OpIMul %int [[b2]] [[a2]]
+// CHECK-NEXT: OpStore %b [[mul0]]
+    b *= a;
+// CHECK-NEXT: [[i2:%\d+]] = OpLoad %uint %i
+// CHECK-NEXT: [[j2:%\d+]] = OpLoad %uint %j
+// CHECK-NEXT: [[mul1:%\d+]] = OpIMul %uint [[j2]] [[i2]]
+// CHECK-NEXT: OpStore %j [[mul1]]
+    j *= i;
+// CHECK-NEXT: [[o2:%\d+]] = OpLoad %float %o
+// CHECK-NEXT: [[p2:%\d+]] = OpLoad %float %p
+// CHECK-NEXT: [[mul2:%\d+]] = OpFMul %float [[p2]] [[o2]]
+// CHECK-NEXT: OpStore %p [[mul2]]
+    p *= o;
+
+// CHECK-NEXT: [[a3:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b3:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[div0:%\d+]] = OpSDiv %int [[b3]] [[a3]]
+// CHECK-NEXT: OpStore %b [[div0]]
+    b /= a;
+// CHECK-NEXT: [[i3:%\d+]] = OpLoad %uint %i
+// CHECK-NEXT: [[j3:%\d+]] = OpLoad %uint %j
+// CHECK-NEXT: [[div1:%\d+]] = OpUDiv %uint [[j3]] [[i3]]
+// CHECK-NEXT: OpStore %j [[div1]]
+    j /= i;
+// CHECK-NEXT: [[o3:%\d+]] = OpLoad %float %o
+// CHECK-NEXT: [[p3:%\d+]] = OpLoad %float %p
+// CHECK-NEXT: [[div2:%\d+]] = OpFDiv %float [[p3]] [[o3]]
+// CHECK-NEXT: OpStore %p [[div2]]
+    p /= o;
+
+// CHECK-NEXT: [[a4:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b4:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[mod0:%\d+]] = OpSRem %int [[b4]] [[a4]]
+// CHECK-NEXT: OpStore %b [[mod0]]
+    b %= a;
+// CHECK-NEXT: [[i4:%\d+]] = OpLoad %uint %i
+// CHECK-NEXT: [[j4:%\d+]] = OpLoad %uint %j
+// CHECK-NEXT: [[mod1:%\d+]] = OpUMod %uint [[j4]] [[i4]]
+// CHECK-NEXT: OpStore %j [[mod1]]
+    j %= i;
+// CHECK-NEXT: [[o4:%\d+]] = OpLoad %float %o
+// CHECK-NEXT: [[p4:%\d+]] = OpLoad %float %p
+// CHECK-NEXT: [[mod2:%\d+]] = OpFRem %float [[p4]] [[o4]]
+// CHECK-NEXT: OpStore %p [[mod2]]
+    p %= o;
+
+    // Spot check that we can use the result on both the left-hand side
+    // and the right-hand side
+// CHECK-NEXT: [[a5:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b5:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[sub3:%\d+]] = OpISub %int [[b5]] [[a5]]
+// CHECK-NEXT: OpStore %b [[sub3]]
+// CHECK-NEXT: [[b6:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[c0:%\d+]] = OpLoad %int %c
+// CHECK-NEXT: [[d0:%\d+]] = OpLoad %int %d
+// CHECK-NEXT: [[add3:%\d+]] = OpIAdd %int [[d0]] [[c0]]
+// CHECK-NEXT: OpStore %d [[add3]]
+// CHECK-NEXT: OpStore %d [[b6]]
+    (d += c) = (b -= a);
+}

tools/clang/test/CodeGenSPIRV/binary-op.arith-assign.vector.hlsl

@@ -0,0 +1,115 @@
+// Run: %dxc -T vs_6_0 -E main
+
+void main() {
+// CHECK-LABEL: %bb_entry = OpLabel
+
+    int1 a, b, c, d;
+    int2 i, j;
+    uint3 o, p;
+    float4 x, y;
+
+// CHECK:      [[a0:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b0:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[add0:%\d+]] = OpIAdd %int [[b0]] [[a0]]
+// CHECK-NEXT: OpStore %b [[add0]]
+    b += a;
+// CHECK-NEXT: [[i0:%\d+]] = OpLoad %v2int %i
+// CHECK-NEXT: [[j0:%\d+]] = OpLoad %v2int %j
+// CHECK-NEXT: [[add1:%\d+]] = OpIAdd %v2int [[j0]] [[i0]]
+// CHECK-NEXT: OpStore %j [[add1]]
+    j += i;
+// CHECK-NEXT: [[o0:%\d+]] = OpLoad %v3uint %o
+// CHECK-NEXT: [[p0:%\d+]] = OpLoad %v3uint %p
+// CHECK-NEXT: [[add2:%\d+]] = OpIAdd %v3uint [[p0]] [[o0]]
+// CHECK-NEXT: OpStore %p [[add2]]
+    p += o;
+// CHECK-NEXT: [[x0:%\d+]] = OpLoad %v4float %x
+// CHECK-NEXT: [[y0:%\d+]] = OpLoad %v4float %y
+// CHECK-NEXT: [[add3:%\d+]] = OpFAdd %v4float [[y0]] [[x0]]
+// CHECK-NEXT: OpStore %y [[add3]]
+    y += x;
+
+// CHECK-NEXT: [[a1:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b1:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[sub0:%\d+]] = OpISub %int [[b1]] [[a1]]
+// CHECK-NEXT: OpStore %b [[sub0]]
+    b -= a;
+// CHECK-NEXT: [[i1:%\d+]] = OpLoad %v2int %i
+// CHECK-NEXT: [[j1:%\d+]] = OpLoad %v2int %j
+// CHECK-NEXT: [[sub1:%\d+]] = OpISub %v2int [[j1]] [[i1]]
+// CHECK-NEXT: OpStore %j [[sub1]]
+    j -= i;
+// CHECK-NEXT: [[o1:%\d+]] = OpLoad %v3uint %o
+// CHECK-NEXT: [[p1:%\d+]] = OpLoad %v3uint %p
+// CHECK-NEXT: [[sub2:%\d+]] = OpISub %v3uint [[p1]] [[o1]]
+// CHECK-NEXT: OpStore %p [[sub2]]
+    p -= o;
+// CHECK-NEXT: [[x1:%\d+]] = OpLoad %v4float %x
+// CHECK-NEXT: [[y1:%\d+]] = OpLoad %v4float %y
+// CHECK-NEXT: [[sub3:%\d+]] = OpFSub %v4float [[y1]] [[x1]]
+// CHECK-NEXT: OpStore %y [[sub3]]
+    y -= x;
+
+// CHECK-NEXT: [[a2:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b2:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[mul0:%\d+]] = OpIMul %int [[b2]] [[a2]]
+// CHECK-NEXT: OpStore %b [[mul0]]
+    b *= a;
+// CHECK-NEXT: [[i2:%\d+]] = OpLoad %v2int %i
+// CHECK-NEXT: [[j2:%\d+]] = OpLoad %v2int %j
+// CHECK-NEXT: [[mul1:%\d+]] = OpIMul %v2int [[j2]] [[i2]]
+// CHECK-NEXT: OpStore %j [[mul1]]
+    j *= i;
+// CHECK-NEXT: [[o2:%\d+]] = OpLoad %v3uint %o
+// CHECK-NEXT: [[p2:%\d+]] = OpLoad %v3uint %p
+// CHECK-NEXT: [[mul2:%\d+]] = OpIMul %v3uint [[p2]] [[o2]]
+// CHECK-NEXT: OpStore %p [[mul2]]
+    p *= o;
+// CHECK-NEXT: [[x2:%\d+]] = OpLoad %v4float %x
+// CHECK-NEXT: [[y2:%\d+]] = OpLoad %v4float %y
+// CHECK-NEXT: [[mul3:%\d+]] = OpFMul %v4float [[y2]] [[x2]]
+// CHECK-NEXT: OpStore %y [[mul3]]
+    y *= x;
+
+// CHECK-NEXT: [[a4:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b4:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[div0:%\d+]] = OpSDiv %int [[b4]] [[a4]]
+// CHECK-NEXT: OpStore %b [[div0]]
+    b /= a;
+// CHECK-NEXT: [[i4:%\d+]] = OpLoad %v2int %i
+// CHECK-NEXT: [[j4:%\d+]] = OpLoad %v2int %j
+// CHECK-NEXT: [[div1:%\d+]] = OpSDiv %v2int [[j4]] [[i4]]
+// CHECK-NEXT: OpStore %j [[div1]]
+    j /= i;
+// CHECK-NEXT: [[o4:%\d+]] = OpLoad %v3uint %o
+// CHECK-NEXT: [[p4:%\d+]] = OpLoad %v3uint %p
+// CHECK-NEXT: [[div2:%\d+]] = OpUDiv %v3uint [[p4]] [[o4]]
+// CHECK-NEXT: OpStore %p [[div2]]
+    p /= o;
+// CHECK-NEXT: [[x4:%\d+]] = OpLoad %v4float %x
+// CHECK-NEXT: [[y4:%\d+]] = OpLoad %v4float %y
+// CHECK-NEXT: [[div3:%\d+]] = OpFDiv %v4float [[y4]] [[x4]]
+// CHECK-NEXT: OpStore %y [[div3]]
+    y /= x;
+
+// CHECK-NEXT: [[a5:%\d+]] = OpLoad %int %a
+// CHECK-NEXT: [[b5:%\d+]] = OpLoad %int %b
+// CHECK-NEXT: [[mod0:%\d+]] = OpSRem %int [[b5]] [[a5]]
+// CHECK-NEXT: OpStore %b [[mod0]]
+    b %= a;
+// CHECK-NEXT: [[i5:%\d+]] = OpLoad %v2int %i
+// CHECK-NEXT: [[j5:%\d+]] = OpLoad %v2int %j
+// CHECK-NEXT: [[mod1:%\d+]] = OpSRem %v2int [[j5]] [[i5]]
+// CHECK-NEXT: OpStore %j [[mod1]]
+    j %= i;
+// CHECK-NEXT: [[o5:%\d+]] = OpLoad %v3uint %o
+// CHECK-NEXT: [[p5:%\d+]] = OpLoad %v3uint %p
+// CHECK-NEXT: [[mod2:%\d+]] = OpUMod %v3uint [[p5]] [[o5]]
+// CHECK-NEXT: OpStore %p [[mod2]]
+    p %= o;
+// CHECK-NEXT: [[x5:%\d+]] = OpLoad %v4float %x
+// CHECK-NEXT: [[y5:%\d+]] = OpLoad %v4float %y
+// CHECK-NEXT: [[mod3:%\d+]] = OpFRem %v4float [[y5]] [[x5]]
+// CHECK-NEXT: OpStore %y [[mod3]]
+    y %= x;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -58,6 +58,16 @@ TEST_F(FileTest, BinaryOpMixedArithmetic) {
   runFileTest("binary-op.arithmetic.mixed.hlsl");
 }
 
+TEST_F(FileTest, BinaryOpScalarArithAssign) {
+  runFileTest("binary-op.arith-assign.scalar.hlsl");
+}
+TEST_F(FileTest, BinaryOpVectorArithAssign) {
+  runFileTest("binary-op.arith-assign.vector.hlsl");
+}
+TEST_F(FileTest, BinaryOpMixedArithAssign) {
+  runFileTest("binary-op.arith-assign.mixed.hlsl");
+}
+
 TEST_F(FileTest, BinaryOpScalarComparison) {
   runFileTest("binary-op.comparison.scalar.hlsl");
 }