Ensure FMA optimizations kick in under embedded broadcast (#116891)

Author: tannergooding

src/coreclr/jit/gentree.cpp

@@ -28308,6 +28308,22 @@ bool GenTree::OperIsVectorConditionalSelect() const
     return false;
 }
 
+//------------------------------------------------------------------------
+// OperIsVectorFusedMultiplyOp: Is this a vector FusedMultiplyOp hwintrinsic
+//
+// Return Value:
+//    true if the node is a vector FusedMultiplyOp hwintrinsic
+//    otherwise; false
+//
+bool GenTree::OperIsVectorFusedMultiplyOp() const
+{
+    if (OperIsHWIntrinsic())
+    {
+        return AsHWIntrinsic()->OperIsVectorFusedMultiplyOp();
+    }
+    return false;
+}
+
 //------------------------------------------------------------------------
 // OperIsMemoryLoad: Does this HWI node have memory load semantics?
 //

src/coreclr/jit/gentree.h

@@ -1682,6 +1682,7 @@ struct GenTree
     bool OperIsConvertMaskToVector() const;
     bool OperIsConvertVectorToMask() const;
     bool OperIsVectorConditionalSelect() const;
+    bool OperIsVectorFusedMultiplyOp() const;
 
     // This is here for cleaner GT_LONG #ifdefs.
     static bool OperIsLong(genTreeOps gtOper)
@@ -6468,14 +6469,63 @@ struct GenTreeHWIntrinsic : public GenTreeJitIntrinsic
 
     bool OperIsVectorConditionalSelect() const
     {
+        switch (GetHWIntrinsicId())
+        {
 #if defined(TARGET_XARCH)
-        return OperIsHWIntrinsic(NI_Vector128_ConditionalSelect) || OperIsHWIntrinsic(NI_Vector256_ConditionalSelect) ||
-               OperIsHWIntrinsic(NI_Vector512_ConditionalSelect);
-#elif defined(TARGET_ARM64)
-        return OperIsHWIntrinsic(NI_AdvSimd_BitwiseSelect) || OperIsHWIntrinsic(NI_Sve_ConditionalSelect);
-#else
-        return false;
-#endif
+            case NI_Vector128_ConditionalSelect:
+            case NI_Vector256_ConditionalSelect:
+            case NI_Vector512_ConditionalSelect:
+            {
+                return true;
+            }
+#endif // TARGET_XARCH
+
+#if defined(TARGET_ARM64)
+            case NI_AdvSimd_BitwiseSelect:
+            case NI_Sve_ConditionalSelect:
+            {
+                return true;
+            }
+#endif // TARGET_ARM64
+
+            default:
+            {
+                return false;
+            }
+        }
+    }
+
+    bool OperIsVectorFusedMultiplyOp() const
+    {
+        switch (GetHWIntrinsicId())
+        {
+#if defined(TARGET_XARCH)
+            case NI_AVX2_MultiplyAdd:
+            case NI_AVX2_MultiplyAddNegated:
+            case NI_AVX2_MultiplyAddNegatedScalar:
+            case NI_AVX2_MultiplyAddScalar:
+            case NI_AVX2_MultiplySubtract:
+            case NI_AVX2_MultiplySubtractNegated:
+            case NI_AVX2_MultiplySubtractNegatedScalar:
+            case NI_AVX2_MultiplySubtractScalar:
+            case NI_AVX512_FusedMultiplyAdd:
+            case NI_AVX512_FusedMultiplyAddNegated:
+            case NI_AVX512_FusedMultiplyAddNegatedScalar:
+            case NI_AVX512_FusedMultiplyAddScalar:
+            case NI_AVX512_FusedMultiplySubtract:
+            case NI_AVX512_FusedMultiplySubtractNegated:
+            case NI_AVX512_FusedMultiplySubtractNegatedScalar:
+            case NI_AVX512_FusedMultiplySubtractScalar:
+            {
+                return true;
+            }
+#endif // TARGET_XARCH
+
+            default:
+            {
+                return false;
+            }
+        }
     }
 
     bool OperRequiresAsgFlag() const;

src/coreclr/jit/lowerxarch.cpp

