Arm64: Improve support for HW_Flag_ReturnsPerElementMask

src/coreclr/jit/emitarm64sve.cpp

@@ -4042,6 +4042,7 @@ void emitter::emitInsSve_R_R_R(instruction     ins,
             assert(isPredicateRegister(reg1)); // MMMM
             assert(isPredicateRegister(reg2)); // gggg
             assert(isPredicateRegister(reg3)); // NNNN
+            opt = INS_OPTS_SCALABLE_B;
             fmt = IF_SVE_DC_3A;
             break;
 
@@ -6336,11 +6337,15 @@ void emitter::emitInsSve_R_R_R_R(instruction     ins,
             }
             else
             {
-                assert(opt == INS_OPTS_SCALABLE_B);
+                assert(insOptsScalable(opt));
                 assert(isPredicateRegister(reg1)); // dddd
                 assert(isPredicateRegister(reg2)); // gggg
                 assert(isPredicateRegister(reg3)); // nnnn
                 assert(isPredicateRegister(reg4)); // mmmm
+                // We support all lane arrangements, although we require byte arrangement for the
+                // encoding as there is only one encoding. This operation is bitwise, so it will
+                // preserve other lane arrangements anyway.
+                opt = INS_OPTS_SCALABLE_B;
                 fmt = IF_SVE_CZ_4A;
             }
             break;

src/coreclr/jit/gentree.cpp

@@ -9435,6 +9435,14 @@ GenTree* Compiler::gtNewZeroConNode(var_types type)
         vecCon->gtSimdVal     = simd_t::Zero();
         return vecCon;
     }
+#ifdef FEATURE_MASKED_HW_INTRINSICS
+    else if (varTypeIsMask(type))
+    {
+        GenTreeMskCon* mskCon = gtNewMskConNode(TYP_MASK);
+        mskCon->gtSimdMaskVal = simdmask_t::Zero();
+        return mskCon;
+    }
+#endif // FEATURE_MASKED_HW_INTRINSICS
 #endif // FEATURE_SIMD
 
     type = genActualType(type);

src/coreclr/jit/gentree.h

@@ -1968,6 +1968,7 @@ struct GenTree
     inline bool IsMaskZero() const;
     inline bool IsMaskAllBitsSet() const;
     inline bool IsTrueMask(var_types simdBaseType) const;
+    inline bool IsSelectZero() const;
 
     inline uint64_t GetIntegralVectorConstElement(size_t index, var_types simdBaseType);
 
@@ -9847,6 +9848,22 @@ inline bool GenTree::IsTrueMask(var_types simdBaseType) const
     return false;
 }
 
+//------------------------------------------------------------------------
+// IsSelectZero: Is the given node a zero value for the purposes of
+//               conditional selection. ConditionalSelect can operate on all
+//               vectors or all masks.
+//
+// Returns true if the node is an all false mask node or a zero vector node.
+//
+// If such a node is used in op3 of ConditionalSelect, it will result in a
+// simple filtering operation on the vector or mask node in op2, using the mask
+// provided in op1.
+//
+inline bool GenTree::IsSelectZero() const
+{
+    return IsVectorZero() || IsMaskZero();
+}
+
 //-------------------------------------------------------------------
 // GetIntegralVectorConstElement: Gets the value of a given element in an integral vector constant
 //

src/coreclr/jit/hwintrinsic.h

@@ -1197,6 +1197,20 @@ struct HWIntrinsicInfo
     }
 #endif // FEATURE_MASKED_HW_INTRINSICS
 
+    // IsSveConditionalSelect: Is this intrinsic a ConditionalSelect intrinsic?
+    //
+    // Arguments:
+    //     id       -- Intrinsic ID to test
+    //
+    // Return value:
+    //     Returns true if the ID is either of the vector or mask variant of
+    //     ConditionalSelect.
+    //
+    static bool IsSveConditionalSelect(NamedIntrinsic id)
+    {
+        return (id == NI_Sve_ConditionalSelect) || (id == NI_Sve_ConditionalSelect_Predicates);
+    }
+
 #endif // TARGET_ARM64
 
     static bool HasSpecialSideEffect(NamedIntrinsic id)

src/coreclr/jit/hwintrinsicarm64.cpp

@@ -2948,7 +2948,7 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
             }
 
             // Was not able to generate a pattern, instead import a truemaskall
-            retNode = gtNewSimdHWIntrinsicNode(TYP_MASK, op1, intrinsic, simdBaseType, simdSize);
+            retNode = gtNewSimdHWIntrinsicNode(retType, op1, intrinsic, simdBaseType, simdSize);
             break;
         }
 

src/coreclr/jit/hwintrinsiccodegenarm64.cpp

@@ -552,7 +552,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
         {
             // Handle case where op2 is operation that needs embedded mask
             GenTree* op2 = intrin.op2;
-            assert(intrin.id == NI_Sve_ConditionalSelect);
+            assert(HWIntrinsicInfo::IsSveConditionalSelect(intrin.id));
             assert(op2->OperIsHWIntrinsic());
             assert(op2->isContained());
 
@@ -597,15 +597,15 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
             // Shared code for setting up embedded mask arg for intrinsics with 3+ operands
 
             auto emitEmbeddedMaskSetupInstrs = [&] {
-                if (intrin.op3->IsVectorZero() || (targetReg != falseReg) || (targetReg != embMaskOp1Reg))
+                if (intrin.op3->IsSelectZero() || (targetReg != falseReg) || (targetReg != embMaskOp1Reg))
                 {
                     return 1;
                 }
                 return 0;
             };
 
             auto emitEmbeddedMaskSetup = [&] {
-                if (intrin.op3->IsVectorZero())
+                if (intrin.op3->IsSelectZero())
                 {
                     // If `falseReg` is zero, then move the first operand of `intrinEmbMask` in the
                     // destination using /Z.
@@ -712,7 +712,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
 
                         if (intrin.op3->isContained())
                         {
-                            assert(intrin.op3->IsVectorZero());
+                            assert(intrin.op3->IsSelectZero());
 
                             if (intrin.op1->isContained() || intrin.op1->IsTrueMask(node->GetSimdBaseType()))
                             {
@@ -818,7 +818,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                             // Predicate functionality is currently not exposed for this API,
                             // but the FADDA instruction only has a predicated variant.
                             // Thus, we expect the JIT to wrap this with CndSel.
-                            assert(intrin.op3->IsVectorZero());
+                            assert(intrin.op3->IsSelectZero());
                             break;
 
                         case NI_Sve2_AddSaturate:
@@ -881,7 +881,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                         }
                     };
 
-                    if (intrin.op3->IsVectorZero())
+                    if ((intrin.op3->IsSelectZero()))
                     {
                         // If `falseReg` is zero, then move the first operand of `intrinEmbMask` in the
                         // destination using /Z.
@@ -1228,7 +1228,7 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                     // This handles optimizations for instructions that have
                     // an implicit 'zero' vector of what would be the second operand.
                     if (HWIntrinsicInfo::SupportsContainment(intrin.id) && intrin.op2->isContained() &&
-                        intrin.op2->IsVectorZero())
+                        intrin.op2->IsSelectZero())
                     {
                         GetEmitter()->emitIns_R_R(ins, emitSize, targetReg, op1Reg, opt);
                     }
@@ -2787,7 +2787,6 @@ 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;