Revert "Vector128.WithElement codegen regression in .NET 9.0

src/coreclr/jit/gentree.cpp

@@ -27867,8 +27867,14 @@ 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)
     {
@@ -27934,20 +27940,6 @@ 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,7 +1832,6 @@ 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();
@@ -1969,56 +1968,6 @@ 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

@@ -4273,6 +4273,34 @@ 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

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

src/coreclr/jit/rationalize.cpp

@@ -385,6 +385,57 @@ 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)