Optimize FMA codegen base on the overwritten

src/coreclr/jit/gentree.cpp

@@ -19791,6 +19791,39 @@ uint16_t GenTreeLclVarCommon::GetLclOffs() const
     }
 }
 
+#if defined(TARGET_XARCH) && defined(FEATURE_HW_INTRINSICS)
+//------------------------------------------------------------------------
+// GetOverwrittenOpNumForFMA: check if the result is written into one of the operands
+//
+// Return Value:
+//     The operand number or 0 if not overwritten.
+//
+unsigned GenTreeHWIntrinsic::GetOverwrittenOpNumForFMA(GenTree* use, GenTree* op1, GenTree* op2, GenTree* op3)
+{
+    // only FMA intrinsic node should call into this function
+    assert(HWIntrinsicInfo::lookupIsa(gtHWIntrinsicId) == InstructionSet_FMA);
+    if (!use->OperIs(GT_STORE_LCL_VAR))
+        return 0;
+    GenTreeLclVarCommon* overwritten = use->AsLclVarCommon();
+    unsigned overwrittenLclNum = overwritten->GetLclNum();
+    unsigned overwrittenOpNum = 0; // 1->op1, 2->op2, 3->op3
+    if (op1->IsLocal() && op1->AsLclVarCommon()->GetLclNum() == overwrittenLclNum)
+    {
+        overwrittenOpNum = 1;
+    }
+    else if (op2->IsLocal() && op2->AsLclVarCommon()->GetLclNum() == overwrittenLclNum)
+    {
+        overwrittenOpNum = 2;
+    }
+    else if (op3->IsLocal() && op3->AsLclVarCommon()->GetLclNum() == overwrittenLclNum)
+    {
+        overwrittenOpNum = 3;
+    }
+
+    return overwrittenOpNum;
+}
+#endif // TARGET_XARCH && FEATURE_HW_INTRINSICS
+
 #ifdef TARGET_ARM
 //------------------------------------------------------------------------
 // IsOffsetMisaligned: check if the field needs a special handling on arm.

src/coreclr/jit/gentree.h

@@ -5187,6 +5187,8 @@ struct GenTreeHWIntrinsic : public GenTreeJitIntrinsic
     bool OperIsMemoryLoadOrStore() const; // Returns true for the HW Intrinsic instructions that have MemoryLoad or
                                           // MemoryStore semantics, false otherwise
 
