Arm64: Implement VectorTableLookup/VectorTableLookupExtension intrinsinsic + Consecutive registers support

src/coreclr/jit/codegenlinear.cpp

@@ -1624,6 +1624,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/compiler.cpp

@@ -6718,6 +6718,10 @@ int Compiler::compCompileHelper(CORINFO_MODULE_HANDLE classPtr,
     compBasicBlockID = 0;
 #endif
 
+#ifdef TARGET_ARM64
+    info.compNeedsConsecutiveRegisters = false;
+#endif
+
     /* Initialize emitter */
 
     if (!compIsForInlining())

src/coreclr/jit/compiler.h

@@ -2809,6 +2809,10 @@ class Compiler
                                               CORINFO_CLASS_HANDLE clsHnd,
                                               CORINFO_SIG_INFO*    sig,
                                               CorInfoType          simdBaseJitType);
+
+#ifdef TARGET_ARM64
+    GenTreeFieldList* gtConvertTableOpToFieldList(GenTree* op, unsigned fieldCount);
+#endif
 #endif // FEATURE_HW_INTRINSICS
 
     GenTree* gtNewMustThrowException(unsigned helper, var_types type, CORINFO_CLASS_HANDLE clsHnd);
@@ -10061,6 +10065,10 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
         // Number of class profile probes in this method
         unsigned compHandleHistogramProbeCount;
 
+#ifdef TARGET_ARM64
+        bool compNeedsConsecutiveRegisters;
+#endif
+
     } info;
 
     ReturnTypeDesc compRetTypeDesc; // ABI return type descriptor for the method

src/coreclr/jit/fginline.cpp

@@ -1447,6 +1447,10 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo)
 
     lvaGenericsContextInUse |= InlineeCompiler->lvaGenericsContextInUse;
 
+#ifdef TARGET_ARM64
+    info.compNeedsConsecutiveRegisters |= InlineeCompiler->info.compNeedsConsecutiveRegisters;
+#endif
+
     // If the inlinee compiler encounters switch tables, disable hot/cold splitting in the root compiler.
     // TODO-CQ: Implement hot/cold splitting of methods with switch tables.
     if (InlineeCompiler->fgHasSwitch && opts.compProcedureSplitting)

src/coreclr/jit/gentree.cpp

@@ -23016,6 +23016,7 @@ GenTree* Compiler::gtNewSimdShuffleNode(var_types   type,
     op2->AsVecCon()->gtSimdVal = vecCns;
 
     return gtNewSimdHWIntrinsicNode(type, op1, op2, lookupIntrinsic, simdBaseJitType, simdSize, isSimdAsHWIntrinsic);
+
 #else
 #error Unsupported platform
 #endif // !TARGET_XARCH && !TARGET_ARM64
@@ -23885,6 +23886,38 @@ GenTree* Compiler::gtNewSimdWithElementNode(var_types   type,
     return gtNewSimdHWIntrinsicNode(type, op1, op2, op3, hwIntrinsicID, simdBaseJitType, simdSize, isSimdAsHWIntrinsic);
 }
 
+#ifdef TARGET_ARM64
+//------------------------------------------------------------------------
+// gtConvertTableOpToFieldList: Convert a operand that represents table of rows into
+//    field list, where each field represents a row in the table.
+//
+// Arguments:
+//    op                  -- Operand to convert.
+//    fieldCount          -- Number of fields or rows present.
+//
+// Return Value:
+//    The GenTreeFieldList node.
+//
+GenTreeFieldList* Compiler::gtConvertTableOpToFieldList(GenTree* op, unsigned fieldCount)
+{
+    LclVarDsc* opVarDsc  = lvaGetDesc(op->AsLclVar());
+    unsigned   lclNum    = lvaGetLclNum(opVarDsc);
+    unsigned   fieldSize = opVarDsc->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;
+    }
+    return fieldList;
+}
+#endif // TARGET_ARM64
+
 GenTree* Compiler::gtNewSimdWithLowerNode(var_types   type,
                                           GenTree*    op1,
                                           GenTree*    op2,

src/coreclr/jit/hwintrinsic.h

@@ -176,6 +176,9 @@ enum HWIntrinsicFlag : unsigned int
 
     // The intrinsic supports some sort of containment analysis
     HW_Flag_SupportsContainment = 0x2000,
+
+    // The intrinsic needs consecutive registers
+    HW_Flag_NeedsConsecutiveRegisters = 0x4000,
 #else
 #error Unsupported platform
 #endif
@@ -751,6 +754,14 @@ struct HWIntrinsicInfo
         return (flags & HW_Flag_SpecialCodeGen) != 0;
     }
 
