Add pattern matching for SVE intrinsics that operate on mask operands (#114438)

* Add pattern matching for SVE intrinsics that operate on mask operands

Introduces `fgMorphTryUseAllMaskVariant` for ARM64 that looks for various
named intrinsics that have operands that look 'mask-like'. E.g. source
operands originating from Sve.CreateTrueMask* may be recognized as
masks, causing the JIT to prefer to use the predicated version of the
instruction as codegen for the intrinsic. It will also inspect
ConditionalSelect intrinsic nodes to match instructions with governing
predicates. The transform runs during morph.

It's possible to emit the following instructions after this patch:
* ZIP{1,2} <Pd>.<T>, <Pn>.<T>, <Pm>.<T> (Sve.ZipLow, Sve.ZipHigh)
* UZP{1,2} <Pd>.<T>, <Pn>.<T>, <Pm>.<T> (Sve.UnzipEven, Sve.UnzipOdd)
* TRN{1,2} <Pd>.<T>, <Pn>.<T>, <Pm>.<T> (Sve.TransposeEven, Sve.TransposeOdd)
* REV <Pd>.<T>, <Pn>.<T>                (Sve.ReverseElement)
* AND <Pd>.B, <Pg>/Z, <Pn>.B, <Pm>.B    (Sve.And)
* BIC <Pd>.B, <Pg>/Z, <Pn>.B, <Pm>.B    (Sve.BitwiseClear)
* EOR <Pd>.B, <Pg>/Z, <Pn>.B, <Pm>.B    (Sve.Xor)
* ORR <Pd>.B, <Pg>/Z, <Pn>.B, <Pm>.B    (Sve.Or)
* SEL <Pd>.B, <Pg>, <Pn>.B, <Pm>.B      (Sve.ConditionalSelect)

Contributes towards #101970

* Fix test failure and add FileCheck tests

* Don't run tests on OSX

* Don't run tests for Mono

* Move the transform later in fgOptimizeHWIntrinsic

* Rename gtNewSimdAllFalseMaskNode

* Re-design using HW_Flag_AllMaskVariant

* Add missing function documentation in hwintrinsic.h

* Fix integer comparison and add assertion

* Refactor to follow similar path to XARCH

* fix the refactoring

* jit formatting

* Move code into morph.cpp

---------

Co-authored-by: Kunal Pathak <[email protected]>

Author: snickolls-arm

src/coreclr/jit/compiler.h

@@ -3132,6 +3132,7 @@ class Compiler
 
 #if defined(TARGET_ARM64)
     GenTree* gtNewSimdAllTrueMaskNode(CorInfoType simdBaseJitType, unsigned simdSize);
+    GenTree* gtNewSimdFalseMaskByteNode(unsigned simdSize);
 #endif
 
     GenTree* gtNewSimdBinOpNode(genTreeOps  op,
@@ -6691,6 +6692,15 @@ class Compiler
     GenTree* fgMorphHWIntrinsic(GenTreeHWIntrinsic* tree);
     GenTree* fgOptimizeHWIntrinsic(GenTreeHWIntrinsic* node);
     GenTree* fgOptimizeHWIntrinsicAssociative(GenTreeHWIntrinsic* node);
+#if defined(FEATURE_MASKED_HW_INTRINSICS)
+    GenTreeHWIntrinsic* fgOptimizeForMaskedIntrinsic(GenTreeHWIntrinsic* node);
+#endif // FEATURE_MASKED_HW_INTRINSICS
+#ifdef TARGET_ARM64
+    bool canMorphVectorOperandToMask(GenTree* node);
+    bool canMorphAllVectorOperandsToMasks(GenTreeHWIntrinsic* node);
+    GenTree* doMorphVectorOperandToMask(GenTree* node, GenTreeHWIntrinsic* parent);
+    GenTreeHWIntrinsic* fgMorphTryUseAllMaskVariant(GenTreeHWIntrinsic* node);
+#endif // TARGET_ARM64
 #endif // FEATURE_HW_INTRINSICS
     GenTree* fgOptimizeCommutativeArithmetic(GenTreeOp* tree);
     GenTree* fgOptimizeRelationalComparisonWithCasts(GenTreeOp* cmp);

src/coreclr/jit/hwintrinsic.h

@@ -233,6 +233,11 @@ enum HWIntrinsicFlag : unsigned int
     // The intrinsic is a reduce operation.
     HW_Flag_ReduceOperation = 0x2000000,
 
+    // This intrinsic could be implemented with another intrinsic when it is operating on operands that are all of
+    // type TYP_MASK, and this other intrinsic will produces a value of this type. Used in morph to convert vector
+    // operations into mask operations when the intrinsic is operating on mask vectors (mainly bitwise operations).
+    HW_Flag_HasAllMaskVariant = 0x4000000,
+
 #else
 #error Unsupported platform
 #endif
@@ -1133,6 +1138,67 @@ struct HWIntrinsicInfo
         }
     }
 
