[RyuJIT Wasm] Handle all codegen cases for PackedSimd.Shuffle

src/coreclr/jit/codegenwasm.cpp

@@ -2272,7 +2272,19 @@ void CodeGen::genRangeCheck(GenTree* tree)
     assert(tree->OperIs(GT_BOUNDS_CHECK));
     GenTreeBoundsChk* boundsCheck = tree->AsBoundsChk();
     genConsumeOperands(boundsCheck);
-    GetEmitter()->emitIns(INS_I_ge_u);
+#ifdef FEATURE_SIMD
+    if (varTypeIsSIMD(boundsCheck->GetIndex()->TypeGet()))
+    {
+        GetEmitter()->emitIns(INS_i8x16_splat);
+        GetEmitter()->emitIns(INS_i8x16_ge_u);
+        GetEmitter()->emitIns(INS_v128_any_true);
+    }
+    else
+#endif
+    {
+        GetEmitter()->emitIns(INS_I_ge_u);
+    }
+
     genJumpToThrowHlpBlk(boundsCheck->gtThrowKind);
 }
 

src/coreclr/jit/hwintrinsiccodegenwasm.cpp

@@ -26,9 +26,6 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 //
 void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
 {
-    // emitIns_Lane
-    // emitIns_Memarg_Lane
-
     const HWIntrinsic info(node);
     genConsumeMultiOpOperands(node);
 
@@ -40,7 +37,12 @@ void CodeGen::genHWIntrinsic(GenTreeHWIntrinsic* node)
         {
             case HW_Category_SIMD:
             {
-                if ((info.id == NI_PackedSimd_Swizzle) && node->Op(2)->isContained())
+                if (info.id == NI_PackedSimd_Shuffle)
+                {
+                    assert(node->Op(3)->isContained());
+                    GetEmitter()->emitIns_V128Imm(ins, node->Op(3)->AsVecCon()->gtSimdVal.u8);
+                }
+                else if ((info.id == NI_PackedSimd_Swizzle) && node->Op(2)->isContained())
                 {
                     // A constant, fully in-range mask was lowered to an immediate i8x16.shuffle.
                     // prior codegen left the source on the value stack once (the mask

src/coreclr/jit/lower.h

@@ -512,6 +512,7 @@ class Lowering final : public Phase
 #elif defined(TARGET_WASM)
     GenTree* LowerHWIntrinsicCompareUnsignedLong(GenTreeHWIntrinsic* node);
     GenTree* LowerHWIntrinsicWithImm(GenTreeHWIntrinsic* node);
+    GenTree* LowerHWIntrinsicNativeShuffle(GenTreeHWIntrinsic* node);
     void     LowerHWIntrinsicSwizzle(GenTreeHWIntrinsic* node);
 #endif // !TARGET_XARCH && !TARGET_ARM64
     GenTree* InsertNewSimdCreateScalarUnsafeNode(var_types type,

src/coreclr/jit/lowerwasm.cpp

@@ -921,6 +921,11 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node)
             return LowerHWIntrinsicWithImm(node);
         }
 
+        case NI_PackedSimd_Shuffle:
+        {
+            return LowerHWIntrinsicNativeShuffle(node);
+        }
+
         case NI_PackedSimd_LoadScalarAndSplatVector128:
         case NI_PackedSimd_LoadScalarVector128:
         case NI_PackedSimd_LoadWideningVector128:
@@ -962,6 +967,8 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node)
 GenTree* Lowering::LowerHWIntrinsicWithImm(GenTreeHWIntrinsic* node)
 {
     GenTree* immOp = node->GetImmOp();
+    assert(varTypeIsIntegral(immOp->TypeGet()) && "Immediate operand must be an integral type");
+
     if (!immOp->IsCnsIntOrI())
     {
         // This node has a non-constant immediate operand, so it will need a jump table
@@ -1316,6 +1323,169 @@ GenTree* Lowering::LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node)
     return LowerNode(node);
 }
 
+// --------------------------------------------------------------------------------
+// LowerHWIntrinsicNativeShuffle: Lowers a PackedSimd Shuffle call with a possibly non-constant mask
+//
+// Arguments:
+//    node - The hardware intrinsic node.
+//
+// Notes:
+//  If the shuffle mask is a constant vector, it can be contained as an immediate and emitted. Otherwise,
+//  the shuffle is rewritten into two swizzles for the upper and lower input vectors and combined into the final result.
+GenTree* Lowering::LowerHWIntrinsicNativeShuffle(GenTreeHWIntrinsic* node)
+{
+    assert(node->GetHWIntrinsicId() == NI_PackedSimd_Shuffle);
+
+    GenTree*  op1         = node->Op(1);
+    GenTree*  shuffleMask = node->Op(3);
+    var_types resultType  = node->TypeGet();
+
+    // No extra work to do if the shuffle is an in-range constant vector, it can be contained as an immediate and
+    // emitted.
+    if (shuffleMask->IsCnsVec())
+    {
+        const simd_t& mask       = shuffleMask->AsVecCon()->gtSimdVal;
+        bool          allInRange = true;
+        for (int i = 0; i < 16; i++)
+        {
+            if (mask.u8[i] >= 32)
+            {
+                allInRange = false;
+                break;
+            }
+        }
+        if (allInRange)
+        {
+            ContainCheckHWIntrinsic(node);
+            return node->gtNext;
+        }
+    }
+
+    // If the shuffle mask is not a constant vector or the mask is not in range, we will need to rewrite the shuffle
+    // into a BOUNDS_CHECK around the mask, and two swizzles: 1 to handle elements from the first vector, and 1 to
+    // handle elements from the second vector. The two swizzles will then be or'd together to produce the final result.
+    // We will be constructing IR like the following:
+    //  op1             = ...
+    //  op2             = ...
+    //                  /--* op2
+    //                  /--* LCL_VAR op2Tmp
+    //                  STORE_LCL_VAR op2Tmp
+    // originalMask     = ...
+    //                  /--* originalMask
+    //                  /--* LCL_VAR originalMaskTmp
+    //                  STORE_LCL_VAR originalMaskTmp
+    //  m0        = *   LCL_VAR originalMaskTmp
+    //  b0        = *   CNS_INT       int    32
+    //                  /--* m0 simd
+    //                  +--* b0  int
+    //                  GT_BOUNDS_CHECK RNG_CHK_FAIL
+    //  m1        = *   LCL_VAR originalMaskTmp
+    //                  /--*  op1 simd
+    //                  +--*  m1  simd
+    //   tmp1      = *  HWINTRINSIC   simd   byte    PackedSimd.Swizzle
+    //   op2Reload = *  LCL_VAR op2tmp
+    //   m2        = *  LCL_VAR originalMaskTmp
+    //   upperBnd  = *  CNS_VEC      simd   byte    <0x10, 0x10, ...>
+    //                   /--*  m2 simd
+    //                   +--*  upperBnd simd
+    //   upperMask = *  HWINTRINSIC   simd    byte    PackedSimd.Subtract
+    //                   /--*  op2Reload simd
+    //                   +--*  upperMask simd
+    //   tmp3      = *  HWINTRINSIC   simd   byte    PackedSimd.Swizzle
+    //                   /--*  tmp1 simd
+    //                   +--*  tmp3 simd
+    //   res       =  *  HWINTRINSIC   simd   byte    PackedSimd.Or
+    //
+    // This is roughly equivalent to the following C#:
+    // ...
+    // if (GreaterThanAny(originalMask, 31)) { throw new ArrayIndexOutOfBoundsException(); }
+    // tmp1 = PackedSimd.Swizzle(op1, originalMask);
+    // tmp2 = PackedSimd.Swizzle(op2, PackedSimd.Subtract(originalMask, PackedSimd.Splat(0x10)));
+    // result  = PackedSimd.Or(tmp1, tmp2);
+    //...
+
+    // op2 needs to be moved, replace with a local, but don't immediately reload
+    LIR::Use op2Use(BlockRange(), &node->Op(2), node);
+    op2Use.ReplaceWithLclVar(m_compiler);
+    GenTree* op2Reload = node->Op(2);
+    BlockRange().Remove(op2Reload);
+
+    // Shuffle mask will be used several times, replace with a local
+    LIR::Use     shuffleMaskUse(BlockRange(), &node->Op(3), node);
+    unsigned int shuffleMaskTmp   = shuffleMaskUse.ReplaceWithLclVar(m_compiler);
+    GenTree*     maskReloadForChk = node->Op(3);
+
+    // Insert a bounds check for the shuffle mask. A bounds check against a simd value will be handled by codegen as a
+    // GreaterThanAny(vec, bound) type operation.
+    GenTree* bound = m_compiler->gtNewIconNode(32);
+    BlockRange().InsertBefore(node, bound);
+    LowerNode(bound);
+
+    GenTree* boundsCheck =
+        new (m_compiler, GT_BOUNDS_CHECK) GenTreeBoundsChk(maskReloadForChk, bound, SCK_ARG_RNG_EXCPN);
+    BlockRange().InsertBefore(node, boundsCheck);
+    LowerNode(boundsCheck);
+
+    // Do a swizzle of the first vector with the original shuffle mask (now loaded from a local), which will produce a
+    // vector with the elements from the first vector in the correct order, and zero's for all elements which correspond
+    // to the second vector.
+    GenTree* maskReload1 = m_compiler->gtNewLclVarNode(shuffleMaskTmp, shuffleMask->TypeGet());
+    BlockRange().InsertBefore(node, maskReload1);
+    LowerNode(maskReload1);
+
+    GenTree* swizzle1 =
+        m_compiler->gtNewSimdHWIntrinsicNode(resultType, op1, maskReload1, NI_PackedSimd_Swizzle, TYP_BYTE, 16);
+    BlockRange().InsertBefore(node, swizzle1);
+    LowerNode(swizzle1);
+
+    // Re-load op2
+    BlockRange().InsertBefore(node, op2Reload);
+    LowerNode(op2Reload);
+
+    // Re-load the original shuffle mask
+    GenTreeLclVar* maskReload2 = m_compiler->gtNewLclVarNode(shuffleMaskTmp, shuffleMask->TypeGet());
+    BlockRange().InsertBefore(node, maskReload2);
+    LowerNode(maskReload2);
+
+    // Create constant upper bound vector of <16, 16, ..., 16>
+    GenTreeVecCon* upperBound = m_compiler->gtNewVconNode(shuffleMask->TypeGet());
+    upperBound->EvaluateBroadcastInPlace(TYP_BYTE, static_cast<int64_t>(16));
+    BlockRange().InsertBefore(node, upperBound);
+    LowerNode(upperBound);
+
+    // Use the above to subtract 16 from each mask element to mark each element which corresponds to the lower vector as
+    // unused, leading to a zero in the result of the swizzle.
+    GenTree* upperMask = m_compiler->gtNewSimdHWIntrinsicNode(shuffleMask->TypeGet(), maskReload2, upperBound,
+                                                              NI_PackedSimd_Subtract, TYP_BYTE, 16);
+    BlockRange().InsertBefore(node, upperMask);
+    LowerNode(upperMask);
+
+    GenTree* swizzle2 =
+        m_compiler->gtNewSimdHWIntrinsicNode(resultType, op2Reload, upperMask, NI_PackedSimd_Swizzle, TYP_BYTE, 16);
+    BlockRange().InsertBefore(node, swizzle2);
+    LowerNode(swizzle2);
+
+    // Since we've left zero's for all the elements which correspond to the upper vector, we can just or the two
+    // swizzles together to get the final result.
+    GenTreeHWIntrinsic* result =
+        m_compiler->gtNewSimdHWIntrinsicNode(resultType, swizzle1, swizzle2, NI_PackedSimd_Or, TYP_BYTE, 16);
+    BlockRange().InsertBefore(node, result);
+
+    LIR::Use use;
+    if (BlockRange().TryGetUse(node, &use))
+    {
+        use.ReplaceWith(result);
+    }
+    else
+    {
+        result->SetUnusedValue();
+    }
+
+    BlockRange().Remove(node);
+
+    return LowerNode(result);
+}
+
 //----------------------------------------------------------------------------------------------
 // ContainCheckHWIntrinsic: Perform containment analysis for a hardware intrinsic node.
 //
@@ -1333,4 +1503,11 @@ void Lowering::ContainCheckHWIntrinsic(GenTreeHWIntrinsic* node)
             MakeSrcContained(node, immOp);
         }
     }
+    else if (intrinsicId == NI_PackedSimd_Shuffle && node->Op(3)->IsCnsVec())
+    {
+        // Shuffle is a special case where the mask is a constant vector immediate
+        // which must be contained. If the mask is non-constant we will re-write the shuffle into a fallback equivalent
+        // operation. (see LowerHWIntrinsicNativeShuffle).
+        MakeSrcContained(node, node->Op(3));
+    }
 }