diff --git a/src/coreclr/jit/lowerxarch.cpp b/src/coreclr/jit/lowerxarch.cpp index 0359b806c6e1bc..c7b0a966838e42 100644 --- a/src/coreclr/jit/lowerxarch.cpp +++ b/src/coreclr/jit/lowerxarch.cpp @@ -2162,6 +2162,81 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node) return nextNode; } + // We can also recognize when the value being inserted is itself a scalar + // extracted from another vector, such as: + // + // Insert(vector, CreateScalarUnsafe(vector2.GetElement(idx1)), idx2) + // + // In this case, insertps can select the source element directly from + // vector2 using count_s (bits 6-7), which lets us fold away the separate + // extraction (which otherwise requires its own movshdup/unpckhps/shufps). + // + // insertps operates purely on the 32-bit lanes of the low 128 bits of its + // source register, so this is valid for any 4-byte element type (float, int, + // or uint) and for a source vector of any width, provided the element being + // extracted is one of the first four. That is all count_s can encode and all + // that resides in the low 128 bits, which is the only part insertps reads. + + if ((count_s == 0) && op2->OperIsHWIntrinsic(NI_Vector_CreateScalarUnsafe)) + { + GenTreeHWIntrinsic* createScalar = op2->AsHWIntrinsic(); + GenTree* scalarOp = createScalar->Op(1); + + if (scalarOp->OperIsHWIntrinsic() && (genTypeSize(scalarOp->AsHWIntrinsic()->GetSimdBaseType()) == 4)) + { + GenTreeHWIntrinsic* extract = scalarOp->AsHWIntrinsic(); + NamedIntrinsic extractId = extract->GetHWIntrinsicId(); + + GenTree* srcVec = nullptr; + GenTree* srcIdx = nullptr; + ssize_t count_s_new = 0; + + if (extractId == NI_Vector_ToScalar) + { + // ToScalar is effectively GetElement with a source index of zero + srcVec = extract->Op(1); + } + else if ((extractId == NI_Vector_GetElement) && extract->Op(2)->IsCnsIntOrI()) + { + srcVec = extract->Op(1); + srcIdx = extract->Op(2); + count_s_new = srcIdx->AsIntConCommon()->IconValue(); + } + + // The source index must fit in count_s, which also guarantees the element + // resides in the low 128 bits of the source register. + // + // We additionally require that the source vector is used from a register. + // When the extraction reads from a contained memory operand, the existing + // lowering already produces an optimal `insertps xmm, m32` (which encodes + // the element offset in the address) and rewriting it to the register form + // would instead force a full vector load. + + if ((srcVec != nullptr) && !srcVec->isContained() && (count_s_new >= 0) && (count_s_new <= 3) && + IsInvariantInRange(srcVec, node)) + { + count_s = count_s_new; + + ival = (count_s << 6) | (count_d << 4) | (zmask); + op3->AsIntConCommon()->SetIconValue(ival); + + // Carry the original source vector directly and remove the now + // unused extraction and scalar creation nodes. + + node->Op(2) = srcVec; + op2 = srcVec; + + if (srcIdx != nullptr) + { + BlockRange().Remove(srcIdx); + } + + BlockRange().Remove(extract); + BlockRange().Remove(createScalar); + } + } + } + if (!op1->OperIsHWIntrinsic()) { // Nothing to do if op1 isn't an intrinsic diff --git a/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.cs b/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.cs new file mode 100644 index 00000000000000..e56db6c15bfb91 --- /dev/null +++ b/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.cs @@ -0,0 +1,179 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +// Validates the JIT folding of Insert(vector, vector2.GetElement(idx1), idx2) +// into a single insertps that selects the source element directly via the +// count_s field of the immediate. Both the register-source form (which the fold +// rewrites) and the memory-source form (already handled) must match a scalar +// reference for every combination of source and destination indices. +// +// The indices must be JIT-time constants for the insertps immediate (and hence +// the fold) to be formed, so each combination is dispatched through a switch of +// literal WithElement/GetElement calls rather than passing the indices as +// arguments. + +using System; +using System.Runtime.CompilerServices; +using System.Runtime.Intrinsics; +using Xunit; + +public static class InsertScalarFromVector +{ + private static readonly float[] s_dst = { 10f, 11f, 12f, 13f }; + private static readonly float[] s_src = { 20f, 21f, 22f, 23f }; + + // src is produced by an arithmetic op so it stays in a register, forcing the + // register GetElement path that the fold collapses into insertps' count_s. + // combo == dstIdx * 4 + srcIdx, with both indices as compile-time constants. + [MethodImpl(MethodImplOptions.NoInlining)] + private static Vector128 FromRegister(Vector128 dst, Vector128 a, Vector128 b, int combo) + { + Vector128 src = a + b; + return combo switch + { + 0 => dst.WithElement(0, src.GetElement(0)), + 1 => dst.WithElement(0, src.GetElement(1)), + 2 => dst.WithElement(0, src.GetElement(2)), + 3 => dst.WithElement(0, src.GetElement(3)), + 4 => dst.WithElement(1, src.GetElement(0)), + 5 => dst.WithElement(1, src.GetElement(1)), + 6 => dst.WithElement(1, src.GetElement(2)), + 7 => dst.WithElement(1, src.GetElement(3)), + 8 => dst.WithElement(2, src.GetElement(0)), + 9 => dst.WithElement(2, src.GetElement(1)), + 10 => dst.WithElement(2, src.GetElement(2)), + 11 => dst.WithElement(2, src.GetElement(3)), + 12 => dst.WithElement(3, src.GetElement(0)), + 13 => dst.WithElement(3, src.GetElement(1)), + 14 => dst.WithElement(3, src.GetElement(2)), + 15 => dst.WithElement(3, src.GetElement(3)), + _ => throw new ArgumentOutOfRangeException(nameof(combo)), + }; + } + + // src is passed directly and is typically consumed from memory, exercising the + // pre-existing insertps-with-memory-operand path. + [MethodImpl(MethodImplOptions.NoInlining)] + private static Vector128 FromMemory(Vector128 dst, Vector128 src, int combo) + { + return combo switch + { + 0 => dst.WithElement(0, src.GetElement(0)), + 1 => dst.WithElement(0, src.GetElement(1)), + 2 => dst.WithElement(0, src.GetElement(2)), + 3 => dst.WithElement(0, src.GetElement(3)), + 4 => dst.WithElement(1, src.GetElement(0)), + 5 => dst.WithElement(1, src.GetElement(1)), + 6 => dst.WithElement(1, src.GetElement(2)), + 7 => dst.WithElement(1, src.GetElement(3)), + 8 => dst.WithElement(2, src.GetElement(0)), + 9 => dst.WithElement(2, src.GetElement(1)), + 10 => dst.WithElement(2, src.GetElement(2)), + 11 => dst.WithElement(2, src.GetElement(3)), + 12 => dst.WithElement(3, src.GetElement(0)), + 13 => dst.WithElement(3, src.GetElement(1)), + 14 => dst.WithElement(3, src.GetElement(2)), + 15 => dst.WithElement(3, src.GetElement(3)), + _ => throw new ArgumentOutOfRangeException(nameof(combo)), + }; + } + + private static Vector128 Reference(int dstIdx, int srcIdx) + { + float[] result = (float[])s_dst.Clone(); + result[dstIdx] = s_src[srcIdx]; + return Vector128.Create(result[0], result[1], result[2], result[3]); + } + + [Fact] + public static void TestAllCombinations() + { + Vector128 dst = Vector128.Create(s_dst[0], s_dst[1], s_dst[2], s_dst[3]); + Vector128 src = Vector128.Create(s_src[0], s_src[1], s_src[2], s_src[3]); + Vector128 zero = Vector128.Zero; + + for (int dstIdx = 0; dstIdx < 4; dstIdx++) + { + for (int srcIdx = 0; srcIdx < 4; srcIdx++) + { + int combo = (dstIdx * 4) + srcIdx; + Vector128 expected = Reference(dstIdx, srcIdx); + + Assert.Equal(expected, FromRegister(dst, src, zero, combo)); + Assert.Equal(expected, FromMemory(dst, src, combo)); + } + } + } + + // Chained WithElement using elements pulled from another vector exercises + // composing the source-selection fold with the existing insert chains. + [MethodImpl(MethodImplOptions.NoInlining)] + private static Vector128 Chained(Vector128 dst, Vector128 a, Vector128 b) + { + Vector128 src = a + b; + return dst.WithElement(0, src.GetElement(3)) + .WithElement(2, src.GetElement(1)); + } + + [Fact] + public static void TestChained() + { + Vector128 dst = Vector128.Create(s_dst[0], s_dst[1], s_dst[2], s_dst[3]); + Vector128 src = Vector128.Create(s_src[0], s_src[1], s_src[2], s_src[3]); + Vector128 zero = Vector128.Zero; + + Vector128 actual = Chained(dst, src, zero); + Vector128 expected = Vector128.Create(s_src[3], s_dst[1], s_src[1], s_dst[3]); + + Assert.Equal(expected, actual); + } + + private static readonly float[] s_src8 = { 20f, 21f, 22f, 23f, 24f, 25f, 26f, 27f }; + + // Extracting from a vector wider than 128 bits is also foldable when the element + // is one of the source's first four (which live in the low 128 bits): the fold + // feeds the wide register straight to insertps for GetElement(0)/ToScalar, while + // the existing GetElement lowering narrows the higher elements via GetLower/GetUpper. + // combo == dstIdx * 8 + srcIdx, with both indices as compile-time constants. + [MethodImpl(MethodImplOptions.NoInlining)] + private static Vector128 FromWideRegister(Vector128 dst, Vector256 a, Vector256 b, int combo) + { + Vector256 src = a + b; + return combo switch + { + 0 => dst.WithElement(0, src.GetElement(0)), + 9 => dst.WithElement(1, src.GetElement(1)), + 18 => dst.WithElement(2, src.GetElement(2)), + 27 => dst.WithElement(3, src.GetElement(3)), + 4 => dst.WithElement(0, src.GetElement(4)), + 13 => dst.WithElement(1, src.GetElement(5)), + 22 => dst.WithElement(2, src.GetElement(6)), + 31 => dst.WithElement(3, src.GetElement(7)), + _ => throw new ArgumentOutOfRangeException(nameof(combo)), + }; + } + + private static Vector128 WideReference(int dstIdx, int srcIdx) + { + float[] result = (float[])s_dst.Clone(); + result[dstIdx] = s_src8[srcIdx]; + return Vector128.Create(result[0], result[1], result[2], result[3]); + } + + [Fact] + public static void TestWideSource() + { + Vector128 dst = Vector128.Create(s_dst[0], s_dst[1], s_dst[2], s_dst[3]); + Vector256 src = Vector256.Create(s_src8[0], s_src8[1], s_src8[2], s_src8[3], + s_src8[4], s_src8[5], s_src8[6], s_src8[7]); + Vector256 zero = Vector256.Zero; + + for (int i = 0; i < 8; i++) + { + int dstIdx = i % 4; + int combo = (dstIdx * 8) + i; + + Assert.Equal(WideReference(dstIdx, i), FromWideRegister(dst, src, zero, combo)); + } + } +} diff --git a/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.csproj b/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.csproj new file mode 100644 index 00000000000000..de6d5e08882e86 --- /dev/null +++ b/src/tests/JIT/opt/Vectorization/InsertScalarFromVector.csproj @@ -0,0 +1,8 @@ + + + True + + + + +