Vector128.WithElement codegen regression in .NET 9.0

src/coreclr/jit/gentree.cpp

@@ -27886,14 +27886,8 @@ GenTree* Compiler::gtNewSimdWithElementNode(
     var_types      simdBaseType  = JitType2PreciseVarType(simdBaseJitType);
 
     assert(varTypeIsArithmetic(simdBaseType));
-    assert(op2->IsCnsIntOrI());
     assert(varTypeIsArithmetic(op3));
 
-    ssize_t imm8  = op2->AsIntCon()->IconValue();
-    ssize_t count = simdSize / genTypeSize(simdBaseType);
-
-    assert((0 <= imm8) && (imm8 < count));
-
 #if defined(TARGET_XARCH)
     switch (simdBaseType)
     {
@@ -27959,6 +27953,20 @@ GenTree* Compiler::gtNewSimdWithElementNode(
 #error Unsupported platform
 #endif // !TARGET_XARCH && !TARGET_ARM64
 
+    int  immUpperBound    = getSIMDVectorLength(simdSize, simdBaseType) - 1;
+    bool rangeCheckNeeded = !op2->OperIsConst();
+
+    if (!rangeCheckNeeded)
+    {
+        ssize_t imm8     = op2->AsIntCon()->IconValue();
+        rangeCheckNeeded = (imm8 < 0) || (imm8 > immUpperBound);
+    }
+
+    if (rangeCheckNeeded)
+    {
+        op2 = addRangeCheckForHWIntrinsic(op2, 0, immUpperBound);
+    }
+
     return gtNewSimdHWIntrinsicNode(type, op1, op2, op3, hwIntrinsicID, simdBaseJitType, simdSize);
 }
 

src/coreclr/jit/hwintrinsiccodegenxarch.cpp

@@ -1832,6 +1832,7 @@ void CodeGen::genBaseIntrinsic(GenTreeHWIntrinsic* node, insOpts instOptions)
 
     GenTree* op1 = (node->GetOperandCount() >= 1) ? node->Op(1) : nullptr;
     GenTree* op2 = (node->GetOperandCount() >= 2) ? node->Op(2) : nullptr;
+    GenTree* op3 = (node->GetOperandCount() >= 3) ? node->Op(3) : nullptr;
 
     genConsumeMultiOpOperands(node);
     regNumber op1Reg = (op1 == nullptr) ? REG_NA : op1->GetRegNum();
@@ -1968,6 +1969,57 @@ void CodeGen::genBaseIntrinsic(GenTreeHWIntrinsic* node, insOpts instOptions)
             break;
         }
 
+        case NI_Vector128_WithElement:
+        case NI_Vector256_WithElement:
+        case NI_Vector512_WithElement:
+        {
+            // Optimize the case where op2 is not a constant.
+            assert(!op2->OperIsConst());
+
+            // We don't have an instruction to implement this intrinsic if the index is not a constant.
+            // So we will use the SIMD temp location to store the vector, set the value and then reload it.
+            // The range check will already have been performed, so at this point we know we have an index
+            // within the bounds of the vector.
+
+            unsigned simdInitTempVarNum = compiler->lvaSIMDInitTempVarNum;
+            noway_assert(simdInitTempVarNum != BAD_VAR_NUM);
+
+            bool     isEBPbased;
+            unsigned offs = compiler->lvaFrameAddress(simdInitTempVarNum, &isEBPbased);
+
+#if !FEATURE_FIXED_OUT_ARGS
+            if (!isEBPbased)
+            {
+                // Adjust the offset by the amount currently pushed on the CPU stack
+                offs += genStackLevel;
+            }
+#else
+            assert(genStackLevel == 0);
+#endif // !FEATURE_FIXED_OUT_ARGS
+
+            regNumber indexReg = op2->GetRegNum();
+            regNumber valueReg = op3->GetRegNum(); // New element value to be stored
+
+            // Store the vector to the temp location.
+            GetEmitter()->emitIns_S_R(ins_Store(simdType, compiler->isSIMDTypeLocalAligned(simdInitTempVarNum)),
+                                        emitTypeSize(simdType), op1Reg, simdInitTempVarNum, 0);
+
+            // Set the desired element.
+            GetEmitter()->emitIns_ARX_R(ins_Move_Extend(op3->TypeGet(), false),   // Store
+                                        emitTypeSize(baseType),                   // Of the vector baseType
+                                        valueReg,                                 // From valueReg
+                                        (isEBPbased) ? REG_EBP : REG_ESP,         // Stack-based
+                                        indexReg,                                 // Indexed
+                                        genTypeSize(baseType),                    // by the size of the baseType
+                                        offs);                                    // Offset
+
+            // Write back the modified vector to the original location.
+            GetEmitter()->emitIns_R_S(ins_Load(simdType, compiler->isSIMDTypeLocalAligned(simdInitTempVarNum)),
+                                        emitTypeSize(simdType), targetReg, simdInitTempVarNum, 0);
+            break;
+        }
+
+
         case NI_Vector128_GetElement:
         case NI_Vector256_GetElement:
         case NI_Vector512_GetElement:

src/coreclr/jit/hwintrinsicxarch.cpp

@@ -3806,34 +3806,6 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
             assert(sig->numArgs == 3);
             GenTree* indexOp = impStackTop(1).val;
 
-            if (!indexOp->OperIsConst())
-            {
-                if (!opts.OptimizationEnabled())
-                {
-                    // Only enable late stage rewriting if optimizations are enabled
-                    // as we won't otherwise encounter a constant at the later point
-                    return nullptr;
-                }
-
-                op3 = impPopStack().val;
-                op2 = impPopStack().val;
-                op1 = impSIMDPopStack();
-
-                retNode = gtNewSimdHWIntrinsicNode(retType, op1, op2, op3, intrinsic, simdBaseJitType, simdSize);
-
-                retNode->AsHWIntrinsic()->SetMethodHandle(this, method R2RARG(*entryPoint));
-                break;
-            }
-
-            ssize_t imm8  = indexOp->AsIntCon()->IconValue();
-            ssize_t count = simdSize / genTypeSize(simdBaseType);
-
-            if ((imm8 >= count) || (imm8 < 0))
-            {
-                // Using software fallback if index is out of range (throw exception)
-                return nullptr;
-            }
-
             switch (simdBaseType)
             {
                 // Using software fallback if simdBaseType is not supported by hardware

src/coreclr/jit/lowerxarch.cpp

@@ -5653,7 +5653,14 @@ GenTree* Lowering::LowerHWIntrinsicWithElement(GenTreeHWIntrinsic* node)
     GenTree* op2 = node->Op(2);
     GenTree* op3 = node->Op(3);
 
-    assert(op2->OperIsConst());
+    if (!op2->OperIsConst())
+    {
+        comp->getSIMDInitTempVarNum(simdType);
+
+        // We will specially handle WithElement in codegen when op2 isn't a constant
+        ContainCheckHWIntrinsic(node);
+        return node->gtNext;
+    }
 
     ssize_t count     = simdSize / genTypeSize(simdBaseType);
     ssize_t imm8      = op2->AsIntCon()->IconValue();

src/coreclr/jit/rationalize.cpp

@@ -385,57 +385,6 @@ void Rationalizer::RewriteHWIntrinsicAsUserCall(GenTree** use, ArrayStack<GenTre
             break;
         }
 
-        case NI_Vector128_WithElement:
-#if defined(TARGET_XARCH)
-        case NI_Vector256_WithElement:
-        case NI_Vector512_WithElement:
-#elif defined(TARGET_ARM64)
-        case NI_Vector64_WithElement:
-#endif
-        {
-            assert(operandCount == 3);
-
-            GenTree* op1 = operands[0];
-            GenTree* op2 = operands[1];
-            GenTree* op3 = operands[2];
-
-            if (op2->OperIsConst())
-            {
-                ssize_t imm8  = op2->AsIntCon()->IconValue();
-                ssize_t count = simdSize / genTypeSize(simdBaseType);
-
-                if ((imm8 >= count) || (imm8 < 0))
-                {
-                    // Using software fallback if index is out of range (throw exception)
-                    break;
-                }
-
-#if defined(TARGET_XARCH)
-                if (varTypeIsIntegral(simdBaseType))
-                {
-                    if (varTypeIsLong(simdBaseType))
-                    {
-                        if (!comp->compOpportunisticallyDependsOn(InstructionSet_SSE41_X64))
-                        {
-                            break;
-                        }
-                    }
-                    else if (!varTypeIsShort(simdBaseType))
-                    {
-                        if (!comp->compOpportunisticallyDependsOn(InstructionSet_SSE41))
-                        {
-                            break;
-                        }
-                    }
-                }
-#endif // TARGET_XARCH
-
-                result = comp->gtNewSimdWithElementNode(retType, op1, op2, op3, simdBaseJitType, simdSize);
-                break;
-            }
-            break;
-        }
-
         default:
         {
             if (sigInfo.numArgs == 0)