Refactor handling of two immediates in hwintrinsic (#102071)

* Refactor handling of two immediates in hwintrinsic

* Add immOp2 assert

* Update comments for getHWIntrinsicImmTypes

* fix arg name

* update function summaries

Author: a74nh

src/coreclr/jit/codegen.h

@@ -1009,7 +1009,7 @@ class CodeGen final : public CodeGenInterface
     class HWIntrinsicImmOpHelper final
     {
     public:
-        HWIntrinsicImmOpHelper(CodeGen* codeGen, GenTree* immOp, GenTreeHWIntrinsic* intrin);
+        HWIntrinsicImmOpHelper(CodeGen* codeGen, GenTree* immOp, GenTreeHWIntrinsic* intrin, int immNum = 1);
 
         void EmitBegin();
         void EmitCaseEnd();

src/coreclr/jit/compiler.h

@@ -4603,6 +4603,34 @@ class Compiler
         NamedIntrinsic intrinsic, GenTree* immOp, bool mustExpand, int immLowerBound, int immUpperBound);
     GenTree* addRangeCheckForHWIntrinsic(GenTree* immOp, int immLowerBound, int immUpperBound);
 
+    void getHWIntrinsicImmOps(NamedIntrinsic    intrinsic,
+                              CORINFO_SIG_INFO* sig,
+                              GenTree**         immOp1Ptr,
+                              GenTree**         immOp2Ptr);
+
+    bool CheckHWIntrinsicImmRange(NamedIntrinsic intrinsic,
+                                  CorInfoType simdBaseJitType,
+                                  GenTree* immOp,
+                                  bool mustExpand,
+                                  int immLowerBound,
+                                  int immUpperBound,
+                                  bool hasFullRangeImm,
+                                  bool *useFallback);
+
+#if defined(TARGET_ARM64)
+
+    void getHWIntrinsicImmTypes(NamedIntrinsic       intrinsic,
+                                CORINFO_SIG_INFO*    sig,
+                                unsigned             immNumber,
+                                var_types            simdBaseType,
+                                CorInfoType          simdBaseJitType,
+                                CORINFO_CLASS_HANDLE op2ClsHnd,
+                                CORINFO_CLASS_HANDLE op3ClsHnd,
+                                unsigned*            immSimdSize,
+                                var_types*           immSimdBaseType);
+
+#endif // TARGET_ARM64
+
 #endif // FEATURE_HW_INTRINSICS
     GenTree* impArrayAccessIntrinsic(CORINFO_CLASS_HANDLE clsHnd,
                                      CORINFO_SIG_INFO*    sig,

src/coreclr/jit/hwintrinsic.cpp

@@ -1032,6 +1032,97 @@ struct HWIntrinsicSignatureReader final
     }
 };
 
