Arm64: re-enable use of predicate variants (#117313)

Fixes #101970

Predicate variants were implemented in #114438 and then turned off in #115566. The code was then removed in #117101 when the AMD64 version was moved to from morph to folding.

This is a simple rework of that code.
Replaces #116854

Author: a74nh

src/coreclr/jit/gentree.cpp

@@ -32106,6 +32106,15 @@ bool GenTree::CanDivOrModPossiblyOverflow(Compiler* comp) const
     return true;
 }
 
+//------------------------------------------------------------------------
+// gtFoldExprHWIntrinsic: Attempt to fold a HWIntrinsic
+//
+// Arguments:
+//    tree - HWIntrinsic to fold
+//
+// Return Value:
+//    folded expression if it could be folded, else the original tree
+//
 #if defined(FEATURE_HW_INTRINSICS)
 GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
 {
@@ -32249,7 +32258,8 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
     // We shouldn't find AND_NOT nodes since it should only be produced in lowering
     assert(oper != GT_AND_NOT);
 
-#if defined(FEATURE_MASKED_HW_INTRINSICS) && defined(TARGET_XARCH)
+#ifdef FEATURE_MASKED_HW_INTRINSICS
+#ifdef TARGET_XARCH
     if (GenTreeHWIntrinsic::OperIsBitwiseHWIntrinsic(oper))
     {
         // Comparisons that produce masks lead to more verbose trees than
@@ -32367,7 +32377,75 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
             }
         }
     }
-#endif // FEATURE_MASKED_HW_INTRINSICS && TARGET_XARCH
+#elif defined(TARGET_ARM64)
+    // Check if the tree can be folded into a mask variant
+    if (HWIntrinsicInfo::HasAllMaskVariant(tree->GetHWIntrinsicId()))
+    {
+        NamedIntrinsic maskVariant = HWIntrinsicInfo::GetMaskVariant(tree->GetHWIntrinsicId());
+
+        assert(opCount == (size_t)HWIntrinsicInfo::lookupNumArgs(maskVariant));
+
+        // Check all operands are valid
+        bool canFold = true;
+        if (ni == NI_Sve_ConditionalSelect)
+        {
+            assert(varTypeIsMask(op1));
+            canFold = (op2->OperIsConvertMaskToVector() && op3->OperIsConvertMaskToVector());
+        }
+        else
+        {
+            for (size_t i = 1; i <= opCount && canFold; i++)
+            {
+                canFold &= tree->Op(i)->OperIsConvertMaskToVector();
+            }
+        }
+
+        if (canFold)
+        {
+            // Convert all the operands to masks
+            for (size_t i = 1; i <= opCount; i++)
+            {
+                if (tree->Op(i)->OperIsConvertMaskToVector())
+                {
+                    // Replace with op1.
+                    tree->Op(i) = tree->Op(i)->AsHWIntrinsic()->Op(1);
+                }
+                else if (tree->Op(i)->IsVectorZero())
+                {
+                    // Replace the vector of zeroes with a mask of zeroes.
+                    tree->Op(i) = gtNewSimdFalseMaskByteNode();
+                    tree->Op(i)->SetMorphed(this);
+                }
+                assert(varTypeIsMask(tree->Op(i)));
+            }
+
+            // Switch to the mask variant
+            switch (opCount)
+            {
+                case 1:
+                    tree->ResetHWIntrinsicId(maskVariant, tree->Op(1));
+                    break;
+                case 2:
+                    tree->ResetHWIntrinsicId(maskVariant, tree->Op(1), tree->Op(2));
+                    break;
+                case 3:
+                    tree->ResetHWIntrinsicId(maskVariant, this, tree->Op(1), tree->Op(2), tree->Op(3));
+                    break;
+                default:
+                    unreached();
+            }
+
+            tree->gtType = TYP_MASK;
+            tree->SetMorphed(this);
+            tree = gtNewSimdCvtMaskToVectorNode(retType, tree, simdBaseJitType, simdSize)->AsHWIntrinsic();
+            tree->SetMorphed(this);
+            op1 = tree->Op(1);
+            op2 = nullptr;
+            op3 = nullptr;
+        }
+    }
+#endif // TARGET_ARM64
+#endif // FEATURE_MASKED_HW_INTRINSICS
 
     GenTree* cnsNode   = nullptr;
     GenTree* otherNode = nullptr;
