arm64: Optimize ARM64 compare mask ExtractMostSignificantBits consumers

src/coreclr/jit/assertionprop.cpp

@@ -78,6 +78,147 @@ static Range GetRange(Compiler* comp, GenTree* tree, BasicBlock* block, ASSERT_V
     return Limit(Limit::keUnknown);
 }
 
+#if defined(FEATURE_HW_INTRINSICS) && defined(TARGET_ARM64)
+//----------------------------------------------------------------------------------------------
+// AllComponentsEitherZeroOrAllBitsSet: Check if a SIMD VN has per-element boolean values.
+//
+// Arguments:
+//    comp     - The compiler instance
+//    vn       - The value number
+//    baseType - The expected SIMD element base type
+//
+// Return Value:
+//    True if every SIMD element is known to be either all-bits-set or zero.
+//
+static bool AllComponentsEitherZeroOrAllBitsSet(Compiler* comp, ValueNum vn, var_types baseType)
+{
+    if (vn == ValueNumStore::NoVN)
+    {
+        return false;
+    }
+
+    vn = comp->vnStore->VNNormalValue(vn);
+
+    if (comp->vnStore->IsVNConstant(vn))
+    {
+        switch (comp->vnStore->TypeOfVN(vn))
+        {
+            case TYP_SIMD8:
+            {
+                simd8_t val = comp->vnStore->GetConstantSimd8(vn);
+                return val.IsAllBitsSet() || val.IsZero();
+            }
+
+            case TYP_SIMD16:
+            {
+                simd16_t val = comp->vnStore->GetConstantSimd16(vn);
+                return val.IsAllBitsSet() || val.IsZero();
+            }
+
+            default:
+            {
+                return false;
+            }
+        }
+    }
+
+    VNFuncApp      funcApp;
+    NamedIntrinsic intrinsicId;
+    unsigned       simdSize;
+    var_types      intrinsicSimdBaseType;
+
+    if (!comp->vnStore->IsVNHWIntrinsicFunc(vn, &funcApp, &intrinsicId, &simdSize, &intrinsicSimdBaseType))
+    {
+        return false;
+    }
+
+    if ((simdSize != 8) && (simdSize != 16))
+    {
+        return false;
+    }
+
+    bool       isScalar = false;
+    genTreeOps oper     = GenTreeHWIntrinsic::GetOperForHWIntrinsicId(intrinsicId, baseType, &isScalar);
+
+    if (isScalar)
+    {
+        return false;
+    }
+
+    switch (oper)
+    {
+        case GT_EQ:
+        case GT_NE:
+        case GT_GT:
+        case GT_GE:
+        case GT_LE:
+        case GT_LT:
+        {
+            // The comparison intrinsic may have used a wider unsigned base type to produce the
+            // per-element mask, e.g. a TYP_BYTE compare can be implemented as TYP_UBYTE.
+            return varTypeIsIntegral(baseType) && (genTypeSize(baseType) <= genTypeSize(intrinsicSimdBaseType));
+        }
+
+        case GT_NOT:
+        {
+            return AllComponentsEitherZeroOrAllBitsSet(comp, funcApp.GetArg(0), baseType);
+        }
+
+        case GT_OR:
+        case GT_AND:
+        case GT_XOR:
+        case GT_AND_NOT:
+        {
+            return AllComponentsEitherZeroOrAllBitsSet(comp, funcApp.GetArg(0), baseType) &&
+                   AllComponentsEitherZeroOrAllBitsSet(comp, funcApp.GetArg(1), baseType);
+        }
+
+        default:
+        {
+            return false;
+        }
+    }
+}
+
+//----------------------------------------------------------------------------------------------
+// optAssertionProp_HWIntrinsic: Propagate VN-derived facts to hwintrinsic tree flags.
+//
+// Arguments:
+//    comp - The compiler instance
+//    tree - The hwintrinsic node
+//
+static void optAssertionProp_HWIntrinsic(Compiler* comp, GenTreeHWIntrinsic* tree)
+{
+    NamedIntrinsic intrinsic = tree->GetHWIntrinsicId();
+
+    if ((intrinsic != NI_Vector64_ExtractMostSignificantBits) && (intrinsic != NI_Vector128_ExtractMostSignificantBits))
+    {
+        return;
+    }
+
+    assert(tree->GetOperandCount() == 1);
+
+    GenTree* op1 = tree->Op(1);
+
+    if (op1->OperIsHWIntrinsic() && !Compiler::IsHWIntrinsicCmpMask(op1->AsHWIntrinsic()->GetHWIntrinsicId()))
+    {
+        return;
+    }
+
+    ValueNum op1VN = comp->vnStore->VNConservativeNormalValue(op1->gtVNPair);
+
+    auto vnVisitor = [comp, tree](ValueNum vn) -> ValueNumStore::VNVisit {
+        return AllComponentsEitherZeroOrAllBitsSet(comp, vn, tree->GetSimdBaseType()) ? ValueNumStore::VNVisit::Continue
+                                                                                      : ValueNumStore::VNVisit::Abort;
+    };
+
+    if (comp->vnStore->VNVisitReachingVNs(op1VN, vnVisitor) == ValueNumStore::VNVisit::Continue)
+    {
+        tree->gtFlags |= GTF_HW_ZERO_OR_ALL_BITS_SET;
+    }
+}
+#endif // FEATURE_HW_INTRINSICS && TARGET_ARM64
+
 //------------------------------------------------------------------------
 // SymbolicToRealValue: Convert a symbolic value to a 64-bit signed integer.
 //