+//------------------------------------------------------------------------
+// CheckHWIntrinsicImmRange: Check if an immediate is within the valid range
+//
+// Arguments:
+//    intrinsicId       -- HW intrinsic id
+//    simdBaseJitType   -- The base JIT type of SIMD type of the intrinsic
+//    immOp             -- Immediate to check is within range
+//    mustExpand        -- true if the intrinsic must expand to a GenTree*; otherwise, false
+//    immLowerBound     -- the lower valid bound of the immediate
+//    immLowerBound     -- the upper valid bound of the immediate
+//    hasFullRangeImm   -- the range has all valid values. The immediate is always within range.
+//    useFallback [OUT] -- Only set if false is returned. A fallback can be used instead.
+//
+// Return Value:
+//    returns true if immOp is within range. Otherwise false.
+//
+bool Compiler::CheckHWIntrinsicImmRange(NamedIntrinsic intrinsic,
+                                        CorInfoType    simdBaseJitType,
+                                        GenTree*       immOp,
+                                        bool           mustExpand,
+                                        int            immLowerBound,
+                                        int            immUpperBound,
+                                        bool           hasFullRangeImm,
+                                        bool*          useFallback)
+{
+    *useFallback = false;
+
+    if (!hasFullRangeImm && immOp->IsCnsIntOrI())
+    {
+        const int ival = (int)immOp->AsIntCon()->IconValue();
+        bool      immOutOfRange;
+#ifdef TARGET_XARCH
+        if (HWIntrinsicInfo::isAVX2GatherIntrinsic(intrinsic))
+        {
+            immOutOfRange = (ival != 1) && (ival != 2) && (ival != 4) && (ival != 8);
+        }
+        else
+#endif
+        {
+            immOutOfRange = (ival < immLowerBound) || (ival > immUpperBound);
+        }
+
+        if (immOutOfRange)
+        {
+            assert(!mustExpand);
+            // The imm-HWintrinsics that do not accept all imm8 values may throw
+            // ArgumentOutOfRangeException when the imm argument is not in the valid range
+            return false;
+        }
+    }
+    else if (!immOp->IsCnsIntOrI())
+    {
+        if (HWIntrinsicInfo::NoJmpTableImm(intrinsic))
+        {
+            *useFallback = true;
+            return false;
+        }
+#if defined(TARGET_XARCH)
+        else if (HWIntrinsicInfo::MaybeNoJmpTableImm(intrinsic))
+        {
+#if defined(TARGET_X86)
+            var_types simdBaseType = JitType2PreciseVarType(simdBaseJitType);
+
+            if (varTypeIsLong(simdBaseType))
+            {
+                if (!mustExpand)
+                {
+                    return false;
+                }
+            }
+            else
+#endif // TARGET_XARCH
+            {
+                *useFallback = true;
+                return false;
+            }
+        }
+#endif // TARGET_XARCH
+        else if (!mustExpand)
+        {
+            // When the imm-argument is not a constant and we are not being forced to expand, we need to
+            // return false so a GT_CALL to the intrinsic method is emitted instead. The
+            // intrinsic method is recursive and will be forced to expand, at which point
+            // we emit some less efficient fallback code.
+            return false;
+        }
+    }
+
+    return true;
+}
+
 //------------------------------------------------------------------------
 // impHWIntrinsic: Import a hardware intrinsic as a GT_HWINTRINSIC node if possible
 //
@@ -1162,190 +1253,64 @@ GenTree* Compiler::impHWIntrinsic(NamedIntrinsic        intrinsic,
     }
 
     var_types simdBaseType = TYP_UNKNOWN;
-    GenTree*  immOp        = nullptr;
 
     if (simdBaseJitType != CORINFO_TYPE_UNDEF)
     {
         simdBaseType = JitType2PreciseVarType(simdBaseJitType);
     }
 
+    const unsigned simdSize = HWIntrinsicInfo::lookupSimdSize(this, intrinsic, sig);
+
     HWIntrinsicSignatureReader sigReader;
     sigReader.Read(info.compCompHnd, sig);
 