@@ -1387,7 +1387,6 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
 {
     assert(node->GetOperandCount() == 3);
 
-    bool isAvx512 = false;
     bool negated  = false;
     bool subtract = false;
     bool isScalar = false;
@@ -1397,111 +1396,59 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
     switch (intrinsic)
     {
         case NI_AVX2_MultiplyAdd:
+        case NI_AVX512_FusedMultiplyAdd:
         {
             break;
         }
 
         case NI_AVX2_MultiplyAddScalar:
+        case NI_AVX512_FusedMultiplyAddScalar:
         {
             isScalar = true;
             break;
         }
 
         case NI_AVX2_MultiplyAddNegated:
+        case NI_AVX512_FusedMultiplyAddNegated:
         {
             negated = true;
             break;
         }
 
         case NI_AVX2_MultiplyAddNegatedScalar:
-        {
-            negated  = true;
-            isScalar = true;
-            break;
-        }
-
-        case NI_AVX2_MultiplySubtract:
-        {
-            subtract = true;
-            break;
-        }
-
-        case NI_AVX2_MultiplySubtractScalar:
-        {
-            subtract = true;
-            isScalar = true;
-            break;
-        }
-
-        case NI_AVX2_MultiplySubtractNegated:
-        {
-            subtract = true;
-            negated  = true;
-            break;
-        }
-
-        case NI_AVX2_MultiplySubtractNegatedScalar:
-        {
-            subtract = true;
-            negated  = true;
-            isScalar = true;
-            break;
-        }
-
-        case NI_AVX512_FusedMultiplyAdd:
-        {
-            isAvx512 = true;
-            break;
-        }
-
-        case NI_AVX512_FusedMultiplyAddScalar:
-        {
-            isAvx512 = true;
-            isScalar = true;
-            break;
-        }
-
-        case NI_AVX512_FusedMultiplyAddNegated:
-        {
-            isAvx512 = true;
-            negated  = true;
-            break;
-        }
-
         case NI_AVX512_FusedMultiplyAddNegatedScalar:
         {
-            isAvx512 = true;
             negated  = true;
             isScalar = true;
             break;
         }
 
+        case NI_AVX2_MultiplySubtract:
         case NI_AVX512_FusedMultiplySubtract:
         {
-            isAvx512 = true;
             subtract = true;
             break;
         }
 
+        case NI_AVX2_MultiplySubtractScalar:
         case NI_AVX512_FusedMultiplySubtractScalar:
         {
-            isAvx512 = true;
             subtract = true;
             isScalar = true;
             break;
         }
 
+        case NI_AVX2_MultiplySubtractNegated:
         case NI_AVX512_FusedMultiplySubtractNegated:
         {
-            isAvx512 = true;
             subtract = true;
             negated  = true;
             break;
         }
 
+        case NI_AVX2_MultiplySubtractNegatedScalar:
         case NI_AVX512_FusedMultiplySubtractNegatedScalar:
         {
-            isAvx512 = true;
             subtract = true;
             negated  = true;
             isScalar = true;
@@ -1543,10 +1490,7 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
                     argOp->ClearContained();
                     ContainCheckHWIntrinsic(arg->AsHWIntrinsic());
 
-                    // We want to toggle the tracking and then check it again,
-                    // which is the simplest way to handle cases like -CreateScalarUnsafe(-x)
                     negatedArgs[i - 1] ^= true;
-                    i -= 1;
                 }
 
                 break;
@@ -1562,11 +1506,18 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
                     break;
                 }
 
-                GenTree* argOp = hwArg->Op(1);
+                GenTree* argOp = hwArg->Op(2);
+
+                if (!argOp->isContained())
+                {
+                    // A constant should have already been contained
+                    break;
+                }
+
+                // xor is bitwise and the actual xor node might be a different base type
+                // from the FMA node, so we check if its negative zero using the FMA base
+                // type since that's what the end negation would end up using
 
-                // xor is bitwise and the actual xor node might not be floating-point
-                // so we check if its negative zero using the FMA base type since that's
-                // what the end negation will end up using
                 if (argOp->IsVectorNegativeZero(node->GetSimdBaseType()))
                 {
                     BlockRange().Remove(hwArg);
@@ -1576,11 +1527,9 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
                     argOp->ClearContained();
                     node->Op(i) = argOp;
 
-                    // We want to toggle the tracking and then check it again,
-                    // which is the simplest way to handle cases like -CreateScalarUnsafe(-x)
                     negatedArgs[i - 1] ^= true;
-                    i -= 1;
                 }
+
                 break;
             }
         }
@@ -1590,7 +1539,7 @@ void Lowering::LowerFusedMultiplyOp(GenTreeHWIntrinsic* node)
     negated ^= negatedArgs[1];
     subtract ^= negatedArgs[2];
 
-    if (isAvx512)
+    if (intrinsic >= FIRST_NI_AVX512)
     {
         if (negated)
         {
@@ -9928,7 +9877,28 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
 
                     if (containedOperand != nullptr)
                     {
-                        if (containedOperand->IsCnsVec() && node->isEmbeddedBroadcastCompatibleHWIntrinsic(comp))
+                        bool isEmbeddedBroadcastCompatible =
+                            containedOperand->IsCnsVec() && node->isEmbeddedBroadcastCompatibleHWIntrinsic(comp);
+
+                        bool       isScalarArg = false;
+                        genTreeOps oper        = node->GetOperForHWIntrinsicId(&isScalarArg);
+
+                        // We want to skip trying to make this an embedded broadcast in certain scenarios
+                        // because it will prevent other transforms that will be better for codegen.
+
+                        LIR::Use use;
+
+                        if ((oper == GT_XOR) && BlockRange().TryGetUse(node, &use) &&
+                            use.User()->OperIsVectorFusedMultiplyOp())
+                        {
+                            // xor is bitwise and the actual xor node might be a different base type
+                            // from the FMA node, so we check if its negative zero using the FMA base
+                            // type since that's what the end negation would end up using
+                            var_types fmaSimdBaseType     = use.User()->AsHWIntrinsic()->GetSimdBaseType();
+                            isEmbeddedBroadcastCompatible = !containedOperand->IsVectorNegativeZero(fmaSimdBaseType);
+                        }
+
+                        if (isEmbeddedBroadcastCompatible)
                         {
                             TryFoldCnsVecForEmbeddedBroadcast(node, containedOperand->AsVecCon());
                         }