@@ -5867,6 +6008,12 @@ GenTree* Compiler::optAssertionProp(ASSERT_VALARG_TP assertions, GenTree* tree,
         case GT_CALL:
             return optAssertionProp_Call(assertions, tree->AsCall(), stmt);
 
+#if defined(FEATURE_HW_INTRINSICS) && defined(TARGET_ARM64)
+        case GT_HWINTRINSIC:
+            optAssertionProp_HWIntrinsic(this, tree->AsHWIntrinsic());
+            return nullptr;
+#endif // FEATURE_HW_INTRINSICS && TARGET_ARM64
+
         case GT_EQ:
         case GT_NE:
         case GT_LT:

src/coreclr/jit/compiler.h

@@ -10481,6 +10481,42 @@ class Compiler
         }
     }
 
+#if defined(FEATURE_HW_INTRINSICS) && defined(TARGET_ARM64)
+    //----------------------------------------------------------------------------------------------
+    // IsHWIntrinsicCmpMask: Checks if the hwintrinsic produces a SIMD comparison mask.
+    //
+    // Arguments:
+    //    intrinsic - The hwintrinsic id
+    //
+    // Return Value:
+    //    True if the hwintrinsic produces a mask where each SIMD element is either all-bits-set or zero.
+    //
+    static bool IsHWIntrinsicCmpMask(NamedIntrinsic intrinsic)
+    {
+        switch (intrinsic)
+        {
+            case NI_AdvSimd_CompareEqual:
+            case NI_AdvSimd_CompareGreaterThan:
+            case NI_AdvSimd_CompareGreaterThanOrEqual:
+            case NI_AdvSimd_CompareLessThan:
+            case NI_AdvSimd_CompareLessThanOrEqual:
+            case NI_AdvSimd_Arm64_CompareEqual:
+            case NI_AdvSimd_Arm64_CompareGreaterThan:
+            case NI_AdvSimd_Arm64_CompareGreaterThanOrEqual:
+            case NI_AdvSimd_Arm64_CompareLessThan:
+            case NI_AdvSimd_Arm64_CompareLessThanOrEqual:
+            {
+                return true;
+            }
+
+            default:
+            {
+                return false;
+            }
+        }
+    }
+#endif // FEATURE_HW_INTRINSICS && TARGET_ARM64
+
 private:
     unsigned getSIMDInitTempVarNum(var_types simdType);
 

src/coreclr/jit/gentree.cpp

@@ -3001,6 +3001,11 @@ bool GenTree::Compare(GenTree* op1, GenTree* op2, bool swapOK)
 
 #ifdef FEATURE_HW_INTRINSICS
         case GT_HWINTRINSIC:
+            if ((op1->gtFlags & GTF_HW_ZERO_OR_ALL_BITS_SET) != (op2->gtFlags & GTF_HW_ZERO_OR_ALL_BITS_SET))
+            {
+                return false;
+            }
+
             return GenTreeHWIntrinsic::Equals(op1->AsHWIntrinsic(), op2->AsHWIntrinsic());
 #endif
 

src/coreclr/jit/gentree.h

@@ -541,6 +541,7 @@ enum GenTreeFlags : unsigned
 #ifdef FEATURE_HW_INTRINSICS
     GTF_HW_EM_OP                  = 0x10000000, // GT_HWINTRINSIC -- node is used as an operand to an embedded mask
     GTF_HW_USER_CALL              = 0x20000000, // GT_HWINTRINSIC -- node is implemented via a user call
+    GTF_HW_ZERO_OR_ALL_BITS_SET   = 0x40000000, // GT_HWINTRINSIC -- each SIMD element is either zero or all-bits-set
 #endif // FEATURE_HW_INTRINSICS
 };