From f1297646ea11df536f7ae634b619ce6fa1054adc Mon Sep 17 00:00:00 2001 From: Tanner Gooding Date: Thu, 9 Jul 2026 10:03:54 -0700 Subject: [PATCH 1/2] JIT: Fold insertps chains that source a scalar extracted from another vector Recognize Insert(vector, CreateScalarUnsafe(vector2.GetElement(idx)), idx2) and rewrite it into a single insertps that selects the source lane directly via count_s, folding away the separate extraction. This is valid for any 4-byte element type (float/int/uint) and any source vector width, provided the extracted element is one of the first four lanes (all that count_s can encode and all that resides in the low 128 bits that insertps reads). Both NI_Vector_ToScalar (index 0) and constant-index NI_Vector_GetElement feeders are handled, keeping the register form only when the source is used from a register so the existing insertps xmm, m32 memory form is preserved. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/coreclr/jit/lowerxarch.cpp | 75 ++++++++ .../Vectorization/InsertScalarFromVector.cs | 179 ++++++++++++++++++ .../InsertScalarFromVector.csproj | 8 + 3 files changed, 262 insertions(+) create mode 100644 src/tests/JIT/opt/Vectorization/InsertScalarFromVector.cs create mode 100644 src/tests/JIT/opt/Vectorization/InsertScalarFromVector.csproj diff --git a/src/coreclr/jit/lowerxarch.cpp b/src/coreclr/jit/lowerxarch.cpp index 0359b806c6e1bc..8dd72360be240c 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 + + + + + From 28de33b89811a91e91b4918fdc1e83fe6ca2a5e7 Mon Sep 17 00:00:00 2001 From: Tanner Gooding Date: Thu, 9 Jul 2026 11:14:02 -0700 Subject: [PATCH 2/2] Apply formatting patch --- src/coreclr/jit/lowerxarch.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/coreclr/jit/lowerxarch.cpp b/src/coreclr/jit/lowerxarch.cpp index 8dd72360be240c..c7b0a966838e42 100644 --- a/src/coreclr/jit/lowerxarch.cpp +++ b/src/coreclr/jit/lowerxarch.cpp @@ -2180,15 +2180,15 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node) if ((count_s == 0) && op2->OperIsHWIntrinsic(NI_Vector_CreateScalarUnsafe)) { GenTreeHWIntrinsic* createScalar = op2->AsHWIntrinsic(); - GenTree* scalarOp = createScalar->Op(1); + 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; + GenTree* srcVec = nullptr; + GenTree* srcIdx = nullptr; ssize_t count_s_new = 0; if (extractId == NI_Vector_ToScalar) @@ -2212,8 +2212,8 @@ GenTree* Lowering::LowerHWIntrinsic(GenTreeHWIntrinsic* node) // 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)) + if ((srcVec != nullptr) && !srcVec->isContained() && (count_s_new >= 0) && (count_s_new <= 3) && + IsInvariantInRange(srcVec, node)) { count_s = count_s_new;