Adding support for the SSE Shuffle intrinsic

Author: tannergooding

src/jit/compiler.h

@@ -2065,6 +2065,13 @@ class Compiler
         var_types type, GenTree* op1, NamedIntrinsic hwIntrinsicID, var_types baseType, unsigned size);
     GenTreeHWIntrinsic* gtNewSimdHWIntrinsicNode(
         var_types type, GenTree* op1, GenTree* op2, NamedIntrinsic hwIntrinsicID, var_types baseType, unsigned size);
+    GenTreeHWIntrinsic* gtNewSimdHWIntrinsicNode(var_types      type,
+                                                 GenTree*       op1,
+                                                 GenTree*       op2,
+                                                 GenTree*       op3,
+                                                 NamedIntrinsic hwIntrinsicID,
+                                                 var_types      baseType,
+                                                 unsigned       size);
     GenTreeHWIntrinsic* gtNewSimdHWIntrinsicNode(var_types      type,
                                                  GenTree*       op1,
                                                  GenTree*       op2,

src/jit/gentree.cpp

@@ -17921,6 +17921,18 @@ GenTreeHWIntrinsic* Compiler::gtNewSimdHWIntrinsicNode(
     return new (this, GT_HWIntrinsic) GenTreeHWIntrinsic(type, op1, op2, hwIntrinsicID, baseType, size);
 }
 
+GenTreeHWIntrinsic* Compiler::gtNewSimdHWIntrinsicNode(var_types      type,
+                                                       GenTree*       op1,
+                                                       GenTree*       op2,
+                                                       GenTree*       op3,
+                                                       NamedIntrinsic hwIntrinsicID,
+                                                       var_types      baseType,
+                                                       unsigned       size)
+{
+    return new (this, GT_HWIntrinsic)
+        GenTreeHWIntrinsic(type, gtNewArgList(op1, op2, op3), hwIntrinsicID, baseType, size);
+}
+
 GenTreeHWIntrinsic* Compiler::gtNewSimdHWIntrinsicNode(var_types      type,
                                                        GenTree*       op1,
                                                        GenTree*       op2,

src/jit/hwintrinsiccodegenxarch.cpp

@@ -353,6 +353,95 @@ void CodeGen::genSSEIntrinsic(GenTreeHWIntrinsic* node)
             emit->emitIns_SIMD_R_R_R(INS_xorps, targetReg, targetReg, targetReg, TYP_SIMD16);
             break;
 
+        case NI_SSE_Shuffle:
+        {
+            GenTreeArgList* argList;
+
+            // Shuffle takes 3 operands, so op1 should be an arg list with two
+            // additional node in the chain.
+            assert(baseType == TYP_FLOAT);
+            assert(op1->OperIsList());
+            assert(op1->AsArgList()->Rest() != nullptr);
+            assert(op1->AsArgList()->Rest()->Rest() != nullptr);
+            assert(op1->AsArgList()->Rest()->Rest()->Rest() == nullptr);
+            assert(op2 == nullptr);
+
+            argList = op1->AsArgList();
+            op1     = argList->Current();
+            op1Reg  = op1->gtRegNum;
+            genConsumeRegs(op1);
+
+            argList = argList->Rest();
+            op2     = argList->Current();
+            op2Reg  = op2->gtRegNum;
+            genConsumeRegs(op2);
+
+            argList = argList->Rest();
+            op3     = argList->Current();
+            genConsumeRegs(op3);
+
+            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;
+        }
+
         case NI_SSE_Sqrt:
             assert(baseType == TYP_FLOAT);
             assert(op2 == nullptr);

src/jit/hwintrinsicxarch.cpp

@@ -468,11 +468,35 @@ 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;
         }
 
+        case NI_SSE_Shuffle:
+            assert(sig->numArgs == 3);
+            assert(getBaseTypeOfSIMDType(sig->retTypeSigClass) == TYP_FLOAT);
+
+            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;
+
         case NI_SSE_Add:
         case NI_SSE_And:
         case NI_SSE_AndNot:

src/jit/lowerxarch.cpp

@@ -2337,6 +2337,18 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
             break;
         }
 
+        case NI_SSE_Shuffle:
+        {
+            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));
             break;

src/jit/lsraxarch.cpp

@@ -2534,13 +2534,34 @@ void LinearScan::TreeNodeInfoInitHWIntrinsic(GenTreeHWIntrinsic* intrinsicTree,
             info->srcCount += GetOperandInfo(op1);
         }
     }
+
     if (op2 != nullptr)
     {
         info->srcCount += GetOperandInfo(op2);
     }
 
     switch (intrinsicID)
     {
+        case NI_SSE_Shuffle:
+        {
+            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;
+        }
+
 #ifdef _TARGET_X86_
         case NI_SSE42_Crc32:
         {