JIT: LSRA never reuses an enregistered 256/512-bit vector constant on windows-x64, even in call-free code #134540

LinearScan::isMatchingConstant (src/coreclr/jit/lsra.cpp) rejects every GT_CNS_VEC whose register type needs a partial callee save:

case GT_CNS_VEC:
{
    return
#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
        !Compiler::varTypeNeedsPartialCalleeSave(physRegRecord->assignedInterval->registerType) &&
#endif
        GenTreeVecCon::Equals(refPosition->treeNode->AsVecCon(), otherTreeNode->AsVecCon());
}

Compiler::varTypeNeedsPartialCalleeSave returns true for every TYP_SIMD32/TYP_SIMD64 on x64 and for TYP_SIMD16/TYP_SIMD12 on ARM64, so this is a pure type test — it never asks whether a clobber actually occurred. The guard was added by #74110 for a genuine hazard: the Windows x64 ABI preserves only the low 128 bits of xmm6–xmm15, and freeKilledRegs clears the constant bit only for the registers in a call's kill set (RBM_FLT_CALLEE_TRASH = xmm0–xmm5), so a wide constant parked in ymm6–ymm15 would survive a call in LSRA's bookkeeping while its upper half was destroyed. But the rejection also fires when no call — and therefore no clobber — can have happened.

Two consequences:

• The identical Vector128 source shape already reuses the register today, so the JIT emits strictly more code for the wider type.
• FEATURE_PARTIAL_SIMD_CALLEE_SAVE is 0 under UNIX_AMD64_ABI (src/coreclr/jit/targetamd64.h), so the whole guard is compiled out on linux-x64 and the reuse already happens there. Windows-x64 emits more instructions than linux-x64 for the same source.

Minimal repro

using System;
using System.Runtime.CompilerServices;
using System.Runtime.Intrinsics;
using System.Runtime.Intrinsics.X86;

public static class VecRepro
{
    [MethodImpl(MethodImplOptions.NoInlining)]
    public static bool TwoEmpty256(Vector256<int> a, Vector256<int> b)
        => Avx.TestZ(a, Vector256<int>.AllBitsSet)
         & Avx.TestZ(b, Vector256<int>.AllBitsSet);

    // Control: the identical 128-bit shape already reuses the register today.
    [MethodImpl(MethodImplOptions.NoInlining)]
    public static bool TwoEmpty128(Vector128<int> a, Vector128<int> b)
        => Sse41.TestZ(a, Vector128<int>.AllBitsSet)
         & Sse41.TestZ(b, Vector128<int>.AllBitsSet);

    [MethodImpl(MethodImplOptions.NoInlining)]
    public static int CountPairs256(Vector256<int>[] a, Vector256<int>[] b)
    {
        int n = 0;
        for (int i = 0; i < a.Length; i++)
        {
            if (Avx.TestZ(a[i], Vector256<int>.AllBitsSet)
              & Avx.TestZ(b[i], Vector256<int>.AllBitsSet))
            {
                n++;
            }
        }
        return n;
    }

    [MethodImpl(MethodImplOptions.NoInlining)]
    private static int Opaque(int v) => v + 1;

    // Negative case: the constant must NOT be reused across a call, because a call
    // destroys the upper half of the partially callee-saved float registers.
    [MethodImpl(MethodImplOptions.NoInlining)]
    public static bool AcrossCall256(Vector256<int> a, Vector256<int> b, out int sink)
    {
        bool first = Avx.TestZ(a, Vector256<int>.AllBitsSet);
        sink = Opaque(first ? 1 : 0);
        return first & Avx.TestZ(b, Vector256<int>.AllBitsSet);
    }

    public static int Main()
    {
        if (!Avx2.IsSupported)
        {
            Console.WriteLine("AVX2 required.");
            return 100;
        }

        var a256 = new Vector256<int>[8];
        var b256 = new Vector256<int>[8];
        var a128 = new Vector128<int>[8];
        a256[3] = Vector256.Create(7);
        a128[3] = Vector128.Create(7);

        Console.WriteLine($"TwoEmpty256   = {TwoEmpty256(a256[0], a256[3])}");
        Console.WriteLine($"TwoEmpty128   = {TwoEmpty128(a128[0], a128[3])}");
        Console.WriteLine($"CountPairs256 = {CountPairs256(a256, b256)}");
        Console.WriteLine($"AcrossCall256 = {AcrossCall256(a256[0], a256[3], out int s)} {s}");
        return 100;
    }
}