+#ifdef TARGET_ARM64
+    static bool NeedsConsecutiveRegisters(NamedIntrinsic id)
+    {
+        HWIntrinsicFlag flags = lookupFlags(id);
+        return (flags & HW_Flag_NeedsConsecutiveRegisters) != 0;
+    }
+#endif
+
     static bool HasRMWSemantics(NamedIntrinsic id)
     {
         HWIntrinsicFlag flags = lookupFlags(id);

src/coreclr/jit/hwintrinsicarm64.cpp

@@ -1900,7 +1900,84 @@ 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->TypeGet() == TYP_STRUCT)
+            {
+                info.compNeedsConsecutiveRegisters = true;
+                unsigned fieldCount                = info.compCompHnd->getClassNumInstanceFields(argClass);
+
+                if (!op1->OperIs(GT_LCL_VAR))
+                {
+                    unsigned tmp = lvaGrabTemp(true DEBUGARG("VectorTableLookup temp tree"));
+
+                    impAssignTempGen(tmp, op1, CHECK_SPILL_NONE);
+                    op1 = gtNewLclvNode(tmp, argType);
+                }
+
+                op1 = gtConvertTableOpToFieldList(op1, fieldCount);
+            }
+            else
+            {
+                assert(varTypeIsSIMD(op1->TypeGet()));
+            }
+
+            retNode = gtNewSimdHWIntrinsicNode(retType, op1, op2, intrinsic, simdBaseJitType, simdSize);
+            break;
+        }
+        case NI_AdvSimd_VectorTableLookupExtension:
+        case NI_AdvSimd_Arm64_VectorTableLookupExtension:
+        {
+            assert(sig->numArgs == 3);
+
+            CORINFO_ARG_LIST_HANDLE arg1     = sig->args;
+            CORINFO_ARG_LIST_HANDLE arg2     = info.compCompHnd->getArgNext(arg1);
+            CORINFO_ARG_LIST_HANDLE arg3     = info.compCompHnd->getArgNext(arg2);
+            var_types               argType  = TYP_UNKNOWN;
+            CORINFO_CLASS_HANDLE    argClass = NO_CLASS_HANDLE;
+
+            argType = JITtype2varType(strip(info.compCompHnd->getArgType(sig, arg3, &argClass)));
+            op3     = getArgForHWIntrinsic(argType, argClass);
+            argType = JITtype2varType(strip(info.compCompHnd->getArgType(sig, arg2, &argClass)));
+            op2     = impPopStack().val;
+            op1     = impPopStack().val;
+
+            if (op2->TypeGet() == TYP_STRUCT)
+            {
+                info.compNeedsConsecutiveRegisters = true;
+                unsigned fieldCount                = info.compCompHnd->getClassNumInstanceFields(argClass);
+
+                if (!op2->OperIs(GT_LCL_VAR))
+                {
+                    unsigned tmp = lvaGrabTemp(true DEBUGARG("VectorTableLookupExtension temp tree"));
+
+                    impAssignTempGen(tmp, op2, CHECK_SPILL_NONE);
+                    op2 = gtNewLclvNode(tmp, argType);
+                }
+
+                op2 = gtConvertTableOpToFieldList(op2, fieldCount);
+            }
+            else
+            {
+                assert(varTypeIsSIMD(op1->TypeGet()));
+            }
+
+            retNode = gtNewSimdHWIntrinsicNode(retType, op1, op2, op3, intrinsic, simdBaseJitType, simdSize);
+            break;
+        }
         default:
         {
             return nullptr;

src/coreclr/jit/hwintrinsiccodegenarm64.cpp

@@ -1002,6 +1002,110 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
                                           (emitSize == EA_8BYTE) ? INS_OPTS_8B : INS_OPTS_16B);
                 break;
 