+#ifdef FEATURE_MASKED_HW_INTRINSICS
+    // HasAllMaskVariant: Does the intrinsic have an intrinsic variant that operates on mask types?
+    //
+    // Arguments:
+    //    id       -- the intrinsic to check for a mask-type variant.
+    //
+    // Return Value:
+    //    true when the intrinsic has a mask-type variant, else false
+    //
+    static bool HasAllMaskVariant(NamedIntrinsic id)
+    {
+        const HWIntrinsicFlag flags = lookupFlags(id);
+        return (flags & HW_Flag_HasAllMaskVariant) != 0;
+    }
+
+    // GetMaskVariant: Given an intrinsic that has a variant that operates on mask types, return the ID of
+    //                 this variant intrinsic. Call HasAllMaskVariant before using this function, as it will
+    //                 assert if no match is found.
+    //
+    // Arguments:
+    //    id       -- the intrinsic with a mask-type variant.
+    //
+    // Return Value:
+    //    The ID of the mask-type variant for the given intrinsic
+    //
+    static NamedIntrinsic GetMaskVariant(NamedIntrinsic id)
+    {
+        assert(HasAllMaskVariant(id));
+        switch (id)
+        {
+            case NI_Sve_And:
+                return NI_Sve_And_Predicates;
+            case NI_Sve_BitwiseClear:
+                return NI_Sve_BitwiseClear_Predicates;
+            case NI_Sve_Xor:
+                return NI_Sve_Xor_Predicates;
+            case NI_Sve_Or:
+                return NI_Sve_Or_Predicates;
+            case NI_Sve_ZipHigh:
+                return NI_Sve_ZipHigh_Predicates;
+            case NI_Sve_ZipLow:
+                return NI_Sve_ZipLow_Predicates;
+            case NI_Sve_UnzipOdd:
+                return NI_Sve_UnzipOdd_Predicates;
+            case NI_Sve_UnzipEven:
+                return NI_Sve_UnzipEven_Predicates;
+            case NI_Sve_TransposeEven:
+                return NI_Sve_TransposeEven_Predicates;
+            case NI_Sve_TransposeOdd:
+                return NI_Sve_TransposeOdd_Predicates;
+            case NI_Sve_ReverseElement:
+                return NI_Sve_ReverseElement_Predicates;
+            case NI_Sve_ConditionalSelect:
+                return NI_Sve_ConditionalSelect_Predicates;
+
+            default:
+                unreached();
+        }
+    }
+#endif // FEATURE_MASKED_HW_INTRINSICS
+
 #endif // TARGET_ARM64
 
     static bool HasSpecialSideEffect(NamedIntrinsic id)

src/coreclr/jit/hwintrinsicarm64.cpp

@@ -3341,7 +3341,7 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
 }
 
 //------------------------------------------------------------------------
-// gtNewSimdEmbeddedMaskNode: Create an embedded mask
+// gtNewSimdAllTrueMaskNode: Create an embedded mask with all bits set to true
 //
 // Arguments:
 //    simdBaseJitType -- the base jit type of the nodes being masked
@@ -3355,4 +3355,18 @@ GenTree* Compiler::gtNewSimdAllTrueMaskNode(CorInfoType simdBaseJitType, unsigne
     return gtNewSimdHWIntrinsicNode(TYP_MASK, NI_Sve_CreateTrueMaskAll, simdBaseJitType, simdSize);
 }
 
+//------------------------------------------------------------------------
+// gtNewSimdFalseMaskByteNode: Create an embedded mask with all bits set to false
+//
+// Arguments:
+//    simdSize        -- the simd size of the nodes being masked
+//
+// Return Value:
+//    The mask
+//
+GenTree* Compiler::gtNewSimdFalseMaskByteNode(unsigned simdSize)
+{
+    return gtNewSimdHWIntrinsicNode(TYP_MASK, NI_Sve_CreateFalseMaskByte, CORINFO_TYPE_UBYTE, simdSize);
+}
+
 #endif // FEATURE_HW_INTRINSICS

src/coreclr/jit/hwintrinsiccodegenarm64.cpp

@@ -702,6 +702,14 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                                 break;
                             }
 
+                            case NI_Sve_And_Predicates:
+                            case NI_Sve_BitwiseClear_Predicates:
+                            case NI_Sve_Or_Predicates:
+                            case NI_Sve_Xor_Predicates:
+                                GetEmitter()->emitIns_R_R_R_R(insEmbMask, emitSize, targetReg, maskReg, embMaskOp1Reg,
+                                                              embMaskOp2Reg, INS_OPTS_SCALABLE_B);
+                                break;
+
                             default:
                             {
                                 GetEmitter()->emitIns_R_R_R_R(insEmbMask, emitSize, targetReg, maskReg, embMaskOp1Reg,
@@ -2478,6 +2486,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
 
             case NI_Sve_CreateBreakAfterPropagateMask:
             case NI_Sve_CreateBreakBeforePropagateMask:
+            case NI_Sve_ConditionalSelect_Predicates:
             {
                 GetEmitter()->emitInsSve_R_R_R_R(ins, emitSize, targetReg, op1Reg, op2Reg, op3Reg, INS_OPTS_SCALABLE_B);
                 break;