Run with DOTNET_TieredCompilation=0, DOTNET_ReadyToRun=0 and DOTNET_JitDisasm=TwoEmpty256 TwoEmpty128 CountPairs256 on windows-x64. Reproduced both with DOTNET_EnableAVX512=0 (x64 + VEX) and with AVX-512 enabled (x64 + VEX + EVEX) — the listings are byte-identical, so this is not an AVX-512 artifact.

Current codegen

TwoEmpty256 — 36 bytes, PerfScore 22.75, 11 instructions:

       vpcmpeqd ymm0, ymm0, ymm0        ; constant #1
       vptest   ymm0, ymmword ptr [rcx]
       sete     al
       movzx    rax, al
       vpcmpeqd ymm0, ymm0, ymm0        ; redundant: ymm0 already holds all-ones
       vptest   ymm0, ymmword ptr [rdx]

The redundancy also lands inside loop bodies — the CountPairs256 loop block is 41 bytes against 37 for the 128-bit control, at bbWeight 3.96.

Expected codegen

What the JIT already produces for the 128-bit control TwoEmpty128 (29 bytes, PerfScore 15.25, 9 instructions), and what a prototype relaxation produces for TwoEmpty256 (32 bytes, PerfScore 22.25, 10 instructions):

       vpcmpeqd ymm0, ymm0, ymm0
       vptest   ymm0, ymmword ptr [rcx]
       sete     al
       movzx    rax, al
       vptest   ymm0, ymmword ptr [rdx]  ; register reused

Impact

Measured on 4856f0c16f89d8cc625a62ee7bab7ed4afe14f9d, windows-x64, with a prototype that replaces the type test with a clobber test (patch below). All numbers are static — code size, instruction count, and PerfScore.

Repro methods (base → prototype):

SuperPMI asm diffs, checked base JIT vs. checked prototype JIT, 12 default collections for JIT-EE fbbaf45f-5b0e-4767-b962-5084c8caae77.windows.x64:

Across those 419 contexts: 411 PerfScore improvements, 0 PerfScore regressions, 8 unchanged. Largest individual wins:

Known regressions and measurement limitations:

• Code size: no size or PerfScore regression appeared in any of the 12 collections in this run. A separate Release-JIT diff pair, run against a baseline clrjit.dll binary built in a different environment, reported 2 regressing contexts (+253 bytes) on libraries_tests.run with asymmetric missing-data counts; the same signature appeared for unrelated patches in the same batch, so it is attributed to baseline build provenance rather than to this change. It has not been independently reproduced against a same-environment baseline.
• JIT throughput: tpdiff (PIN, and against that same externally built baseline) shows MinOpts +0.40% to +0.94% on all 12 collections, with overall −0.05% to +0.18% and FullOpts −0.10% to +0.12%. The prototype does add MinOpts work for no MinOpts benefit — processKills and updateAssignedInterval are shared with allocateRegistersMinimal, while the reuse itself is gated on opts.OptimizationEnabled() — so this should not be dismissed as baseline noise without a same-environment re-run. A real fix should probably gate the new bookkeeping on OptimizationEnabled(); that variant was not measured.
• No runtime measurement. No BenchmarkDotNet or other timed benchmark was run. vpcmpeqd r,r,r and vxorps r,r,r are dependency-breaking idioms, so the wall-clock effect is expected to be small even where a hot loop loses an instruction per iteration. Nothing here supports a speedup claim.
• windows-x64 only. ARM64 and linux-x64 were not measured.
• No JitStress / JitStressRegs / GC-stress runs; the assert-enabled checked SuperPMI replay over all 12 collections is the only failure gate that was exercised.

Notes

• Scope. ARM64 is the larger exposure, not x64: varTypeNeedsPartialCalleeSave is true there for the far more common TYP_SIMD16/TYP_SIMD12, so today ARM64 never reuses any enregistered vector constant. Relaxing the guard would enable it there too, and that was not measured. An #ifdef TARGET_XARCH scoping, or a separate ARM64 evaluation, is probably the right first step.
• Correctness. Reusing a register that still holds the value is a no-op provided the register really is unchanged. GenTreeVecCon::Equals still proves bitwise identity; vpcmpeqd ymm,ymm,ymm has no memory operand, no faults, and no GC-visible state; and vector constants are non-GC, so the SetReuseRegVal GC-tracking concern of [JIT] Fix re-use val zero on GC tracking #84051 does not apply. The only hazard is the partial callee save that Disable matching constants for vectors that needs upper half to be save/restore #74110 guards against, which needs an actual clobber test.
• Residual risk in the prototype's invalidation rule. It keys on "the kill set intersects RBM_FLT_CALLEE_TRASH". CORINFO_HELP_ASSIGN_REF / CHECKED_ASSIGN_REF use RBM_CALLEE_TRASH_WRITEBARRIER, which contains no float register on x64, so a write-barrier call does not clear the mask — safe only because those helpers are hand-written assembly that leaves xmm/ymm alone. That is an assumption, not a proof. Likewise, the claim that genVzeroupperIfNeeded can never run between a wide-constant definition and its reuse without a coinciding kill is reasoned from source, not tested.
• Related work, none of which closes this gap. JIT: Reuse enregistered FP/SIMD/mask constants across non-last uses #130312 (open, "Reuse enregistered FP/SIMD/mask constants across non-last uses") adds a build-time coalescing path that explicitly continues on varTypeNeedsPartialCalleeSave and does not touch isMatchingConstant; it edits the same LSRA constant machinery, so the two would need to be sequenced. JIT: Rematerialize spilled FP/SIMD/mask constants instead of stack spill/reload #130313 (open) rematerializes spilled constants from the codegen side. arm64: Reuse SVE mask constants in LSRA #131309 (open) is ARM64 SVE GT_CNS_MSK reuse in the register selector. Disable matching constants for vectors that needs upper half to be save/restore #74110 added the guard; Ensure that GT_CNS_VEC is handled in LinearScan::isMatchingConstant #70171 and [JIT] Fix re-use val zero on GC tracking #84051 are adjacent merged work. Open issues whose symptoms this partly explains: Suboptimal ASM emmited for Vector256<T>.Zero and Vector128<T>.Zero #76067 (Vector256<T>.Zero / Vector128<T>.Zero re-materialized instead of reused), Improve handling of reused constants in the register allocator #70182, Jit creates multiple zeroes in registers #37079.
• Alternative upstream fix for the repro shape specifically: have lowering recognize TestZ(x, AllBitsSet) and emit vptest x, x, removing the constant entirely. That is a narrower, separate change and would not help the Zero/AllBitsSet operand cases that dominate the corpus wins.
• Any change here should re-run the test that motivated Disable matching constants for vectors that needs upper half to be save/restore #74110 and add a call-containing negative test (wide constant live in ymm6–ymm15 across a call).

Prototype patch

Experimental only — it has not been through stress testing, it is windows-x64-measured only, and the MinOpts throughput question above is unresolved.

