[Arm64] SIMDIntrinsicGetItem

Fix srcCount
Fix item containment
Implement [Arm64] genSIMDIntrinsicGetItem w/ non const index
Handle out of range const index

3 files changed, 48 insertions, 7 deletions

diff --git a/src/jit/codegenarm64.cpp b/src/jit/codegenarm64.cpp
index 89fe0c92e3..b9e29011ca 100644
--- a/src/jit/codegenarm64.cpp
+++ b/src/jit/codegenarm64.cpp
@@ -4499,19 +4499,43 @@ void CodeGen::genSIMDIntrinsicGetItem(GenTreeSIMD* simdNode)
     // Optimize the case of op1 is in memory and trying to access ith element.
     assert(op1->isUsedFromReg());
 
+    emitAttr baseTypeSize = emitTypeSize(baseType);
+
     if (op2->IsCnsIntOrI())
     {
         assert(op2->isContained());
 
-        emitAttr     attr  = emitTypeSize(baseType);
-        unsigned int index = (unsigned int)op2->gtIntCon.gtIconVal;
+        ssize_t index = op2->gtIntCon.gtIconVal;
 
-        getEmitter()->emitIns_R_R_I(INS_mov, attr, targetReg, srcReg, index);
+        if (getEmitter()->isValidVectorIndex(emitTypeSize(simdType), baseTypeSize, index))
+        {
+            // Only generate code for the get if the index is valid
+            // Otherwise generated code will throw
+            getEmitter()->emitIns_R_R_I(INS_mov, baseTypeSize, targetReg, srcReg, index);
+        }
     }
     else
     {
-        NYI("getItem() with non const index");
-        assert(op2->IsCnsIntOrI());
+        unsigned simdInitTempVarNum = compiler->lvaSIMDInitTempVarNum;
+        noway_assert(compiler->lvaSIMDInitTempVarNum != BAD_VAR_NUM);
+
+        regNumber indexReg = op2->gtRegNum;
+        regNumber tmpReg   = simdNode->ExtractTempReg();
+
+        assert(genIsValidIntReg(tmpReg));
+        assert(tmpReg != indexReg);
+
+        unsigned baseTypeScale = genLog2(EA_SIZE_IN_BYTES(baseTypeSize));
+
+        // Load the address of simdInitTempVarNum
+        getEmitter()->emitIns_R_S(INS_lea, EA_PTRSIZE, tmpReg, simdInitTempVarNum, 0);
+
+        // Store the vector to simdInitTempVarNum
+        getEmitter()->emitIns_R_R(INS_str, emitTypeSize(simdType), srcReg, tmpReg);
+
+        // Load item at simdInitTempVarNum[index]
+        getEmitter()->emitIns_R_R_R_Ext(ins_Load(baseType), baseTypeSize, targetReg, tmpReg, indexReg, INS_OPTS_LSL,
+                                        baseTypeScale);
     }
 
     genProduceReg(simdNode);
diff --git a/src/jit/lowerarmarch.cpp b/src/jit/lowerarmarch.cpp
index 9f73fb47e1..2213960c11 100644
--- a/src/jit/lowerarmarch.cpp
+++ b/src/jit/lowerarmarch.cpp
@@ -777,7 +777,10 @@ void Lowering::ContainCheckSIMD(GenTreeSIMD* simdNode)
             op2 = simdNode->gtOp.gtOp2;
 
             // If the index is a constant, mark it as contained.
-            CheckImmedAndMakeContained(simdNode, op2);
+            if (op2->IsCnsIntOrI())
+            {
+                MakeSrcContained(simdNode, op2);
+            }
             break;
         }
 
diff --git a/src/jit/lsraarm64.cpp b/src/jit/lsraarm64.cpp
index cf00f6d17c..346115fba9 100644
--- a/src/jit/lsraarm64.cpp
+++ b/src/jit/lsraarm64.cpp
@@ -812,7 +812,21 @@ void LinearScan::TreeNodeInfoInitSIMD(GenTreeSIMD* simdTree)
             break;
 
         case SIMDIntrinsicGetItem:
-            info->srcCount = 1;
+            // We have an object and an item, which may be contained.
+            info->srcCount = simdTree->gtGetOp2()->isContained() ? 1 : 2;
+
+            if (!simdTree->gtGetOp2()->IsCnsIntOrI())
+            {
+                // If the index is not a constant, we will need a general purpose register
+                info->internalIntCount = 1;
+
+                // If the index is not a constant, we will use the SIMD temp location to store the vector.
+                compiler->getSIMDInitTempVarNum();
+
+                // internal register must not clobber input index
+                simdTree->gtOp.gtOp2->gtLsraInfo.isDelayFree = true;
+                info->hasDelayFreeSrc                        = true;
+            }
             break;
 
         case SIMDIntrinsicAdd: