JIT: Materialize partially-constant Vector.Create as CNS_VEC + inserts

src/coreclr/jit/gentree.h

@@ -6935,15 +6935,18 @@ struct GenTreeVecCon : public GenTree
     static unsigned ElementCount(unsigned simdSize, var_types simdBaseType);
 
     template <typename simdTypename>
-    static bool IsHWIntrinsicCreateConstant(GenTreeHWIntrinsic* node, simdTypename& simdVal)
+    static bool IsHWIntrinsicCreateConstant(GenTreeHWIntrinsic* node,
+                                            simdTypename&       simdVal,
+                                            simdmask_t*         cnsMask = nullptr)
     {
         NamedIntrinsic intrinsic    = node->GetHWIntrinsicId();
         var_types      simdType     = node->TypeGet();
         var_types      simdBaseType = node->GetSimdBaseType();
         unsigned       simdSize     = node->GetSimdSize();
 
-        size_t argCnt    = node->GetOperandCount();
-        size_t cnsArgCnt = 0;
+        size_t     argCnt    = node->GetOperandCount();
+        size_t     cnsArgCnt = 0;
+        simdmask_t mask      = {};
 
         switch (intrinsic)
         {
@@ -6958,13 +6961,15 @@ struct GenTreeVecCon : public GenTree
                 if ((argCnt == 1) && HandleArgForHWIntrinsicCreate<simdTypename>(node->Op(1), 0, simdVal, simdBaseType))
                 {
                     // CreateScalar leaves the upper bits as zero
+                    mask.u64[0] |= 1;
 
                     if (intrinsic != NI_Vector_CreateScalar)
                     {
                         // Now assign the rest of the arguments.
                         for (unsigned i = 1; i < ElementCount(simdSize, simdBaseType); i++)
                         {
                             HandleArgForHWIntrinsicCreate<simdTypename>(node->Op(1), i, simdVal, simdBaseType);
+                            mask.u64[0] |= (1ULL << i);
                         }
                     }
 
@@ -6976,17 +6981,27 @@ struct GenTreeVecCon : public GenTree
                     {
                         if (HandleArgForHWIntrinsicCreate<simdTypename>(node->Op(i), i - 1, simdVal, simdBaseType))
                         {
+                            mask.u64[0] |= (1ULL << (i - 1));
                             cnsArgCnt++;
                         }
                     }
                 }
 
                 assert((argCnt == 1) || (argCnt == ElementCount(simdSize, simdBaseType)));
+
+                if (cnsMask != nullptr)
+                {
+                    *cnsMask = mask;
+                }
                 return argCnt == cnsArgCnt;
             }
 
             default:
             {
+                if (cnsMask != nullptr)
+                {
+                    *cnsMask = mask;
+                }
                 return false;
             }
         }

src/coreclr/jit/lower.cpp

@@ -13081,6 +13081,151 @@ GenTree* Lowering::InsertNewSimdCreateScalarUnsafeNode(var_types simdType,
     return result;
 }
 
