From 6319bc1f081fd4079d904ad7f988b7c4c6a4d59c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 7 Oct 2026 13:51:35 -0400 Subject: [PATCH 1/3] Fence the AVX2 complex K steps so their broadcasts stop spilling At 16 ymm registers LLVM's haswell/znver2 scheduler hoists every B broadcast of a planar or fmaddsub K step to the top of the loop body, holds them all live and spills them. `memory_fence()`, an empty `asm sideeffect` with a memory clobber (no instruction), after each column group pins the loads next to their FMAs; the FMAs stay free to move. On AVX2, fmaddsub also runs OpenBLAS's Zen zgemm order: the swap(a)*bi pass, a fence, then the a*br pass on A loaded again, so a and swap(a) are never live together. Only 256-bit kernels (`W * sizeof(R) == 32`) take the fenced order; every other shape gets the same statements in the same order as before. code_llvm/code_native of `add_tile` are identical to ad94b73 for all 57 non-256-bit complex and mixed-domain menu shapes (B packed and B in place) under -C icelake-server, znver4, znver2 and haswell. Hot-loop stack references per K step, Julia 1.12.7, -C znver2 (haswell identical), ad94b73 -> this commit, B packed / B read in place (the #22 unpacked-B codegen does not change the picture): planar 4x5/W4, 8x5/W8 4 / 4 -> 0 / 0 fmaddsub 4x6/W4, 8x6/W8 24 / 24 -> 0 / 0 fmaddsub 4x5/W4, 8x5/W8 0 / 0 -> 0 / 0 planar 4x6/W4, 8x6/W8 14 / 12 -> 0 / 13 (not in any menu) The mixed-domain kernels run the real SIMDKernel step and do not spill. Co-Authored-By: Claude Opus 5.5 --- src/microkernels/steps.jl | 139 ++++++++++------------ src/microkernels/vecops.jl | 4 + test/microkernels/helpers.jl | 27 +++++ test/microkernels/test_fmaddsub_kernel.jl | 43 ++----- test/microkernels/test_planar_kernel.jl | 7 ++ 5 files changed, 110 insertions(+), 110 deletions(-) diff --git a/src/microkernels/steps.jl b/src/microkernels/steps.jl index d17d4fa..b290b1f 100644 --- a/src/microkernels/steps.jl +++ b/src/microkernels/steps.jl @@ -131,56 +131,37 @@ end naiv = [Symbol(:nai, v) for v in 0:(MV - 1)] brv = [Symbol(:br, j) for j in 1:NR] biv = [Symbol(:bi, j) for j in 1:NR] + outs = [Symbol(:c, i) for i in 1:(2NV)] - load_a = Any[] + body = Any[] for v in 0:(MV - 1) - push!( - load_a, - :( - $(arv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) - ) - ) - ) - push!( - load_a, - :( - $(aiv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p, 1) - ) - ) - ) - push!(load_a, :($(naiv[v + 1]) = -$(aiv[v + 1]))) + push!(body, :($(arv[v + 1]) = panel_vload(Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p)))) + push!(body, :($(aiv[v + 1]) = panel_vload(Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p, 1)))) + push!(body, :($(naiv[v + 1]) = -$(aiv[v + 1]))) end - - load_b = Any[] - for j in 1:NR - push!(load_b, :(($(brv[j]), $(biv[j])) = b_complex(packed_b, kernel, $j, p))) - end - - acc_exprs = Vector{Any}(undef, NA) + acc_exprs = Vector{Any}(undef, 2NV) for j in 1:NR, v in 0:(MV - 1) - idx = acc_index(MV, v, j) - acc_exprs[idx] = :( # re: ar*br - ai*bi - muladd( - $(naiv[v + 1]), $(biv[j]), - muladd($(arv[v + 1]), $(brv[j]), acc[$idx]) - ) - ) - acc_exprs[NV + idx] = :( # im: ar*bi + ai*br - muladd( - $(aiv[v + 1]), $(brv[j]), - muladd($(arv[v + 1]), $(biv[j]), acc[$(NV + idx)]) - ) - ) + i = acc_index(MV, v, j) + acc_exprs[i] = :(muladd($(naiv[v + 1]), $(biv[j]), muladd($(arv[v + 1]), $(brv[j]), acc[$i]))) # re: ar*br - ai*bi + acc_exprs[NV + i] = :(muladd($(aiv[v + 1]), $(brv[j]), muladd($(arv[v + 1]), $(biv[j]), acc[$(NV + i)]))) # im: ar*bi + ai*br + end + # At AVX2's 16 registers LLVM hoists every broadcast of the K step and + # spills them; a fence per column keeps them next to their FMAs. + fence = W * sizeof(R) == 32 + for j in 1:NR + push!(body, :(($(brv[j]), $(biv[j])) = b_complex(packed_b, kernel, $j, p))) + fence || continue + for i in acc_index(MV, 0, j):acc_index(MV, MV - 1, j) + push!(body, :($(outs[i]) = $(acc_exprs[i])), :($(outs[NV + i]) = $(acc_exprs[NV + i]))) + end + push!(body, :(memory_fence())) end return quote Base.@_inline_meta @inbounds begin - $(load_a...) - $(load_b...) - return $(Expr(:tuple, acc_exprs...)) + $(body...) + return $(Expr(:tuple, (fence ? outs : acc_exprs)...)) end end end @@ -210,49 +191,55 @@ end ) where {MR, NR, T, W, R, NA, PA, PB} MV = (2 * MR) ÷ W check_acc(:accumulate_step, R, T, NA, MV * NR) - - av = [Symbol(:a, v) for v in 0:(MV - 1)] - sv = [Symbol(:s, v) for v in 0:(MV - 1)] - brv = [Symbol(:br, j) for j in 1:NR] - biv = [Symbol(:bi, j) for j in 1:NR] - - load_a = Any[] + loads = Any[] for v in 0:(MV - 1) - push!( - load_a, - :( - $(av[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) - ) - ) - ) - push!(load_a, :($(sv[v + 1]) = swap_pairs($(av[v + 1])))) + push!(loads, :(panel_vload(Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p)))) end - - load_b = Any[] - for j in 1:NR - push!(load_b, :((br_s, bi_s) = b_complex(packed_b, kernel, $j, p))) - push!(load_b, :($(brv[j]) = Vec{$W, $R}(br_s))) - push!(load_b, :($(biv[j]) = Vec{$W, $R}(bi_s))) - end - - acc_exprs = Vector{Any}(undef, NA) + # Per accumulator, the `swap(a)*bi` term into `m`, then the `a*br` term into `c`. + swapped, direct = Vector{Any}(undef, NA), Vector{Any}(undef, NA) for j in 1:NR, v in 0:(MV - 1) - idx = acc_index(MV, v, j) - acc_exprs[idx] = :( - fmaddsub( - $(av[v + 1]), $(brv[j]), - fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) - ) - ) + i = acc_index(MV, v, j) + swapped[i] = :($(Symbol(:m, i)) = fmaddsub($(Symbol(:s, v)), $(Symbol(:bi, j)), acc[$i])) + direct[i] = :($(Symbol(:c, i)) = fmaddsub($(Symbol(:a, v)), $(Symbol(:br, j)), $(Symbol(:m, i)))) + end + body = Any[] + if W * sizeof(R) == 32 + # On AVX2 the two passes are fenced apart, the second on A loaded again, + # so `a` and `swap(a)` are never live together. + for v in 0:(MV - 1) + push!(body, :($(Symbol(:s, v)) = swap_pairs($(loads[v + 1])))) + end + for j in 1:NR + push!(body, :($(Symbol(:bi, j)) = Vec{$W, $R}(b_complex(packed_b, kernel, $j, p)[2]))) + append!(body, swapped[acc_index(MV, 0, j):acc_index(MV, MV - 1, j)]) + end + push!(body, :(memory_fence())) + for v in 0:(MV - 1) + push!(body, :($(Symbol(:a, v)) = $(loads[v + 1]))) + end + for j in 1:NR + push!(body, :($(Symbol(:br, j)) = Vec{$W, $R}(b_complex(packed_b, kernel, $j, p)[1]))) + append!(body, direct[acc_index(MV, 0, j):acc_index(MV, MV - 1, j)]) + end + push!(body, :(memory_fence())) + else + for v in 0:(MV - 1) + push!(body, :($(Symbol(:a, v)) = $(loads[v + 1])), :($(Symbol(:s, v)) = swap_pairs($(Symbol(:a, v))))) + end + for j in 1:NR + push!(body, :((br_s, bi_s) = b_complex(packed_b, kernel, $j, p))) + push!(body, :($(Symbol(:br, j)) = Vec{$W, $R}(br_s)), :($(Symbol(:bi, j)) = Vec{$W, $R}(bi_s))) + end + for i in 1:NA + push!(body, swapped[i], direct[i]) + end end return quote Base.@_inline_meta @inbounds begin - $(load_a...) - $(load_b...) - return $(Expr(:tuple, acc_exprs...)) + $(body...) + return $(Expr(:tuple, (Symbol(:c, i) for i in 1:NA)...)) end end end diff --git a/src/microkernels/vecops.jl b/src/microkernels/vecops.jl index f60175a..b71034b 100644 --- a/src/microkernels/vecops.jl +++ b/src/microkernels/vecops.jl @@ -37,6 +37,10 @@ using SIMD: Vec, shufflevector end end +# Emits no instruction: a `memory` clobber that stops LLVM moving (or merging) +# loads across it, while the register-only FMAs stay free to move. +@inline memory_fence() = Base.llvmcall("call void asm sideeffect \"\", \"~{memory}\"()\nret void", Cvoid, Tuple{}) + # `[x1, x0, x3, x2, ...]`: one in-lane `vshufpd`/`vpermilps`. @generated function swap_pairs(x::Vec{N, R}) where {N, R} iseven(N) || return :(throw(ArgumentError("swap_pairs: expected even N, got $N"))) diff --git a/test/microkernels/helpers.jl b/test/microkernels/helpers.jl index f6ddcbf..d6f1592 100644 --- a/test/microkernels/helpers.jl +++ b/test/microkernels/helpers.jl @@ -7,8 +7,10 @@ include("../helpers.jl") using QuasiStrided: SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel +using QuasiStrided: PackedPanel using SIMD: Vec using StridedViews: StridedView +using InteractiveUtils: code_native # `[x0, y0, x1, y1, ...]` mk_ilv(x, y) = vec(permutedims(hcat(x, y))) @@ -207,3 +209,28 @@ function mk_contract_full(k) end return nothing end + +# The innermost loop of `add_tile`'s native code over packed panels, for exact +# instruction counts. Only on an FMA3 x86 host and not on hosted CI, whose +# virtualized CPU feature sets do not reliably match. +const MK_HAS_FMA = Sys.ARCH === :x86_64 && target_profile().isa in (:avx2, :avx512) && get(ENV, "CI", "false") != "true" +function mk_hot_loop(k) + R = realtype(k) + asm = sprint() do io + code_native( + io, QuasiStrided.add_tile, + (typeof(k), typeof(zero_accumulator(k)), PackedPanel{R}, PackedPanel{R}, Int); + debuginfo = :none, syntax = :intel + ) + end + lines = split(asm, '\n') + labels = Dict{String, Int}() + best = nothing + for (n, l) in enumerate(lines) + m = match(r"^(\.LBB\w+):", l) + m === nothing || (labels[m[1]] = n) + b = match(r"^\s+j\w+\s+(\.LBB\w+)", l) + b !== nothing && haskey(labels, b[1]) && (best = (labels[b[1]], n)) + end + return best === nothing ? "" : join(lines[best[1]:best[2]], '\n') +end diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 39d25a5..62b7203 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -1,9 +1,8 @@ include("helpers.jl") using QuasiStrided: FMAddSubKernel, lanewidth, kernel_shapes, packed_panel, - accumulator_planes, target_profile, PackedPanel + accumulator_planes, target_profile using SIMD: Vec -using InteractiveUtils: code_native const QSF = QuasiStrided @@ -71,40 +70,16 @@ const QSF = QuasiStrided # Exact instruction counts in the hot loop of `add_tile`, so only on # an FMA3 x86 host and not on hosted CI, whose virtualized CPU feature # sets do not reliably match. - isa = target_profile().isa - has_fma = Sys.ARCH === :x86_64 && isa in (:avx2, :avx512) && get(ENV, "CI", "false") != "true" - function hot_loop(asm) - lines = split(asm, '\n') - labels = Dict{String, Int}() - best = nothing - for (n, l) in enumerate(lines) - m = match(r"^(\.LBB\w+):", l) - m === nothing || (labels[m[1]] = n) - b = match(r"^\s+j\w+\s+(\.LBB\w+)", l) - b !== nothing && haskey(labels, b[1]) && (best = (labels[b[1]], n)) - end - return best === nothing ? "" : join(lines[best[1]:best[2]], '\n') - end - shapes = isa === :avx512 ? - ((ComplexF64, (8, 8, 8)), (ComplexF32, (16, 8, 16)), (ComplexF64, (4, 5, 4))) : - ((ComplexF64, (4, 5, 4)), (ComplexF32, (8, 5, 8))) + avx2 = ((ComplexF64, (4, 5, 4)), (ComplexF32, (8, 5, 8)), (ComplexF64, (4, 6, 4)), (ComplexF32, (8, 6, 8))) + shapes = target_profile().isa === :avx512 ? ((ComplexF64, (8, 8, 8)), (ComplexF32, (16, 8, 16)), avx2...) : avx2 for (T, (MR, NR, W)) in shapes - k = FMAddSubKernel(Val(MR), Val(NR), T, Val(W)) - R = real(T) - asm = sprint() do io - code_native( - io, QuasiStrided.add_tile, - (typeof(k), typeof(zero_accumulator(k)), PackedPanel{R}, PackedPanel{R}, Int); - debuginfo = :none, syntax = :intel - ) - end - loop = hot_loop(asm) + loop = mk_hot_loop(FMAddSubKernel(Val(MR), Val(NR), T, Val(W))) MV = (2 * MR) ÷ W - @test count(r"vfmaddsub\d+p", loop) == 2 * MV * NR skip = !has_fma - @test count(r"vf(n?madd|n?msub)\d+p[sd]", loop) == 0 skip = !has_fma - @test count(r"v(mul|add|sub)p[sd]", loop) == 0 skip = !has_fma - @test count(r"v(shufp|permilp)", loop) == MV skip = !has_fma - @test count(r"\[r[sb]p", loop) == 0 skip = !has_fma + @test count(r"vfmaddsub\d+p", loop) == 2 * MV * NR skip = !MK_HAS_FMA + @test count(r"vf(n?madd|n?msub)\d+p[sd]", loop) == 0 skip = !MK_HAS_FMA + @test count(r"v(mul|add|sub)p[sd]", loop) == 0 skip = !MK_HAS_FMA + @test count(r"v(shufp|permilp)", loop) == MV skip = !MK_HAS_FMA + @test count(r"\[r[sb]p", loop) == 0 skip = !MK_HAS_FMA end end end diff --git a/test/microkernels/test_planar_kernel.jl b/test/microkernels/test_planar_kernel.jl index ed26043..427c7c2 100644 --- a/test/microkernels/test_planar_kernel.jl +++ b/test/microkernels/test_planar_kernel.jl @@ -18,4 +18,11 @@ using QuasiStrided: PlanarKernel, lanewidth, @test_throws ArgumentError PlanarKernel(Val(6), Val(4), ComplexF64, Val(4)) @test_throws ArgumentError PlanarKernel(Val(8), Val(4), ComplexF64, Val(0)) @test_throws ArgumentError PlanarKernel(Val(8), Val(4), Float64) + + # The fenced AVX2 tiles keep every broadcast in registers. + for (T, (MR, NR, W)) in ((ComplexF64, (4, 5, 4)), (ComplexF32, (8, 5, 8))) + loop = mk_hot_loop(PlanarKernel(Val(MR), Val(NR), T, Val(W))) + @test count(r"vfn?madd\d+p", loop) == 4 * NR skip = !MK_HAS_FMA + @test count(r"\[r[sb]p", loop) == 0 skip = !MK_HAS_FMA + end end From c3f46d6f01e3d751ada76ab35e3ece587b5ca6d5 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 7 Oct 2026 13:51:35 -0400 Subject: [PATCH 2/3] Open the complex pack and store fast paths to AVX2 `complex_fastpath_isa_eligible` was AVX-512 only, so on AVX2 every complex sliver packed through the scalar loop and every complex tile (planar, fmaddsub and the mixed-domain kernels, which share those stores) stored element by element. The AVX2 codegen is the AVX-512 path at half width. NEON and unknown targets keep the scalar loops. Measured on the pre-refactor branch (Cascade Lake, QS_FAKE_ISA=avx2 -C znver2, gate off -> on, GF/s): kernel-only planar 8x6/W8 85.4 -> 108.5; engine ComplexF32 2000^3 94.2 -> 105.1, shallowK 47.0 -> 85.5, smallN 47.8 -> 70.4. Co-Authored-By: Claude Opus 5.5 --- src/packing/pack_contiguous.jl | 6 +++--- test/packing/test_pack_contiguous.jl | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index d2410ff..d3925fb 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -20,9 +20,9 @@ lane_convertible(::Type, ::Type) = false lane_convertible(::Type{<:LaneFloat}, ::Type{<:LaneFloat}) = true lane_convertible(::Type{Complex{S}}, ::Type{Complex{T}}) where {S <: LaneFloat, T <: LaneFloat} = true -# The deinterleaving complex packer and the planar store pay off only with -# 512-bit vector registers. -@inline complex_fastpath_isa_eligible(profile::TargetProfile) = profile.isa === :avx512 +# The deinterleaving complex packer and the complex vector stores (unmeasured +# on NEON). +@inline complex_fastpath_isa_eligible(profile::TargetProfile) = profile.isa in (:avx512, :avx2) @inline complex_fastpath_isa_eligible() = complex_fastpath_isa_eligible(target_profile()) # --- Real --- diff --git a/test/packing/test_pack_contiguous.jl b/test/packing/test_pack_contiguous.jl index f89d7c2..c02f694 100644 --- a/test/packing/test_pack_contiguous.jl +++ b/test/packing/test_pack_contiguous.jl @@ -95,11 +95,11 @@ end end @testset "complex pack fast path: ISA gate is a register-width question" begin - @test complex_fastpath_isa_eligible(synthetic(:avx512)) - for key in (:avx2, :neon, :unknown) + @test complex_fastpath_isa_eligible(synthetic(:avx512)) && complex_fastpath_isa_eligible(synthetic(:avx2)) + for key in (:neon, :unknown) @test !complex_fastpath_isa_eligible(synthetic(key)) end - @test FASTPATH_ON == (target_profile().isa === :avx512) + @test FASTPATH_ON == (target_profile().isa in (:avx512, :avx2)) end @testset "pack! contiguous fast path: fires exactly when eligible ($T, L=$L, side $i)" for From 27439f66f0acb19cf7fc8fa9c2bdfb7a1a7e4d9e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 7 Oct 2026 13:51:35 -0400 Subject: [PATCH 3/3] AVX2 complex default: fmaddsub 4x6/W4, 8x6/W8 where C's rows take its vector store On AVX2, `select_shape` hands a complex contraction to FMAddSub 4x6/W4 (ComplexF64) or 8x6/W8 (ComplexF32) when C's unit run along M fills whole `MR`-row slivers (`run == m_length || run % MR == 0`); otherwise it keeps planar 4x5/8x5. FMAddSub's single accumulator plane holds 12 accumulators + 2 A + 1 broadcast, 15 of 16 registers, and stays spill-free with B packed or in place. Planar NR = 6 fills all 16 and spills 13 per K step with B in place, so it is not a candidate. Where C's rows defeat the vector store, fmaddsub's lane-pair scalar store costs about 1.3x planar's per element, and with a short K the store dominates (ccsd_t_2/3/4, ao2mo_2/3), hence the run condition. `plan_contract` now computes the run for complex `T` too. The FMAddSub NR = 5 entries (4,5,4)/(8,5,8) leave the menus: nothing selects them, and the fenced NR = 6 tiles are spill-free. One `shape_override` row now gives both lane-pair kernels (1m, FMAddSub) their AVX2 shape; for OneMKernel ComplexF32 that is (8,6,8) where the fit used to hand it the AVX-512 menu head (24,8,16). Plans over the 490 suite :contract cases with uniform legs, same-eltype and both mixed-domain orientations, icelake-server/znver4/znver2 profiles, ad94b73 vs this branch: identical everywhere except 419 same-eltype plans on znver2, which change kernel (planar -> fmaddsub) and the blocking derived from it; the execution path is unchanged. Rome, same-node A/B of the same change on 4a77066, before the cleanup (#22) (jobs 7143376/7143377, two rounds each; GF/s of the default kernel, main planar 4x5 -> fmaddsub 4x6, OpenBLAS in brackets): ComplexF64 64^3 35.9 -> 40.6 [42.3], 256^3 42.1 -> 44.6 [49.3], 512^3 41.2 -> 48.2 [50.3], 2048^3 40.5 -> 45.1 [51.2], smallMN 29.3 -> 37.0 [30.2]; time geomean 0.87, no case slower ComplexF32 64^3 61.1 -> 76.7 [72.5], 256^3 80.7 -> 90.4 [95.2], 2048^3 81.5 -> 99.3 [102.8], smallN 44.0 -> 73.0 [53.1]; time geomean 0.79, no case slower Suite (complex :contract and :network): QuasiStrided time geomean 0.78 (ComplexF32) / 0.83 (ComplexF64) of main. Consistently slower: the dim32/96_1_3_1 contract_scrambled cases (1.23-1.47), and in ComplexF64 intensli_7_dim24 and dim8-16 0_2_3 cases (1.06-1.14). A diagnostic run traced the scrambled cases to the kernel method (fmaddsub), not the fence. Co-Authored-By: Claude Opus 5.5 --- src/planning/kernel_selection.jl | 27 ++++++++++++++--------- src/planning/plan.jl | 4 ++-- test/execution/test_unpackedb.jl | 2 +- test/microkernels/test_fmaddsub_kernel.jl | 6 ++--- test/planning/test_kernel_selection.jl | 24 ++++++++++++++++---- 5 files changed, 42 insertions(+), 21 deletions(-) diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index 5cb222e..51a0984 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -2,8 +2,8 @@ # not name one, a pure function of a `TargetProfile` and `T`. # Kernel types, unparameterised, name the complex-arithmetic schemes here. -# `OneMKernel`/`FMAddSubKernel` are chosen only by naming the kernel (FMAddSub -# also by the AVX-512 small-M demotion, `small_m_shape`). +# `OneMKernel` is chosen only by naming the kernel; `FMAddSubKernel` also on +# AVX2 and by the AVX-512 small-M demotion (`select_shape`). default_kernel_type(::Type{<:Real}) = SIMDKernel default_kernel_type(::Type{<:Complex}) = PlanarKernel @@ -38,10 +38,9 @@ const KERNEL_SHAPES_C32_PLANAR = ( (48, 3, 16), (32, 6, 16), (16, 8, 16), (8, 5, 8), (8, 6, 4), (4, 6, 4), ) const KERNEL_SHAPES_C32_ONEM = ((24, 8, 16), (32, 6, 16), (16, 8, 16), (8, 6, 8)) -# FMAddSub starts from 1m's shapes (same accumulator layout); the `NR = 5` AVX2 -# entries fit 16 registers, where `NR = 6` spills. -const KERNEL_SHAPES_C64_FMADDSUB = ((12, 8, 8), (8, 8, 8), (4, 6, 4), (4, 5, 4)) -const KERNEL_SHAPES_C32_FMADDSUB = ((24, 8, 16), (16, 8, 16), (8, 6, 8), (8, 5, 8)) +# FMAddSub starts from 1m's shapes (same accumulator layout). +const KERNEL_SHAPES_C64_FMADDSUB = ((12, 8, 8), (8, 8, 8), (4, 6, 4)) +const KERNEL_SHAPES_C32_FMADDSUB = ((24, 8, 16), (16, 8, 16), (8, 6, 8)) """ kernel_shapes(T, K = default_kernel_type(T)) -> NTuple{<:Any,NTuple{3,Int}} @@ -131,10 +130,9 @@ function shape_override(isa::Symbol, ::Type{T}, ::Type{K}) where {T, K} elseif K <: PlanarKernel && T === ComplexF32 isa === :avx512 && return (48, 3, 16) isa === :avx2 && return (8, 5, 8) - elseif K <: OneMKernel && T === ComplexF64 - # The AVX2-sized 1m shape, for a caller asking for `OneMKernel`'s shape - # (the fit would hand it an AVX-512 shape). - isa === :avx2 && return (4, 6, 4) + elseif K <: Union{OneMKernel, FMAddSubKernel} + # The AVX2-sized lane-pair shapes (the fit would hand them an AVX-512 one). + isa === :avx2 && return T === ComplexF64 ? (4, 6, 4) : (8, 6, 8) end return nothing end @@ -218,7 +216,8 @@ with, for compute type `T` and kernel type `K = default_kernel_type(T, TA, TB)`, extents `m_length`/`k_length` and C's unit-stride run along M (`run`; `m_length` when no layout is known). The steps, in order: - 1. The host's shape for `T` (`derived_shape`). + 1. The host's shape for `T` (`derived_shape`); for complex `T` on AVX2 the + FMAddSub one, where C's rows take its vector store. 2. Extent (real): the MV = 4 shape steps down to MV = 2 where `m_length` pads less (`extent_shape`). 3. Small M: an `m_length` below one tile takes the fitted shape, or for @@ -232,6 +231,12 @@ tuple to `UnionAll`. """ @inline function select_shape(::Type{T}, ::Type{K}, m_length::Int, k_length::Int, run::Int) where {T, K} profile = target_profile() + if K <: PlanarKernel && profile.isa === :avx2 + # FMAddSub's one accumulator plane fits `NR = 6` with a register to spare + # (planar's two spill), but its scalar store is the slower one. + shape = derived_shape(profile, T, FMAddSubKernel) + (run == m_length || run % shape[1] == 0) && return (shape, Val(FMAddSubKernel)) + end shape = extent_shape(derived_shape(profile, T, K), T, K, m_length) if m_length > 0 && m_length < shape[1] # Static, so a real `T`'s kernel type stays concrete. diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 6aabae2..c63aa86 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -225,8 +225,8 @@ function planned( morder = order_free_labels(mlabels, indC, C) norder = order_free_labels(nlabels, indC, C) - # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. - run_m = T <: Real || K isa MixedKernel ? leading_unit_run(morder, indC, C) : 0 + # Only real `T` swaps, so `run_n` is a placeholder for complex. + run_m = leading_unit_run(morder, indC, C) run_n = T <: Real ? leading_unit_run(norder, indC, C) : 0 mgroup = AxisGroup(morder, (indA, A), (indC, C)) # maps: (A, C) diff --git a/test/execution/test_unpackedb.jl b/test/execution/test_unpackedb.jl index 6d59547..3a7cf2d 100644 --- a/test/execution/test_unpackedb.jl +++ b/test/execution/test_unpackedb.jl @@ -55,7 +55,7 @@ end k0 = plan_contract(_mm_maker(T, 64, 8, 8, 2)()...).kernel W = lanewidth(k0) kernels = T <: Real ? (k0, SIMDKernel(Val(8), Val(6), T)) : - (k0, QuasiStrided.PlanarKernel(Val(W), Val(5), T, Val(W)), QuasiStrided.FMAddSubKernel(Val(W), Val(5), T, Val(W))) + (k0, QuasiStrided.PlanarKernel(Val(W), Val(5), T, Val(W)), QuasiStrided.FMAddSubKernel(Val(W), Val(6), T, Val(W))) for k in kernels, (m_block, k_block, n_block) in ((8, 3, 6), (40, 1, 100)) _ub_check(_mm_maker(T, 2 * tile_size(k, 1) + 1, 13, 2 * tile_size(k, 2) + 1, 3), 1.5, 0.5; kernel = k, m_block, k_block, n_block) end diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 62b7203..f892873 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -7,7 +7,7 @@ using SIMD: Vec const QSF = QuasiStrided @testset "FMAddSubKernel" begin - full = ((ComplexF64, (12, 8, 8)), (ComplexF32, (8, 5, 8))) + full = ((ComplexF64, (12, 8, 8)), (ComplexF32, (8, 6, 8))) for (T, (MR, NR, W)) in full mk_contract(FMAddSubKernel(Val(MR), Val(NR), T, Val(W))) end @@ -57,7 +57,7 @@ const QSF = QuasiStrided end @testset "PackedPanel operands (the driver's form) give the same bits as Vectors" begin - for (T, (MR, NR, W)) in ((ComplexF64, (4, 5, 4)), (ComplexF32, (16, 8, 16))) + for (T, (MR, NR, W)) in ((ComplexF64, (4, 6, 4)), (ComplexF32, (16, 8, 16))) k = FMAddSubKernel(Val(MR), Val(NR), T, Val(W)) rng = MersenneTwister(3) pa, pb = mk_pack(k, rand(rng, T, MR, 9), rand(rng, T, 9, NR)) @@ -70,7 +70,7 @@ const QSF = QuasiStrided # Exact instruction counts in the hot loop of `add_tile`, so only on # an FMA3 x86 host and not on hosted CI, whose virtualized CPU feature # sets do not reliably match. - avx2 = ((ComplexF64, (4, 5, 4)), (ComplexF32, (8, 5, 8)), (ComplexF64, (4, 6, 4)), (ComplexF32, (8, 6, 8))) + avx2 = ((ComplexF64, (4, 6, 4)), (ComplexF32, (8, 6, 8))) shapes = target_profile().isa === :avx512 ? ((ComplexF64, (8, 8, 8)), (ComplexF32, (16, 8, 16)), avx2...) : avx2 for (T, (MR, NR, W)) in shapes loop = mk_hot_loop(FMAddSubKernel(Val(MR), Val(NR), T, Val(W))) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index c2bd1f7..816064c 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -235,6 +235,22 @@ end end end +@testset "AVX2 complex: fmaddsub where C's rows take its vector store, else planar" begin + saved = QuasiStrided.TARGET[] + try + QuasiStrided.TARGET[] = synthetic(:avx2) + for T in (ComplexF64, ComplexF32) + W = 32 ÷ sizeof(real(T)) + fms, planar = ((W, NR_DEFAULT, W), Val(FMAddSubKernel)), ((W, 5, W), Val(PlanarKernel)) + select(m_length, run) = QuasiStrided.select_shape(T, PlanarKernel, m_length, 64, run) + @test select(64, 64) === select(64, 2W) === select(1, 1) === fms + @test select(64, W + 2) === select(64, 1) === planar + end + finally + QuasiStrided.TARGET[] = saved + end +end + @testset "mixed-domain selection: the real default of real(T), mapped" begin CR, RC = QuasiStrided.ComplexRealKernel, QuasiStrided.RealComplexKernel mapped(m, (MR, NR, W)) = m === CR ? (MR ÷ 2, NR, W) : (MR, NR ÷ 2, W) @@ -307,10 +323,10 @@ end profile = target_profile() @test planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarKernel)) - @test auto_kernel(T, 1024) isa QuasiStrided.PlanarKernel - # The small-M demotion: FMAddSub on AVX-512, planar elsewhere. - small = auto_kernel(T, 1) - @test small isa (profile.isa === :avx512 ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) + fms = profile.isa in (:avx512, :avx2) + @test auto_kernel(T, 1024) isa (profile.isa === :avx2 ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) + # The small-M demotion: FMAddSub on AVX-512 (and AVX2's default), planar elsewhere. + @test auto_kernel(T, 1) isa (fms ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) shape1m = first(kernel_shapes(T, OneMKernel)) k1m = QuasiStrided.kernel_from_shape(shape1m, T, OneMKernel)