+            case NI_AdvSimd_VectorTableLookup:
+            case NI_AdvSimd_Arm64_VectorTableLookup:
+            {
+                unsigned regCount = 0;
+                if (intrin.op1->OperIsFieldList())
+                {
+                    GenTreeFieldList* fieldList  = intrin.op1->AsFieldList();
+                    GenTree*          firstField = fieldList->Uses().GetHead()->GetNode();
+                    op1Reg                       = firstField->GetRegNum();
+                    INDEBUG(regNumber argReg = op1Reg);
+                    for (GenTreeFieldList::Use& use : fieldList->Uses())
+                    {
+                        regCount++;
+#ifdef DEBUG
+
+                        GenTree* argNode = use.GetNode();
+                        assert(argReg == argNode->GetRegNum());
+                        argReg = REG_NEXT(argReg);
+#endif
+                    }
+                }
+                else
+                {
+                    regCount = 1;
+                    op1Reg   = intrin.op1->GetRegNum();
+                }
+
+                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_VectorTableLookupExtension:
+            case NI_AdvSimd_Arm64_VectorTableLookupExtension:
+            {
+                assert(isRMW);
+                unsigned regCount = 0;
+                op1Reg            = intrin.op1->GetRegNum();
+                op3Reg            = intrin.op3->GetRegNum();
+                assert(targetReg != op3Reg);
+                if (intrin.op2->OperIsFieldList())
+                {
+                    GenTreeFieldList* fieldList  = intrin.op2->AsFieldList();
+                    GenTree*          firstField = fieldList->Uses().GetHead()->GetNode();
+                    op2Reg                       = firstField->GetRegNum();
+                    INDEBUG(regNumber argReg = op2Reg);
+                    for (GenTreeFieldList::Use& use : fieldList->Uses())
+                    {
+                        regCount++;
+#ifdef DEBUG
+
+                        GenTree* argNode = use.GetNode();
+
+                        // registers should be consecutive
+                        assert(argReg == argNode->GetRegNum());
+                        // and they should not interfere with targetReg
+                        assert(targetReg != argReg);
+                        argReg = REG_NEXT(argReg);
+#endif
+                    }
+                }
+                else
+                {
+                    regCount = 1;
+                    op2Reg   = intrin.op2->GetRegNum();
+                }
+
+                switch (regCount)
+                {
+                    case 2:
+                        ins = INS_tbx_2regs;
+                        break;
+                    case 3:
+                        ins = INS_tbx_3regs;
+                        break;
+                    case 4:
+                        ins = INS_tbx_4regs;
+                        break;
+                    default:
+                        assert(regCount == 1);
+                        assert(ins == INS_tbx);
+                        break;
+                }
+
+                GetEmitter()->emitIns_Mov(INS_mov, emitTypeSize(node), targetReg, op1Reg, /* canSkip */ true);
+                GetEmitter()->emitIns_R_R_R(ins, emitSize, targetReg, op2Reg, op3Reg, opt);
+                break;
+            }
             default:
                 unreached();
         }

src/coreclr/jit/hwintrinsiclistarm64.h

@@ -477,8 +477,8 @@ 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,       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,       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|HW_Flag_NeedsConsecutiveRegisters)
+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_SpecialImport|HW_Flag_SpecialCodeGen|HW_Flag_HasRMWSemantics|HW_Flag_NeedsConsecutiveRegisters)
 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)
 HARDWARE_INTRINSIC(AdvSimd,       ZeroExtendWideningUpper,                                          16,      1,     {INS_uxtl2,          INS_uxtl2,          INS_uxtl2,          INS_uxtl2,          INS_uxtl2,          INS_uxtl2,          INS_invalid,        INS_invalid,        INS_invalid,        INS_invalid},     HW_Category_SIMD,                  HW_Flag_BaseTypeFromFirstArg)
@@ -651,8 +651,8 @@ 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, 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, 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|HW_Flag_NeedsConsecutiveRegisters)
+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_SpecialImport|HW_Flag_SpecialCodeGen|HW_Flag_HasRMWSemantics|HW_Flag_NeedsConsecutiveRegisters)
 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)