+//----------------------------------------------------------------------------------------------
+// Lowering::NonZeroConstantElementCount: Counts how many of the constant elements of a partially
+//    constant Vector Create node have a bit pattern that is not all-bits-zero.
+//
+//  Arguments:
+//    simdVal      - The constant value, with the constant elements populated and the non-constant
+//                   elements left as zero.
+//    cnsMask      - A mask with a bit set for each element that is a constant.
+//    simdBaseType - The base type of the vector.
+//
+// Returns:
+//    The number of constant elements whose bit pattern is not all-bits-zero.
+//
+// Remarks:
+//    This checks the raw bytes rather than the numeric value on purpose. An all-bits-zero lane is
+//    effectively free to produce (for example, insertps can zero lanes as part of another insert,
+//    and CreateScalarUnsafe zero-extends the upper elements), so such lanes should not count towards
+//    the profitability of materializing a vector constant.
+//
+//    A floating-point -0.0 has its sign bit set, so it is not all-bits-zero and must be counted:
+//    the free zeroing paths would produce +0.0, which is a different value, so a -0.0 lane genuinely
+//    requires a materialized constant. Do not "simplify" this to a numeric == 0 check.
+//
+unsigned Lowering::NonZeroConstantElementCount(const simd_t* simdVal, simdmask_t cnsMask, var_types simdBaseType)
+{
+    unsigned       elementSize = genTypeSize(simdBaseType);
+    uint64_t       maskBits    = cnsMask.GetRawBits();
+    const uint8_t* bytes       = reinterpret_cast<const uint8_t*>(simdVal);
+    unsigned       count       = 0;
+
+    while (maskBits != 0)
+    {
+        unsigned index = BitOperations::BitScanForward(maskBits);
+        maskBits &= (maskBits - 1);
+
+        unsigned base = index * elementSize;
+
+        // Treat the lane as non-zero if any byte is set. This keeps -0.0 (sign bit only) counted,
+        // since the free zeroing paths would otherwise turn it into +0.0.
+        for (unsigned i = 0; i < elementSize; i++)
+        {
+            if (bytes[base + i] != 0)
+            {
+                count++;
+                break;
+            }
+        }
+    }
+
+    return count;
+}
+
+//----------------------------------------------------------------------------------------------
+// Lowering::LowerHWIntrinsicCreateWithInserts: Lowers a Vector Create node whose operands are
+//    partially constant by materializing the constant operands as a vector constant and inserting
+//    the remaining non-constant operands into it.
+//
+//  Arguments:
+//    node    - The Create intrinsic node being lowered. This must represent a single 128-bit (or
+//              smaller) lane, as larger vectors are split into per-lane Create nodes before reaching
+//              here.
+//    simdVal - The constant value, with the constant elements populated and the non-constant
+//              elements left as zero.
+//    cnsMask - A mask with a bit set for each operand that is a constant. At least two bits must be
+//              set (otherwise materializing a constant is not worthwhile).
+//
+// Returns:
+//    The next node to lower.
+//
+// Remarks:
+//    For example, Vector128.Create(1, 2, 3, x) is turned into Vector128.Create(1, 2, 3, 0) (a single
+//    CNS_VEC) with a single WithElement inserting x, rather than a chain of four inserts.
+//
+GenTree* Lowering::LowerHWIntrinsicCreateWithInserts(GenTreeHWIntrinsic* node,
+                                                     const simd_t*       simdVal,
+                                                     simdmask_t          cnsMask)
+{
+    var_types simdType     = node->TypeGet();
+    var_types simdBaseType = node->GetSimdBaseType();
+    unsigned  simdSize     = node->GetSimdSize();
+    size_t    argCnt       = node->GetOperandCount();
+
+    if ((simdSize == 8) && (simdType == TYP_DOUBLE))
+    {
+        // TODO-Cleanup: Struct retyping means we have the wrong type here. We need to
+        //               manually fix it up so the simdType checks below are correct.
+        simdType = TYP_SIMD8;
+    }
+
+    uint64_t maskBits = cnsMask.GetRawBits();
+    assert(BitOperations::PopCount(maskBits) >= 2);
+
+    // Materialize the constant elements as a vector constant. The non-constant operands
+    // are then inserted into it below.
+    GenTreeVecCon* vecCon = m_compiler->gtNewVconNode(simdType);
+    memcpy(&vecCon->gtSimdVal, simdVal, simdSize);
+    BlockRange().InsertBefore(node, vecCon);
+
+    GenTree* result = vecCon;
+
+    for (size_t i = 1; i <= argCnt; i++)
+    {
+        GenTree* opN = node->Op(i);
+
+        if ((maskBits & (1ULL << (i - 1))) != 0)
+        {
+            // This operand is a constant and is already part of the vector constant.
+#if !defined(TARGET_64BIT)
+            if (opN->OperIsLong())
+            {
+                BlockRange().Remove(opN->gtGetOp1());
+                BlockRange().Remove(opN->gtGetOp2());
+            }
+#endif // !TARGET_64BIT
+            BlockRange().Remove(opN);
+            continue;
+        }
+
+        GenTree* idx = m_compiler->gtNewIconNode(i - 1, TYP_INT);
+
+        // Place the insert as early as possible to avoid creating a lot of long lifetimes.
+        GenTree* insertionPoint = LIR::LastNode(result, opN);
+
+        GenTree* insert = m_compiler->gtNewSimdWithElementNode(simdType, result, idx, opN, simdBaseType, simdSize);
+        BlockRange().InsertAfter(insertionPoint, idx, insert);
+
+        result = insert;
+        LowerNode(insert);
+    }
+
+    LIR::Use use;
+    if (BlockRange().TryGetUse(node, &use))
+    {
+        use.ReplaceWith(result);
+    }
+    else
+    {
+        result->SetUnusedValue();
+    }
+
+    GenTree* next = node->gtNext;
+    BlockRange().Remove(node);
+    return next;
+}
+
 //----------------------------------------------------------------------------------------------
 // Lowering::NormalizeIndexToNativeSized:
 //   Prepare to use an index for address calculations by ensuring it is native sized.

src/coreclr/jit/lower.h

@@ -498,8 +498,10 @@ class Lowering final : public Phase
     void     LowerHWIntrinsicCC(GenTreeHWIntrinsic* node, NamedIntrinsic newIntrinsicId, GenCondition condition);
     GenTree* LowerHWIntrinsicCmpOp(GenTreeHWIntrinsic* node, genTreeOps cmpOp);
     GenTree* LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node);