+    unsigned GetOverwrittenOpNumForFMA(GenTree* use, GenTree* op1, GenTree* op2, GenTree* op3);
+
 #if DEBUGGABLE_GENTREE
     GenTreeHWIntrinsic() : GenTreeJitIntrinsic()
     {

src/coreclr/jit/hwintrinsiccodegenxarch.cpp

@@ -2133,47 +2133,74 @@ void CodeGen::genFMAIntrinsic(GenTreeHWIntrinsic* node)
     // Intrinsics with CopyUpperBits semantics cannot have op1 be contained
     assert(!copiesUpperBits || !op1->isContained());
 
-    if (op2->isContained() || op2->isUsedFromSpillTemp())
+    unsigned overwrittenOpNum = 0;
+    LIR::Use use;
+    if (LIR::AsRange(compiler->compCurBB).TryGetUse(node, &use))
     {
-        // 132 form: op1 = (op1 * op3) + [op2]
+        overwrittenOpNum = node->GetOverwrittenOpNumForFMA(use.User(), op1, op2, op3);
+    }
 
-        ins    = (instruction)(ins - 1);
-        op1Reg = op1->GetRegNum();
-        op2Reg = op3->GetRegNum();
-        op3    = op2;
+    if (overwrittenOpNum == 1)
+    {
+        if (op2->isContained())
+        {
+            // op1 = (op1 * [op2]) + op3
+            // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+            ins    = (instruction)(ins - 1);
+            op1Reg = op1->GetRegNum();
+            op2Reg = op3->GetRegNum();
+            op3    = op2;
+        }
+        else
+        {
+            // op1 = (op1 * op2) + [op3]
+            // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+            op1Reg = op1->GetRegNum();
+            op2Reg = op2->GetRegNum();
+        }
     }
-    else if (op1->isContained() || op1->isUsedFromSpillTemp())
+    else if (overwrittenOpNum == 3)
     {
-        // 231 form: op3 = (op2 * op3) + [op1]
-
+        // 231 form: XMM1 = (XMM2 * [XMM3]) + XMM1
+        // One of the following:
+        // op3 = ([op1] * op2) + op3
+        // op3 = (op1 * [op2]) + op3
         ins    = (instruction)(ins + 1);
         op1Reg = op3->GetRegNum();
-        op2Reg = op2->GetRegNum();
-        op3    = op1;
+        if (op1->isContained())
+        {
+            // op3 = ([op1] * op2) + op3
+            op2Reg = op2->GetRegNum();
+            op3    = op1;
+        }
+        else
+        {
+            // op3 = (op1 * [op2]) + op3
+            op2Reg = op1->GetRegNum();
+            op3    = op2;
+        }
     }
     else
     {
-        // 213 form: op1 = (op2 * op1) + [op3]
-
-        op1Reg = op1->GetRegNum();
-        op2Reg = op2->GetRegNum();
-
-        isCommutative = !copiesUpperBits;
-    }
-
-    if (isCommutative && (op1Reg != targetReg) && (op2Reg == targetReg))
-    {
-        assert(node->isRMWHWIntrinsic(compiler));
+        assert(overwrittenOpNum == 2 || overwrittenOpNum == 0);
+        if (op1->isContained())
+        {
+            // op2 = ([op1] * op2) + op3
+            // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+            ins    = (instruction)(ins - 1);
+            op1Reg = op2->GetRegNum();
+            op2Reg = op3->GetRegNum();
+            op3    = op1;
 
-        // We have "reg2 = (reg1 * reg2) +/- op3" where "reg1 != reg2" on a RMW intrinsic.
-        //
-        // For non-commutative intrinsics, we should have ensured that op2 was marked
-        // delay free in order to prevent it from getting assigned the same register
-        // as target. However, for commutative intrinsics, we can just swap the operands
-        // in order to have "reg2 = reg2 op reg1" which will end up producing the right code.
+        }
+        else
+        {
+            // op2 = (op1 * op2) + [op3]
+            // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+            op1Reg = op2->GetRegNum();
+            op2Reg = op1->GetRegNum();
+        }
 
-        op2Reg = op1Reg;
-        op1Reg = targetReg;
     }
 
     genHWIntrinsic_R_R_R_RM(ins, attr, targetReg, op1Reg, op2Reg, op3);

src/coreclr/jit/lowerxarch.cpp

@@ -6314,38 +6314,92 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
                     if ((intrinsicId >= NI_FMA_MultiplyAdd) && (intrinsicId <= NI_FMA_MultiplySubtractNegatedScalar))
                     {
                         bool supportsRegOptional = false;
+                        unsigned overwrittenOpNum    = 0;
+                        LIR::Use use;
 
-                        if (IsContainableHWIntrinsicOp(node, op3, &supportsRegOptional))
+                        if (BlockRange().TryGetUse(node, &use))
                         {
-                            // 213 form: op1 = (op2 * op1) + [op3]
-                            MakeSrcContained(node, op3);
+                            overwrittenOpNum = node->GetOverwrittenOpNumForFMA(use.User(), op1, op2, op3);
                         }
-                        else if (IsContainableHWIntrinsicOp(node, op2, &supportsRegOptional))
-                        {
-                            // 132 form: op1 = (op1 * op3) + [op2]
-                            MakeSrcContained(node, op2);
-                        }
-                        else if (IsContainableHWIntrinsicOp(node, op1, &supportsRegOptional))
+
+                        switch (overwrittenOpNum)
                         {
-                            // Intrinsics with CopyUpperBits semantics cannot have op1 be contained
+                            case 1:
+                            {
+                                if (IsContainableHWIntrinsicOp(node, op2, &supportsRegOptional))
+                                {
+                                    // op1 = (op1 * [op2]) + op3
+                                    // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+                                    MakeSrcContained(node, op2);
+                                }
+                                else if (IsContainableHWIntrinsicOp(node, op3, &supportsRegOptional))
+                                {
+                                    // op1 = (op1 * op2) + [op3]
+                                    // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+                                    MakeSrcContained(node, op3);
 
-                            if (!HWIntrinsicInfo::CopiesUpperBits(intrinsicId))
+                                }
+                                else
+                                {
+                                    assert(supportsRegOptional);
+                                    // op1 = (op1 * op2) + [op3]
+                                    // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+                                    op3->SetRegOptional();
+                                }
+                                break;
+                            }
+                            case 3:
                             {
-                                // 231 form: op3 = (op2 * op3) + [op1]
-                                MakeSrcContained(node, op1);
+                                // 231 form: XMM1 = (XMM2 * [XMM3]) + XMM1
+                                // One of the following:
+                                // op3 = ([op1] * op2) + op3
+                                // op3 = (op1 * [op2]) + op3
+                                if (IsContainableHWIntrinsicOp(node, op1, &supportsRegOptional) && !HWIntrinsicInfo::CopiesUpperBits(intrinsicId))
+                                {
+                                    // Intrinsics with CopyUpperBits semantics cannot have op1 be contained
+
+                                    MakeSrcContained(node, op1);
+                                }
+                                else if (IsContainableHWIntrinsicOp(node, op2, &supportsRegOptional))
+                                {
+                                    MakeSrcContained(node, op2);
+                                }
+                                else
+                                {
+                                    assert(supportsRegOptional);
+                                    op2->SetRegOptional();
+                                }
+                                break;
                             }
-                        }
-                        else
-                        {
-                            assert(supportsRegOptional);
+                            default:
+                            {
+                                assert(overwrittenOpNum == 2 || overwrittenOpNum == 0);
+                                if (IsContainableHWIntrinsicOp(node, op1, &supportsRegOptional) && !HWIntrinsicInfo::CopiesUpperBits(intrinsicId))
+                                {
+                                    // op2 = ([op1] * op2) + op3
+                                    // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+                                    MakeSrcContained(node, op1);
+                                }
+                                else if (IsContainableHWIntrinsicOp(node, op3, &supportsRegOptional))
+                                {
+                                    // op2 = (op1 * op2) + [op3]
+                                    // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+                                    MakeSrcContained(node, op3);
+                                }
+                                else
+                                {
+                                    assert(supportsRegOptional);
 
-                            // TODO-XArch-CQ: Technically any one of the three operands can
-                            //                be reg-optional. With a limitation on op1 where
-                            //                it can only be so if CopyUpperBits is off.
-                            //                https://github.com/dotnet/runtime/issues/6358
+                                    // TODO-XArch-CQ: Technically any one of the three operands can
+                                    //                be reg-optional. With a limitation on op1 where
+                                    //                it can only be so if CopyUpperBits is off.
+                                    //                https://github.com/dotnet/runtime/issues/6358
 
-                            // 213 form: op1 = (op2 * op1) + op3
-                            op3->SetRegOptional();
+                                    // op2 = (op1 * op2) + [op3]
+                                    // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+                                    op3->SetRegOptional();
+                                }
+                            }
                         }
                     }
                     else

src/coreclr/jit/lsraxarch.cpp

@@ -2334,47 +2334,85 @@ int LinearScan::BuildHWIntrinsic(GenTreeHWIntrinsic* intrinsicTree)
 
                 const bool copiesUpperBits = HWIntrinsicInfo::CopiesUpperBits(intrinsicId);
 
+                unsigned overwrittenOpNum = 0;
+                LIR::Use use;
+                if (LIR::AsRange(blockSequence[curBBSeqNum]).TryGetUse(intrinsicTree, &use))
+                {
+                    overwrittenOpNum = intrinsicTree->GetOverwrittenOpNumForFMA(use.User(), op1, op2, op3);
+                }
+
                 // Intrinsics with CopyUpperBits semantics cannot have op1 be contained
                 assert(!copiesUpperBits || !op1->isContained());
 
-                if (op2->isContained())
+                if (overwrittenOpNum == 1)
                 {
-                    // 132 form: op1 = (op1 * op3) + [op2]
+                    if (op2->isContained())
+                    {
+                        // op1 = (op1 * [op2]) + op3
+                        // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+                        tgtPrefUse = BuildUse(op1);
 
-                    tgtPrefUse = BuildUse(op1);
+                        srcCount += 1;
+                        srcCount += BuildOperandUses(op2);
+                        srcCount += BuildDelayFreeUses(op3, op1);
+                    }
+                    else
+                    {
+                        //assert(op3->isContained());
+                        // op1 = (op1 * op2) + [op3]
+                        // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+                        tgtPrefUse = BuildUse(op1);
 
-                    srcCount += 1;
-                    srcCount += BuildOperandUses(op2);
-                    srcCount += BuildDelayFreeUses(op3, op1);
+                        srcCount += 1;
+                        srcCount += op3->isContained() ? BuildOperandUses(op3) : BuildDelayFreeUses(op3, op1);
+                        srcCount += BuildDelayFreeUses(op2, op1);
+
+                    }
                 }
-                else if (op1->isContained())
+                else if (overwrittenOpNum == 3)
                 {
-                    // 231 form: op3 = (op2 * op3) + [op1]
-
+                    // 231 form: XMM1 = (XMM2 * [XMM3]) + XMM1
+                    // One of the following:
+                    // op3 = ([op1] * op2) + op3
+                    // op3 = (op1 * [op2]) + op3
                     tgtPrefUse = BuildUse(op3);
-
-                    srcCount += BuildOperandUses(op1);
-                    srcCount += BuildDelayFreeUses(op2, op1);
                     srcCount += 1;
+                    if (op1->isContained())
+                    {
+                        srcCount += BuildOperandUses(op1);
+                        srcCount += BuildDelayFreeUses(op2, op3);
+                    }
+                    else
+                    {
+                        //assert(op2->isContained());
+                        srcCount += op2->isContained() ? BuildOperandUses(op2) : BuildDelayFreeUses(op2, op3);
+                        srcCount += BuildDelayFreeUses(op1, op3);
+                    }
+
                 }
                 else
                 {
-                    // 213 form: op1 = (op2 * op1) + [op3]
+                    assert(overwrittenOpNum == 2 || overwrittenOpNum == 0);
 
-                    tgtPrefUse = BuildUse(op1);
-                    srcCount += 1;
-
-                    if (copiesUpperBits)
+                    if (op1->isContained())
                     {
-                        srcCount += BuildDelayFreeUses(op2, op1);
+                        // op2 = ([op1] * op2) + op3
+                        // 132 form: XMM1 = (XMM1 * [XMM3]) + XMM2
+                        tgtPrefUse = BuildUse(op2);
+                        srcCount += 1;
+                        srcCount += BuildOperandUses(op1);
+                        srcCount += BuildDelayFreeUses(op3, op2);
                     }
                     else
                     {
-                        tgtPrefUse2 = BuildUse(op2);
+                        // op2 = (op1 * op2) + [op3]
+                        // 213 form: XMM1 = (XMM2 * XMM1) + [XMM3]
+
+                        tgtPrefUse = BuildUse(op2);
                         srcCount += 1;
+                        srcCount += op3->isContained() ? BuildOperandUses(op3) : BuildDelayFreeUses(op3, op1);
+                        srcCount += BuildDelayFreeUses(op1, op2);
                     }
-
-                    srcCount += op3->isContained() ? BuildOperandUses(op3) : BuildDelayFreeUses(op3, op1);
                 }
 
                 buildUses = false;