Updating Sse.Shuffle to have a compiler fallback for indirect calls

Author: tannergooding

src/jit/hwintrinsiccodegenxarch.cpp

@@ -517,7 +517,6 @@ void CodeGen::genSSEIntrinsic(GenTreeHWIntrinsic* node)
         case NI_SSE_Shuffle:
         {
             GenTreeArgList* argList;
-            ssize_t         ival;
 
             // Shuffle takes 3 operands, so op1 should be an arg list with two
             // additional node in the chain.
@@ -540,10 +539,67 @@ void CodeGen::genSSEIntrinsic(GenTreeHWIntrinsic* node)
 
             argList = argList->Rest();
             op3     = argList->Current();
-            ival    = op3->AsIntConCommon()->IconValue();
             genConsumeRegs(op3);
 
-            emit->emitIns_SIMD_R_R_R_I(INS_shufps, targetReg, op1Reg, op2Reg, (int)ival, TYP_SIMD16);
+            if (op3->IsCnsIntOrI())
+            {
+                ssize_t ival = op3->AsIntConCommon()->IconValue();
+                emit->emitIns_SIMD_R_R_R_I(INS_shufps, targetReg, op1Reg, op2Reg, (int)ival, TYP_SIMD16);
+            }
+            else
+            {
+                // We emit a fallback case for the scenario when op3 is not a constant. This should normally
+                // happen when the intrinsic is called indirectly, such as via Reflection. However, it can
+                // also occur if the consumer calls it directly and just doesn't pass a constant value.
+
+                const unsigned jmpCount = 256;
+                BasicBlock*    jmpTable[jmpCount];
+
+                unsigned jmpTableBase = emit->emitBBTableDataGenBeg(jmpCount, true);
+                unsigned jmpTableOffs = 0;
+
+                // Emit the jump table
+
+                JITDUMP("\n      J_M%03u_DS%02u LABEL   DWORD\n", Compiler::s_compMethodsCount, jmpTableBase);
+
+                for (unsigned i = 0; i < jmpCount; i++)
+                {
+                    jmpTable[i] = genCreateTempLabel();
+                    JITDUMP("            DD      L_M%03u_BB%02u\n", Compiler::s_compMethodsCount, jmpTable[i]->bbNum);
+                    emit->emitDataGenData(i, jmpTable[i]);
+                }
+
+                emit->emitDataGenEnd();
+
+                // Compute and jump to the appropriate offset in the switch table
+
+                regNumber baseReg = node->ExtractTempReg();   // the start of the switch table
+                regNumber offsReg = node->GetSingleTempReg(); // the offset into the switch table
+
+                emit->emitIns_R_C(INS_lea, emitTypeSize(TYP_I_IMPL), offsReg, compiler->eeFindJitDataOffs(jmpTableBase),
+                                  0);
+
+                emit->emitIns_R_ARX(INS_mov, EA_4BYTE, offsReg, offsReg, op3->gtRegNum, 4, 0);
+                emit->emitIns_R_L(INS_lea, EA_PTR_DSP_RELOC, compiler->fgFirstBB, baseReg);
+                emit->emitIns_R_R(INS_add, EA_PTRSIZE, offsReg, baseReg);
+                emit->emitIns_R(INS_i_jmp, emitTypeSize(TYP_I_IMPL), offsReg);
+
+                // Emit the switch table entries
+
+                BasicBlock* switchTableBeg = genCreateTempLabel();
+                BasicBlock* switchTableEnd = genCreateTempLabel();
+
+                genDefineTempLabel(switchTableBeg);
+
+                for (unsigned i = 0; i < jmpCount; i++)
+                {
+                    genDefineTempLabel(jmpTable[i]);
+                    emit->emitIns_SIMD_R_R_R_I(INS_shufps, targetReg, op1Reg, op2Reg, i, TYP_SIMD16);
+                    emit->emitIns_J(INS_jmp, switchTableEnd);
+                }
+
+                genDefineTempLabel(switchTableEnd);
+            }
             break;
         }
 

src/jit/hwintrinsicxarch.cpp

@@ -489,7 +489,7 @@ GenTree* Compiler::impSSEIntrinsic(NamedIntrinsic        intrinsic,
             GenTree* left    = gtNewSimdHWIntrinsicNode(TYP_SIMD16, op4, op3, NI_SSE_UnpackLow, TYP_FLOAT, 16);
             GenTree* right   = gtNewSimdHWIntrinsicNode(TYP_SIMD16, op2, op1, NI_SSE_UnpackLow, TYP_FLOAT, 16);
             GenTree* control = gtNewIconNode(68, TYP_UBYTE);
-            
+
             retNode = gtNewSimdHWIntrinsicNode(TYP_SIMD16, left, right, control, NI_SSE_Shuffle, TYP_FLOAT, 16);
             break;
         }