-#ifdef TARGET_ARM64
-    if ((intrinsic == NI_AdvSimd_Insert) || (intrinsic == NI_AdvSimd_InsertScalar) ||
-        ((intrinsic >= NI_AdvSimd_LoadAndInsertScalar) && (intrinsic <= NI_AdvSimd_LoadAndInsertScalarVector64x4)) ||
-        ((intrinsic >= NI_AdvSimd_Arm64_LoadAndInsertScalar) &&
-         (intrinsic <= NI_AdvSimd_Arm64_LoadAndInsertScalarVector128x4)))
-    {
-        assert(sig->numArgs == 3);
-        immOp = impStackTop(1).val;
-        assert(HWIntrinsicInfo::isImmOp(intrinsic, immOp));
-    }
-    else if (intrinsic == NI_AdvSimd_Arm64_InsertSelectedScalar)
-    {
-        // InsertSelectedScalar intrinsic has two immediate operands.
-        // Since all the remaining intrinsics on both platforms have only one immediate
-        // operand, in order to not complicate the shared logic even further we ensure here that
-        // 1) The second immediate operand immOp2 is constant and
-        // 2) its value belongs to [0, sizeof(op3) / sizeof(op3.BaseType)).
-        // If either is false, we should fallback to the managed implementation Insert(dst, dstIdx, Extract(src,
-        // srcIdx)).
-        // The check for the first immediate operand immOp will use the same logic as other intrinsics that have an
-        // immediate operand.
-
-        GenTree* immOp2 = nullptr;
-
-        assert(sig->numArgs == 4);
-
-        immOp  = impStackTop(2).val;
-        immOp2 = impStackTop().val;
-
-        assert(HWIntrinsicInfo::isImmOp(intrinsic, immOp));
-        assert(HWIntrinsicInfo::isImmOp(intrinsic, immOp2));
-
-        if (!immOp2->IsCnsIntOrI())
-        {
-            assert(HWIntrinsicInfo::NoJmpTableImm(intrinsic));
-            return impNonConstFallback(intrinsic, retType, simdBaseJitType);
-        }
-
-        unsigned int otherSimdSize    = 0;
-        CorInfoType  otherBaseJitType = getBaseJitTypeAndSizeOfSIMDType(sigReader.op3ClsHnd, &otherSimdSize);
-        var_types    otherBaseType    = JitType2PreciseVarType(otherBaseJitType);
-
-        assert(otherBaseJitType == simdBaseJitType);
-
-        int immLowerBound2 = 0;
-        int immUpperBound2 = 0;
-
-        HWIntrinsicInfo::lookupImmBounds(intrinsic, otherSimdSize, otherBaseType, &immLowerBound2, &immUpperBound2);
+    GenTree* immOp1          = nullptr;
+    GenTree* immOp2          = nullptr;
+    int      immLowerBound   = 0;
+    int      immUpperBound   = 0;
+    bool     hasFullRangeImm = false;
+    bool     useFallback     = false;
 
-        const int immVal2 = (int)immOp2->AsIntCon()->IconValue();
+    getHWIntrinsicImmOps(intrinsic, sig, &immOp1, &immOp2);
 
-        if ((immVal2 < immLowerBound2) || (immVal2 > immUpperBound2))
+    // Validate the second immediate
+#ifdef TARGET_ARM64
+    if (immOp2 != nullptr)
+    {
+        unsigned  immSimdSize     = simdSize;
+        var_types immSimdBaseType = simdBaseType;
+        getHWIntrinsicImmTypes(intrinsic, sig, 2, simdBaseType, simdBaseJitType, sigReader.op2ClsHnd,
+                               sigReader.op3ClsHnd, &immSimdSize, &immSimdBaseType);
+        HWIntrinsicInfo::lookupImmBounds(intrinsic, immSimdSize, immSimdBaseType, 2, &immLowerBound, &immUpperBound);
+
+        if (!CheckHWIntrinsicImmRange(intrinsic, simdBaseJitType, immOp2, mustExpand, immLowerBound, immUpperBound,
+                                      false, &useFallback))
         {
-            assert(!mustExpand);
-            return nullptr;
+            return useFallback ? impNonConstFallback(intrinsic, retType, simdBaseJitType) : nullptr;
         }
     }
-    else
+#else
+    assert(immOp2 == nullptr);
 #endif
-        if ((sig->numArgs > 0) && HWIntrinsicInfo::isImmOp(intrinsic, impStackTop().val))
-    {
-        // NOTE: The following code assumes that for all intrinsics
-        // taking an immediate operand, that operand will be last.
-        immOp = impStackTop().val;
-    }
-
-    const unsigned simdSize = HWIntrinsicInfo::lookupSimdSize(this, intrinsic, sig);
-
-    int  immLowerBound   = 0;
-    int  immUpperBound   = 0;
-    bool hasFullRangeImm = false;
 
-    if (immOp != nullptr)
+    // Validate the first immediate
+    if (immOp1 != nullptr)
     {
-#ifdef TARGET_XARCH
+#ifdef TARGET_ARM64
+        unsigned  immSimdSize     = simdSize;
+        var_types immSimdBaseType = simdBaseType;
+        getHWIntrinsicImmTypes(intrinsic, sig, 1, simdBaseType, simdBaseJitType, sigReader.op2ClsHnd,
+                               sigReader.op3ClsHnd, &immSimdSize, &immSimdBaseType);
+        HWIntrinsicInfo::lookupImmBounds(intrinsic, immSimdSize, immSimdBaseType, 1, &immLowerBound, &immUpperBound);
+#else
         immUpperBound   = HWIntrinsicInfo::lookupImmUpperBound(intrinsic);
         hasFullRangeImm = HWIntrinsicInfo::HasFullRangeImm(intrinsic);
-#elif defined(TARGET_ARM64)
-        if (category == HW_Category_SIMDByIndexedElement)
-        {
-            CorInfoType  indexedElementBaseJitType;
-            var_types    indexedElementBaseType;
-            unsigned int indexedElementSimdSize = 0;
-
-            if (numArgs == 3)
-            {
-                indexedElementBaseJitType =
-                    getBaseJitTypeAndSizeOfSIMDType(sigReader.op2ClsHnd, &indexedElementSimdSize);
-                indexedElementBaseType = JitType2PreciseVarType(indexedElementBaseJitType);
-            }
-            else
-            {
-                assert(numArgs == 4);
-                indexedElementBaseJitType =
-                    getBaseJitTypeAndSizeOfSIMDType(sigReader.op3ClsHnd, &indexedElementSimdSize);
-                indexedElementBaseType = JitType2PreciseVarType(indexedElementBaseJitType);
-
-                if (intrinsic == NI_Dp_DotProductBySelectedQuadruplet)
-                {
-                    assert(((simdBaseType == TYP_INT) && (indexedElementBaseType == TYP_BYTE)) ||
-                           ((simdBaseType == TYP_UINT) && (indexedElementBaseType == TYP_UBYTE)));
-                    // The second source operand of sdot, udot instructions is an indexed 32-bit element.
-                    indexedElementBaseJitType = simdBaseJitType;
-                    indexedElementBaseType    = simdBaseType;
-                }
-            }
-
-            assert(indexedElementBaseType == simdBaseType);
-            HWIntrinsicInfo::lookupImmBounds(intrinsic, indexedElementSimdSize, simdBaseType, &immLowerBound,
-                                             &immUpperBound);
-        }
-        else
-        {
-            HWIntrinsicInfo::lookupImmBounds(intrinsic, simdSize, simdBaseType, &immLowerBound, &immUpperBound);
-        }
 #endif
 
-        if (!hasFullRangeImm && immOp->IsCnsIntOrI())
+        if (!CheckHWIntrinsicImmRange(intrinsic, simdBaseJitType, immOp1, mustExpand, immLowerBound, immUpperBound,
+                                      hasFullRangeImm, &useFallback))
         {
-            const int ival = (int)immOp->AsIntCon()->IconValue();
-            bool      immOutOfRange;
-#ifdef TARGET_XARCH
-            if (HWIntrinsicInfo::isAVX2GatherIntrinsic(intrinsic))
-            {
-                immOutOfRange = (ival != 1) && (ival != 2) && (ival != 4) && (ival != 8);
-            }
-            else
-#endif
-            {
-                immOutOfRange = (ival < immLowerBound) || (ival > immUpperBound);
-            }
-
-            if (immOutOfRange)
-            {
-                assert(!mustExpand);
-                // The imm-HWintrinsics that do not accept all imm8 values may throw
-                // ArgumentOutOfRangeException when the imm argument is not in the valid range
-                return nullptr;
-            }
-        }
-        else if (!immOp->IsCnsIntOrI())
-        {
-            if (HWIntrinsicInfo::NoJmpTableImm(intrinsic))
-            {
-                return impNonConstFallback(intrinsic, retType, simdBaseJitType);
-            }
-#if defined(TARGET_XARCH)
-            else if (HWIntrinsicInfo::MaybeNoJmpTableImm(intrinsic))
-            {
-#if defined(TARGET_X86)
-                var_types simdBaseType = JitType2PreciseVarType(simdBaseJitType);
-
-                if (varTypeIsLong(simdBaseType))
-                {
-                    if (!mustExpand)
-                    {
-                        return nullptr;
-                    }
-                }
-                else
-#endif // TARGET_XARCH
-                {
-                    return impNonConstFallback(intrinsic, retType, simdBaseJitType);
-                }
-            }
-#endif // TARGET_XARCH
-            else if (!mustExpand)
-            {
-                // When the imm-argument is not a constant and we are not being forced to expand, we need to
-                // return nullptr so a GT_CALL to the intrinsic method is emitted instead. The
-                // intrinsic method is recursive and will be forced to expand, at which point
-                // we emit some less efficient fallback code.
-                return nullptr;
-            }
+            return useFallback ? impNonConstFallback(intrinsic, retType, simdBaseJitType) : nullptr;
         }
     }
 

src/coreclr/jit/hwintrinsic.h

@@ -514,7 +514,7 @@ struct HWIntrinsicInfo
     static int lookupImmUpperBound(NamedIntrinsic intrinsic);
 #elif defined(TARGET_ARM64)
     static void lookupImmBounds(
-        NamedIntrinsic intrinsic, int simdSize, var_types baseType, int* lowerBound, int* upperBound);
+        NamedIntrinsic intrinsic, int simdSize, var_types baseType, int immNumber, int* lowerBound, int* upperBound);
 #else
 #error Unsupported platform
 #endif