-    GenTree* LowerHWIntrinsicDot(GenTreeHWIntrinsic* node);
-    GenTree* LowerHWIntrinsicCndSel(GenTreeHWIntrinsic* node);
+    GenTree* LowerHWIntrinsicCreateWithInserts(GenTreeHWIntrinsic* node, const simd_t* simdVal, simdmask_t cnsMask);
+    static unsigned NonZeroConstantElementCount(const simd_t* simdVal, simdmask_t cnsMask, var_types simdBaseType);
+    GenTree*        LowerHWIntrinsicDot(GenTreeHWIntrinsic* node);
+    GenTree*        LowerHWIntrinsicCndSel(GenTreeHWIntrinsic* node);
 #if defined(TARGET_XARCH)
     void     LowerFusedMultiplyOp(GenTreeHWIntrinsic* node);
     GenTree* LowerHWIntrinsicToScalar(GenTreeHWIntrinsic* node);

src/coreclr/jit/lowerarmarch.cpp

@@ -2247,9 +2247,10 @@ GenTree* Lowering::LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node)
     assert(varTypeIsArithmetic(simdBaseType));
     assert(simdSize != 0);
 
-    bool   isConstant     = GenTreeVecCon::IsHWIntrinsicCreateConstant<simd_t>(node, simdVal);
-    bool   isCreateScalar = HWIntrinsicInfo::IsVectorCreateScalar(intrinsicId);
-    size_t argCnt         = node->GetOperandCount();
+    simdmask_t cnsMask        = {};
+    bool       isConstant     = GenTreeVecCon::IsHWIntrinsicCreateConstant<simd_t>(node, simdVal, &cnsMask);
+    bool       isCreateScalar = HWIntrinsicInfo::IsVectorCreateScalar(intrinsicId);
+    size_t     argCnt         = node->GetOperandCount();
 
     // Check if we have a cast that we can remove. Note that "IsValidConstForMovImm"
     // will reset Op(1) if it finds such a cast, so we do not need to handle it here.
@@ -2338,6 +2339,16 @@ GenTree* Lowering::LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node)
     //          +--*  opN T
     //   node = *  HWINTRINSIC   simd   T Create
 
+    // If two or more of the operands are constants that are not all-bits-zero, we can materialize
+    // them as a vector constant and insert the remaining non-constant operands into it. This is both
+    // fewer nodes and typically cheaper than a chain of inserts starting from CreateScalarUnsafe. We
+    // only consider such constants because all-bits-zero lanes are effectively free to produce, so
+    // materializing them into a constant can regress.
+    if (NonZeroConstantElementCount(&simdVal, cnsMask, simdBaseType) >= 2)
+    {
+        return LowerHWIntrinsicCreateWithInserts(node, &simdVal, cnsMask);
+    }
+
     // We will be constructing the following parts:
     //          /--*  op1  T
     //   tmp1 = *  HWINTRINSIC   simd8  T CreateScalarUnsafe

src/coreclr/jit/lowerxarch.cpp

@@ -4031,9 +4031,10 @@ GenTree* Lowering::LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node)
     GenTree* tmp2 = nullptr;
     GenTree* tmp3 = nullptr;
 
-    bool   isConstant     = GenTreeVecCon::IsHWIntrinsicCreateConstant<simd_t>(node, simdVal);
-    bool   isCreateScalar = HWIntrinsicInfo::IsVectorCreateScalar(intrinsicId);
-    size_t argCnt         = node->GetOperandCount();
+    simdmask_t cnsMask        = {};
+    bool       isConstant     = GenTreeVecCon::IsHWIntrinsicCreateConstant<simd_t>(node, simdVal, &cnsMask);
+    bool       isCreateScalar = HWIntrinsicInfo::IsVectorCreateScalar(intrinsicId);
+    size_t     argCnt         = node->GetOperandCount();
 
     if (isConstant)
     {
@@ -4512,6 +4513,17 @@ GenTree* Lowering::LowerHWIntrinsicCreate(GenTreeHWIntrinsic* node)
 
     assert(simdType == TYP_SIMD16);
 
+    // If two or more of the operands are constants that are not all-bits-zero, we can materialize
+    // them as a vector constant and insert the remaining non-constant operands into it. This is both
+    // fewer nodes and typically cheaper than a chain of inserts starting from CreateScalarUnsafe. We
+    // only consider such constants because all-bits-zero lanes are effectively free to produce (for
+    // example, insertps can zero lanes as part of another insert), so materializing them into a
+    // constant can regress.
+    if (NonZeroConstantElementCount(&simdVal, cnsMask, simdBaseType) >= 2)
+    {
+        return LowerHWIntrinsicCreateWithInserts(node, &simdVal, cnsMask);
+    }
+
     // We will be constructing the following parts:
     //          /--*  op1  T
     //   tmp1 = *  HWINTRINSIC   simd16 T CreateScalarUnsafe