[arm64] JIT: Fold "A * B + C" to MADD/MSUB (#61037)

Author: EgorBo

src/coreclr/jit/codegenlinear.cpp

@@ -1608,9 +1608,14 @@ void CodeGen::genConsumeRegs(GenTree* tree)
         }
 #endif // FEATURE_HW_INTRINSICS
 #endif // TARGET_XARCH
-        else if (tree->OperIs(GT_BITCAST))
+        else if (tree->OperIs(GT_BITCAST, GT_NEG))
         {
-            genConsumeReg(tree->gtGetOp1());
+            genConsumeRegs(tree->gtGetOp1());
+        }
+        else if (tree->OperIs(GT_MUL))
+        {
+            genConsumeRegs(tree->gtGetOp1());
+            genConsumeRegs(tree->gtGetOp2());
         }
         else
         {

src/coreclr/jit/emitarm64.cpp

@@ -13597,6 +13597,44 @@ regNumber emitter::emitInsTernary(instruction ins, emitAttr attr, GenTree* dst,
         // src2 can only be a reg
         assert(!src2->isContained());
     }
+    else if ((src1->OperIs(GT_MUL) && src1->isContained()) || (src2->OperIs(GT_MUL) && src2->isContained()))
+    {
+        assert(ins == INS_add);
+
+        GenTree* mul;
+        GenTree* c;
+        if (src1->OperIs(GT_MUL))
+        {
+            mul = src1;
+            c   = src2;
+        }
+        else
+        {
+            mul = src2;
+            c   = src1;
+        }
+
+        GenTree* a = mul->gtGetOp1();
+        GenTree* b = mul->gtGetOp2();
+
+        assert(varTypeIsIntegral(mul) && !mul->gtOverflow());
+
+        bool msub = false;
+        if (a->OperIs(GT_NEG) && a->isContained())
+        {
+            a    = a->gtGetOp1();
+            msub = true;
+        }
+        if (b->OperIs(GT_NEG) && b->isContained())
+        {
+            b    = b->gtGetOp1();
+            msub = !msub; // it's either "a * -b" or "-a * -b" which is the same as "a * b"
+        }
+
+        emitIns_R_R_R_R(msub ? INS_msub : INS_madd, attr, dst->GetRegNum(), a->GetRegNum(), b->GetRegNum(),
+                        c->GetRegNum());
+        return dst->GetRegNum();
+    }
     else // not floating point
     {
         // src2 can be immed or reg

src/coreclr/jit/lowerarmarch.cpp

@@ -1568,6 +1568,47 @@ void Lowering::ContainCheckBinary(GenTreeOp* node)
 {
     // Check and make op2 contained (if it is a containable immediate)
     CheckImmedAndMakeContained(node, node->gtOp2);
+
+#ifdef TARGET_ARM64
+    // Find "a * b + c" or "c + a * b" in order to emit MADD/MSUB
+    if (comp->opts.OptimizationEnabled() && varTypeIsIntegral(node) && !node->isContained() && node->OperIs(GT_ADD) &&
+        !node->gtOverflow() && (node->gtGetOp1()->OperIs(GT_MUL) || node->gtGetOp2()->OperIs(GT_MUL)))
+    {
+        GenTree* mul;
+        GenTree* c;
+        if (node->gtGetOp1()->OperIs(GT_MUL))
+        {
+            mul = node->gtGetOp1();
+            c   = node->gtGetOp2();
+        }
+        else
+        {
+            mul = node->gtGetOp2();
+            c   = node->gtGetOp1();
+        }
+
+        GenTree* a = mul->gtGetOp1();
+        GenTree* b = mul->gtGetOp2();
+
+        if (!mul->isContained() && !mul->gtOverflow() && !a->isContained() && !b->isContained() && !c->isContained() &&
+            varTypeIsIntegral(mul))
+        {
+            if (a->OperIs(GT_NEG) && !a->gtGetOp1()->isContained() && !a->gtGetOp1()->IsRegOptional())
+            {
+                // "-a * b + c" to MSUB
+                MakeSrcContained(mul, a);
+            }
+            if (b->OperIs(GT_NEG) && !b->gtGetOp1()->isContained())
+            {
+                // "a * -b + c" to MSUB
+                MakeSrcContained(mul, b);
+            }
+            // If both 'a' and 'b' are GT_NEG - MADD will be emitted.
+
+            MakeSrcContained(node, mul);
+        }
+    }
+#endif
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/lsraarm64.cpp

@@ -265,7 +265,6 @@ int LinearScan::BuildNode(GenTree* tree)
                 // everything is made explicit by adding casts.
                 assert(tree->gtGetOp1()->TypeGet() == tree->gtGetOp2()->TypeGet());
             }
-
             FALLTHROUGH;
 
         case GT_AND:

src/coreclr/jit/lsrabuild.cpp

@@ -3071,6 +3071,18 @@ int LinearScan::BuildOperandUses(GenTree* node, regMaskTP candidates)
         return 1;
     }
 #endif // FEATURE_HW_INTRINSICS
+#ifdef TARGET_ARM64
+    if (node->OperIs(GT_MUL))
+    {
+        // Can be contained for MultiplyAdd on arm64
+        return BuildBinaryUses(node->AsOp(), candidates);
+    }
+    if (node->OperIs(GT_NEG))
+    {
+        // Can be contained for MultiplyAdd on arm64
+        return BuildOperandUses(node->gtGetOp1(), candidates);
+    }
+#endif
 
     return 0;
 }