AVX2 complex kernels: fenced K step, fmaddsub 4x6 default, vector pack/store

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

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")))

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 ---

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.

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)

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

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

test/microkernels/test_fmaddsub_kernel.jl

@@ -1,14 +1,13 @@
 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
 
 @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
@@ -58,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))
@@ -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, 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