JIT: propagate vector constants through assertion prop

src/coreclr/jit/assertionprop.cpp

@@ -887,6 +887,10 @@ void Compiler::optPrintAssertion(const AssertionDsc& curAssertion, AssertionInde
             }
             break;
 
+        case O2K_CONST_VEC:
+            printf("VecCns");
+            break;
+
         case O2K_ZEROOBJ:
             printf("ZeroObj");
             break;
@@ -1138,6 +1142,32 @@ AssertionIndex Compiler::optCreateAssertion(GenTree* op1, GenTree* op2, bool equ
                 return optAddAssertion(dsc);
             }
 
+#if defined(FEATURE_HW_INTRINSICS)
+            case GT_CNS_VEC:
+            {
+                // For now, only support SIMD constants up to 16 bytes (SIMD8/12/16).
+                if (!op1->TypeIs(TYP_SIMD8, TYP_SIMD12, TYP_SIMD16) || (op1->TypeGet() != op2->TypeGet()))
+                {
+                    return NO_ASSERTION_INDEX;
+                }
+
+                ValueNum op1VN = optConservativeNormalVN(op1);
+                ValueNum op2VN = optConservativeNormalVN(op2);
+                if (!optLocalAssertionProp && (op1VN == ValueNumStore::NoVN || op2VN == ValueNumStore::NoVN))
+                {
+                    // GlobalAP requires valid VNs.
+                    return NO_ASSERTION_INDEX;
+                }
+
+                simd16_t simdVal = {};
+                memcpy(&simdVal, &op2->AsVecCon()->gtSimdVal, genTypeSize(op2->TypeGet()));
+
+                AssertionDsc dsc =
+                    AssertionDsc::CreateConstLclVarAssertion(this, lclNum, op1VN, simdVal, op2VN, equals);
+                return optAddAssertion(dsc);
+            }
+#endif // FEATURE_HW_INTRINSICS
+
             case GT_CNS_INT:
             {
                 ValueNum op1VN = optConservativeNormalVN(op1);
@@ -1827,6 +1857,61 @@ AssertionInfo Compiler::optAssertionGenJtrue(GenTree* tree)
     GenTree* op1 = relop->AsOp()->gtOp1->gtCommaStoreVal();
     GenTree* op2 = relop->AsOp()->gtOp2->gtCommaStoreVal();
 
+#if defined(FEATURE_HW_INTRINSICS)
+    if (op1->OperIsHWIntrinsic() && (op2->IsIntegralConst(0) || op2->IsIntegralConst(1)))
+    {
+        // We have ==/!= 0/1, we need to normalize it to ==/!= 1
+        if (op2->IsIntegralConst(0))
+        {
+            equals = !equals;
+        }
+
+        GenTreeHWIntrinsic* hwi = op1->AsHWIntrinsic();
+        switch (hwi->GetHWIntrinsicId())
+        {
+#if defined(TARGET_XARCH)
+            case NI_Vector128_op_Equality:
+#elif defined(TARGET_ARM64)
+            case NI_Vector64_op_Equality:
+            case NI_Vector128_op_Equality:
+#endif
+                break;
+#if defined(TARGET_XARCH)
+            case NI_Vector128_op_Inequality:
+#elif defined(TARGET_ARM64)
+            case NI_Vector64_op_Inequality:
+            case NI_Vector128_op_Inequality:
+#endif
+                equals = !equals;
+                break;
+
+            default:
+                return NO_ASSERTION_INDEX;
+        }
+
+        // SIMD floating-point equality is not bitwise equality (+0 == -0, NaN != NaN),
+        // Only enable for integral SIMD base types.
+        if (varTypeIsIntegral(hwi->GetSimdBaseType()))
+        {
+            assert(hwi->GetOperandCount() == 2);
+            op1 = hwi->Op(1);
+            op2 = hwi->Op(2);
+
+            if (!op2->IsCnsVec())
+            {
+                return NO_ASSERTION_INDEX;
+            }
+
+            assert(op1->TypeIs(TYP_SIMD8, TYP_SIMD12, TYP_SIMD16));
+            assert(op1->TypeIs(op2->TypeGet()));
+        }
+        else
+        {
+            return NO_ASSERTION_INDEX;
+        }
+    }
+#endif // FEATURE_HW_INTRINSICS
+
     // Avoid creating local assertions for float types.
     //
     if (optLocalAssertionProp && varTypeIsFloating(op1))
@@ -3174,6 +3259,24 @@ GenTree* Compiler::optConstantAssertionProp(const AssertionDsc&  curAssertion,
             newTree->BashToConst(curAssertion.GetOp2().GetDoubleConstant(), tree->TypeGet());
             break;
 
+#if defined(FEATURE_HW_INTRINSICS)
+        case O2K_CONST_VEC:
+        {
+            // The assertion was created from a LCL_VAR == CNS_VEC where types matched.
+            // For now, only support SIMD constants up to 16 bytes (SIMD8/12/16).
+            if (!tree->TypeIs(TYP_SIMD8, TYP_SIMD12, TYP_SIMD16) || !tree->TypeIs(lvaGetDesc(lclNum)->TypeGet()))
+            {
+                return nullptr;
+            }
+
+            // We can't bash a LCL_VAR into a GenTreeVecCon (different node size), so allocate a fresh node.
+            GenTreeVecCon* vecCon = gtNewVconNode(tree->TypeGet());
+            memcpy(&vecCon->gtSimdVal, &curAssertion.GetOp2().GetSimdConstant(), genTypeSize(tree->TypeGet()));
+            newTree = vecCon;
+            break;
+        }
+#endif // FEATURE_HW_INTRINSICS
+
         case O2K_CONST_INT:
 
             // Don't propagate non-nulll non-static handles if we need to report relocs.
@@ -3493,9 +3596,9 @@ GenTree* Compiler::optAssertionProp_LclVar(ASSERT_VALARG_TP assertions, GenTreeL
         return nullptr;
     }
 
-    // There are no constant assertions for structs in global propagation.
+    // There are no constant assertions for structs in global propagation (except SIMD vector constants).
     //
-    if ((!optLocalAssertionProp && varTypeIsStruct(tree)) || !optCanPropLclVar)
+    if ((!optLocalAssertionProp && varTypeIsStruct(tree) && !varTypeIsSIMD(tree)) || !optCanPropLclVar)
     {
         return nullptr;
     }
@@ -3552,9 +3655,9 @@ GenTree* Compiler::optAssertionProp_LclVar(ASSERT_VALARG_TP assertions, GenTreeL
             continue;
         }
 
-        // There are no constant assertions for structs.
+        // There are no constant assertions for structs (except SIMD vector constants).
         //
-        if (varTypeIsStruct(tree))
+        if (varTypeIsStruct(tree) && !varTypeIsSIMD(tree))
         {
             continue;
         }

src/coreclr/jit/compiler.h

@@ -7967,7 +7967,8 @@ class Compiler
                                    // "Checked bound" alone doesn't mean anything,
                                    // nor it implies that it's never negative.
         O2K_ZEROOBJ,
-        O2K_SUBRANGE
+        O2K_SUBRANGE,
+        O2K_CONST_VEC
     };
 
     struct AssertionDsc
@@ -8034,6 +8035,8 @@ class Compiler
                 unsigned      m_lclNum;
                 double        m_dconVal;
                 IntegralRange m_range;
+                simd16_t m_simdVal; // for O2K_CONST_VEC (TYP_SIMD8/12/16 only). TODO-CQ: support wider SIMD via heap
+                                    // allocation.
                 struct
                 {
                     ssize_t   m_iconVal;
@@ -8066,6 +8069,12 @@ class Compiler
                 return m_dconVal;
             }
 
+            const simd16_t& GetSimdConstant() const
+            {
+                assert(KindIs(O2K_CONST_VEC));
+                return m_simdVal;
+            }
+
             ssize_t GetIntConstant() const
             {
                 assert(KindIs(O2K_CONST_INT));
@@ -8089,7 +8098,7 @@ class Compiler
             ValueNum GetVN() const
             {
                 assert(!m_compiler->optLocalAssertionProp);
-                assert(KindIs(O2K_CONST_INT, O2K_CONST_DOUBLE, O2K_ZEROOBJ));
+                assert(KindIs(O2K_CONST_INT, O2K_CONST_DOUBLE, O2K_ZEROOBJ, O2K_CONST_VEC));
                 assert(m_vn != ValueNumStore::NoVN);
                 return m_vn;
             }
@@ -8417,6 +8426,10 @@ class Compiler
                     // exact match because of positive and negative zero.
                     return (memcmp(&GetOp2().m_dconVal, &that.GetOp2().m_dconVal, sizeof(double)) == 0);
 
+                case O2K_CONST_VEC:
+                    // memcmp the full stored payload; GenTreeVecCon zero-inits gtSimdVal before populating.
+                    return (memcmp(&GetOp2().m_simdVal, &that.GetOp2().m_simdVal, sizeof(simd16_t)) == 0);
+
                 case O2K_ZEROOBJ:
                     return true;
 
@@ -8459,7 +8472,7 @@ class Compiler
         static AssertionDsc CreateConstLclVarAssertion(const Compiler* comp,
                                                        unsigned        lclNum,
                                                        ValueNum        vn,
-                                                       T               cns,
+                                                       const T&        cns,
                                                        ValueNum        cnsVN,
                                                        bool            equals,
                                                        GenTreeFlags    iconFlags = GTF_EMPTY,
@@ -8508,6 +8521,14 @@ class Compiler
                 assert(iconFlags == GTF_EMPTY); // no flags expected for double constants
                 assert(fldSeq == nullptr);      // no fieldSeq expected for double constants
             }
+            else if constexpr (std::is_same_v<T, simd16_t>)
+            {
+                dsc.m_op2.m_kind    = O2K_CONST_VEC;
+                dsc.m_op2.m_simdVal = {};
+                memcpy(&dsc.m_op2.m_simdVal, &cns, sizeof(simd16_t));
+                assert(iconFlags == GTF_EMPTY); // no flags expected for vector constants
+                assert(fldSeq == nullptr);      // no fieldSeq expected for vector constants
+            }
             else if constexpr (std::is_same_v<T, optOp2Kind>)
             {
                 dsc.m_op2.m_kind           = static_cast<optOp2Kind>(cns);