Arm64: Implement VectorTableLookup/VectorTableLookupExtension intrinsinsic + Consecutive registers support

src/coreclr/jit/codegenlinear.cpp

@@ -1617,6 +1617,14 @@ void CodeGen::genConsumeRegs(GenTree* tree)
             genConsumeRegs(tree->gtGetOp1());
             genConsumeRegs(tree->gtGetOp2());
         }
+        else if (tree->OperIsFieldList())
+        {
+            for (GenTreeFieldList::Use& use : tree->AsFieldList()->Uses())
+            {
+                GenTree* fieldNode = use.GetNode();
+                genConsumeRegs(fieldNode);
+            }
+        }
 #endif
         else if (tree->OperIsLocalRead())
         {

src/coreclr/jit/hwintrinsicarm64.cpp

@@ -1876,6 +1876,60 @@ GenTree* Compiler::impSpecialIntrinsic(NamedIntrinsic        intrinsic,
             retNode = impAssignMultiRegTypeToVar(op1, sig->retTypeSigClass DEBUGARG(CorInfoCallConvExtension::Managed));
             break;
         }
+        case NI_AdvSimd_VectorTableLookup:
+        case NI_AdvSimd_Arm64_VectorTableLookup:
+        {
+            assert(sig->numArgs == 2);
+
+            CORINFO_ARG_LIST_HANDLE arg1     = sig->args;
+            CORINFO_ARG_LIST_HANDLE arg2     = info.compCompHnd->getArgNext(arg1);
+            var_types               argType  = TYP_UNKNOWN;
+            CORINFO_CLASS_HANDLE    argClass = NO_CLASS_HANDLE;
+
+            argType = JITtype2varType(strip(info.compCompHnd->getArgType(sig, arg2, &argClass)));
+            op2     = getArgForHWIntrinsic(argType, argClass);
+
+            argType = JITtype2varType(strip(info.compCompHnd->getArgType(sig, arg1, &argClass)));
+            op1     = impPopStack().val;
+
+            if (!op1->OperIsLocal())
+            {
+                unsigned tmp = lvaGrabTemp(true DEBUGARG("VectorTableLookup temp tree"));
+
+                impAssignTempGen(tmp, op1, CHECK_SPILL_NONE);
+                op1 = gtNewLclvNode(tmp, argType);
+            }
+
+            if (argType == TYP_STRUCT)
+            {
+                assert(op1->OperIsLocal());
+
+                LclVarDsc* op1VarDsc  = lvaGetDesc(op1->AsLclVar());
+                unsigned   lclNum     = lvaGetLclNum(op1VarDsc);
+                unsigned   fieldCount = info.compCompHnd->getClassNumInstanceFields(argClass);
+                unsigned   fieldSize  = op1VarDsc->lvSize() / fieldCount;
+                var_types  fieldType  = TYP_SIMD16;
+
+                GenTreeFieldList* fieldList = new (this, GT_FIELD_LIST) GenTreeFieldList();
+                int               offset    = 0;
+                for (unsigned fieldId = 0; fieldId < fieldCount; fieldId++)
+                {
+                    GenTreeLclFld* fldNode = gtNewLclFldNode(lclNum, fieldType, offset);
+                    fieldList->AddField(this, fldNode, offset, fieldType);
+
+                    offset += fieldSize;
+                }
+                op1 = fieldList;
+
+                retNode = gtNewSimdHWIntrinsicNode(retType, op1, op2, intrinsic, simdBaseJitType, simdSize);
+            }
+            else
+            {
+                assert(argType == TYP_SIMD16);
+                retNode = gtNewSimdHWIntrinsicNode(retType, op1, op2, intrinsic, simdBaseJitType, simdSize);
+            }
+            break;
+        }
 
         default:
         {

src/coreclr/jit/hwintrinsiccodegenarm64.cpp

@@ -424,6 +424,11 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
         instruction ins = INS_invalid;
         switch (intrin.id)
         {
+            case NI_AdvSimd_VectorTableLookup:
+            case NI_AdvSimd_Arm64_VectorTableLookup:
+                ins = INS_tbl;
+                break;
+
             case NI_AdvSimd_AddWideningLower:
                 assert(varTypeIsIntegral(intrin.baseType));
                 if (intrin.op1->TypeGet() == TYP_SIMD8)
@@ -489,6 +494,54 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
 
         switch (intrin.id)
         {
+            case NI_AdvSimd_VectorTableLookup:
+            case NI_AdvSimd_Arm64_VectorTableLookup:
+            {
+
+                if (!intrin.op1->OperIsFieldList())
+                {
+                    assert(!"Expect the first operand of VectorTableLookup to be FIELD_LIST");
+                }
+
+                GenTreeFieldList* fieldList  = intrin.op1->AsFieldList();
+                GenTree*          firstField = fieldList->Uses().GetHead()->GetNode();
+                op1Reg                       = firstField->GetRegNum();
+#ifdef DEBUG
+                regNumber argReg = op1Reg;
+#endif
+                unsigned regCount = 0;
+                for (GenTreeFieldList::Use& use : fieldList->Uses())
+                {
+                    regCount++;
+#ifdef DEBUG
+
+                    GenTree* argNode = use.GetNode();
+                    assert(argReg == argNode->GetRegNum());
+                    argReg = (regNumber)(argReg + 1);
+#endif
+                }
+
+                switch (regCount)
+                {
+                    case 2:
+                        ins = INS_tbl_2regs;
+                        break;
+                    case 3:
+                        ins = INS_tbl_3regs;
+                        break;
+                    case 4:
+                        ins = INS_tbl_4regs;
+                        break;
+                    default:
+                        assert(regCount == 1);
+                        assert(ins == INS_tbl);
+                        break;
+                }
+
+                GetEmitter()->emitIns_R_R_R(ins, emitSize, targetReg, op1Reg, op2Reg, opt);
+                break;
+            }
+
             case NI_AdvSimd_BitwiseSelect:
                 // Even though BitwiseSelect is an RMW intrinsic per se, we don't want to mark it as such
                 // since we can handle all possible allocation decisions for targetReg.

src/coreclr/jit/hwintrinsiclistarm64.h

@@ -475,7 +475,7 @@ HARDWARE_INTRINSIC(AdvSimd,       SubtractSaturateScalar,
 HARDWARE_INTRINSIC(AdvSimd,       SubtractScalar,                                                    8,      2,     {INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_sub,            INS_sub,            INS_fsub,           INS_fsub},        HW_Category_SIMD,                  HW_Flag_SIMDScalar)
 HARDWARE_INTRINSIC(AdvSimd,       SubtractWideningLower,                                             8,      2,     {INS_ssubl,          INS_usubl,          INS_ssubl,          INS_usubl,          INS_ssubl,          INS_usubl,          INS_ssubw,          INS_usubw,          INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_BaseTypeFromSecondArg|HW_Flag_SpecialCodeGen)
 HARDWARE_INTRINSIC(AdvSimd,       SubtractWideningUpper,                                            16,      2,     {INS_ssubl2,         INS_usubl2,         INS_ssubl2,         INS_usubl2,         INS_ssubl2,         INS_usubl2,         INS_ssubw2,         INS_usubw2,         INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_BaseTypeFromSecondArg|HW_Flag_SpecialCodeGen)
-HARDWARE_INTRINSIC(AdvSimd,       VectorTableLookup,                                                 8,      2,     {INS_tbl,            INS_tbl,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_NoFlag)
+HARDWARE_INTRINSIC(AdvSimd,       VectorTableLookup,                                                 8,      2,     {INS_tbl,            INS_tbl,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_SpecialImport|HW_Flag_SpecialCodeGen)
 HARDWARE_INTRINSIC(AdvSimd,       VectorTableLookupExtension,                                        8,      3,     {INS_tbx,            INS_tbx,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_HasRMWSemantics)
 HARDWARE_INTRINSIC(AdvSimd,       Xor,                                                              -1,      2,     {INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor,            INS_eor},         HW_Category_SIMD,                  HW_Flag_Commutative)
 HARDWARE_INTRINSIC(AdvSimd,       ZeroExtendWideningLower,                                           8,      1,     {INS_uxtl,           INS_uxtl,           INS_uxtl,           INS_uxtl,           INS_uxtl,           INS_uxtl,           INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_BaseTypeFromFirstArg)
@@ -649,7 +649,7 @@ HARDWARE_INTRINSIC(AdvSimd_Arm64, TransposeEven,
 HARDWARE_INTRINSIC(AdvSimd_Arm64, TransposeOdd,                                                     -1,      2,     {INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2,           INS_trn2},        HW_Category_SIMD,                  HW_Flag_NoFlag)
 HARDWARE_INTRINSIC(AdvSimd_Arm64, UnzipEven,                                                        -1,      2,     {INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1,           INS_uzp1},        HW_Category_SIMD,                  HW_Flag_NoFlag)
 HARDWARE_INTRINSIC(AdvSimd_Arm64, UnzipOdd,                                                         -1,      2,     {INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2,           INS_uzp2},        HW_Category_SIMD,                  HW_Flag_NoFlag)
-HARDWARE_INTRINSIC(AdvSimd_Arm64, VectorTableLookup,                                                16,      2,     {INS_tbl,            INS_tbl,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_NoFlag)
+HARDWARE_INTRINSIC(AdvSimd_Arm64, VectorTableLookup,                                                16,      2,     {INS_tbl,            INS_tbl,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_SpecialImport|HW_Flag_SpecialCodeGen)
 HARDWARE_INTRINSIC(AdvSimd_Arm64, VectorTableLookupExtension,                                       16,      3,     {INS_tbx,            INS_tbx,            INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_HasRMWSemantics)
 HARDWARE_INTRINSIC(AdvSimd_Arm64, ZipHigh,                                                          -1,      2,     {INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2,           INS_zip2},        HW_Category_SIMD,                  HW_Flag_NoFlag)
 HARDWARE_INTRINSIC(AdvSimd_Arm64, ZipLow,                                                           -1,      2,     {INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1,           INS_zip1},        HW_Category_SIMD,                  HW_Flag_NoFlag)

src/coreclr/jit/lsra.cpp

@@ -1243,6 +1243,10 @@ PhaseStatus LinearScan::doLinearScan()
 
     splitBBNumToTargetBBNumMap = nullptr;
 
+#ifdef TARGET_ARM64
+    nextConsecutiveRefPositionMap = nullptr;
+#endif
+
     // This is complicated by the fact that physical registers have refs associated
     // with locations where they are killed (e.g. calls), but we don't want to
     // count these as being touched.
@@ -2218,6 +2222,8 @@ void LinearScan::checkLastUses(BasicBlock* block)
                             loc);
                     foundDiff = true;
                 }
+
+                JITDUMP("++ V%02u in computedLive\n", compiler->lvaTrackedIndexToLclNum(varIndex));
                 VarSetOps::AddElemD(compiler, computedLive, varIndex);
             }
             else if (currentRefPosition->lastUse)
@@ -2233,6 +2239,7 @@ void LinearScan::checkLastUses(BasicBlock* block)
 
             if (currentRefPosition->refType == RefTypeDef || currentRefPosition->refType == RefTypeDummyDef)
             {
+                JITDUMP("-- V%02u from computedLive\n", compiler->lvaTrackedIndexToLclNum(varIndex));
                 VarSetOps::RemoveElemD(compiler, computedLive, varIndex);
             }
         }
@@ -2837,8 +2844,17 @@ regNumber LinearScan::allocateReg(Interval*    currentInterval,
             }
         }
     }
+
     assignPhysReg(availablePhysRegRecord, currentInterval);
     refPosition->registerAssignment = foundRegBit;
+
+#ifdef TARGET_ARM64
+    if (refPosition->isFirstRefPositionOfConsecutiveRegisters())
+    {
+        setNextConsecutiveRegisterAssignment(refPosition, foundRegBit);
+    }
+#endif // TARGET_ARM64
+
     return foundReg;
 }
 
@@ -5297,6 +5313,19 @@ void LinearScan::allocateRegisters()
                     lastAllocatedRefPosition  = &currentRefPosition;
                     regMaskTP copyRegMask     = getRegMask(copyReg, currentInterval->registerType);
                     regMaskTP assignedRegMask = getRegMask(assignedRegister, currentInterval->registerType);
+
+#ifdef TARGET_ARM64
+                    if (currentRefPosition.needsConsecutive)
+                    {
+                        // For consecutive register, it doesn't matter what the assigned register was.
+                        // We have just assigned it `copyRegMask` and that's the one in-use, and not the
+                        // one that was assigned previously.
+                        assignedRegMask = REG_NA;
+
+                        // This should never be the first refposition of the series.
+                        assert(currentRefPosition.regCount == 0);
+                    }
+#endif
                     regsInUseThisLocation |= copyRegMask | assignedRegMask;
                     if (currentRefPosition.lastUse)
                     {
@@ -5351,6 +5380,43 @@ void LinearScan::allocateRegisters()
             }
         }
 
+#ifdef TARGET_ARM64
+        if (currentRefPosition.needsConsecutive)
+        {
+            // For consecutive register, we would like to assign a register (if not already assigned)
+            // to the 1st refPosition and the subsequent refPositions will just get the consecutive register.
+            if (currentRefPosition.regCount > 0)
+            {
+                // 1st refPosition of the series...
+                if (assignedRegister != REG_NA)
+                {
+                    // For the 1st refPosition, if it already has a register assigned, then just assign
+                    // subsequent registers to the remaining position and skip the allocation for the
+                    // 1st refPosition altogether.
+
+                    setNextConsecutiveRegisterAssignment(&currentRefPosition, assignedRegBit);
+                }
+            }
+            else
+            {
+                // remaining refPosition of the series...
+                if (assignedRegBit == currentRefPosition.registerAssignment)
+                {
+                    // For the subsequent position, if they already have the subsequent register assigned, then
+                    // no need to find register to assign.
+                    allocate = false;
+                }
+                else
+                {
+                    // If the subsequent refPosition is not assigned to the consecutive register, then reassign the
+                    // right
+                    // consecutive register.
+                    assignedRegister = REG_NA;
+                }
+            }
+        }
+#endif // TARGET_ARM64
+
         if (assignedRegister == REG_NA)
         {
             if (currentRefPosition.RegOptional())
@@ -11942,7 +12008,7 @@ regMaskTP LinearScan::RegisterSelection::select(Interval*    currentInterval,
     reverseSelect = linearScan->doReverseSelect();
 #endif // DEBUG
 
-    freeCandidates = linearScan->getFreeCandidates(candidates, regType);
+    freeCandidates = linearScan->getFreeCandidates(candidates ARM_ARG(regType));
 
     // If no free candidates, then double check if refPosition is an actual ref.
     if (freeCandidates == RBM_NONE)

src/coreclr/jit/lsra.h

@@ -1179,10 +1179,15 @@ class LinearScan : public LinearScanInterface
 
     void spillGCRefs(RefPosition* killRefPosition);
 
-    /*****************************************************************************
-    * Register selection
-    ****************************************************************************/
-    regMaskTP getFreeCandidates(regMaskTP candidates, var_types regType)
+/*****************************************************************************
+* Register selection
+****************************************************************************/
+
+#if defined(TARGET_ARM64)
+    void setNextConsecutiveRegisterAssignment(RefPosition* firstRefPosition, regMaskTP firstRegAssigned);
+#endif // TARGET_ARM64
+
+    regMaskTP getFreeCandidates(regMaskTP candidates ARM_ARG(var_types regType))
     {
         regMaskTP result = candidates & m_AvailableRegs;
 #ifdef TARGET_ARM
@@ -1354,6 +1359,21 @@ class LinearScan : public LinearScanInterface
 
     regNumber getTempRegForResolution(BasicBlock* fromBlock, BasicBlock* toBlock, var_types type);
 
+#ifdef TARGET_ARM64
+    typedef JitHashTable<RefPosition*, JitPtrKeyFuncs<RefPosition>, RefPosition*> NextConsecutiveRefPositionsMap;
+    NextConsecutiveRefPositionsMap* nextConsecutiveRefPositionMap;
+    NextConsecutiveRefPositionsMap* getNextConsecutiveRefPositionsMap()
+    {
+        if (nextConsecutiveRefPositionMap == nullptr)
+        {
+            nextConsecutiveRefPositionMap =
+                new (getAllocator(compiler)) NextConsecutiveRefPositionsMap(getAllocator(compiler));
+        }
+        return nextConsecutiveRefPositionMap;
+    }
+    RefPosition* getNextConsecutiveRefPosition(RefPosition* refPosition);
+#endif
+
 #ifdef DEBUG
     void dumpVarToRegMap(VarToRegMap map);
     void dumpInVarToRegMap(BasicBlock* block);
@@ -2259,6 +2279,14 @@ class RefPosition
     // would be 0..MAX_RET_REG_COUNT-1.
     unsigned char multiRegIdx : 2;
 
+#ifdef TARGET_ARM64
+    // If this refposition needs consecutive register assignment
+    bool needsConsecutive;
+
+    // How many consecutive registers does this and subsequent refPositions need
+    unsigned char regCount : 3;
+#endif // TARGET_ARM64
+
     // Last Use - this may be true for multiple RefPositions in the same Interval
     unsigned char lastUse : 1;
 
@@ -2344,6 +2372,10 @@ class RefPosition
         , registerAssignment(RBM_NONE)
         , refType(refType)
         , multiRegIdx(0)
+#ifdef TARGET_ARM64
+        , needsConsecutive(false)
+        , regCount(0)
+#endif
         , lastUse(false)
         , reload(false)
         , spillAfter(false)
@@ -2497,6 +2529,19 @@ class RefPosition
         return (isFixedRefOfRegMask(genRegMask(regNum)));
     }
 
+#ifdef TARGET_ARM64
+    /// For consecutive registers, returns true if this RefPosition is
+    /// the first of the series.
+    FORCEINLINE bool isFirstRefPositionOfConsecutiveRegisters()
+    {
+        if (needsConsecutive)
+        {
+            return regCount != 0;
+        }
+        return false;
+    }
+#endif // TARGET_ARM64
+
 #ifdef DEBUG
     // operator= copies everything except 'rpNum', which must remain unique
     RefPosition& operator=(const RefPosition& rp)