diff --git a/src/coreclr/jit/lsra.cpp b/src/coreclr/jit/lsra.cpp
index faf1b4b10c8..37889cf9403 100644
--- a/src/coreclr/jit/lsra.cpp
+++ b/src/coreclr/jit/lsra.cpp
@@ -2757,11 +2757,18 @@ bool LinearScan::isMatchingConstant(RegRecord* physRegRecord, RefPosition* refPo
 #if defined(FEATURE_SIMD)
         case GT_CNS_VEC:
         {
-            return
 #if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
-                !Compiler::varTypeNeedsPartialCalleeSave(physRegRecord->assignedInterval->registerType) &&
+            // Only part of a wide vector register is preserved across a call, so such a register
+            // may only be reused for an identical constant if no call-like kill has occurred
+            // since the constant was materialized.
+            if ((Compiler::varTypeNeedsPartialCalleeSave(physRegRecord->assignedInterval->registerType) ||
+                 Compiler::varTypeNeedsPartialCalleeSave(interval->registerType)) &&
+                !isWideConstantRegValid(physRegRecord->regNum, physRegRecord->assignedInterval->registerType))
+            {
+                return false;
+            }
 #endif
-                GenTreeVecCon::Equals(refPosition->treeNode->AsVecCon(), otherTreeNode->AsVecCon());
+            return GenTreeVecCon::Equals(refPosition->treeNode->AsVecCon(), otherTreeNode->AsVecCon());
         }
 #endif // FEATURE_SIMD
 
@@ -3831,6 +3838,17 @@ void LinearScan::processKills(RefPosition* killRefPosition)
 #endif
 
     regsBusyUntilKill &= ~killRefPosition->getKilledRegisters();
+
+#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+    // A call trashes the non-preserved part of every float register, including the ones that are
+    // not in its kill set (only the low part of the callee-saved float registers is preserved).
+    // Any kill that trashes float registers therefore invalidates every wide constant register.
+    if (!(killedRegs & RBM_FLT_CALLEE_TRASH).IsEmpty())
+    {
+        m_WideConstantsValid = RBM_NONE;
+    }
+#endif // FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+
     INDEBUG(dumpLsraAllocationEvent(LSRA_EVENT_KILL_REGS, nullptr, REG_NA, nullptr, NONE,
                                     killRefPosition->getKilledRegisters()));
 }
@@ -6850,6 +6868,16 @@ void LinearScan::updateAssignedInterval(RegRecord* reg, Interval* interval ARM_A
     if (interval->isConstant)
     {
         setConstantReg(reg->regNum, interval->registerType);
+#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+        if (Compiler::varTypeNeedsPartialCalleeSave(interval->registerType))
+        {
+            setWideConstantReg(reg->regNum, interval->registerType);
+        }
+        else
+        {
+            clearWideConstantReg(reg->regNum, interval->registerType);
+        }
+#endif
     }
     else
     {
diff --git a/src/coreclr/jit/lsra.h b/src/coreclr/jit/lsra.h
index 3541e1a2309..03ee7b2aae0 100644
--- a/src/coreclr/jit/lsra.h
+++ b/src/coreclr/jit/lsra.h
@@ -1763,6 +1763,9 @@ private:
     {
         m_AvailableRegs          = allAvailableRegs;
         m_RegistersWithConstants = RBM_NONE;
+#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+        m_WideConstantsValid = RBM_NONE;
+#endif
     }
 
     bool isRegAvailable(regNumber reg, var_types regType)
@@ -1802,9 +1805,30 @@ private:
                                                  DEBUG_ARG(regNumber assignedReg));
 
     regMaskTP m_RegistersWithConstants;
-    void      clearConstantReg(regNumber reg, var_types regType)
+#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+    // The subset of `m_RegistersWithConstants` that holds a constant whose register is only
+    // partially preserved across a call. Such a register may only be reused for an identical
+    // constant while no call-like kill has occurred since the constant was materialized.
+    regMaskTP m_WideConstantsValid;
+    void      setWideConstantReg(regNumber reg, var_types regType)
+    {
+        m_WideConstantsValid.AddRegNum(reg, regType);
+    }
+    void clearWideConstantReg(regNumber reg, var_types regType)
+    {
+        m_WideConstantsValid.RemoveRegNum(reg, regType);
+    }
+    bool isWideConstantRegValid(regNumber reg, var_types regType)
+    {
+        return m_WideConstantsValid.IsRegNumPresent(reg, regType);
+    }
+#endif // FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+    void clearConstantReg(regNumber reg, var_types regType)
     {
         m_RegistersWithConstants.RemoveRegNum(reg, regType);
+#if FEATURE_PARTIAL_SIMD_CALLEE_SAVE
+        m_WideConstantsValid.RemoveRegNum(reg, regType);
+#endif
     }
     void setConstantReg(regNumber reg, var_types regType)
     {