@@ -33754,7 +33832,7 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
                     // op2 = op2 & op1
                     op2->AsVecCon()->EvaluateBinaryInPlace(GT_AND, false, simdBaseType, op1->AsVecCon());
 
-                    // op3 = op2 & ~op1
+                    // op3 = op3 & ~op1
                     op3->AsVecCon()->EvaluateBinaryInPlace(GT_AND_NOT, false, simdBaseType, op1->AsVecCon());
 
                     // op2 = op2 | op3
@@ -33767,8 +33845,8 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
 
 #if defined(TARGET_ARM64)
             case NI_Sve_ConditionalSelect:
+            case NI_Sve_ConditionalSelect_Predicates:
             {
-                assert(!varTypeIsMask(retType));
                 assert(varTypeIsMask(op1));
 
                 if (cnsNode != op1)
@@ -33797,10 +33875,11 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
 
                 if (op2->IsCnsVec() && op3->IsCnsVec())
                 {
+                    assert(ni == NI_Sve_ConditionalSelect);
                     assert(op2->gtType == TYP_SIMD16);
                     assert(op3->gtType == TYP_SIMD16);
 
-                    simd16_t op1SimdVal;
+                    simd16_t op1SimdVal = {};
                     EvaluateSimdCvtMaskToVector<simd16_t>(simdBaseType, &op1SimdVal, op1->AsMskCon()->gtSimdMaskVal);
 
                     // op2 = op2 & op1
@@ -33809,7 +33888,7 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
                                                  op1SimdVal);
                     op2->AsVecCon()->gtSimd16Val = result;
 
-                    // op3 = op2 & ~op1
+                    // op3 = op3 & ~op1
                     result = {};
                     EvaluateBinarySimd<simd16_t>(GT_AND_NOT, false, simdBaseType, &result, op3->AsVecCon()->gtSimd16Val,
                                                  op1SimdVal);
@@ -33820,6 +33899,30 @@ GenTree* Compiler::gtFoldExprHWIntrinsic(GenTreeHWIntrinsic* tree)
 
                     resultNode = op2;
                 }
+                else if (op2->IsCnsMsk() && op3->IsCnsMsk())
+                {
+                    assert(ni == NI_Sve_ConditionalSelect_Predicates);
+
+                    // op2 = op2 & op1
+                    simdmask_t result = {};
+                    EvaluateBinaryMask<simd16_t>(GT_AND, false, simdBaseType, &result, op2->AsMskCon()->gtSimdMaskVal,
+                                                 op1->AsMskCon()->gtSimdMaskVal);
+                    op2->AsMskCon()->gtSimdMaskVal = result;
+
+                    // op3 = op3 & ~op1
+                    result = {};
+                    EvaluateBinaryMask<simd16_t>(GT_AND_NOT, false, simdBaseType, &result,
+                                                 op3->AsMskCon()->gtSimdMaskVal, op1->AsMskCon()->gtSimdMaskVal);
+                    op3->AsMskCon()->gtSimdMaskVal = result;
+
+                    // op2 = op2 | op3
+                    result = {};
+                    EvaluateBinaryMask<simd16_t>(GT_OR, false, simdBaseType, &result, op2->AsMskCon()->gtSimdMaskVal,
+                                                 op3->AsMskCon()->gtSimdMaskVal);
+                    op2->AsMskCon()->gtSimdMaskVal = result;
+
+                    resultNode = op2;
+                }
                 break;
             }
 #endif // TARGET_ARM64

src/coreclr/jit/simd.h

@@ -1164,6 +1164,12 @@ void EvaluateBinaryMask(
             break;
         }
 
+        case 16:
+        {
+            bitMask = 0x0001000100010001;
+            break;
+        }
+
         default:
         {
             unreached();