@@ -498,10 +498,25 @@ GenTree* Compiler::impSSEIntrinsic(NamedIntrinsic        intrinsic,
         {
             assert(sig->numArgs == 3);
             assert(getBaseTypeOfSIMDType(sig->retTypeSigClass) == TYP_FLOAT);
-            op3     = impPopStack().val;
-            op2     = impSIMDPopStack(TYP_SIMD16);
-            op1     = impSIMDPopStack(TYP_SIMD16);
-            retNode = gtNewSimdHWIntrinsicNode(TYP_SIMD16, op1, op2, op3, intrinsic, TYP_FLOAT, 16);
+
+            op3 = impStackTop().val;
+
+            if (op3->IsCnsIntOrI() || mustExpand)
+            {
+                impPopStack(); // Pop the value we peeked at
+                op2     = impSIMDPopStack(TYP_SIMD16);
+                op1     = impSIMDPopStack(TYP_SIMD16);
+                retNode = gtNewSimdHWIntrinsicNode(TYP_SIMD16, op1, op2, op3, intrinsic, TYP_FLOAT, 16);
+            }
+            else
+            {
+                // When op3 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;
+            }
             break;
         }
 
@@ -586,7 +601,8 @@ GenTree* Compiler::impSSEIntrinsic(NamedIntrinsic        intrinsic,
             CORINFO_CLASS_HANDLE argClass;
 
             CORINFO_ARG_LIST_HANDLE argLst = info.compCompHnd->getArgNext(sig->args);
-            CorInfoType             corType = strip(info.compCompHnd->getArgType(sig, argLst, &argClass)); // type of the second argument
+            CorInfoType             corType =
+                strip(info.compCompHnd->getArgType(sig, argLst, &argClass)); // type of the second argument
 
             if (varTypeIsLong(corType))
             {

src/jit/lowerxarch.cpp

@@ -2307,10 +2307,16 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
     switch (node->gtHWIntrinsicId)
     {
         case NI_SSE_Shuffle:
-            // third operand is an integer constant and marked as contained.
-            assert(op1->AsArgList()->Rest()->Rest()->Current()->IsCnsIntOrI());
-            MakeSrcContained(node, op1->AsArgList()->Rest()->Rest()->Current());
+        {
+            assert(op1->OperIsList());
+            GenTree* op3 = op1->AsArgList()->Rest()->Rest()->Current();
+
+            if (op3->IsCnsIntOrI())
+            {
+                MakeSrcContained(node, op3);
+            }
             break;
+        }
 
         default:
             assert((intrinsicID > NI_HW_INTRINSIC_START) && (intrinsicID < NI_HW_INTRINSIC_END));

src/jit/lsraxarch.cpp

@@ -2550,10 +2550,24 @@ void LinearScan::TreeNodeInfoInitHWIntrinsic(GenTreeHWIntrinsic* intrinsicTree,
             break;
 
         case NI_SSE_Shuffle:
-            // Third operand should be a contained integer constant
-            assert(info->srcCount == 2);
-            assert(op1->AsArgList()->Rest()->Rest()->Current()->isContainedIntOrIImmed());
+        {
+            assert(op1->OperIsList());
+            GenTree* op3 = op1->AsArgList()->Rest()->Rest()->Current();
+
+            if (!op3->isContainedIntOrIImmed())
+            {
+                assert(!op3->IsCnsIntOrI());
+
+                // We need two extra reg when op3 isn't a constant so
+                // the offset into the jump table for the fallback path
+                // can be computed.
+
+                info->internalIntCount = 2;
+                info->setInternalCandidates(this, allRegs(TYP_INT));
+                break;
+            }
             break;
+        }
 
         case NI_SSE_ConvertToSingle:
         case NI_SSE_StaticCast:

tests/src/JIT/HardwareIntrinsics/X86/Sse/Shuffle.cs

@@ -91,6 +91,22 @@ static unsafe int Main(string[] args)
                         Console.WriteLine();
                         testResult = Fail;
                     }
+                    
+                    // XYZW
+                    vf3 = (Vector128<float>)typeof(Sse).GetMethod(nameof(Sse.Shuffle), new Type[] { vf1.GetType(), vf2.GetType(), typeof(byte) }).Invoke(null, new object[] { vf1, vf2, (byte)(27) });
+                    Unsafe.Write(floatTable.outArrayPtr, vf3);
+
+                    if (!floatTable.CheckResult((x, y, z) => (z[0] == x[3]) && (z[1] == x[2]) &&
+                                                             (z[2] == y[1]) && (z[3] == y[0])))
+                    {
+                        Console.WriteLine("SSE Shuffle failed on float:");
+                        foreach (var item in floatTable.outArray)
+                        {
+                            Console.Write(item + ", ");
+                        }
+                        Console.WriteLine();
+                        testResult = Fail;
+                    }
                 }
             }