From 8b1644321b96773c0772ce0e100a7e8c47c4317b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Mon, 28 Sep 2026 18:16:17 -0400 Subject: [PATCH 01/27] Convert storage lanes in the contiguous pack and vector store fast paths The contiguous packers (real A, complex planar/interleaved/1e) and the real, planar and fmaddsub vector stores now accept dense rank-1 storage of any supported type in the kernel's domain (Float32/Float64 for real kernels, ComplexF32/ComplexF64 for complex ones). Lanes load in the storage real type, convert to the kernel real type, and C rounds back once on store. Eligibility folds statically on the storage type. Same-eltype specializations emit identical LLVM IR to 42f962b (59 specializations of store_tile!, _store_tile_{vector,planar_vector, fmaddsub_vector}!, _pack_a_contiguous!, _pack_complex_contiguous!, unsafe_pack_a!/unsafe_pack_b! for F64/F32/C64/C32, diffed after normalizing specialization suffixes). Local pack_a! into a PackedPanel, kc=256, best of 15 (workstation, indicative): F64 <- F32 (MR=16): 1370 -> 615 ns (same-type F64: 633 ns) F32 <- F64 (MR=32): 3083 -> 903 ns (same-type F32: 627 ns) C64 <- C32 planar: 1817 -> 639 ns (same-type C64: 661 ns) C32 <- C64 planar: 3245 -> 909 ns (same-type C32: 662 ns) Full Pkg.test: 778 s wall. probe_ttfx (3 runs, after/base): load 0.5/0.5 s, first call F64 3.0/2.9 s, C64 2.0/2.1 s, K=1 outer 0.5/0.5 s. Co-Authored-By: Claude Opus 5.5 --- src/microkernels/fmaddsub.jl | 19 ++++---- src/microkernels/planar.jl | 27 ++++++----- src/microkernels/simd.jl | 21 +++++--- src/packing/pack_contiguous.jl | 48 +++++++++++-------- .../test_planar_store_fastpath.jl | 17 ++++--- test/microkernels/test_simd_kernel.jl | 21 +++++++- test/packing/test_pack_complex_contiguous.jl | 15 +++--- test/packing/test_pack_real.jl | 3 +- 8 files changed, 109 insertions(+), 62 deletions(-) diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 534e769..7bd939f 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -174,25 +174,25 @@ end end # One full `W÷2`-row block, already in `Complex`'s memory order, so no -# interleave shuffle. Base's `Complex` expression trees in lanes, as planar's +# interleave shuffle. `C` converts from and to its storage type `RC`. Base's `Complex` expression trees in lanes, as planar's # `_planar_store_block!`: # beta == 0: addsub(ar*r, ai*swap(r)) # beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) # otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) @inline function _fmaddsub_store_block!( - sp::Ptr{R}, at::Int, r::Vec{W, R}, + sp::Ptr{RC}, at::Int, r::Vec{W, R}, ar::Vec{W, R}, ai::Vec{W, R}, br::Vec{W, R}, bi::Vec{W, R}, beta::Complex{R} - ) where {R, W} + ) where {RC, R, W} s = _swap_pairs(r) if iszero(beta) new = _addsub(ar * r, ai * s) else - old = vload(Vec{W, R}, sp + sizeof(R) * at) + old = convert(Vec{W, R}, vload(Vec{W, RC}, sp + sizeof(RC) * at)) x = isone(beta) ? old : _addsub(br * old, bi * _swap_pairs(old)) new = _fmaddsub(ar, r, _fmaddsub(ai, s, x)) end - vstore(new, sp + sizeof(R) * at) + vstore(convert(Vec{W, RC}, new), sp + sizeof(RC) * at) return nothing end @@ -202,9 +202,10 @@ end alpha::T, beta::T, kernel::FMAddSubKernel{MR, NR, T, W}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NV} - # The pointer reinterpretation is only sound on dense rank-1 storage of `T`. - S <: DenseVector{T} || - throw(ArgumentError("_store_tile_fmaddsub_vector!: storage $S is not a DenseVector{$T}")) + # The pointer reinterpretation is only sound on dense rank-1 complex storage. + S <: DenseVector && _lane_convertible(eltype(S), T) || + throw(ArgumentError("_store_tile_fmaddsub_vector!: storage $S is not a dense vector convertible to $T")) + RC = real(eltype(S)) iseven(W) || throw(ArgumentError("_store_tile_fmaddsub_vector!: requires an even W, got $W")) MV = (2 * MR) ÷ W _check_acc(:_store_tile_fmaddsub_vector!, R, T, NV, MV * NR) @@ -254,7 +255,7 @@ end br = Vec{$W, $R}(real(beta)) bi = Vec{$W, $R}(imag(beta)) GC.@preserve storage begin - sp = reinterpret(Ptr{$R}, pointer(storage)) + sp = reinterpret(Ptr{$RC}, pointer(storage)) @inbounds begin $(blocks...) end diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 2da99d2..2207fda 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -173,10 +173,11 @@ end end # Vector store: unit-stride rows into rank-1 dense `Complex` storage (so the -# storage can be reinterpreted as `2W` consecutive reals per `W` rows), on an -# ISA the complex fast paths ship for (shared with the complex pack fast path). +# storage can be reinterpreted as `2W` consecutive reals per `W` rows) of `T` or +# another complex type the lanes convert to and from, on an ISA the complex +# fast paths ship for (shared with the complex pack fast path). @inline _complex_vector_eligible(tile::QSTile, ::Type{T}) where {T} = - _unit_stride_rows(tile.rows) && tile.storage isa DenseVector{T} && + _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) && _complex_fastpath_isa_eligible() # Shuffle patterns built from `W` at specialization time, never hardcoded to @@ -198,7 +199,8 @@ end return :(shufflevector(re, im, Val($idx))) end -# One full `W`-row block; `at` is the zero-based index of its first real. +# One full `W`-row block; `at` is the zero-based index of its first real. `C` +# is `RC` in memory and converts to and from `R` around the arithmetic. # The arithmetic transcribes Base's `Complex` expression trees (the ones # `_axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, # `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, @@ -206,15 +208,15 @@ end # `beta == 0/1` and ~1 ULP otherwise, because LLVM contracts Base's scalar # complex `muladd` depending on inlining context. @inline function _planar_store_block!( - sp::Ptr{R}, at::Int, rev::Vec{W, R}, imv::Vec{W, R}, + sp::Ptr{RC}, at::Int, rev::Vec{W, R}, imv::Vec{W, R}, ar::Vec{W, R}, ai::Vec{W, R}, br::Vec{W, R}, bi::Vec{W, R}, beta::Complex{R}, ::Val{W} - ) where {R, W} + ) where {RC, R, W} if iszero(beta) newre = ar * rev - ai * imv newim = ar * imv + ai * rev else - old = vload(Vec{2 * W, R}, sp + sizeof(R) * at) + old = convert(Vec{2 * W, R}, vload(Vec{2 * W, RC}, sp + sizeof(RC) * at)) orv = _deinterleave_re(old, Val(W)) oiv = _deinterleave_im(old, Val(W)) if isone(beta) @@ -226,7 +228,7 @@ end newre = muladd(ar, rev, -muladd(ai, imv, -xr)) newim = muladd(ar, imv, muladd(ai, rev, xi)) end - vstore(_interleave_planes(newre, newim, Val(W)), sp + sizeof(R) * at) + vstore(convert(Vec{2 * W, RC}, _interleave_planes(newre, newim, Val(W))), sp + sizeof(RC) * at) return nothing end @@ -237,9 +239,10 @@ end alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NA} - # The pointer reinterpretation is only sound on dense rank-1 storage of `T`. - S <: DenseVector{T} || - throw(ArgumentError("_store_tile_planar_vector!: storage $S is not a DenseVector{$T}")) + # The pointer reinterpretation is only sound on dense rank-1 complex storage. + S <: DenseVector && _lane_convertible(eltype(S), T) || + throw(ArgumentError("_store_tile_planar_vector!: storage $S is not a dense vector convertible to $T")) + RC = real(eltype(S)) MV = MR ÷ W NV = MV * NR _check_acc(:_store_tile_planar_vector!, R, T, NA, 2NV) @@ -291,7 +294,7 @@ end br = Vec{$W, $R}(real(beta)) bi = Vec{$W, $R}(imag(beta)) GC.@preserve storage begin - sp = reinterpret(Ptr{$R}, pointer(storage)) + sp = reinterpret(Ptr{$RC}, pointer(storage)) @inbounds begin $(blocks...) end diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index f052b37..2584fe4 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -89,11 +89,12 @@ end return acc end -# Vector store eligibility: unit-stride rows into rank-1 dense storage of `T`, -# exactly what SIMD.jl's array `vload`/`vstore` accept. Must admit `Memory{T}`: -# that is the `parent` of an Array-backed `StridedView` on Julia >= 1.11. +# Vector store eligibility: unit-stride rows into rank-1 dense storage of `T` +# or of another real type the lanes convert to and from, exactly what SIMD.jl's +# array `vload`/`vstore` accept. Must admit `Memory`: that is the `parent` of an +# Array-backed `StridedView` on Julia >= 1.11. @inline _vector_store_eligible(tile::QSTile, ::Type{T}) where {T} = - _unit_stride_rows(tile.rows) && tile.storage isa DenseVector{T} + _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) @generated function _store_tile_scattered!( destination::QSTile, acc::NTuple{NV, Vec{W, T}}, @@ -128,12 +129,15 @@ end # Whole `W`-row blocks are one vector load/store; a block straddling `m` is # stored lane by lane, so nothing outside the valid rectangle is touched. # `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. +# `C` is loaded into and rounded back from `T` lanes. @generated function _store_tile_vector!( destination::QSTile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, m::Int, n::Int ) where {S, MR, NR, T, W, NV} NVECA = MR ÷ W + RC = eltype(S) + old = :(convert(Vec{$W, $T}, vload(Vec{$W, $RC}, storage, at))) blocks = Any[] for j in 0:(NR - 1) vblocks = Any[] @@ -145,9 +149,12 @@ end if $((v + 1) * W) <= m at = colbase + $(v * W) + 1 vstore( - iszero(beta) ? alpha * vec : - isone(beta) ? muladd(alpha, vec, vload(Vec{$W, $T}, storage, at)) : - muladd(alpha, vec, beta * vload(Vec{$W, $T}, storage, at)), + convert( + Vec{$W, $RC}, + iszero(beta) ? alpha * vec : + isone(beta) ? muladd(alpha, vec, $old) : + muladd(alpha, vec, beta * $old) + ), storage, at ) elseif $(v * W) < m diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 27f4347..6b60eb3 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -10,27 +10,37 @@ using SIMD: shufflevector @inline _copies_unchanged(::typeof(conj), ::Type{T}) where {T <: Real} = true @inline _copies_unchanged(::Any, ::Type) = false +# Dense rank-1 storage whose elements load as SIMD lanes and convert lane-wise +# to `T`: `T` itself, or another supported type of the same domain. Shared with +# the vector stores. +const _LaneFloat = Union{Float32, Float64} +@inline _dense_lanes(storage::S, ::Type{T}) where {S, T} = + storage isa DenseVector{T} || (storage isa DenseVector && _lane_convertible(eltype(S), T)) +_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 + # Everything but `m == MR` and the stride test folds at compile time. # `_unit_stride_rows` deliberately has no fallback method: an unknown axis type # must be a MethodError. @inline function _pack_a_contiguous_eligible( packed::V, source::QSTile, transform::F, m::Int, ::Val{MR}, ::Type{T} ) where {V, F, MR, T} - return packed isa PackedPanel{T} && source.storage isa DenseVector{T} && + return packed isa PackedPanel{T} && _dense_lanes(source.storage, T) && _copies_unchanged(transform, T) && m == MR && _unit_stride_rows(source.rows) end -# Real A: one `Vec{MR,T}` load/store per K step. +# Real A: one `Vec{MR}` load/convert/store per K step. @inline function _pack_a_contiguous!( - packed::PackedPanel{T}, storage::DenseVector{T}, rowbase::Int, cols::C, + packed::PackedPanel{T}, storage::DenseVector{S}, rowbase::Int, cols::C, ::Val{MR}, kc::Int - ) where {T, C, MR} + ) where {T, S, C, MR} GC.@preserve storage begin sp = pointer(storage) dp = packed.ptr for p in 0:(kc - 1) - v = vload(Vec{MR, T}, sp + sizeof(T) * (rowbase + axis_offset(cols, p))) - vstore(v, dp + sizeof(T) * (MR * p)) + v = vload(Vec{MR, S}, sp + sizeof(S) * (rowbase + axis_offset(cols, p))) + vstore(convert(Vec{MR, T}, v), dp + sizeof(T) * (MR * p)) end end return packed @@ -51,7 +61,7 @@ end valid::Int, ::Val{PD}, ::Type{T} ) where {V, S, AX, F, FMT, PD, T} return !(format isa RealFormat) && - packed isa PackedPanel{real(T)} && storage isa DenseVector{T} && + packed isa PackedPanel{real(T)} && _dense_lanes(storage, T) && _complex_pack_transform_eligible(transform) && valid == PD && _unit_stride_rows(lane_axis) && _complex_fastpath_isa_eligible() @@ -102,21 +112,21 @@ _single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a # Lane `t` of K step `p` is element `elembase + axis_offset(steps, p) + t`; # only the lane axis must be unit-stride, `steps` may be scattered. Pinning -# `Complex{R}` in the signature makes an eltype mismatch a MethodError rather -# than a bitcast to the wrong width. +# `Complex{RS}` in the signature keeps the bitcast to `RS` lanes sound; the +# lanes convert to `R` before the shuffle. @inline function _pack_complex_contiguous!( format::Union{PlanarFormat, InterleavedFormat}, packed::PackedPanel{R}, - storage::DenseVector{Complex{R}}, elembase::Int, steps::C, ::Val{PD}, kc::Int, + storage::DenseVector{Complex{RS}}, elembase::Int, steps::C, ::Val{PD}, kc::Int, transform::F - ) where {R, C, PD, F} + ) where {R, RS, C, PD, F} shuffle = _single_region_shuffle(format) GC.@preserve storage begin - sp = reinterpret(Ptr{R}, pointer(storage)) + sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr for p in 0:(kc - 1) - src = vload( + src = convert( Vec{2 * PD, R}, - sp + sizeof(R) * (2 * (elembase + axis_offset(steps, p))) + vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) ) vstore( shuffle(src, _pack_alt(src, transform), Val(PD)), @@ -128,16 +138,16 @@ _single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a end @inline function _pack_complex_contiguous!( - ::OneEFormat, packed::PackedPanel{R}, storage::DenseVector{Complex{R}}, + ::OneEFormat, packed::PackedPanel{R}, storage::DenseVector{Complex{RS}}, elembase::Int, steps::C, ::Val{PD}, kc::Int, transform::F - ) where {R, C, PD, F} + ) where {R, RS, C, PD, F} GC.@preserve storage begin - sp = reinterpret(Ptr{R}, pointer(storage)) + sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr for p in 0:(kc - 1) - src = vload( + src = convert( Vec{2 * PD, R}, - sp + sizeof(R) * (2 * (elembase + axis_offset(steps, p))) + vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) ) at = 4 * PD * p vstore( diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 342c749..e9a4441 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -43,10 +43,11 @@ mk_store_ab(T) = ( ) # `beta0_exact`: the fast path also matches the scalar store bitwise at beta == 0. -function mk_store_fastpath(K, T, shapes; beta0_exact = false) +# `S` is the storage eltype: C converts to `T` on load and rounds once on store. +function mk_store_fastpath(K, T, shapes; beta0_exact = false, S = T) R = real(T) - tol = 8 * eps(R) - return @testset "$K store fast path $T $s" for s in shapes + tol = 8 * max(eps(R), eps(real(S))) + return @testset "$K store fast path $T into $S $s" for s in shapes MR, NR, W = s k = K(Val(MR), Val(NR), T, Val(W)) blk = k isa PlanarKernel ? W : W ÷ 2 # complex rows per vector @@ -60,7 +61,7 @@ function mk_store_fastpath(K, T, shapes; beta0_exact = false) ) for (m, n) in extents acc = map(v -> typeof(v)(ntuple(_ -> R(2rand(rng) - 1), W)), zero_accumulator(k)) - cold = [Complex(R(2rand(rng) - 1), R(2rand(rng) - 1)) for _ in 1:(m * n)] + cold = [S(2rand(rng) - 1, 2rand(rng) - 1) for _ in 1:(m * n)] for (alpha, beta) in mk_store_ab(T) fast = mk_dense(cold) dfast = DestinationTile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) @@ -68,7 +69,7 @@ function mk_store_fastpath(K, T, shapes; beta0_exact = false) store_tile!(dfast, acc, alpha, beta, k) scal = copy(cold) # scattered rows: always the scalar store store_tile!(DestinationTile(scal, 0, ScatterAxis(collect(0:(m - 1)), m), AffineAxis(0, m, n)), acc, alpha, beta, k) - want = [ref_axpby(alpha, reim(mk_read(k, acc, i, j))..., beta, cold[i + j * m + 1]) for i in 0:(m - 1), j in 0:(n - 1)] + want = [S(ref_axpby(alpha, reim(mk_read(k, acc, i, j))..., beta, T(cold[i + j * m + 1]))) for i in 0:(m - 1), j in 0:(n - 1)] got = reshape(collect(fast), m, n) vectorized = STORE_FASTPATH_ON ? (1:((m ÷ blk) * blk)) : (1:0) @test isequal(got[vectorized, :], want[vectorized, :]) @@ -83,6 +84,8 @@ end @testset "planar store fast path" begin mk_store_fastpath(PlanarKernel, ComplexF64, KERNEL_SHAPES_C64_PLANAR) mk_store_fastpath(PlanarKernel, ComplexF32, KERNEL_SHAPES_C32_PLANAR) + mk_store_fastpath(PlanarKernel, ComplexF64, KERNEL_SHAPES_C64_PLANAR[[1, end]]; S = ComplexF32) + mk_store_fastpath(PlanarKernel, ComplexF32, KERNEL_SHAPES_C32_PLANAR[[1, end]]; S = ComplexF64) @testset "eligibility gate: one violated condition at a time" begin T = ComplexF64 @@ -98,7 +101,7 @@ end GC.@preserve ptr_rows @test !eligible(storage, PtrScatterAxis(pointer(ptr_rows), m)) @test !eligible(view(storage, 1:(m * n)), AffineAxis(0, 1, m)) @test !eligible(reshape(storage, 4m, n), AffineAxis(0, 1, m)) - @test !eligible(zeros(ComplexF32, m * n), AffineAxis(0, 1, m)) + @test eligible(zeros(ComplexF32, m * n), AffineAxis(0, 1, m)) == STORE_FASTPATH_ON @test !eligible(zeros(Float64, m * n), AffineAxis(0, 1, m)) end @@ -117,4 +120,6 @@ end @testset "fmaddsub store fast path" begin mk_store_fastpath(FMAddSubKernel, ComplexF64, KERNEL_SHAPES_C64_FMADDSUB; beta0_exact = true) mk_store_fastpath(FMAddSubKernel, ComplexF32, KERNEL_SHAPES_C32_FMADDSUB; beta0_exact = true) + mk_store_fastpath(FMAddSubKernel, ComplexF64, KERNEL_SHAPES_C64_FMADDSUB[[1, end]]; beta0_exact = true, S = ComplexF32) + mk_store_fastpath(FMAddSubKernel, ComplexF32, KERNEL_SHAPES_C32_FMADDSUB[[1, end]]; beta0_exact = true, S = ComplexF64) end diff --git a/test/microkernels/test_simd_kernel.jl b/test/microkernels/test_simd_kernel.jl index a9f4e53..7feecb6 100644 --- a/test/microkernels/test_simd_kernel.jl +++ b/test/microkernels/test_simd_kernel.jl @@ -18,7 +18,7 @@ using StridedViews: StridedView @test_throws ArgumentError SIMDKernel(Val(6), Val(4), Float64, Val(4)) @test_throws ArgumentError SIMDKernel(Val(6), Val(4), Float64, Val(0)) - @testset "_vector_store_eligible: unit-stride rows into 1-D dense storage of T" begin + @testset "_vector_store_eligible: unit-stride rows into 1-D dense real storage" begin for T in (Float64, Float32) m, n = 8, 8 rows, cols = AffineAxis(0, 1, m), AffineAxis(0, m, n) @@ -29,7 +29,24 @@ using StridedViews: StridedView @test !_vector_store_eligible(DestinationTile(zeros(T, m, n), 0, rows, cols), T) @test !_vector_store_eligible(DestinationTile(zeros(T, 2m * n), 0, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)), T) @test !_vector_store_eligible(DestinationTile(mem, 0, ScatterAxis(collect(0:(m - 1)), m), cols), T) - @test !_vector_store_eligible(DestinationTile(mem, 0, rows, cols), T === Float64 ? Float32 : Float64) + @test _vector_store_eligible(DestinationTile(mem, 0, rows, cols), T === Float64 ? Float32 : Float64) + @test !_vector_store_eligible(DestinationTile(zeros(complex(T), m * n), 0, rows, cols), T) + end + end + + @testset "vector store into converted storage: $T kernel, $S C" for (T, S) in ((Float64, Float32), (Float32, Float64)) + k = SIMDKernel(Val(8), Val(3), T) + rng = MersenneTwister(5) + for (m, n) in ((8, 3), (7, 3), (3, 2)), (alpha, beta) in mk_alphabeta(T) + acc = map(v -> typeof(v)(ntuple(_ -> T(2rand(rng) - 1), lanewidth(k))), zero_accumulator(k)) + cold = S.(2 .* rand(rng, m * n) .- 1) + fast = mk_dense(cold) + dfast = DestinationTile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) + @test _vector_store_eligible(dfast, T) + store_tile!(dfast, acc, alpha, beta, k) + scal = copy(cold) + store_tile!(DestinationTile(scal, 0, ScatterAxis(collect(0:(m - 1)), m), AffineAxis(0, m, n)), acc, alpha, beta, k) + @test mk_close(fast, scal, Float32) end end end diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index 3f3bbba..c6b2575 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -23,8 +23,9 @@ const FAST_A_CASES = [ const FAST_B_NRS = (1, 3, 7, 8) # Packs into a `PackedPanel` (the only destination the gate admits). -function check_panel(side, kernel, T, PD, fmt, lane, step, f) - storage = complex_storage(T, 40000) +# `S` is the storage eltype; a mixed one holds values `T` must round. +function check_panel(side, kernel, T, PD, fmt, lane, step, f; S = T) + storage = S === T ? complex_storage(T, 40000) : S.(complex_storage(ComplexF64, 40000) ./ 3) src, g = pack_fixture(side, storage, 17, lane, step) pack! = side === :a ? pack_a! : pack_b! len = _ref_rpe(fmt) * PD * axis_length(step) @@ -36,8 +37,8 @@ end @testset "complex pack fast path: A panel, $T / $(typeof(fa)) / MR=$MR" for (T, fa, MR) in FAST_A_CASES kernel = ComplexKernelDescriptor(Val(MR), Val(3), T, fa, PlanarFormat()) - for f in (identity, conj) - storage, src = check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), AffineAxis(0, 997, 5), f) + for f in (identity, conj), S in (T, T === ComplexF64 ? ComplexF32 : ComplexF64) + storage, src = check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), AffineAxis(0, 997, 5), f; S) pp = packed_panel(zeros(real(T), 1), 1, 1) @test QS._pack_complex_contiguous_eligible(pp, storage, src.rows, f, fa, MR, Val(MR), T) == FASTPATH_ON @@ -46,8 +47,8 @@ end @testset "complex pack fast path: B panel, $T / NR=$NR" for T in (ComplexF64, ComplexF32), NR in FAST_B_NRS kernel = ComplexKernelDescriptor(Val(4), Val(NR), T, PlanarFormat(), PlanarFormat()) - for f in (identity, conj) - storage, src = check_panel(:b, kernel, T, NR, PlanarFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, 5), f) + for f in (identity, conj), S in (T, T === ComplexF64 ? ComplexF32 : ComplexF64) + storage, src = check_panel(:b, kernel, T, NR, PlanarFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, 5), f; S) pp = packed_panel(zeros(real(T), 1), 1, 1) @test QS._pack_complex_contiguous_eligible(pp, storage, src.cols, f, PlanarFormat(), NR, Val(NR), T) == FASTPATH_ON @@ -90,6 +91,8 @@ end @test !elig(packed = packed_panel(Float32[0.0f0], 1, 1)) # wrong real type @test !elig(store = view(storage, 1:100)) # not a DenseVector @test !elig(store = collect(reshape(storage, 40, 100))) # rank 2 + @test elig(store = ComplexF32.(storage)) == FASTPATH_ON # converted lanes + @test !elig(store = real.(storage)) # real storage @test !elig(ax = AffineAxis(0, 2, MR)) @test !elig(ax = AffineAxis(0, -1, MR)) @test !elig(ax = ScatterAxis(offs, MR)) diff --git a/test/packing/test_pack_real.jl b/test/packing/test_pack_real.jl index 71efd76..0d315a2 100644 --- a/test/packing/test_pack_real.jl +++ b/test/packing/test_pack_real.jl @@ -104,7 +104,8 @@ end kernel = KernelDescriptor(Val(MR), Val(3), T) vals = T.(collect(1.0:2000.0)) - storages = @static isdefined(Base, :Memory) ? (vals, copyto!(Memory{T}(undef, 2000), vals)) : (vals,) + mixed = (T === Float64 ? Float32 : Float64).(collect(1.0:2000.0) ./ 3) + storages = @static isdefined(Base, :Memory) ? (vals, copyto!(Memory{T}(undef, 2000), vals), mixed) : (vals, mixed) koffs = [7, 900, 300, 1500] contig = collect(0:(MR - 1)) GC.@preserve koffs contig for storage in storages From 096ab7f85a4368e0963e9928b7c6e4fbfa25d83e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 08:19:43 -0400 Subject: [PATCH 02/27] Mixed operand eltypes with a compute type and an optional accumulator plan_contract, contract! and QuasiStridedBackend accept any combination of Float32/Float64/ComplexF32/ComplexF64 operands, except a complex input into a real C. The compute type is promote_type(TA, TB, TC), or with `accumulator = Float32/Float64` that precision in the promoted domain; the plan, workspace and kernel are keyed on it. The scalar store helpers, the unpacked-B loads, the dot gather and the outer-product B load convert to the compute type explicitly; a real operand of a complex contraction is promoted by the generic packer. Same-eltype code is unchanged: code_llvm of _axpby_tile!/_axpby_at!/ scale_tile!/execute! (18 F64/C64 specializations) is identical to 8b16443 modulo name suffixes, and return types of plan_contract/_planned (automatic and named kernel) and execute! are unchanged. Partial sums still round-trip through C between kc blocks of K, so a C narrower than the compute type rounds once per K block, not once overall. Workstation under load (loadavg ~16), indicative only: Full Pkg.test, run concurrently with 8b16443: 1281 s vs 1250 s (+2.5%). probe_ttfx, 3 runs (new / 8b16443): load 0.5-0.6 / 0.5-0.6 s, first F64 3.7-4.4 / 3.8-4.2 s, C64 2.3-3.0 / 2.7-3.0 s, K=1 outer 0.7 / 0.6-0.8 s. Fresh-process first @tensor, 3 runs: F64xF64 3.2-5.8 s, F32xF64->F64 3.3-3.6 s, C64xC64 3.2-3.5 s, F64xC64->C64 3.1-3.6 s. Co-Authored-By: Claude Opus 5.5 --- README.md | 19 ++++--- src/execution/dot.jl | 2 +- src/execution/execute.jl | 14 +++-- src/execution/outer.jl | 2 +- src/execution/unpackedb.jl | 5 +- src/execution/workspace.jl | 14 ++--- src/integrations/tensoroperations.jl | 61 ++++++++++++-------- src/microkernels/interface.jl | 17 +++--- src/planning/plan.jl | 65 ++++++++++++++++------ test/execution/test_mixed.jl | 62 +++++++++++++++++++++ test/integrations/test_tensoroperations.jl | 24 ++++++-- test/planning/test_plan_contract.jl | 13 +++++ test/runtests.jl | 1 + 13 files changed, 221 insertions(+), 78 deletions(-) create mode 100644 test/execution/test_mixed.jl diff --git a/README.md b/README.md index af1eab4..2ed41a6 100644 --- a/README.md +++ b/README.md @@ -117,8 +117,13 @@ Implemented: from an analytical model of the detected cache sizes (fixed constants when they are undetected); plus automatically selected paths for dot-product, outer-product and read-B-in-place shapes. -- `Float32`/`Float64`/`ComplexF32`/`ComplexF64`; `conjA`/`conjB` and each - operand's `StridedView.op` are applied during packing. +- `Float32`/`Float64`/`ComplexF32`/`ComplexF64`, mixed freely (a complex input + needs a complex output). The compute type is the promoted eltype, or the + precision chosen with `accumulator = Float32`/`Float64` + (`plan_contract(...; accumulator)`, `QuasiStridedBackend(; accumulator)`); + `C` is rounded on store, once per `kc` block of K. A real operand of a complex contraction is + promoted at pack time. `conjA`/`conjB` and each operand's `StridedView.op` + are applied during packing. - Zero steady-state allocation on Julia >= 1.11 for `execute!` on a reused plan and for `tensorcontract!` through the backend (which pools workspaces per task). On Julia 1.10 `SIMDKernel`'s accumulator is not kept in @@ -127,11 +132,11 @@ Implemented: Not implemented: - Anything but contraction: `tensoradd!`/`tensortrace!` under this backend are - forwarded to `StridedNative()`, so mixed `@tensor` networks still run. -- Mixed eltypes (including real with complex; promotion is TensorOperations' - `promote_contract` job), other eltypes, non-strided operands, an output - aliased with an input, and a conjugated output view: `tensorcontract!` throws - an `ArgumentError` instead of falling back to another backend. + forwarded to `StridedNative()`, so whole `@tensor` networks still run. +- Other eltypes, a complex input with a real output, non-strided operands, an + output aliased with an input, and a conjugated output view: `tensorcontract!` + throws an `ArgumentError` instead of falling back to another backend. +- Dedicated real-times-complex kernels. - Being the default backend: `StridedBLAS()` is faster on most shapes. - Batch axes, traces/diagonals, threading, GPU execution, the 3m complex method, autotuning. diff --git a/src/execution/dot.jl b/src/execution/dot.jl index 44b5c2c..3e9dd0b 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -174,7 +174,7 @@ end gptr::Ptr{T}, vstorage::SV, vbase::Int, vkbuf::Vector{Int}, kblock::Int, transform::F ) where {T, SV, F} @inbounds for t in 0:(kblock - 1) - z = transform(vstorage[vbase + vkbuf[t + 1] + 1])::T + z = convert(T, transform(vstorage[vbase + vkbuf[t + 1] + 1])) unsafe_store!(gptr + sizeof(T) * t, z) if T <: Complex unsafe_store!(gptr + sizeof(T) * (kblock + t), Complex(imag(z), real(z))) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 6c4e0de..f8171e2 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -352,25 +352,27 @@ end A::StridedView, indA::NTuple{NA,Int}, B::StridedView, indB::NTuple{NB,Int}, beta::Number, - indC::NTuple{NC,Int}) where {NA,NB,NC} + indC::NTuple{NC,Int}; accumulator = nothing) where {NA,NB,NC} Compute `C[indC] = alpha * sum_K A[indA] * B[indB] + beta * C[indC]`, with one `Int` label per axis: a label in `indA` and `indB` but not `indC` is contracted (K), and a label in `indC` and exactly one of `indA`/`indB` is free (M or N). Any other label pattern, or a label repeated within one tuple, throws an `ArgumentError`; matched labels of unequal axis length throw a -`DimensionMismatch`. Equivalent to -`execute!(plan_contract(C, A, indA, B, indB, indC), alpha, beta)` — use -those directly to reuse a plan across calls. Returns `C`. +`DimensionMismatch`. Eltypes and `accumulator` are as in +[`plan_contract`](@ref). Equivalent to +`execute!(plan_contract(C, A, indA, B, indB, indC; accumulator), alpha, beta)` +— use those directly to reuse a plan across calls. Returns `C`. """ function contract!( C::StridedView, alpha::Number, A::StridedView, indA::NTuple{NA, Int}, B::StridedView, indB::NTuple{NB, Int}, beta::Number, - indC::NTuple{NC, Int} + indC::NTuple{NC, Int}; + accumulator::Union{Nothing, Type{Float32}, Type{Float64}} = nothing ) where {NA, NB, NC} - plan = plan_contract(C, A, indA, B, indB, indC) + plan = plan_contract(C, A, indA, B, indB, indC; accumulator) execute!(plan, alpha, beta) return C end diff --git a/src/execution/outer.jl b/src/execution/outer.jl index 4212e1f..7df91f4 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -63,7 +63,7 @@ end sz = sizeof(T) mmain = (Qm ÷ W) * W @inbounds for n in 1:nblock - b = btransform(Bstorage[Bbase + bufB[n] + 1])::T + b = convert(T, btransform(Bstorage[Bbase + bufB[n] + 1])) bv = Vec{W, T}(b) cpn = cp + sz * bufC[n] m = 0 diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 414a713..65c69dc 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -26,11 +26,12 @@ Base.length(u::UnpackedBView) = u.per_k * axis_length(u.ksteps) return u.transform(z) end -@inline _b_step_load(u::UnpackedBView, kernel, j::Int, p::Int) = _unpacked_b_element(u, j, p) +@inline _b_step_load(u::UnpackedBView, kernel, j::Int, p::Int) = + convert(scalartype(kernel), _unpacked_b_element(u, j, p)) # Complex kernels: the planar `(re, im)` pair. @inline function _b_step_load2(u::UnpackedBView, kernel, j::Int, p::Int) - z = _unpacked_b_element(u, j, p) + z = convert(scalartype(kernel), _unpacked_b_element(u, j, p)) return (real(z), imag(z)) end diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 5a4f49e..550f71b 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -13,7 +13,7 @@ _SlotCache() = _SlotCache(nothing, IdDict{Any, Any}()) The buffers [`execute!`](@ref) needs. Pass one back as `plan_contract(...; workspace = ws)` to reuse (and grow) it across contractions; one from a non-default allocator is single-use and must be -[`release!`](@ref)d. `T` is the storage element type; the packed panels (`VT`) +[`release!`](@ref)d. `T` is the plan's compute type; the packed panels (`VT`) hold `real(T)`. Field layout is an implementation detail. """ # A `mutable struct` with `const` fields so a plan and a barrier slot hold one @@ -66,7 +66,7 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} eltype(VT) === real(T) || throw( ArgumentError( "ContractWorkspace{$T,$VT}: the packed panels must hold $(real(T)) " * - "(the real type of the storage element type $T), got $(eltype(VT))" + "(the real type of the compute type $T), got $(eltype(VT))" ) ) return new{T, VT}( @@ -134,7 +134,7 @@ end ContractWorkspace(T, kernel, blocking::Blocking; oracle = true, allocator = TensorOperations.DefaultAllocator()) -A workspace for storage type `T`, `kernel` and an effective `blocking`. +A workspace for compute type `T`, `kernel` and an effective `blocking`. `oracle = false` leaves the `execute_tilewise!` buffers empty. Under a non-default `allocator` the packed panels come from `tensoralloc`, are never resized, and must be handed back with [`release!`](@ref). @@ -259,7 +259,7 @@ function _resolve_workspace( return ContractWorkspace(T, kernel, blocking, oracle, allocator) end -# The pool is keyed by `eltype(C)` alone, so also check the packed type +# The pool is keyed by the compute type alone, so also check the packed type # (folds at compile time). @inline function _reuse_workspace( ::Type{T}, ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, @@ -283,9 +283,9 @@ end ) where {T} throw( ArgumentError( - "plan_contract: cannot reuse a $(typeof(ws)) for an eltype-$T contraction " * - "on the default allocator; only a " * - "ContractWorkspace{$T,Vector{$(realtype(kernel))}} -- storage element " * + "plan_contract: cannot reuse a $(typeof(ws)) for a contraction of compute " * + "type $T on the default allocator; only a " * + "ContractWorkspace{$T,Vector{$(realtype(kernel))}} -- compute " * "type $T, packed panels of $(realtype(kernel)) -- is `reserve!`-able" ) ) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index ec30a40..4960b58 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -5,25 +5,35 @@ import TupleTools using StridedViews: StridedView, isstrided """ - QuasiStridedBackend() + QuasiStridedBackend(; accumulator = nothing) TensorOperations backend running contractions on QuasiStrided's engine: @tensor backend = QuasiStridedBackend() C[i, j] := A[i, k] * B[k, j] -`tensorcontract!` accepts strided operands sharing one element type out of -`Float32`/`Float64`/`ComplexF32`/`ComplexF64`, and honours `conjA`/`conjB`. Any -other input (mixed or other eltypes, a non-strided operand, an output aliased +`tensorcontract!` accepts strided operands with element types out of +`Float32`/`Float64`/`ComplexF32`/`ComplexF64`, possibly mixed, and honours +`conjA`/`conjB`. The compute type and `accumulator` (`nothing`, `Float32` or +`Float64`) are as in [`plan_contract`](@ref). Any other input (other eltypes, a +complex input with a real output, a non-strided operand, an output aliased with an input, a conjugated output view) throws an `ArgumentError`; it never falls back to another backend. `tensoradd!` and `tensortrace!` are forwarded to -`TensorOperations.StridedNative()`, so mixed `@tensor` networks still run. +`TensorOperations.StridedNative()`. The backend is not registered with `TensorOperations.select_backend`. """ -struct QuasiStridedBackend <: TO.AbstractBackend end +struct QuasiStridedBackend{A} <: TO.AbstractBackend end -# Task-local pool of workspaces, keyed by `eltype(C)` only: all complex methods -# pack into `Vector{real(T)}` and `reserve!` is grow-only, so they can share one. +function QuasiStridedBackend(; accumulator = nothing) + accumulator in (nothing, Float32, Float64) || throw( + ArgumentError("QuasiStridedBackend: accumulator must be nothing, Float32 or Float64, got $accumulator") + ) + return QuasiStridedBackend{accumulator}() +end + +# Task-local pool of workspaces, keyed by the compute type only: all complex +# methods pack into `Vector{real(T)}` and `reserve!` is grow-only, so they can +# share one. const _QS_WORKSPACE_KEY = :quasistrided_contract_workspaces @inline function _qs_workspace_pool() @@ -63,16 +73,19 @@ function _qs_labels(pA::Index2Tuple, pB::Index2Tuple, pAB::Index2Tuple) return indA, indB, linearize(pAB) end -const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) - @noinline _qs_throw(msg::AbstractString) = throw(ArgumentError(msg)) @noinline function _qs_check_eligible(C, A, B) f = TO.tensorcontract! - eltype(A) === eltype(B) === eltype(C) && eltype(C) ∈ _QS_ELTYPES || _qs_throw( - "QuasiStridedBackend requires all tensors of $f to share a single " * - "element type out of Float32, Float64, ComplexF32 and ComplexF64, got " * - join(map(eltype, (C, A, B)), ", ") + TA, TB, TC = eltype(A), eltype(B), eltype(C) + all(in(_QS_ELTYPES), (TA, TB, TC)) || _qs_throw( + "QuasiStridedBackend requires every tensor of $f to have an element type " * + "out of Float32, Float64, ComplexF32 and ComplexF64, got " * + join((TC, TA, TB), ", ") + ) + (TC <: Complex || (TA <: Real && TB <: Real)) || _qs_throw( + "QuasiStridedBackend: a complex input of $f needs a complex output, got " * + join((TC, TA, TB), ", ") ) all(isstrided, (A, B, C)) || _qs_throw( "QuasiStridedBackend requires strided arrays for $f, got " * @@ -84,7 +97,7 @@ end # Checks shared by both `tensorcontract!` methods. The conjugated-C check # duplicates `plan_contract`'s so that a rejected call never acquires or grows # a pooled workspace (the workspace is an argument to `plan_contract`). -@inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β) +@inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β, accumulator) _qs_check_eligible(C, A, B) TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) @@ -102,7 +115,7 @@ end ) # Dropping `Zero()`/`One()` is safe: the kernels branch on `iszero(alpha/beta)`. - T = eltype(C) + T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) indA, indB, indC = _qs_labels(pA, pB, pAB) return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end @@ -116,14 +129,14 @@ function TO.tensorcontract!( B::AbstractArray, pB::Index2Tuple, conjB::Bool, pAB::Index2Tuple, α::Number, β::Number, - backend::QuasiStridedBackend, + backend::QuasiStridedBackend{AC}, allocator::TO.DefaultAllocator = TO.DefaultAllocator() - ) - Cv, Av, Bv, indA, indB, indC, α′, β′ = _qs_prepare(C, A, pA, B, pB, pAB, α, β) + ) where {AC} + Cv, Av, Bv, indA, indB, indC, α′, β′ = _qs_prepare(C, A, pA, B, pB, pAB, α, β, AC) _planned( _Execute(α′, β′), Cv, Av, indA, Bv, indB, indC, nothing, conjA, conjB, nothing, nothing, nothing, - _qs_task_workspace(eltype(C)), allocator, false + _qs_task_workspace(typeof(α′)), allocator, false, AC # α′ has the compute type ) return C end @@ -137,14 +150,14 @@ function TO.tensorcontract!( B::AbstractArray, pB::Index2Tuple, conjB::Bool, pAB::Index2Tuple, α::Number, β::Number, - backend::QuasiStridedBackend, allocator - ) - Cv, Av, Bv, indA, indB, indC, α′, β′ = _qs_prepare(C, A, pA, B, pB, pAB, α, β) + backend::QuasiStridedBackend{AC}, allocator + ) where {AC} + Cv, Av, Bv, indA, indB, indC, α′, β′ = _qs_prepare(C, A, pA, B, pB, pAB, α, β, AC) checkpoint = TO.allocator_checkpoint!(allocator) plan = plan_contract( Cv, Av, indA, Bv, indB, indC; conjA = conjA, conjB = conjB, - workspace = nothing, allocator = allocator, oracle = false + workspace = nothing, allocator = allocator, oracle = false, accumulator = AC ) try execute!(plan, α′, β′) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 787acbb..0f92e87 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -105,7 +105,8 @@ end _default_lanewidth(::Type{Float64}) = 4 _default_lanewidth(::Type{Float32}) = 8 -# `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. +# `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. These +# helpers load C as `T`, compute in `T` and round on store. function scale_tile!(destination::QSTile, beta::T) where {T} m = nrows(destination) n = ncols(destination) @@ -118,24 +119,24 @@ function scale_tile!(destination::QSTile, beta::T) where {T} end else @inbounds for j in 0:(n - 1), i in 0:(m - 1) - tile_store!(destination, i, j, tile_load(destination, i, j) * beta) + tile_store!(destination, i, j, convert(T, tile_load(destination, i, j)) * beta) end end return destination end # `C = alpha*r + beta*C` at one element. Ternaries, so `beta == 0` never reads C. -@inline _axpby_tile!(dest, i::Int, j::Int, alpha, r, beta) = tile_store!( +@inline _axpby_tile!(dest, i::Int, j::Int, alpha, r, beta::T) where {T} = tile_store!( dest, i, j, iszero(beta) ? alpha * r : - isone(beta) ? muladd(alpha, r, tile_load(dest, i, j)) : - muladd(alpha, r, beta * tile_load(dest, i, j)) + isone(beta) ? muladd(alpha, r, convert(T, tile_load(dest, i, j))) : + muladd(alpha, r, beta * convert(T, tile_load(dest, i, j))) ) -@inline _axpby_at!(storage, idx::Int, alpha, r, beta) = @inbounds storage[idx] = +@inline _axpby_at!(storage, idx::Int, alpha, r, beta::T) where {T} = @inbounds storage[idx] = iszero(beta) ? alpha * r : - isone(beta) ? muladd(alpha, r, storage[idx]) : - muladd(alpha, r, beta * storage[idx]) + isone(beta) ? muladd(alpha, r, convert(T, storage[idx])) : + muladd(alpha, r, beta * convert(T, storage[idx])) # The `(m, n)` to store over, or `(0, 0)` when already done (empty destination, # or `alpha == 0` handled by `scale_tile!`). diff --git a/src/planning/plan.jl b/src/planning/plan.jl index ce22dcc..a82ef9b 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -44,7 +44,7 @@ end mc = nothing, kc = nothing, nc = nothing, workspace = nothing, allocator = TensorOperations.DefaultAllocator(), - oracle = true) -> ContractPlan + oracle = true, accumulator = nothing) -> ContractPlan Plan `C[indC] = A[indA] * B[indB]` (every label in exactly two operands): resolve the labels into M/N/K `AxisGroup`s, validate axis lengths and eltypes, @@ -52,22 +52,30 @@ choose the kernel and blocking, and preallocate every buffer [`execute!`](@ref) needs. Throws `ArgumentError`/`DimensionMismatch` on invalid input. + * Each operand's eltype is one of `Float32`, `Float64`, `ComplexF32`, + `ComplexF64`; a complex `A` or `B` needs a complex `C`. The compute type + `T` is `promote_type` of the three, or, for `accumulator = Float32` or + `Float64`, that precision in the domain (real or complex) of the promoted + type. Operands are converted to `T` on load and `alpha*AB + beta*C` is + evaluated in `T`, rounding to `eltype(C)` on store. `C` holds the partial + sums between `kc` blocks of K, so a K longer than `kc` rounds once per + block. * `kernel = nothing` picks one from the hardware profile and the extents: a - [`SIMDKernel`](@ref) for a real type, a [`PlanarKernel`](@ref) for a + [`SIMDKernel`](@ref) for a real `T`, a [`PlanarKernel`](@ref) for a complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512). [`OneMKernel`](@ref) is used only when named. * Labels within the M and N composites are ordered by their stride in `C`; - the K order follows a cost model of the two packs. For a real type the + the K order follows a cost model of the two packs. For a real `T` the operand roles are swapped (B feeds M) when only the N side gives `C` a unit-stride run long enough for the kernel's register tile; the result is the same either way. * `mc`/`kc`/`nc` override the fields of `default_blocking(kernel)` (each `>= 1`); `mc`/`nc` are rounded up to `mr`/`nr` multiples and all three are clamped to the extents. - * `workspace` reuses an existing [`ContractWorkspace`](@ref), grown by - [`reserve!`](@ref) as needed. A non-default TensorOperations `allocator` - sizes the buffers once via `tensoralloc`, forbids `workspace`, and leaves - [`release!`](@ref) to the caller. + * `workspace` reuses an existing [`ContractWorkspace`](@ref) for compute + type `T`, grown by [`reserve!`](@ref) as needed. A non-default + TensorOperations `allocator` sizes the buffers once via `tensoralloc`, + forbids `workspace`, and leaves [`release!`](@ref) to the caller. * `oracle = false` skips `execute_tilewise!`'s buffers. * `conjA`/`conjB` conjugate A's/B's elements (never `alpha`/`beta`) and compose by XOR with a view's own `conj`/`adjoint` `op`. A conjugated `C` @@ -85,17 +93,18 @@ function plan_contract( nc::Union{Int, Nothing} = nothing, workspace::Union{Nothing, ContractWorkspace} = nothing, allocator = TO.DefaultAllocator(), - oracle::Bool = true + oracle::Bool = true, + accumulator::Union{Nothing, Type{Float32}, Type{Float64}} = nothing ) where {NA, NB, NC} return _planned( identity, C, A, indA, B, indB, indC, - kernel, conjA, conjB, mc, kc, nc, workspace, allocator, oracle + kernel, conjA, conjB, mc, kc, nc, workspace, allocator, oracle, accumulator ) end # Everything `_plan_contract` needs that is concretely typed before the kernel # is known (`run` is the chosen M composite's unit-stride run in C), as one -# value that crosses the kernel barrier. `T` is the phantom eltype. +# value that crosses the kernel barrier. `T` is the phantom compute type. struct _PlanRequest{ T, F, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, WS <: Union{Nothing, ContractWorkspace}, AL, @@ -138,13 +147,10 @@ function _planned( B::StridedView, indB::NTuple{NB, Int}, indC::NTuple{NC, Int}, kernel, conjA::Bool, conjB::Bool, mc::Union{Int, Nothing}, kc::Union{Int, Nothing}, nc::Union{Int, Nothing}, - workspace::Union{Nothing, ContractWorkspace}, allocator, oracle::Bool - ) where {F, NA, NB, NC} - T = eltype(C) - eltype(A) === T || - throw(ArgumentError("eltype(A) = $(eltype(A)) does not match eltype(C) = $T")) - eltype(B) === T || - throw(ArgumentError("eltype(B) = $(eltype(B)) does not match eltype(C) = $T")) + workspace::Union{Nothing, ContractWorkspace}, allocator, oracle::Bool, + accumulator::AC + ) where {F, NA, NB, NC, AC} + T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) # GUARDRAIL: a conjugated `C` is rejected; the engine writes through to # its parent, so there is nowhere to absorb its `op`. @@ -198,6 +204,29 @@ function _planned( return _plan_with_kernel(kernel, atransform, btransform, req) end +const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) + +# Fold to the compute type, or a throw, at compile time. +@inline function _compute_type(::Type{TA}, ::Type{TB}, ::Type{TC}, ::Nothing) where {TA, TB, TC} + (TA in _QS_ELTYPES && TB in _QS_ELTYPES && TC in _QS_ELTYPES) || + _throw_eltypes(TA, TB, TC) + (TC <: Real && !(TA <: Real && TB <: Real)) && _throw_complex_into_real(TA, TB, TC) + return promote_type(TA, TB, TC) +end +@inline function _compute_type(::Type{TA}, ::Type{TB}, ::Type{TC}, ::Type{R}) where {TA, TB, TC, R} + return _compute_type(TA, TB, TC, nothing) <: Complex ? Complex{R} : R +end + +@noinline _throw_eltypes(TA, TB, TC) = throw( + ArgumentError( + "plan_contract: eltypes (A, B, C) = ($TA, $TB, $TC); each must be one of " * + "Float32, Float64, ComplexF32, ComplexF64" + ) +) +@noinline _throw_complex_into_real(TA, TB, TC) = throw( + ArgumentError("plan_contract: a complex operand (A: $TA, B: $TB) needs a complex C, got $TC") +) + # `mr` of the kernel each orientation would run, for the swap decision: a named # kernel either way, else `_default_shape`'s pick at that orientation's extents # and C run. @@ -264,7 +293,7 @@ function _plan_contract( kernel::K, atransform::TA, btransform::TB, req::_PlanRequest{T}, hint::H ) where {K, TA, TB, T, H} scalartype(kernel) === T || - throw(ArgumentError("kernel scalar type $(scalartype(kernel)) does not match eltype(C) = $T")) + throw(ArgumentError("kernel scalar type $(scalartype(kernel)) does not match the compute type $T")) Qm = axis_length(req.mgroup) Qn = axis_length(req.ngroup) diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl new file mode 100644 index 0000000..664ae1a --- /dev/null +++ b/test/execution/test_mixed.jl @@ -0,0 +1,62 @@ +# Mixed operand eltypes: compute type, conversion on load, one rounding on store. + +using LinearAlgebra: norm + +# `(Cv, Av, indA, Bv, indB, indC)` of shape `(M, K, N)`; `:scattered` takes each +# operand from `scattered_fixture` of its own eltype. +function _mixed_fixture(TA, TB, TC, M, K, N, layout, seed) + layout === :scattered && + return (scattered_fixture(TC)[1], scattered_fixture(TA)[2:3]..., scattered_fixture(TB)[4:6]...) + rng = MersenneTwister(seed) + A, B, C = randn(rng, TA, M, K), randn(rng, TB, K, N), randn(rng, TC, M, N) + return (StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3)) +end + +# The reference in Float64/ComplexF64 from the exactly promoted inputs. +function _mixed_ref(fx, alpha, beta; kw...) + Cv, Av, iA, Bv, iB, iC = fx + R = promote_type(eltype(Av), eltype(Bv), eltype(Cv)) <: Complex ? ComplexF64 : Float64 + return _brute_ref(R.(Array(Av)), iA, R.(Array(Bv)), iB, R.(Array(Cv)), iC, alpha, beta; kw...) +end + +@testset "mixed eltypes: $TA x $TB -> $TC, accumulator = $acc, $layout" for (TA, TB, TC, acc, layout, MKN, beta, conjA, conjB, path) in ( + (Float32, Float64, Float64, nothing, :dense, (23, 37, 19), 0, false, false, QuasiStrided._NestPath), + (Float64, Float32, Float32, nothing, :dense, (32, 40, 12), 1, false, false, QuasiStrided._NestPath), + (Float64, ComplexF64, ComplexF64, nothing, :dense, (21, 33, 17), 0.3 - 0.7im, false, true, QuasiStrided._NestPath), + (ComplexF32, Float64, ComplexF64, nothing, :scattered, (32, 32, 32), 0.5, true, false, QuasiStrided._NestPath), + (ComplexF32, ComplexF64, ComplexF64, nothing, :dense, (19, 29, 23), 1, true, true, QuasiStrided._NestPath), + (Float64, Float64, ComplexF64, nothing, :dense, (18, 25, 14), 0.25 + 0.5im, false, false, QuasiStrided._NestPath), + (ComplexF64, Float32, ComplexF64, Float32, :dense, (17, 31, 13), 0, false, false, QuasiStrided._NestPath), + (Float32, Float64, Float64, nothing, :dense, (1, 70, 29), 0.5, false, false, QuasiStrided._DotPath), + (Float64, Float32, Float64, nothing, :dense, (37, 1, 21), 1, false, false, QuasiStrided._OuterPath), + ) + fx = _mixed_fixture(TA, TB, TC, MKN..., layout, 7) + alpha = 1.5 + ref = _mixed_ref(fx, alpha, beta; conjA, conjB) + plan = plan_contract(fx...; conjA, conjB, accumulator = acc) + T = acc === nothing ? promote_type(TA, TB, TC) : (TC <: Complex ? Complex{acc} : acc) + @test plan isa ContractPlan{T} + @test _path_of(plan) isa path + execute!(plan, alpha, beta) + K = layout === :scattered ? 32 : MKN[2] + rtol = 10 * sqrt(K) * max(eps(real(T)), eps(real(TC))) + @test norm(Array(fx[1]) - ref) <= rtol * norm(ref) +end + +@testset "mixed eltypes: steady-state execute! allocates nothing" begin + rng = MersenneTwister(3) + Amat, Bmat, Cmat = randn(rng, Float32, 40, 30), randn(rng, ComplexF64, 30, 20), zeros(ComplexF64, 40, 20) + @test _steady_allocs!(execute!, _mm_plan(Cmat, Amat, Bmat), Cmat) == 0 skip = (VERSION < v"1.11") +end + +# One K panel: between panels the partial sums round-trip through C. +@testset "accumulator = Float64: Float32 operands within one ulp" begin + rng = MersenneTwister(5) + A, B, C0 = randn(rng, Float32, 64, 512), randn(rng, Float32, 512, 48), randn(rng, Float32, 64, 48) + ref = Float32.(Float64.(A) * Float64.(B) .+ 0.5 .* Float64.(C0)) + C64, C32 = copy(C0), copy(C0) + execute!(_mm_plan(C64, A, B; accumulator = Float64, kc = 512), 1, 0.5) + execute!(_mm_plan(C32, A, B), 1, 0.5) + @test all(abs.(C64 .- ref) .<= eps.(ref)) + @test maximum(abs.(C64 .- ref)) < maximum(abs.(C32 .- ref)) +end diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 5d3cb9c..39f9703 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -129,13 +129,11 @@ end @testset "hard-reject: ineligible eltypes and non-strided operands" begin pA, pB, pAB = _MATMUL_PAB - # (eltype A, eltype B, eltype C): types outside the four, and mixtures of - # eligible ones (promotion is TO's `promote_contract` job). + # (eltype A, eltype B, eltype C): types outside the four, and complex into real. for (TA, TB, TC) in ( (Float16, Float16, Float16), (Complex{Float16}, Complex{Float16}, Complex{Float16}), (Complex{Int}, Complex{Int}, Complex{Int}), (Complex{BigFloat}, Complex{BigFloat}, Complex{BigFloat}), - (Float32, Float64, Float64), (ComplexF32, ComplexF64, ComplexF64), - (Float64, ComplexF64, ComplexF64), (ComplexF64, Float64, ComplexF64), + (Float64, Float32, Float16), (ComplexF64, Float64, Float64), (Float32, ComplexF32, Float32), ) A, B, C = ones(TA, (3, 4)), ones(TB, (4, 5)), zeros(TC, (3, 5)) @test_throws ArgumentError tensorcontract!(C, A, pA, false, B, pB, false, pAB, 1, 0, qsbackend) @@ -143,6 +141,24 @@ end D4, D5 = Diagonal(randn(4)), Diagonal(randn(5)) @test_throws ArgumentError tensorcontract!(zeros(4, 5), D4, pA, false, randn(4, 5), pB, false, pAB, 1, 0, qsbackend) @test_throws ArgumentError tensorcontract!(zeros(4, 5), randn(4, 5), pA, false, D5, pB, false, pAB, 1, 0, qsbackend) + @test_throws ArgumentError QuasiStrided.QuasiStridedBackend(accumulator = Float16) +end + +@testset "@tensor with mixed eltypes and an accumulator" begin + Random.seed!(4321) + A32, B64, Bc = randn(Float32, 6, 7, 5), randn(7, 8), randn(ComplexF64, 7, 8) + A64 = Float64.(A32) + @tensor backend = qsbackend D[a, b, c] := A32[a, k, c] * B64[k, b] + @tensor Dref[a, b, c] := A64[a, k, c] * B64[k, b] + @test eltype(D) === Float64 && D ≈ Dref + @tensor backend = qsbackend E[a, b, c] := conj(A32[a, k, c]) * Bc[k, b] + @tensor Eref[a, b, c] := A64[a, k, c] * Bc[k, b] + @test eltype(E) === ComplexF64 && E ≈ Eref + B32 = randn(Float32, 7, 8) + @tensor backend = QuasiStrided.QuasiStridedBackend(accumulator = Float64) F[a, b, c] := A32[a, k, c] * B32[k, b] + B32w = Float64.(B32) + @tensor Fref[a, b, c] := A64[a, k, c] * B32w[k, b] + @test eltype(F) === Float32 && all(abs.(F .- Float32.(Fref)) .<= eps.(Float32.(Fref))) end @testset "hard-reject: C aliasing an input (eltype = $T)" for T in (Float64, ComplexF64) diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 9ab1342..4a96f21 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -25,6 +25,19 @@ @test_throws ArgumentError contract!(Cv, 1.0, Av, (1, 2), Bv, (2, 3), 0.0, (2, 3)) end +@testset "driver: eltype validation errors" begin + mk(TA, TB, TC) = (StridedView(zeros(TC, 3, 5)), StridedView(ones(TA, 3, 4)), (1, 2), StridedView(ones(TB, 4, 5)), (2, 3), (1, 3)) + for eltypes in ( + (Float16, Float16, Float16), (Int, Float64, Float64), (Float64, Float64, BigFloat), + (ComplexF64, Float64, Float64), (Float32, ComplexF32, Float32), + ) + @test_throws ArgumentError plan_contract(mk(eltypes...)...) + end + ws = plan_contract(mk(Float32, Float32, Float32)...).workspace + @test_throws ArgumentError plan_contract(mk(Float32, Float32, Float32)...; accumulator = Float64, workspace = ws) + @test plan_contract(mk(Float32, Float64, Float32)...; accumulator = Float32, workspace = ws).workspace === ws +end + @testset "Blocking: field validation" begin b = Blocking(4, 8, 16) diff --git a/test/runtests.jl b/test/runtests.jl index 5c9d966..ce45612 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -55,6 +55,7 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip include("execution/test_unpackedb.jl") include("execution/test_dot.jl") include("execution/test_outer.jl") + include("execution/test_mixed.jl") # Last: its `using TensorOperations` makes `scalartype` ambiguous for later files. include("integrations/test_tensoroperations.jl") From b9ad2b1fd31e651964a4a042af99d742c5e55c66 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 09:57:45 -0400 Subject: [PATCH 03/27] Mixed-domain kernels: complex x real and real x complex on the real SIMD kernel ComplexRealKernel runs SIMDKernel{2MR,NR,real(T),W} over interleaved A and real B (accumulator in the fmaddsub layout, reusing its stores, unpacked B eligible). RealComplexKernel runs SIMDKernel{MR,2NR,real(T),W} over real A and interleaved B, storing through the planar generators with a column-pair index map. The descriptor admits (Interleaved, Real) and (Real, Interleaved) for a complex T; a RealFormat operand packs real(T). Named kernels run through plan_contract unchanged. Codegen (code_llvm, debuginfo stripped, name suffixes normalized) against 096ab7f: identical for store_tile!/execute_tile! of SIMD/planar/fmaddsub/1m, _store_tile_planar[_vector]!, _store_tile_fmaddsub_vector!, _store_tile_lanepair!, and unsafe_pack_a!/unsafe_pack_b! (dense/scattered, identity/conj, PackedPanel/Vector, same and mixed storage precision): 264 functions. The mixed kernels' _accumulate_step is identical to the equivalent SIMDKernel's. Pkg.test() run concurrently with 096ab7f on a Cascade Lake workstation: 856 s vs 888 s wall. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 2 + src/execution/unpackedb.jl | 3 +- src/microkernels/fmaddsub.jl | 2 +- src/microkernels/interface.jl | 8 ++ src/microkernels/mixed.jl | 124 ++++++++++++++++++++++++ src/microkernels/planar.jl | 23 +++-- src/packing/format.jl | 8 +- src/packing/pack.jl | 11 ++- test/microkernels/test_mixed_kernel.jl | 113 +++++++++++++++++++++ test/microkernels/test_scalar_kernel.jl | 29 +++--- test/packing/test_kernel_descriptor.jl | 1 + test/runtests.jl | 1 + 12 files changed, 298 insertions(+), 27 deletions(-) create mode 100644 src/microkernels/mixed.jl create mode 100644 test/microkernels/test_mixed_kernel.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 453d153..51c21a9 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -32,6 +32,7 @@ include("microkernels/simd.jl") include("microkernels/planar.jl") include("microkernels/onem.jl") include("microkernels/fmaddsub.jl") +include("microkernels/mixed.jl") # --- Planning: labels, conjugation, kernel and blocking choice, the plan --- include("planning/labels.jl") @@ -70,6 +71,7 @@ end :public, :contract!, :plan_contract, :execute!, :ContractPlan, :ContractWorkspace, :Blocking, :default_blocking, :ScalarKernel, :SIMDKernel, :PlanarKernel, :OneMKernel, :FMAddSubKernel, + :ComplexRealKernel, :RealComplexKernel, :target_profile, :cache_topology, :TargetProfile, :CacheLevel ) diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 65c69dc..952a84b 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -40,10 +40,11 @@ end @inline _unpacked_b_kernel_eligible(::SIMDKernel) = true @inline _unpacked_b_kernel_eligible(::PlanarKernel) = true @inline _unpacked_b_kernel_eligible(::FMAddSubKernel) = true +@inline _unpacked_b_kernel_eligible(::ComplexRealKernel) = true @inline _unpacked_b_kernel_eligible(::Any) = false # The same by method, to predict the path before the kernel exists. -@inline _unpacked_b_method_eligible(::Union{RealMethod, PlanarMethod, FMAddSubMethod}) = true +@inline _unpacked_b_method_eligible(::Union{RealMethod, PlanarMethod, FMAddSubMethod, ComplexRealMethod}) = true @inline _unpacked_b_method_eligible(::Any) = false # `:always`/`:never` override the rule below, for benchmarks and tests. diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 7bd939f..78ee65d 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -199,7 +199,7 @@ end # Same unroll and full-block / row-tail split as `_store_tile_planar_vector!`. @generated function _store_tile_fmaddsub_vector!( destination::QSTile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::FMAddSubKernel{MR, NR, T, W}, + alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NV} # The pointer reinterpretation is only sound on dense rank-1 complex storage. diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 0f92e87..8597eae 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -20,6 +20,8 @@ struct RealMethod <: ComplexMethod end # what a real kernel reports struct PlanarMethod <: ComplexMethod end # split re/im planes, 4 real FMAs per MAC struct OneMMethod <: ComplexMethod end # Van Zee's 1m: a real 2mr x nr kernel struct FMAddSubMethod <: ComplexMethod end # interleaved A, x86 `vfmaddsub` +struct ComplexRealMethod <: ComplexMethod end # complex A, real B: the real kernel on 2MR rows +struct RealComplexMethod <: ComplexMethod end # real A, complex B: the real kernel on 2NR columns # Reals per packed A / B element. Blocking divides the real `mc`/`nc` by these, # so every method gets the same packed byte budget. @@ -31,12 +33,18 @@ a_reals(::OneMMethod) = 4 b_reals(::OneMMethod) = 2 a_reals(::FMAddSubMethod) = 2 b_reals(::FMAddSubMethod) = 2 +a_reals(::ComplexRealMethod) = 2 +b_reals(::ComplexRealMethod) = 1 +a_reals(::RealComplexMethod) = 1 +b_reals(::RealComplexMethod) = 2 # Accumulator planes held live: planar keeps separate re/im planes. accumulator_planes(::RealMethod) = 1 accumulator_planes(::PlanarMethod) = 2 accumulator_planes(::OneMMethod) = 1 accumulator_planes(::FMAddSubMethod) = 1 +accumulator_planes(::ComplexRealMethod) = 1 +accumulator_planes(::RealComplexMethod) = 1 complex_method(::Any) = RealMethod() diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl new file mode 100644 index 0000000..01e95a7 --- /dev/null +++ b/src/microkernels/mixed.jl @@ -0,0 +1,124 @@ +# Mixed-domain kernels (BLIS's mixed-domain method): a complex operand times a +# real one is a real product once the complex operand is viewed as real, so +# both run the real `SIMDKernel` verbatim, at 2 real FMAs per complex MAC. +# * complex A x real B: interleaved A is a real `2MR x kc` panel; the +# accumulator is 1m/fmaddsub's interleaved layout. +# * real A x complex B: interleaved B is a real `kc x 2NR` panel; accumulator +# column `2j` holds the real and `2j+1` the imaginary part of column `j`. + +using SIMD: Vec + +""" + ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Mixed-domain microkernel for complex `A` times real `B`, computing in complex +`T`: a real `SIMDKernel` of `2MR x NR` over `SIMD.Vec{W,real(T)}` lanes. `2MR` +must be a multiple of `W`, and `W` must be even. +""" +struct ComplexRealKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, NR, T} + descriptor::ComplexKernelDescriptor{MR, NR, T, InterleavedFormat, RealFormat} + inner::KI + + function ComplexRealKernel{MR, NR, T, W, KI}( + descriptor::ComplexKernelDescriptor{MR, NR, T, InterleavedFormat, RealFormat}, + inner::KI + ) where {MR, NR, T, W, KI <: SIMDKernel} + _check_vector_shape("ComplexRealKernel", 2 * MR, W, true) + KI === SIMDKernel{2 * MR, NR, real(T), W} || throw( + ArgumentError("ComplexRealKernel's inner kernel must be SIMDKernel{$(2 * MR),$NR,$(real(T)),$W}, got $KI") + ) + return new{MR, NR, T, W, KI}(descriptor, inner) + end +end + +""" + RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Mixed-domain microkernel for real `A` times complex `B`, computing in complex +`T`: a real `SIMDKernel` of `MR x 2NR` over `SIMD.Vec{W,real(T)}` lanes. `MR` +must be a multiple of `W`. +""" +struct RealComplexKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, NR, T} + descriptor::ComplexKernelDescriptor{MR, NR, T, RealFormat, InterleavedFormat} + inner::KI + + function RealComplexKernel{MR, NR, T, W, KI}( + descriptor::ComplexKernelDescriptor{MR, NR, T, RealFormat, InterleavedFormat}, + inner::KI + ) where {MR, NR, T, W, KI <: SIMDKernel} + _check_vector_shape("RealComplexKernel", MR, W) + KI === SIMDKernel{MR, 2 * NR, real(T), W} || throw( + ArgumentError("RealComplexKernel's inner kernel must be SIMDKernel{$MR,$(2 * NR),$(real(T)),$W}, got $KI") + ) + return new{MR, NR, T, W, KI}(descriptor, inner) + end +end + +const _MixedKernel = Union{ComplexRealKernel, RealComplexKernel} + +function ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} + T <: Complex || + throw(ArgumentError("ComplexRealKernel requires a complex element type, got $T")) + descriptor = ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), RealFormat()) + inner = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) + return ComplexRealKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) +end + +function RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} + T <: Complex || + throw(ArgumentError("RealComplexKernel requires a complex element type, got $T")) + descriptor = ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()) + inner = SIMDKernel(Val(MR), Val(2 * NR), real(T), Val(W)) + return RealComplexKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) +end + +function ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} + T <: Complex || + throw(ArgumentError("ComplexRealKernel requires a complex element type, got $T")) + return ComplexRealKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) +end +function RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} + T <: Complex || + throw(ArgumentError("RealComplexKernel requires a complex element type, got $T")) + return RealComplexKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) +end + +complex_method(::ComplexRealKernel) = ComplexRealMethod() +complex_method(::RealComplexKernel) = RealComplexMethod() +lanewidth(kernel::_MixedKernel) = lanewidth(kernel.inner) + +zero_accumulator(kernel::_MixedKernel) = zero_accumulator(kernel.inner) + +@inline Base.accumulate( + kernel::K, acc::NTuple{NV, Vec{W, R}}, packed_a::PA, packed_b::PB, kc::Int +) where {K <: _MixedKernel, NV, W, R, PA, PB} = + accumulate(kernel.inner, acc, packed_a, packed_b, kc) + +# Inlined or not as the fmaddsub and planar stores each reuses. +function store_tile!( + destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + alpha::T, beta::T, kernel::ComplexRealKernel{MR, NR, T, W} + ) where {MR, NR, T, W, R, NV} + m, n = _store_prologue!(destination, alpha, beta) + (m == 0 || n == 0) && return destination + + if _complex_vector_eligible(destination, T) + return _store_tile_fmaddsub_vector!(destination, acc, alpha, beta, kernel, m, n) + end + + return _store_tile_lanepair!(destination, acc, alpha, beta, kernel, m, n) +end + +@inline function store_tile!( + destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + alpha::T, beta::T, kernel::RealComplexKernel{MR, NR, T, W} + ) where {MR, NR, T, W, R, NV} + m, n = _store_prologue!(destination, alpha, beta) + (m == 0 || n == 0) && return destination + + if _complex_vector_eligible(destination, T) + return _store_tile_planar_vector!(destination, acc, alpha, beta, kernel, m, n) + end + + return _store_tile_planar!(destination, acc, alpha, beta, kernel, m, n) +end diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 2207fda..a30b417 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -134,10 +134,17 @@ function Base.accumulate( return acc end +# Generator-time accumulator indices of the re/im vectors of block `v`, column +# `j`: planar's two planes, or `RealComplexKernel`'s column pairs. +function _planar_acc_index(kernel::Type, MV::Int, NR::Int, v::Int, j::Int) + kernel <: PlanarKernel && return (v + MV * j + 1, MV * NR + v + MV * j + 1) + return (v + MV * 2j + 1, v + MV * (2j + 1) + 1) +end + # Scalar fallback store, for every destination the vector store cannot take. @generated function _store_tile_planar!( destination::QSTile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W}, + alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {MR, NR, T, W, R, NA} MV = MR ÷ W @@ -146,12 +153,12 @@ end blocks = Any[] for j in 0:(NR - 1), v in 0:(MV - 1) - idx = v + MV * j + 1 + ire, iim = _planar_acc_index(kernel, MV, NR, v, j) push!( blocks, quote if $j < n - revec = acc[$idx] - imvec = acc[$(NV + idx)] + revec = acc[$ire] + imvec = acc[$iim] for lane in 1:$W i = $(v * W) + lane - 1 i < m || break @@ -236,7 +243,7 @@ end # block of `2W` reals. The raw pointer is only dereferenced inside `GC.@preserve`. @generated function _store_tile_planar_vector!( destination::QSTile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W}, + alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NA} # The pointer reinterpretation is only sound on dense rank-1 complex storage. @@ -251,11 +258,11 @@ end for j in 0:(NR - 1) vblocks = Any[] for v in 0:(MV - 1) - idx = v + MV * j + 1 + ire, iim = _planar_acc_index(kernel, MV, NR, v, j) push!( vblocks, quote - revec = acc[$idx] - imvec = acc[$(NV + idx)] + revec = acc[$ire] + imvec = acc[$iim] if $((v + 1) * W) <= m _planar_store_block!( sp, 2 * (colbase + $(v * W)), revec, imvec, diff --git a/src/packing/format.jl b/src/packing/format.jl index 491e716..07d1a9f 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -55,9 +55,13 @@ function _check_descriptor(MR, NR, T, FA, FB) T === Float32 || T === Float64 || throw(ArgumentError("KernelDescriptor requires T ∈ (Float32, Float64), got $T")) else - (FA === RealFormat || FB === RealFormat) && throw( + # A real operand of a complex `T` only in the mixed-domain pairings. + mixed = (FA, FB) === (InterleavedFormat, RealFormat) || + (FA, FB) === (RealFormat, InterleavedFormat) + !mixed && (FA === RealFormat || FB === RealFormat) && throw( ArgumentError( - "ComplexKernelDescriptor requires complex formats on both operands, got ($FA, $FB)" + "ComplexKernelDescriptor requires complex formats on both operands, " * + "or (InterleavedFormat, RealFormat) / (RealFormat, InterleavedFormat), got ($FA, $FB)" ) ) T === ComplexF32 || T === ComplexF64 || throw( diff --git a/src/packing/pack.jl b/src/packing/pack.jl index ab78faf..dcfde8e 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -98,7 +98,7 @@ end transform::F, m::Int, kc::Int ) where {FMT, V, MR, NR, T2, F} if format isa RealFormat - if _pack_a_contiguous_eligible(packed, source, transform, m, Val(MR), T2) + if _pack_a_contiguous_eligible(packed, source, transform, m, Val(MR), real(T2)) rowbase = source.base + source.rows.base return _pack_a_contiguous!(packed, source.storage, rowbase, source.cols, Val(MR), kc) end @@ -112,7 +112,7 @@ end end load = (i, p) -> tile_load(source, i, p) plane_offset = (plane, i, p) -> packed_a_plane_offset(kernel, plane, i, p) - return _pack_panel!(packed, T2, format, Val(MR), kc, m, transform, load, plane_offset) + return _pack_panel!(packed, _element_type(format, T2), format, Val(MR), kc, m, transform, load, plane_offset) end @inline function _pack_b_sliver!( @@ -129,9 +129,14 @@ end end load = (j, p) -> tile_load(source, p, j) plane_offset = (plane, j, p) -> packed_b_plane_offset(kernel, plane, j, p) - return _pack_panel!(packed, T2, format, Val(NR), kc, n, transform, load, plane_offset) + return _pack_panel!(packed, _element_type(format, T2), format, Val(NR), kc, n, transform, load, plane_offset) end +# What a packed element converts to: the real operand of a mixed-domain kernel +# packs `real(T)`. +@inline _element_type(::RealFormat, ::Type{T}) where {T} = real(T) +@inline _element_type(::PackFormat, ::Type{T}) where {T} = T + # One sliver over `kc` K steps; `PD` (MR or NR) is a compile-time constant. A # full sliver gets a constant-trip inner loop that LLVM fully unrolls; a tail # writes its valid lanes and its padding as two loops. A per-lane diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl new file mode 100644 index 0000000..cfdd951 --- /dev/null +++ b/test/microkernels/test_mixed_kernel.jl @@ -0,0 +1,113 @@ +using QuasiStrided: ComplexRealKernel, RealComplexKernel, ComplexRealMethod, RealComplexMethod, packed_panel + +mk_optypes(k::ComplexRealKernel) = (scalartype(k), real(scalartype(k))) +mk_optypes(k::RealComplexKernel) = (real(scalartype(k)), scalartype(k)) +mk_pack_a(::ComplexRealKernel, A) = mk_cols(c -> mk_ilv(real(c), imag(c)), A) +mk_pack_b(::ComplexRealKernel, B) = vec(permutedims(B)) +mk_pack_a(::RealComplexKernel, A) = vec(A) +mk_pack_b(::RealComplexKernel, B) = mk_cols(c -> mk_ilv(real(c), imag(c)), permutedims(B)) +function mk_read(k::ComplexRealKernel, acc, i, j) + W = lanewidth(k) + v, u = divrem(i, W ÷ 2) + vec = acc[v + (2 * mr(k) ÷ W) * j + 1] + return Complex(vec[2u + 1], vec[2u + 2]) +end +function mk_read(k::RealComplexKernel, acc, i, j) + W = lanewidth(k) + MV = mr(k) ÷ W + v, u = divrem(i, W) + return Complex(acc[v + MV * 2j + 1][u + 1], acc[v + MV * (2j + 1) + 1][u + 1]) +end + +@testset "mixed-domain kernels" begin + for (K, s, full) in ( + (ComplexRealKernel, (ComplexF64, 12, 8, 8), true), (ComplexRealKernel, (ComplexF32, 8, 5, 8), true), + (ComplexRealKernel, (ComplexF64, 3, 2, 2), false), + (RealComplexKernel, (ComplexF64, 24, 4, 8), true), (RealComplexKernel, (ComplexF32, 16, 3, 8), true), + (RealComplexKernel, (ComplexF64, 2, 3, 2), false), + ) + T, MR, NR, W = s + mk_contract(K(Val(MR), Val(NR), T, Val(W)); full) + end + + @testset "construction" begin + cr = ComplexRealKernel(Val(8), Val(4), ComplexF32) + rc = RealComplexKernel(Val(8), Val(4), ComplexF64) + @test complex_method(cr) === ComplexRealMethod() + @test complex_method(rc) === RealComplexMethod() + @test (a_reals(ComplexRealMethod()), b_reals(ComplexRealMethod())) == (2, 1) + @test (a_reals(RealComplexMethod()), b_reals(RealComplexMethod())) == (1, 2) + @test accumulator_planes(ComplexRealMethod()) == accumulator_planes(RealComplexMethod()) == 1 + @test cr.inner isa SIMDKernel{16, 4, Float32, 8} + @test rc.inner isa SIMDKernel{8, 8, Float64, 4} + @test (packed_a_per_k(cr), packed_b_per_k(cr), packed_a_per_k(rc), packed_b_per_k(rc)) == (16, 4, 8, 8) + @test_throws ArgumentError ComplexRealKernel(Val(3), Val(4), ComplexF64, Val(3)) # odd W + @test_throws ArgumentError RealComplexKernel(Val(6), Val(4), ComplexF64, Val(4)) + @test_throws ArgumentError ComplexRealKernel(Val(8), Val(4), Float64) + @test_throws ArgumentError RealComplexKernel(Val(8), Val(4), Float64, Val(4)) + end + + # Packing from storage of another precision, `conj` on both sides. + @testset "pack -> accumulate -> store: $(nameof(K)) $T from $SA x $SB" for (K, T, SA, SB, MR, NR, W) in ( + (ComplexRealKernel, ComplexF32, ComplexF64, Float64, 8, 5, 8), + (RealComplexKernel, ComplexF32, Float64, ComplexF64, 16, 3, 8), + (ComplexRealKernel, ComplexF64, ComplexF64, Float32, 4, 3, 4), + (RealComplexKernel, ComplexF64, Float32, ComplexF32, 8, 2, 4), + ) + k = K(Val(MR), Val(NR), T, Val(W)) + R = real(T) + rng = MersenneTwister(7) + kc, lda = 6, MR + 2 + Av, Bv = rand(rng, SA, lda * kc + 1), rand(rng, SB, kc * NR) + for f in (identity, conj), (m, n) in ((MR, NR), (MR - 1, NR - 1)), panel in (false, true) + srcA = SourceTile(Av, 1, AffineAxis(0, 1, m), AffineAxis(0, lda, kc)) + srcB = SourceTile(Bv, 0, AffineAxis(0, NR, kc), AffineAxis(0, 1, n)) + A = zeros(T, MR, kc) + B = zeros(T, kc, NR) + A[1:m, :] = f.(reshape(Av[2:(1 + lda * kc)], lda, kc)[1:m, :]) + B[:, 1:n] = f.(permutedims(reshape(Bv, NR, kc))[:, 1:n]) + TA, TB = mk_optypes(k) + pa, pb = zeros(R, packed_a_length(k, kc)), zeros(R, packed_b_length(k, kc)) + GC.@preserve pa pb begin + dpa = panel ? packed_panel(pa, 1, length(pa)) : pa + dpb = panel ? packed_panel(pb, 1, length(pb)) : pb + pack_a!(dpa, srcA, k, f) + pack_b!(dpb, srcB, k, f) + end + @test isequal((pa, pb), mk_pack(k, TA.(A), TB.(B))) + alpha, beta = T(1.5, -0.5), T(0.25, 1) + C0 = rand(rng, T, m * n) + got = mk_dense(copy(C0)) + execute_tile!(k, DestinationTile(got, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)), pa, pb, kc, alpha, beta) + @test mk_close(got, alpha .* vec((A * B)[1:m, 1:n]) .+ beta .* C0, T) + end + end + + @testset "end to end: $(nameof(typeof(k))) $TA x $TB -> $TC, B $mode" for (k, TA, TB, TC, acc, mode) in ( + (ComplexRealKernel(Val(12), Val(8), ComplexF64, Val(8)), ComplexF64, Float64, ComplexF64, nothing, :always), + (ComplexRealKernel(Val(12), Val(8), ComplexF64, Val(8)), ComplexF64, Float64, ComplexF64, nothing, :never), + (ComplexRealKernel(Val(8), Val(5), ComplexF32, Val(16)), ComplexF32, Float64, ComplexF32, Float32, :auto), + (RealComplexKernel(Val(24), Val(4), ComplexF64, Val(8)), Float64, ComplexF64, ComplexF64, nothing, :auto), + (RealComplexKernel(Val(16), Val(3), ComplexF32, Val(16)), Float64, ComplexF32, ComplexF32, Float32, :auto), + ) + rng = MersenneTwister(4242) + A, B, C0 = rand(rng, TA, 37, 41), rand(rng, TB, 41, 23), rand(rng, TC, 37, 23) + alpha, beta = 1.5 - 0.25im, -0.75 + 0.5im + saved = QuasiStrided._UNPACKED_B_MODE[] + QuasiStrided._UNPACKED_B_MODE[] = mode + try + for (cA, cB) in ((false, false), (true, false), (false, true)) + C = copy(C0) + plan = plan_contract( + StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3); + kernel = k, kc = 16, conjA = cA, conjB = cB, accumulator = acc + ) + execute!(plan, alpha, beta) + want = alpha .* ((cA ? conj.(A) : A) * (cB ? conj.(B) : B)) .+ beta .* C0 + @test maximum(abs, C .- want) <= 64 * mk_tol(TC) * maximum(abs, want) + end + finally + QuasiStrided._UNPACKED_B_MODE[] = saved + end + end +end diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index 3c1f45e..5c47933 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -54,6 +54,9 @@ mk_alphabeta(T) = T <: Complex ? # The storage the driver hands `store_tile!`: `Memory{T}` on Julia >= 1.11. mk_dense(v::AbstractVector{T}) where {T} = copyto!(parent(StridedView(zeros(T, length(v)))), v) +# Element types of the test operands A and B; the mixed-domain kernels differ. +mk_optypes(k) = (scalartype(k), scalartype(k)) + mk_run_acc(k, pa, pb, kc) = accumulate(k, zero_accumulator(k), pa, pb, kc) mk_run_exec(k, dst, pa, pb, kc, alpha, beta) = (execute_tile!(k, dst, pa, pb, kc, alpha, beta); nothing) @@ -68,11 +71,12 @@ end function mk_contract_light(k) T = scalartype(k) + TA, TB = mk_optypes(k) MR, NR = mr(k), nr(k) rng = MersenneTwister(100MR + NR) kc = 5 - A = rand(rng, T, MR, kc) - B = rand(rng, T, kc, NR) + A = rand(rng, TA, MR, kc) + B = rand(rng, TB, kc, NR) pa, pb = mk_pack(k, A, B) @test (length(pa), length(pb)) == (packed_a_length(k, kc), packed_b_length(k, kc)) acc = mk_run_acc(k, pa, pb, kc) @@ -99,13 +103,14 @@ end function mk_contract_full(k) T = scalartype(k) + TA, TB = mk_optypes(k) R = real(T) MR, NR = mr(k), nr(k) rng = MersenneTwister(100MR + NR + 1) m = max(1, MR - 1) kc = 5 - A = rand(rng, T, MR, kc) - B = rand(rng, T, kc, NR) + A = rand(rng, TA, MR, kc) + B = rand(rng, TB, kc, NR) pa, pb = mk_pack(k, A, B) acc = mk_run_acc(k, pa, pb, kc) @@ -126,8 +131,8 @@ function mk_contract_full(k) pad, sentinel = 2, T(-77) extents = unique(((MR, NR), (1, 1), (m, NR), (MR, max(1, NR - 1)), (MR ÷ 2 + 1, min(2, NR)), (0, 0))) for kc in (0, 1, 5), (alpha, beta) in mk_alphabeta(T), (m, n) in extents, scattered in (false, true) - A = rand(rng, T, MR, kc) - B = rand(rng, T, kc, NR) + A = rand(rng, TA, MR, kc) + B = rand(rng, TB, kc, NR) pa, pb = kc == 0 ? (R[], R[]) : mk_pack(k, A, B) Cold = iszero(beta) ? fill(mk_nan(T), m * n) : rand(rng, T, m * n) storage = [fill(sentinel, pad); Cold; fill(sentinel, pad)] @@ -142,10 +147,10 @@ function mk_contract_full(k) # Nonfinite padding lanes (rows >= m, columns >= n) never reach C. m, n = max(1, MR - 1), max(1, NR - 1) - A = rand(rng, T, MR, 2) - B = rand(rng, T, 2, NR) - A[(m + 1):end, :] .= mk_inf(T) - B[:, (n + 1):end] .= mk_inf(T) + A = rand(rng, TA, MR, 2) + B = rand(rng, TB, 2, NR) + A[(m + 1):end, :] .= mk_inf(TA) + B[:, (n + 1):end] .= mk_inf(TB) pa, pb = mk_pack(k, A, B) acc = mk_run_acc(k, pa, pb, 2) @test any(!isfinite(mk_read(k, acc, i, j)) for i in 0:(MR - 1), j in 0:(NR - 1)) @@ -174,8 +179,8 @@ function mk_contract_full(k) @test_throws DimensionMismatch execute_tile!(k, fulltile(st), pa, pb[1:(end - 1)], 2, one(T), zero(T)) # Destination layouts: each lands on exactly its own addresses. - A = rand(rng, T, MR, kc) - B = rand(rng, T, kc, NR) + A = rand(rng, TA, MR, kc) + B = rand(rng, TB, kc, NR) pa, pb = mk_pack(k, A, B) AB = A * B ld = m + 3 diff --git a/test/packing/test_kernel_descriptor.jl b/test/packing/test_kernel_descriptor.jl index 7502ef3..ee171bb 100644 --- a/test/packing/test_kernel_descriptor.jl +++ b/test/packing/test_kernel_descriptor.jl @@ -30,6 +30,7 @@ end @test packed_b_length(planar, 7) == 2NR * 7 @test_throws ArgumentError ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), PlanarFormat()) + @test packed_b_per_k(ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat())) == 2NR @test_throws ArgumentError ComplexKernelDescriptor(Val(MR), Val(NR), Float64, PlanarFormat(), PlanarFormat()) @test_throws ArgumentError ComplexKernelDescriptor(Val(0), Val(NR), T, PlanarFormat(), PlanarFormat()) end diff --git a/test/runtests.jl b/test/runtests.jl index ce45612..bb45471 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,6 +41,7 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip include("microkernels/test_planar_store_fastpath.jl") include("microkernels/test_onem_kernel.jl") include("microkernels/test_fmaddsub_kernel.jl") + include("microkernels/test_mixed_kernel.jl") include("planning/test_kernel_selection.jl") include("planning/test_plan_contract.jl") From 88826e2146c35e3635fc3b53470557c5e5c2f1d5 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 10:01:26 -0400 Subject: [PATCH 04/27] Add benchmark/bench_mixed.jl: mixed-domain kernels vs promoted vs real GEMM Indicative single-core timings on a Cascade Lake workstation (noisy), inner tile = default real shape (F64 32x6/W8, F32 64x6/W16): dtype case n promoted/mixed real/mixed ComplexF64 CR 512 1.84 0.93 ComplexF64 RC 512 1.94 0.95 ComplexF64 CR 2048 1.76 1.00 ComplexF64 RC 2048 2.01 1.05 ComplexF32 CR 512 1.92 1.00 ComplexF32 RC 512 2.07 0.95 ComplexF32 CR 2048 1.91 1.00 ComplexF32 RC 2048 1.99 1.02 probe_ttfx.jl 3x against 096ab7f: Float64 first call 3.2-3.3 s vs 3.1-3.2 s, ComplexF64 2.1-2.2 s vs 2.1-2.2 s (within noise). Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_mixed.jl | 62 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 benchmark/bench_mixed.jl diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl new file mode 100644 index 0000000..cc1181b --- /dev/null +++ b/benchmark/bench_mixed.jl @@ -0,0 +1,62 @@ +# Mixed-domain GEMM, single core: complex x real (CR) and real x complex (RC) +# through the named mixed-domain kernel, the default plan (which promotes the +# real operand to complex), and the real GEMM of the same real FMA count +# (2M x N x K for CR, M x 2N x K for RC). +# +# julia --project=benchmark benchmark/bench_mixed.jl [--dtypes ComplexF64,ComplexF32] +# [--sizes 512,2048] [--reps 5] [--outdir DIR] +# +# The mixed kernels' inner real tile is the host's default real shape, so +# `real/mixed` is the fraction of real efficiency reached. + +include(joinpath(@__DIR__, "harness.jl")) + +using QuasiStrided: ComplexRealKernel, RealComplexKernel, target_profile, _derived_shape + +const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) +const SIZES = parse_ints(argopt("sizes", "512,2048")) +const REPS = argopt("reps", 5) +const OUTDIR = outdir() +const CSV_PATH = joinpath(OUTDIR, "bench_mixed.csv") + +function time_plan(C, A, B; kw...) + plan = plan_contract(StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3); kw...) + return median_time_s(() -> execute!(plan, 1, 0); reps = REPS) +end + +function run(io) + print_env_header(stdout, "bench_mixed.jl") + println(io, "dtype,case,n,kernel,t_mixed,t_promoted,t_real") + rng = MersenneTwister(1) + for T in RUN_DTYPES, n in SIZES + R = real(T) + MRr, NRr, W = _derived_shape(target_profile(), R) + cases = ( + ( + "CR", ComplexRealKernel(Val(MRr ÷ 2), Val(NRr), T, Val(W)), + randn(rng, T, n, n), randn(rng, R, n, n), (2n, n), + ), + ( + "RC", RealComplexKernel(Val(MRr), Val(NRr ÷ 2), T, Val(W)), + randn(rng, R, n, n), randn(rng, T, n, n), (n, 2n), + ), + ) + for (case, kernel, A, B, (Mr, Nr)) in cases + C = zeros(T, n, n) + tm = time_plan(C, A, B; kernel) + tp = time_plan(C, A, B) + tr = time_plan(zeros(R, Mr, Nr), randn(rng, R, Mr, n), randn(rng, R, n, Nr)) + tag = "$(mr(kernel))x$(nr(kernel))/W$(lanewidth(kernel))" + @printf( + "%-10s %s %5d %-10s mixed %8.4f s promoted %8.4f s real %8.4f s promoted/mixed %.2f real/mixed %.2f\n", + T, case, n, tag, tm, tp, tr, tp / tm, tr / tm + ) + println(io, join((T, case, n, tag, tm, tp, tr), ',')) + flush(stdout) + end + end + return +end + +open(run, CSV_PATH, "w") +println("wrote ", CSV_PATH) From 42ad4bdae90e1594f6a5529ca98b44164b3c3eaf Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 09:45:22 -0400 Subject: [PATCH 05/27] Accumulate a narrow C in a compute-type panel across kc blocks When real(eltype(C)) is narrower than real(T) and K spans more than one kc block, the nest runs each jc column block against a dense M x min(N, nc) panel of T in the workspace (new grow-only slot c_panel): the block of C is loaded into the panel (skipped for beta == 0), all pc/ic/jr/ir iterations store into it, and it is rounded back to C once. The choice is a _PanelPath{P} singleton from _select_path, predicted exactly by the planning hint (the default kc does not depend on the complex method), so the plain nest is never compiled for such a plan. The dot path is skipped in that case; the outer path has one K block. execute_tilewise! keeps rounding per kc block. Float32 operands, accumulator = Float64, Float32 C, 64x512x48 at the default kc = 341: every element is now within 1 ulp of the correctly rounded Float64 result (was up to ~4 ulp). Same-eltype (F64, C64; packed, unpacked B) LLVM for execute!, _execute_path! and _execute_nest! is identical to 096ab7f modulo function-name suffixes and the workspace's dereferenceable size (168 -> 176 bytes, the new field); return types unchanged. Local measurements (noisy workstation, indicative): Pkg.test() 879.6 s vs 873.2 s for 096ab7f run concurrently; probe_ttfx.jl x3 identical (Float64 3.0-3.1 s, ComplexF64 2.0-2.1 s, outer 0.5 s); first @tensor call with Float32 operands and accumulator = Float64 over two kc blocks 3.35-3.42 s vs 3.17-3.28 s same-eltype Float64 (~1.05x). 1024^3, execute! only: narrow C with accumulator = Float64 27.9-28.6 ms vs all-Float64 26.6-28.0 ms (~1.04x); 096ab7f's per-block-rounding narrow C was 23.2-24.1 ms. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- src/execution/barrier.jl | 19 +++++- src/execution/execute.jl | 112 +++++++++++++++++++++++++++++++---- src/execution/oracle.jl | 3 +- src/execution/workspace.jl | 55 ++++++++++------- src/planning/plan.jl | 9 +-- test/execution/test_mixed.jl | 40 +++++++++---- 7 files changed, 188 insertions(+), 52 deletions(-) diff --git a/README.md b/README.md index 2ed41a6..41ca46c 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ Implemented: needs a complex output). The compute type is the promoted eltype, or the precision chosen with `accumulator = Float32`/`Float64` (`plan_contract(...; accumulator)`, `QuasiStridedBackend(; accumulator)`); - `C` is rounded on store, once per `kc` block of K. A real operand of a complex contraction is + the result is rounded to `eltype(C)` once. A real operand of a complex contraction is promoted at pack time. `conjA`/`conjB` and each operand's `StridedView.op` are applied during packing. - Zero steady-state allocation on Julia >= 1.11 for `execute!` on a reused diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index a6ecb7a..e85ddce 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -49,20 +49,31 @@ struct _NestPath{UNPACKED_B, AFF} end # sliver axis, so the flag would only split specialisations. @inline _nest_path(unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup) = _nest_path_from_flags( - unpacked_b, + _NEST_PATHS, unpacked_b, _is_ramp_map(mgroup, 1), _is_ramp_map(mgroup, 2), !unpacked_b && _is_ramp_map(ngroup, 1), _is_ramp_map(ngroup, 2), _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2) ) +# The nest run against a dense panel of C in the compute type, whose M and N +# maps are ramps. +struct _PanelPath{P <: _NestPath} end + +@inline _panel_path(unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup) = + _nest_path_from_flags( + _PANEL_PATHS, unpacked_b, + _is_ramp_map(mgroup, 1), true, !unpacked_b && _is_ramp_map(ngroup, 1), true, + _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2) +) + # Table lookup rather than `_NestPath{u, aff}()` built at runtime (a dynamic # type application) or a 128-leaf branch tree (slow to infer). -@inline function _nest_path_from_flags(flags::Vararg{Bool, 7}) +@inline function _nest_path_from_flags(table::Vector{Any}, flags::Vararg{Bool, 7}) index = 1 for i in 1:7 index += flags[i] << (i - 1) end - return @inbounds _NEST_PATHS[index] + return @inbounds table[index] end const _NEST_PATHS = let @@ -74,6 +85,8 @@ const _NEST_PATHS = let table end +const _PANEL_PATHS = Any[_PanelPath{typeof(p)}() for p in _NEST_PATHS] + # The dot path at lane width `W`; `MATB`: the matrix operand is B (`Qm == 1`). struct _DotPath{MATB, W} end diff --git a/src/execution/execute.jl b/src/execution/execute.jl index f8171e2..1451d02 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -57,28 +57,41 @@ end # outer-product path, else the nest (B packed or read in place). @inline _select_path(plan::ContractPlan{T}) where {T} = _select_path( T, _unpacked_b_kernel_eligible(plan.kernel), - plan.Astorage, plan.Bstorage, plan.Cstorage, plan.mgroup, plan.ngroup, plan.kgroup, plan + plan.Astorage, plan.Bstorage, plan.Cstorage, plan.mgroup, plan.ngroup, plan.kgroup, plan, + plan.blocking.kc ) # On the plan's parts, so `plan_contract` can predict the path before the # plan exists (`_path_hint`). `unpack_ok`: the kernel admits unpacked B. # `capacity`: the plan whose workspace must hold the dot path's vector, or -# `nothing` to assume it does. +# `nothing` to assume it does. `kc`: the requested K block, `nothing` for +# the default. @inline function _select_path( ::Type{T}, unpack_ok::Bool, Astorage, Bstorage, Cstorage, - mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, capacity + mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, capacity, kc ) where {T} Qm = axis_length(mgroup) Qn = axis_length(ngroup) Qk = axis_length(kgroup) - if (Qm == 1 || Qn == 1) && _dot_applicable(T, Astorage, Bstorage, kgroup, Qm, Qn, Qk) && + panel = _c_panel_needed(T, Cstorage, Qk, kc) + if (Qm == 1 || Qn == 1) && !panel && _dot_applicable(T, Astorage, Bstorage, kgroup, Qm, Qn, Qk) && _dot_capacity_ok(capacity) W = _dot_lanewidth(T) return Qm == 1 ? _lane_path(_DotPath{true}, W) : _lane_path(_DotPath{false}, W) end Qk == 1 && _outer_applicable(T, Astorage, Cstorage, mgroup, Qm) && return _lane_path(_OuterPath, _dot_lanewidth(T)) - return _nest_path(unpack_ok && _unpacked_b_rule(mgroup, kgroup), mgroup, ngroup, kgroup) + unpacked_b = unpack_ok && _unpacked_b_rule(mgroup, kgroup) + panel && return _panel_path(unpacked_b, mgroup, ngroup, kgroup) + return _nest_path(unpacked_b, mgroup, ngroup, kgroup) +end + +# Whether the partial sums between K blocks must live in a compute-type panel +# rather than in a C of narrower eltype. Static `false` unless C is narrower, +# so the default blocking is looked up only then. +@inline function _c_panel_needed(::Type{T}, Cstorage, Qk::Int, kc) where {T} + sizeof(real(eltype(Cstorage))) < sizeof(real(T)) || return false + return Qk > (kc === nothing ? _resolved_defaults(T).real_row.kc : kc) end # Run `path` behind a dynamic call. The plan crosses in the workspace slot @@ -125,7 +138,8 @@ end # on the path as well as the kernel. @inline _path_hint(::_Execute, req::_PlanRequest{T}, unpack_ok::Bool) where {T} = _select_path( T, unpack_ok, - req.Astorage, req.Bstorage, req.Cstorage, req.mgroup, req.ngroup, req.kgroup, nothing + req.Astorage, req.Bstorage, req.Cstorage, req.mgroup, req.ngroup, req.kgroup, nothing, + req.kc ) @inline _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} = @@ -153,6 +167,7 @@ end _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) @inline _hint_holds(plan::ContractPlan, ::_DotPath) = _dot_capacity_ok(plan) @inline _hint_holds(plan::ContractPlan, ::_OuterPath) = true +@inline _hint_holds(plan::ContractPlan, ::_PanelPath{P}) where {P} = _hint_holds(plan, P()) _execute_path!(plan::ContractPlan, alphaT, betaT, ::_DotPath{MATB, W}) where {MATB, W} = (_execute_dot!(plan, alphaT, betaT, MATB, Val(W)); nothing) @@ -173,16 +188,89 @@ function _execute_path!( _execute_nest!( plan, ws, kernel, mr(kernel), nr(kernel), axis_length(plan.mgroup), axis_length(plan.ngroup), axis_length(plan.kgroup), - plan.blocking.mc, plan.blocking.kc, plan.blocking.nc, alphaT, betaT, path + plan.blocking.mc, plan.blocking.kc, plan.blocking.nc, alphaT, betaT, path, nothing + ) + end + return nothing +end + +# The nest over `pplan`, a copy of `plan` whose C is the workspace panel with +# dense M/N maps, one `jc` block at a time: `_panel_enter!` loads the block of +# C into the panel and `_panel_exit!` rounds it back. +function _execute_path!( + plan::ContractPlan{T}, alphaT::T, betaT::T, ::_PanelPath{P} + ) where {T, P} + kernel = plan.kernel + ws = plan.workspace + Qm = axis_length(plan.mgroup) + pplan = ContractPlan( + kernel, _dense_second_map(plan.mgroup, 1), _dense_second_map(plan.ngroup, Qm), + plan.kgroup, plan.blocking, plan.Astorage, plan.Abase, plan.Bstorage, plan.Bbase, + ws.c_panel, 0, plan.atransform, plan.btransform, ws + ) + GC.@preserve ws begin + _execute_nest!( + pplan, ws, kernel, mr(kernel), nr(kernel), + Qm, axis_length(plan.ngroup), axis_length(plan.kgroup), + plan.blocking.mc, plan.blocking.kc, plan.blocking.nc, alphaT, betaT, P(), plan ) end return nothing end +# `g` with its second map replaced by the column-major one scaled by `step`. +@inline function _dense_second_map(g::AxisGroup{D, 2}, step::Int) where {D} + dense = ntuple(d -> step * _unchecked_axis_length(ntuple(i -> i < d ? g.lengths[i] : 1, Val(D))), Val(D)) + return AxisGroup(g.lengths, (g.strides[1], dense)) +end + +# `target`: `nothing`, or the plan whose C the panel stands in for. +@inline _panel_enter!(::Nothing, plan, ws, jc, nblock, betaT) = plan +@inline _panel_exit!(::Nothing, plan, ws, jc, nblock) = nothing + +function _panel_enter!(target::ContractPlan, plan::ContractPlan, ws, jc::Int, nblock::Int, betaT) + iszero(betaT) || _panel_copy!(target, plan.Cstorage, ws, jc, nblock, true) + Qm = axis_length(plan.mgroup) + return ContractPlan( + plan.kernel, plan.mgroup, plan.ngroup, plan.kgroup, plan.blocking, + plan.Astorage, plan.Abase, plan.Bstorage, plan.Bbase, plan.Cstorage, -jc * Qm, + plan.atransform, plan.btransform, plan.workspace + ) +end + +_panel_exit!(target::ContractPlan, plan::ContractPlan, ws, jc::Int, nblock::Int) = + _panel_copy!(target, plan.Cstorage, ws, jc, nblock, false) + +# Columns `jc .+ (0:nblock-1)` of `target`'s C into (`load`) or out of the +# panel, converting to the destination's eltype. Borrows the M/N offset +# buffers, which the nest refills before reading them again. +function _panel_copy!(target::ContractPlan, panel::Vector, ws, jc::Int, nblock::Int, load::Bool) + C = target.Cstorage + Qm = axis_length(target.mgroup) + mc = target.blocking.mc + fill_offsets!((ws.n_buf_B, ws.n_buf_C), target.ngroup, jc, nblock) + ic = 0 + while ic < Qm + mblock = min(mc, Qm - ic) + fill_offsets!((ws.m_buf_A, ws.m_buf_C), target.mgroup, ic, mblock) + for j in 1:nblock, i in 1:mblock + c = target.Cbase + ws.m_buf_C[i] + ws.n_buf_C[j] + 1 + q = ic + i + (j - 1) * Qm + if load + panel[q] = C[c] + else + C[c] = panel[q] + end + end + ic += mblock + end + return nothing +end + function _execute_nest!( plan::ContractPlan{T}, ws, kernel::K, MRk::Int, NRk::Int, Qm::Int, Qn::Int, Qk::Int, mc_eff::Int, kc_eff::Int, nc_eff::Int, - alphaT::T, betaT::T, ::_NestPath{UNPACKED_B, AFF} + alphaT::T, betaT::T, ::_NestPath{UNPACKED_B, AFF}, target ) where {T, K, UNPACKED_B, AFF} # GUARDRAIL: reals per sliver per K step address the packed panels; # `MRk`/`NRk` count register-tile rows. They differ for complex kernels. @@ -207,6 +295,7 @@ function _execute_nest!( jc = 0 while jc < Qn nblock = min(nc_eff, Qn - jc) + cplan = _panel_enter!(target, plan, ws, jc, nblock, betaT) n_slivers = cld(nblock, NRk) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( @@ -284,7 +373,7 @@ function _execute_nest!( end checked_span_bounds(plan.Abase, rng_mA, rng_kA, lenA) - checked_span_bounds(plan.Cbase, rng_mC, rng_nC, lenC) + checked_span_bounds(cplan.Cbase, rng_mC, rng_nC, lenC) for r in 0:(m_slivers - 1) rfirst = r * MRk @@ -299,12 +388,12 @@ function _execute_nest!( # --- loop 2: jr over N-slivers; loop 1: ir over M-slivers --- if UNPACKED_B _micro_tiles_unpacked_b!( - kernel, plan, ws, rowsB_k, m_slivers, n_slivers, + kernel, cplan, ws, rowsB_k, m_slivers, n_slivers, MRk, NRk, MRp, kblock, alphaT, beta_eff, aff_mC, aff_nC ) else _micro_tiles_packed_b!( - kernel, plan, ws, m_slivers, n_slivers, + kernel, cplan, ws, m_slivers, n_slivers, MRk, NRk, MRp, NRp, kblock, alphaT, beta_eff, aff_mC, aff_nC ) end @@ -316,6 +405,7 @@ function _execute_nest!( pc += kblock end + _panel_exit!(target, plan, ws, jc, nblock) jc += nblock end diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl index b3cc3ef..2ddf32f 100644 --- a/src/execution/oracle.jl +++ b/src/execution/oracle.jl @@ -1,7 +1,8 @@ # The tile-by-tile correctness oracle for `execute!`: per register tile, K in # `kc` panels, pack one sliver each and call the checked `execute_tile!`. Uses # only its own `tw_*` buffers, so it shares no mutable state with `execute!`. -# Needs a plan built with `oracle = true`. +# Needs a plan built with `oracle = true`. Unlike `execute!`, rounds a C narrower +# than the compute type once per `kc` block. function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) where {T} ws = plan.workspace # `tw_packed_a` is never empty for a legal blocking unless `oracle = false`. diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 550f71b..5805ab0 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -17,7 +17,7 @@ contractions; one from a non-default allocator is single-use and must be hold `real(T)`. Field layout is an implementation detail. """ # A `mutable struct` with `const` fields so a plan and a barrier slot hold one -# pointer to it instead of copying twenty-one GC-tracked fields. +# pointer to it instead of copying twenty-two GC-tracked fields. mutable struct ContractWorkspace{T, VT <: AbstractVector} # Macro-block offset buffers and their per-sliver descriptors. const m_buf_A::Vector{Int} @@ -50,6 +50,10 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} const slots::_SlotCache + # The compute-type panel of C for an eltype of C narrower than `T` (see + # `_PanelPath`); empty otherwise. + const c_panel::Vector{T} + # GUARDRAIL: the packed panels hold `real(T)`, never `T`. function ContractWorkspace{T, VT}( m_buf_A::Vector{Int}, m_buf_C::Vector{Int}, @@ -61,7 +65,7 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} tile_m_buf_A::Vector{Int}, tile_m_buf_C::Vector{Int}, tile_n_buf_B::Vector{Int}, tile_n_buf_C::Vector{Int}, tw_k_buf_A::Vector{Int}, tw_k_buf_B::Vector{Int}, - tw_packed_a::VT, tw_packed_b::VT + tw_packed_a::VT, tw_packed_b::VT, c_panel::Vector{T} ) where {T, VT <: AbstractVector} eltype(VT) === real(T) || throw( ArgumentError( @@ -74,7 +78,7 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} m_desc_A, m_desc_C, n_desc_B, n_desc_C, packed_a, packed_b, tile_m_buf_A, tile_m_buf_C, tile_n_buf_B, tile_n_buf_C, - tw_k_buf_A, tw_k_buf_B, tw_packed_a, tw_packed_b, _SlotCache(), + tw_k_buf_A, tw_k_buf_B, tw_packed_a, tw_packed_b, _SlotCache(), c_panel, ) end end @@ -114,7 +118,7 @@ end @inline function _build_workspace( ::Type{T}, s, ntw::Int, ints::F, - packed_a::VT, packed_b::VT, tw_packed_a::VT, tw_packed_b::VT + packed_a::VT, packed_b::VT, tw_packed_a::VT, tw_packed_b::VT, panel::Int ) where {T, F, VT <: AbstractVector} return ContractWorkspace{T, VT}( ints(s.mc), ints(s.mc), @@ -126,21 +130,24 @@ end ints(s.mr), ints(s.mr), ints(s.nr), ints(s.nr), ints(ntw), ints(ntw), - tw_packed_a, tw_packed_b, + tw_packed_a, tw_packed_b, Vector{T}(undef, panel), ) end """ ContractWorkspace(T, kernel, blocking::Blocking; - oracle = true, allocator = TensorOperations.DefaultAllocator()) + oracle = true, allocator = TensorOperations.DefaultAllocator(), + panel = 0) A workspace for compute type `T`, `kernel` and an effective `blocking`. -`oracle = false` leaves the `execute_tilewise!` buffers empty. Under a +`oracle = false` leaves the `execute_tilewise!` buffers empty. `panel` is the +length of the compute-type panel of C (see [`plan_contract`](@ref)). Under a non-default `allocator` the packed panels come from `tensoralloc`, are never resized, and must be handed back with [`release!`](@ref). """ function ContractWorkspace( - ::Type{T}, kernel, blocking::Blocking, oracle::Bool, ::TO.DefaultAllocator + ::Type{T}, kernel, blocking::Blocking, oracle::Bool, ::TO.DefaultAllocator, + panel::Int = 0 ) where {T} s = _workspace_sizes(kernel, blocking) R = realtype(kernel) @@ -148,12 +155,12 @@ function ContractWorkspace( T, s, oracle ? s.kc : 0, n -> Vector{Int}(undef, n), Vector{R}(undef, s.packed_a), Vector{R}(undef, s.packed_b), Vector{R}(undef, oracle ? s.tw_packed_a : 0), - Vector{R}(undef, oracle ? s.tw_packed_b : 0), + Vector{R}(undef, oracle ? s.tw_packed_b : 0), panel, ) end function ContractWorkspace( - ::Type{T}, kernel, blocking::Blocking, oracle::Bool, allocator + ::Type{T}, kernel, blocking::Blocking, oracle::Bool, allocator, panel::Int = 0 ) where {T} s = _workspace_sizes(kernel, blocking) R = realtype(kernel) @@ -166,22 +173,23 @@ function ContractWorkspace( return _build_workspace( T, s, oracle ? s.kc : 0, n -> _alloc_offsets(n, allocator), - packed_a, packed_b, tw_packed_a, tw_packed_b + packed_a, packed_b, tw_packed_a, tw_packed_b, panel ) end function ContractWorkspace( ::Type{T}, kernel, blocking::Blocking; - oracle::Bool = true, allocator = TO.DefaultAllocator() + oracle::Bool = true, allocator = TO.DefaultAllocator(), panel::Int = 0 ) where {T} - return ContractWorkspace(T, kernel, blocking, oracle, allocator) + return ContractWorkspace(T, kernel, blocking, oracle, allocator, panel) end # Grow `ws` in place (never shrink, never reallocate what is large enough) # for `kernel` at `blocking`. Only for the GC-owned path: an allocator's # temporaries must never be resized. function reserve!( - ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, oracle::Bool + ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, oracle::Bool, + panel::Int = 0 ) where {T, R} s = _workspace_sizes(kernel, blocking) @@ -204,6 +212,7 @@ function reserve!( _grow!(ws.tile_m_buf_C, s.mr) _grow!(ws.tile_n_buf_B, s.nr) _grow!(ws.tile_n_buf_C, s.nr) + _grow!(ws.c_panel, panel) if oracle _grow!(ws.tw_k_buf_A, s.kc) @@ -238,16 +247,17 @@ end # Build or reuse the plan's workspace, dispatching on the allocator type. function _resolve_workspace( ::Type{T}, workspace, kernel, blocking::Blocking, oracle::Bool, - allocator::TO.DefaultAllocator + allocator::TO.DefaultAllocator, panel::Int ) where {T} workspace === nothing && - return ContractWorkspace(T, kernel, blocking, oracle, allocator) - return _reuse_workspace(T, workspace, kernel, blocking, oracle) + return ContractWorkspace(T, kernel, blocking, oracle, allocator, panel) + return _reuse_workspace(T, workspace, kernel, blocking, oracle, panel) end # Explicit allocator: sized once, never reused; the caller owns `release!`. function _resolve_workspace( - ::Type{T}, workspace, kernel, blocking::Blocking, oracle::Bool, allocator + ::Type{T}, workspace, kernel, blocking::Blocking, oracle::Bool, allocator, + panel::Int ) where {T} workspace === nothing || throw( ArgumentError( @@ -256,17 +266,17 @@ function _resolve_workspace( "sized once at construction and must not be resized or reused" ) ) - return ContractWorkspace(T, kernel, blocking, oracle, allocator) + return ContractWorkspace(T, kernel, blocking, oracle, allocator, panel) end # The pool is keyed by the compute type alone, so also check the packed type # (folds at compile time). @inline function _reuse_workspace( ::Type{T}, ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, - oracle::Bool + oracle::Bool, panel::Int ) where {T, R} R === realtype(kernel) || _throw_packed_eltype_mismatch(T, ws, kernel) - return reserve!(ws, kernel, blocking, oracle) + return reserve!(ws, kernel, blocking, oracle, panel) end @noinline function _throw_packed_eltype_mismatch(::Type{T}, ws, kernel) where {T} @@ -279,7 +289,8 @@ end end @noinline function _reuse_workspace( - ::Type{T}, ws::ContractWorkspace, kernel, blocking::Blocking, oracle::Bool + ::Type{T}, ws::ContractWorkspace, kernel, blocking::Blocking, oracle::Bool, + panel::Int ) where {T} throw( ArgumentError( diff --git a/src/planning/plan.jl b/src/planning/plan.jl index a82ef9b..c859214 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -57,9 +57,9 @@ invalid input. `T` is `promote_type` of the three, or, for `accumulator = Float32` or `Float64`, that precision in the domain (real or complex) of the promoted type. Operands are converted to `T` on load and `alpha*AB + beta*C` is - evaluated in `T`, rounding to `eltype(C)` on store. `C` holds the partial - sums between `kc` blocks of K, so a K longer than `kc` rounds once per - block. + evaluated in `T`, rounding to `eltype(C)` once. For an `eltype(C)` of + lower precision than `T`, a K longer than `kc` accumulates in a panel of + `T` in the workspace, of `M * min(N, nc)` elements. * `kernel = nothing` picks one from the hardware profile and the extents: a [`SIMDKernel`](@ref) for a real `T`, a [`PlanarKernel`](@ref) for a complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512). @@ -315,7 +315,8 @@ function _plan_contract( nc_eff = Qn == 0 ? NRk : min(_roundup(requested.nc, NRk), _roundup(Qn, NRk)) kc_eff = Qk == 0 ? 1 : min(requested.kc, Qk) blocking = Blocking(mc_eff, kc_eff, nc_eff) - ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator) + panel = _c_panel_needed(T, req.Cstorage, Qk, kc_eff) ? Qm * min(nc_eff, Qn) : 0 + ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel) plan = ContractPlan( kernel, req.mgroup, req.ngroup, req.kgroup, blocking, diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 664ae1a..6ad78bc 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -49,14 +49,34 @@ end @test _steady_allocs!(execute!, _mm_plan(Cmat, Amat, Bmat), Cmat) == 0 skip = (VERSION < v"1.11") end -# One K panel: between panels the partial sums round-trip through C. -@testset "accumulator = Float64: Float32 operands within one ulp" begin - rng = MersenneTwister(5) - A, B, C0 = randn(rng, Float32, 64, 512), randn(rng, Float32, 512, 48), randn(rng, Float32, 64, 48) - ref = Float32.(Float64.(A) * Float64.(B) .+ 0.5 .* Float64.(C0)) - C64, C32 = copy(C0), copy(C0) - execute!(_mm_plan(C64, A, B; accumulator = Float64, kc = 512), 1, 0.5) - execute!(_mm_plan(C32, A, B), 1, 0.5) - @test all(abs.(C64 .- ref) .<= eps.(ref)) - @test maximum(abs.(C64 .- ref)) < maximum(abs.(C32 .- ref)) +@testset "Float32 C, Float64 compute type, K over several kc blocks: within one ulp ($TA, $layout, beta = $beta)" for (TA, acc, layout) in ( + (Float32, Float64, :dense), (Float64, nothing, :dense), (Float32, Float64, :scattered), + ), beta in (0, 0.5) + fx = if layout === :scattered + Cv = scattered_fixture(Float32, 32, 512)[1] + copyto!(Cv, randn(MersenneTwister(4), Float32, size(Cv))) + (Cv, scattered_fixture(TA, 32, 512)[2:end]...) + else + rng = MersenneTwister(5) + A, B, C = randn(rng, TA, 64, 512), randn(rng, TA, 512, 48), randn(rng, Float32, 64, 48) + (StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3)) + end + ref = Float32.(_mixed_ref(fx, 1, beta)) + C32 = copy(fx[1]) + execute!(plan_contract(C32, map(x -> x isa StridedView ? StridedView(Float32.(x)) : x, fx[2:end])...), 1, beta) + plan = plan_contract(fx...; accumulator = acc) + @test _path_of(plan) isa QuasiStrided._PanelPath + @test axis_length(plan.kgroup) > plan.blocking.kc + execute!(plan, 1, beta) + err = abs.(Array(fx[1]) .- ref) + @test all(err .<= eps.(ref)) + @test maximum(err) < maximum(abs.(Array(C32) .- ref)) +end + +@testset "Float32 C, Float64 compute type: steady-state execute! allocates nothing" begin + rng = MersenneTwister(6) + Amat, Bmat, Cmat = randn(rng, Float32, 64, 700), randn(rng, Float32, 700, 40), zeros(Float32, 64, 40) + plan = _mm_plan(Cmat, Amat, Bmat; accumulator = Float64) + @test _path_of(plan) isa QuasiStrided._PanelPath + @test _steady_allocs!(execute!, plan, Cmat) == 0 skip = (VERSION < v"1.11") end From a3d4189d988d92b2ca236db63b0cabb8d08193a8 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 10:31:35 -0400 Subject: [PATCH 06/27] Select the mixed-domain kernels automatically for a complex T with one real operand A complex T with a complex A and a real B now plans a ComplexRealKernel, a real A and a complex B a RealComplexKernel, statically from the operand eltypes (no A/B swap for a complex T). Both run the real default shape logic of real(T) on the real problem, (2Qm, Qn, 2run) for CR and (Qm, 2Qn, run) for RC, and map the shape to (MR/2, NR, W) and (MR, NR/2, W): the extent, store, small-M and run demotions carry over and reuse the real cached defaults. A named CR/RC kernel whose RealFormat side is complex throws ArgumentError at plan time. Same-eltype selection and blocking for F32/F64/C32/C64 are unchanged over 7 shapes plus a swapped layout; code_llvm of _planned and the plan_contract kwcall is identical to 293e212 modulo SSA/label names, type tag addresses and GC-frame slot order (same frame size). Co-Authored-By: Claude Opus 5.5 --- src/planning/defaults.jl | 11 ++++++ src/planning/kernel_selection.jl | 35 ++++++++++++++++++ src/planning/plan.jl | 45 +++++++++++++++-------- test/execution/test_mixed.jl | 35 +++++++++++------- test/planning/test_kernel_selection.jl | 49 ++++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 26 deletions(-) diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index bc1fad0..8ac74f0 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -76,6 +76,17 @@ _default_kernel(::Type{T}) where {T} = return (d.fitted, method) end +# The automatic `(shape, method)` under `method`, `_default_method(T, TA, TB)`. +@inline _default_shape(::Type{T}, ::ComplexMethod, Qm::Int, Qn::Int, run::Int) where {T} = + _default_shape(T, Qm, Qn, run) + +# The real default shape of the real problem, mapped: the real extent, store and +# small-M demotions carry over, on the real type's cached defaults. +@inline function _default_shape(::Type{T}, method::_MixedMethod, Qm::Int, Qn::Int, run::Int) where {T} + shape, _ = _default_shape(real(T), _real_problem(method, Qm, Qn, run)...) + return (_mixed_shape(method, shape), method) +end + # `@noinline`: the return type is the Union of `T`'s menu kernels. @noinline function _default_kernel(::Type{T}, Qm::Int, Qn::Int) where {T} shape, method = _default_shape(T, Qm, Qn) diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index d37d658..b1bfd9f 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -8,6 +8,17 @@ _default_method(::Type{<:Real}) = RealMethod() _default_method(::Type{<:Complex}) = PlanarMethod() +const _MixedMethod = Union{ComplexRealMethod, RealComplexMethod} + +# The method for compute type `T` and (oriented) operand eltypes: a mixed-domain +# method when exactly one operand is real. +_default_method(::Type{T}, ::Type{TA}, ::Type{TB}) where {T, TA, TB} = + _default_method(T) +_default_method(::Type{T}, ::Type{<:Complex}, ::Type{<:Real}) where {T <: Complex} = + ComplexRealMethod() +_default_method(::Type{T}, ::Type{<:Real}, ::Type{<:Complex}) where {T <: Complex} = + RealComplexMethod() + const NR_DEFAULT = 6 # An (8, 6) tile at one 256-bit register's lane width: the shape wherever no rule applies. @@ -48,6 +59,20 @@ kernel_shapes(::Type{ComplexF32}, ::PlanarMethod) = KERNEL_SHAPES_C32_PLANAR kernel_shapes(::Type{ComplexF32}, ::OneMMethod) = KERNEL_SHAPES_C32_ONEM kernel_shapes(::Type{ComplexF64}, ::FMAddSubMethod) = KERNEL_SHAPES_C64_FMADDSUB kernel_shapes(::Type{ComplexF32}, ::FMAddSubMethod) = KERNEL_SHAPES_C32_FMADDSUB +# The mixed menus are the real menu of `real(T)`, mapped by `_mixed_shape`. +kernel_shapes(::Type{T}, method::_MixedMethod) where {T <: Union{ComplexF32, ComplexF64}} = + map(s -> _mixed_shape(method, s), kernel_shapes(real(T), RealMethod())) + +# A real inner shape to the mixed kernel's `(MR, NR, W)`, and back: complex A +# rows and complex B columns each span two reals of the inner kernel. +_mixed_shape(::ComplexRealMethod, (MR, NR, W)::NTuple{3, Int}) = (MR ÷ 2, NR, W) +_mixed_shape(::RealComplexMethod, (MR, NR, W)::NTuple{3, Int}) = (MR, NR ÷ 2, W) +_real_shape(::ComplexRealMethod, (MR, NR, W)::NTuple{3, Int}) = (2 * MR, NR, W) +_real_shape(::RealComplexMethod, (MR, NR, W)::NTuple{3, Int}) = (MR, 2 * NR, W) + +# The real problem the inner kernel sees: extents and C's unit-stride M run in reals. +_real_problem(::ComplexRealMethod, Qm::Int, Qn::Int, run::Int) = (2 * Qm, Qn, 2 * run) +_real_problem(::RealComplexMethod, Qm::Int, Qn::Int, run::Int) = (Qm, 2 * Qn, run) # The kernel type implementing `method` for `T`, or `nothing` (exactly the # pairs with a menu). @@ -55,6 +80,8 @@ _kernel_type(::RealMethod, ::Type{<:Union{Float32, Float64}}) = SIMDKernel _kernel_type(::PlanarMethod, ::Type{<:Union{ComplexF32, ComplexF64}}) = PlanarKernel _kernel_type(::OneMMethod, ::Type{<:Union{ComplexF32, ComplexF64}}) = OneMKernel _kernel_type(::FMAddSubMethod, ::Type{<:Union{ComplexF32, ComplexF64}}) = FMAddSubKernel +_kernel_type(::ComplexRealMethod, ::Type{<:Union{ComplexF32, ComplexF64}}) = ComplexRealKernel +_kernel_type(::RealComplexMethod, ::Type{<:Union{ComplexF32, ComplexF64}}) = RealComplexKernel _kernel_type(::Any, ::Type) = nothing # An unrolled `shape === menu[i] ? ... :` ladder over the menu whose branches @@ -263,6 +290,14 @@ end return target === nothing ? shape : target end +@inline function _demote_shape_for_run( + ::Type{T}, shape::NTuple{3, Int}, method::_MixedMethod, run::Int, Qm::Int, Qk::Int + ) where {T} + m, _, r = _real_problem(method, Qm, 0, run) + real_shape = _demote_shape_for_run(real(T), _real_shape(method, shape), RealMethod(), r, m, Qk) + return _mixed_shape(method, real_shape) +end + function _run_demotion_target(::Type{T}, mrk::Int, method, run::Int, Qm::Int, Qk::Int) where {T} T <: Real || return nothing kmax = T === Float64 ? _RUN_DEMOTE_KMAX_F64 : _RUN_DEMOTE_KMAX_F32 diff --git a/src/planning/plan.jl b/src/planning/plan.jl index c859214..cd10d06 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -62,8 +62,9 @@ invalid input. `T` in the workspace, of `M * min(N, nc)` elements. * `kernel = nothing` picks one from the hardware profile and the extents: a [`SIMDKernel`](@ref) for a real `T`, a [`PlanarKernel`](@ref) for a - complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512). - [`OneMKernel`](@ref) is used only when named. + complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512), a + [`ComplexRealKernel`](@ref)/[`RealComplexKernel`](@ref) for a complex `T` + with a real `B`/`A`. [`OneMKernel`](@ref) is used only when named. * Labels within the M and N composites are ordered by their stride in `C`; the K order follows a cost model of the two packs. For a real `T` the operand roles are swapped (B feeds M) when only the N side gives `C` a @@ -151,6 +152,8 @@ function _planned( accumulator::AC ) where {F, NA, NB, NC, AC} T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) + method = _default_method(T, eltype(A), eltype(B)) + _check_kernel_domain(kernel, eltype(A), eltype(B)) # GUARDRAIL: a conjugated `C` is rejected; the engine writes through to # its parent, so there is nowhere to absorb its `op`. @@ -168,8 +171,9 @@ function _planned( morder = _order_free_labels(mlabels, indC, C) norder = _order_free_labels(nlabels, indC, C) - # Only real `T` swaps or run-demotes; `0` is a placeholder. - run_m = T <: Real ? _leading_unit_run(morder, indC, C) : 0 + # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. + runs = T <: Real || method isa _MixedMethod + run_m = runs ? _leading_unit_run(morder, indC, C) : 0 run_n = T <: Real ? _leading_unit_run(norder, indC, C) : 0 mgroup = _build_pair_group(morder, indA, A, indC, C) # maps: (A, C) @@ -183,7 +187,7 @@ function _planned( Qk = axis_length(kgroup) - mr_asis, mr_swapped = _candidate_mrs(T, kernel, Qm, Qn, run_m, run_n) + mr_asis, mr_swapped = _candidate_mrs(T, method, kernel, Qm, Qn, run_m, run_n) if T <: Real && _prefer_swap(run_m, run_n, mr_asis, mr_swapped) # B takes the M role: groups, K maps, storages, run and transforms move @@ -194,18 +198,31 @@ function _planned( parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), run_n, mc, kc, nc, workspace, allocator, oracle ) - return _plan_with_kernel(kernel, btransform, atransform, req_swapped) + return _plan_with_kernel(kernel, method, btransform, atransform, req_swapped) end req = _plan_request( T, f, mgroup, ngroup, kgroup, parent(A), offset(A), parent(B), offset(B), parent(C), offset(C), run_m, mc, kc, nc, workspace, allocator, oracle ) - return _plan_with_kernel(kernel, atransform, btransform, req) + return _plan_with_kernel(kernel, method, atransform, btransform, req) end const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) +# A named mixed-domain kernel's `RealFormat` side packs a real operand only. +@inline _check_kernel_domain(kernel, ::Type{TA}, ::Type{TB}) where {TA, TB} = + _check_kernel_domain(complex_method(kernel), kernel, TA, TB) +@inline _check_kernel_domain(::ComplexMethod, kernel, ::Type, ::Type) = nothing +@inline _check_kernel_domain(::ComplexRealMethod, kernel, ::Type, ::Type{TB}) where {TB} = + TB <: Real || _throw_kernel_domain(kernel, "B", TB) +@inline _check_kernel_domain(::RealComplexMethod, kernel, ::Type{TA}, ::Type) where {TA} = + TA <: Real || _throw_kernel_domain(kernel, "A", TA) + +@noinline _throw_kernel_domain(kernel, side, T) = throw( + ArgumentError("plan_contract: $(typeof(kernel)) needs a real $side, got eltype $T") +) + # Fold to the compute type, or a throw, at compile time. @inline function _compute_type(::Type{TA}, ::Type{TB}, ::Type{TC}, ::Nothing) where {TA, TB, TC} (TA in _QS_ELTYPES && TB in _QS_ELTYPES && TC in _QS_ELTYPES) || @@ -230,11 +247,11 @@ end # `mr` of the kernel each orientation would run, for the swap decision: a named # kernel either way, else `_default_shape`'s pick at that orientation's extents # and C run. -@inline _candidate_mrs(::Type{T}, kernel, Qm::Int, Qn::Int, run_m::Int, run_n::Int) where {T} = +@inline _candidate_mrs(::Type{T}, method, kernel, Qm::Int, Qn::Int, run_m::Int, run_n::Int) where {T} = (mr(kernel), mr(kernel)) -@inline function _candidate_mrs(::Type{T}, ::Nothing, Qm::Int, Qn::Int, run_m::Int, run_n::Int) where {T} - mr_asis = _default_shape(T, Qm, Qn, run_m)[1][1] - mr_swapped = T <: Real ? _default_shape(T, Qn, Qm, run_n)[1][1] : mr_asis +@inline function _candidate_mrs(::Type{T}, method, ::Nothing, Qm::Int, Qn::Int, run_m::Int, run_n::Int) where {T} + mr_asis = _default_shape(T, method, Qm, Qn, run_m)[1][1] + mr_swapped = T <: Real ? _default_shape(T, method, Qn, Qm, run_n)[1][1] : mr_asis return mr_asis, mr_swapped end @@ -244,11 +261,11 @@ end # kernel's code is compiled (holding the menu-wide kernel Union would box the # request; a static ladder would compile every menu kernel). All barrier # arguments are singletons or heap objects: 0 B, one method-cache hit. -@inline _plan_with_kernel(kernel, atransform, btransform, req::_PlanRequest) = +@inline _plan_with_kernel(kernel, method, atransform, btransform, req::_PlanRequest) = _plan_contract(kernel, atransform, btransform, req, nothing) -@inline function _plan_with_kernel(::Nothing, atransform, btransform, req::_PlanRequest{T}) where {T} +@inline function _plan_with_kernel(::Nothing, method, atransform, btransform, req::_PlanRequest{T}) where {T} Qm = axis_length(req.mgroup) - shape, method = _default_shape(T, Qm, axis_length(req.ngroup), req.run) + shape, method = _default_shape(T, method, Qm, axis_length(req.ngroup), req.run) shape = _demote_shape_for_run(T, shape, method, req.run, Qm, axis_length(req.kgroup)) vshape = _menu_val(shape, T, method) # The execution path, predicted so the callee is specialised on it. Only a diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 6ad78bc..cc47506 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -19,16 +19,19 @@ function _mixed_ref(fx, alpha, beta; kw...) return _brute_ref(R.(Array(Av)), iA, R.(Array(Bv)), iB, R.(Array(Cv)), iC, alpha, beta; kw...) end -@testset "mixed eltypes: $TA x $TB -> $TC, accumulator = $acc, $layout" for (TA, TB, TC, acc, layout, MKN, beta, conjA, conjB, path) in ( - (Float32, Float64, Float64, nothing, :dense, (23, 37, 19), 0, false, false, QuasiStrided._NestPath), - (Float64, Float32, Float32, nothing, :dense, (32, 40, 12), 1, false, false, QuasiStrided._NestPath), - (Float64, ComplexF64, ComplexF64, nothing, :dense, (21, 33, 17), 0.3 - 0.7im, false, true, QuasiStrided._NestPath), - (ComplexF32, Float64, ComplexF64, nothing, :scattered, (32, 32, 32), 0.5, true, false, QuasiStrided._NestPath), - (ComplexF32, ComplexF64, ComplexF64, nothing, :dense, (19, 29, 23), 1, true, true, QuasiStrided._NestPath), - (Float64, Float64, ComplexF64, nothing, :dense, (18, 25, 14), 0.25 + 0.5im, false, false, QuasiStrided._NestPath), - (ComplexF64, Float32, ComplexF64, Float32, :dense, (17, 31, 13), 0, false, false, QuasiStrided._NestPath), - (Float32, Float64, Float64, nothing, :dense, (1, 70, 29), 0.5, false, false, QuasiStrided._DotPath), - (Float64, Float32, Float64, nothing, :dense, (37, 1, 21), 1, false, false, QuasiStrided._OuterPath), +const _CR, _RC = QuasiStrided.ComplexRealKernel, QuasiStrided.RealComplexKernel +const _PROMOTED = Union{QuasiStrided.PlanarKernel, QuasiStrided.FMAddSubKernel} + +@testset "mixed eltypes: $TA x $TB -> $TC, accumulator = $acc, $layout" for (TA, TB, TC, acc, layout, MKN, beta, conjA, conjB, path, K) in ( + (Float32, Float64, Float64, nothing, :dense, (23, 37, 19), 0, false, false, QuasiStrided._NestPath, SIMDKernel), + (Float64, Float32, Float32, nothing, :dense, (32, 40, 12), 1, false, false, QuasiStrided._NestPath, SIMDKernel), + (Float64, ComplexF64, ComplexF64, nothing, :dense, (21, 33, 17), 0.3 - 0.7im, false, true, QuasiStrided._NestPath, _RC), + (ComplexF32, Float64, ComplexF64, nothing, :scattered, (32, 32, 32), 0.5, true, false, QuasiStrided._NestPath, _CR), + (ComplexF32, ComplexF64, ComplexF64, nothing, :dense, (19, 29, 23), 1, true, true, QuasiStrided._NestPath, _PROMOTED), + (Float64, Float64, ComplexF64, nothing, :dense, (18, 25, 14), 0.25 + 0.5im, false, false, QuasiStrided._NestPath, _PROMOTED), + (ComplexF64, Float32, ComplexF64, Float32, :dense, (17, 31, 13), 0, false, false, QuasiStrided._NestPath, _CR), + (Float32, Float64, Float64, nothing, :dense, (1, 70, 29), 0.5, false, false, QuasiStrided._DotPath, SIMDKernel), + (Float64, Float32, Float64, nothing, :dense, (37, 1, 21), 1, false, false, QuasiStrided._OuterPath, SIMDKernel), ) fx = _mixed_fixture(TA, TB, TC, MKN..., layout, 7) alpha = 1.5 @@ -36,13 +39,21 @@ end plan = plan_contract(fx...; conjA, conjB, accumulator = acc) T = acc === nothing ? promote_type(TA, TB, TC) : (TC <: Complex ? Complex{acc} : acc) @test plan isa ContractPlan{T} + @test plan.kernel isa K @test _path_of(plan) isa path execute!(plan, alpha, beta) - K = layout === :scattered ? 32 : MKN[2] - rtol = 10 * sqrt(K) * max(eps(real(T)), eps(real(TC))) + Qk = layout === :scattered ? 32 : MKN[2] + rtol = 10 * sqrt(Qk) * max(eps(real(T)), eps(real(TC))) @test norm(Array(fx[1]) - ref) <= rtol * norm(ref) end +@testset "a named mixed-domain kernel with a complex RealFormat side throws at plan time" begin + Cv, Av, Bv = (StridedView(zeros(ComplexF64, 8, 8)) for _ in 1:3) + for kernel in (_CR(Val(4), Val(6), ComplexF64, Val(4)), _RC(Val(8), Val(3), ComplexF64, Val(4))) + @test_throws ArgumentError plan_contract(Cv, Av, (1, 2), Bv, (2, 3), (1, 3); kernel) + end +end + @testset "mixed eltypes: steady-state execute! allocates nothing" begin rng = MersenneTwister(3) Amat, Bmat, Cmat = randn(rng, Float32, 40, 30), randn(rng, ComplexF64, 30, 20), zeros(ComplexF64, 40, 20) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index b87b59f..9a8eb8a 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -229,6 +229,55 @@ end end end +@testset "mixed-domain selection: the real default of real(T), mapped" begin + CR, RC = QuasiStrided.ComplexRealMethod(), QuasiStrided.RealComplexMethod() + mapped(m, (MR, NR, W)) = m === CR ? (MR ÷ 2, NR, W) : (MR, NR ÷ 2, W) + dshape = QuasiStrided._default_shape + dmethod = QuasiStrided._default_method + for T in (ComplexF64, ComplexF32) + R = real(T) + @test all(((MR, NR, W),) -> iseven(NR) && iseven(W), kernel_shapes(R)) + for m in (CR, RC) + @test kernel_shapes(T, m) === map(s -> mapped(m, s), kernel_shapes(R)) + for shape in kernel_shapes(T, m) + k = QuasiStrided._kernel_from_shape(shape, T, m) + @test complex_method(k) === m && (mr(k), nr(k), lanewidth(k)) === shape + end + end + @test dmethod(T, T, R) === dmethod(T, ComplexF32, Float64) === CR + @test dmethod(T, R, T) === dmethod(T, Float32, ComplexF64) === RC + @test dmethod(T, R, R) === dmethod(T, T, T) === dmethod(T, ComplexF32, T) === PlanarMethod() + @test dmethod(R, Float32, Float64) === RealMethod() + end + # (isa, vb, nregisters) => (F64, F32, C64, C32) default shapes at a large extent. + expected = ( + (:avx512, 64, 32) => ((32, 6, 8), (64, 6, 16), (24, 3, 8), (48, 3, 16)), + (:avx2, 32, 16) => ((8, 6, 4), (16, 6, 8), (4, 5, 4), (8, 5, 8)), + (:neon, 16, 32) => ((8, 6, 4), (8, 6, 8), (4, 6, 2), (8, 6, 4)), + ) + saved = QuasiStrided._TARGET[] + try + for ((key, vb, nreg), shapes) in expected + QuasiStrided._TARGET[] = synthetic(key, vb; nregisters = nreg) + Q = 4096 + for (T, shape) in zip((Float64, Float32, ComplexF64, ComplexF32), shapes) + @test dshape(T, Q, Q, Q) === (shape, QuasiStrided._default_method(T)) + @test dshape(T, QuasiStrided._default_method(T), Q, Q, Q) === dshape(T, Q, Q, Q) + end + for (T, rshape) in zip((ComplexF64, ComplexF32), shapes[1:2]), m in (CR, RC) + @test dshape(T, m, Q, Q, Q) === (mapped(m, rshape), m) + # The real extent, store and small-M demotions carry over. + for Qm in (1, 3, 17, 40, 100), run in (0, 1, 8, Qm) + real_problem = m === CR ? (2 * Qm, 50, 2 * run) : (Qm, 100, run) + @test dshape(T, m, Qm, 50, run) === (mapped(m, dshape(real(T), real_problem...)[1]), m) + end + end + end + finally + QuasiStrided._TARGET[] = saved + end +end + @testset "_derived_shape is always a constructible member of the method's menu" begin methods = ( (Float64, RealMethod()), (Float32, RealMethod()), (ComplexF64, PlanarMethod()), From 565ad27c7f4386043d604b483bf5ef144c517ec6 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 10:54:26 -0400 Subject: [PATCH 07/27] bench_mixed.jl: time the promoted path through the named planar kernel The default plan now selects the mixed-domain kernel, so the promoted column names the host's planar kernel instead. End-to-end, @tensor with QuasiStridedBackend(), median, workstation (indicative only), this branch vs 293e212 (promoted): F64 x C64 -> C64 512^3: 6.2 ms vs 12.9 ms (2.08x); 2048^3: 0.415 s vs 0.804 s (1.94x) C64 x F64 -> C64 512^3: 6.9 ms vs 12.4 ms (1.80x); 2048^3: 0.456 s vs 0.798 s (1.75x) First @tensor call, fresh process (3 runs): F64 x C64 3.17-3.30 s, C64 x F64 3.40-3.54 s, C64 x C64 3.06-3.14 s (1.04x / 1.11x). probe_ttfx.jl 3x: F64 3.0-3.1 s, C64 2.0-2.1 s, both trees. Pkg.test() run concurrently with 293e212: 850.9 s vs 862.2 s wall. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_mixed.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index cc1181b..5a981a6 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -1,6 +1,6 @@ # Mixed-domain GEMM, single core: complex x real (CR) and real x complex (RC) -# through the named mixed-domain kernel, the default plan (which promotes the -# real operand to complex), and the real GEMM of the same real FMA count +# through the named mixed-domain kernel, the default planar kernel (which +# promotes the real operand to complex), and the real GEMM of the same real FMA count # (2M x N x K for CR, M x 2N x K for RC). # # julia --project=benchmark benchmark/bench_mixed.jl [--dtypes ComplexF64,ComplexF32] @@ -11,7 +11,8 @@ include(joinpath(@__DIR__, "harness.jl")) -using QuasiStrided: ComplexRealKernel, RealComplexKernel, target_profile, _derived_shape +using QuasiStrided: ComplexRealKernel, RealComplexKernel, PlanarMethod, target_profile, + _derived_shape, _kernel_from_shape const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) const SIZES = parse_ints(argopt("sizes", "512,2048")) @@ -31,6 +32,7 @@ function run(io) for T in RUN_DTYPES, n in SIZES R = real(T) MRr, NRr, W = _derived_shape(target_profile(), R) + planar = _kernel_from_shape(_derived_shape(target_profile(), T), T, PlanarMethod()) cases = ( ( "CR", ComplexRealKernel(Val(MRr ÷ 2), Val(NRr), T, Val(W)), @@ -44,7 +46,7 @@ function run(io) for (case, kernel, A, B, (Mr, Nr)) in cases C = zeros(T, n, n) tm = time_plan(C, A, B; kernel) - tp = time_plan(C, A, B) + tp = time_plan(C, A, B; kernel = planar) tr = time_plan(zeros(R, Mr, Nr), randn(rng, R, Mr, n), randn(rng, R, n, Nr)) tag = "$(mr(kernel))x$(nr(kernel))/W$(lanewidth(kernel))" @printf( From 87abcd8611beff0b5e4d6c0fcd1462401ae91808 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:43:45 -0400 Subject: [PATCH 08/27] Allocate the compute-type panel of C through the workspace allocator The panel was a plain Vector{T} whatever the allocator. It now comes from tensoralloc as a temporary like the packed panels (a PtrArray under ManualAllocator, so no GC allocation) and release! frees it first. ContractWorkspace and ContractPlan gain a type parameter for its storage. Same-eltype code_llvm of execute!, _execute_path!, _execute_nest! and the micro-tile specialisations (F64, F32, C64, C32, permuted C, dot, outer) is identical to 565ad27 after normalising symbol numbering; code_typed differs only in the printed workspace type. Co-Authored-By: Claude Opus 5.5 --- src/execution/execute.jl | 2 +- src/execution/workspace.jl | 52 +++++++++++++++------------- src/integrations/tensoroperations.jl | 4 +-- src/planning/plan.jl | 4 +-- test/execution/test_execute.jl | 2 +- 5 files changed, 34 insertions(+), 30 deletions(-) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 1451d02..bf08a69 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -244,7 +244,7 @@ _panel_exit!(target::ContractPlan, plan::ContractPlan, ws, jc::Int, nblock::Int) # Columns `jc .+ (0:nblock-1)` of `target`'s C into (`load`) or out of the # panel, converting to the destination's eltype. Borrows the M/N offset # buffers, which the nest refills before reading them again. -function _panel_copy!(target::ContractPlan, panel::Vector, ws, jc::Int, nblock::Int, load::Bool) +function _panel_copy!(target::ContractPlan, panel::AbstractVector, ws, jc::Int, nblock::Int, load::Bool) C = target.Cstorage Qm = axis_length(target.mgroup) mc = target.blocking.mc diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 5805ab0..9080ab6 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -8,17 +8,18 @@ end _SlotCache() = _SlotCache(nothing, IdDict{Any, Any}()) """ - ContractWorkspace{T,VT<:AbstractVector} + ContractWorkspace{T,VT<:AbstractVector,PT<:AbstractVector} The buffers [`execute!`](@ref) needs. Pass one back as `plan_contract(...; workspace = ws)` to reuse (and grow) it across contractions; one from a non-default allocator is single-use and must be [`release!`](@ref)d. `T` is the plan's compute type; the packed panels (`VT`) -hold `real(T)`. Field layout is an implementation detail. +hold `real(T)`, the panel of C (`PT`) holds `T`. Field layout is an +implementation detail. """ # A `mutable struct` with `const` fields so a plan and a barrier slot hold one # pointer to it instead of copying twenty-two GC-tracked fields. -mutable struct ContractWorkspace{T, VT <: AbstractVector} +mutable struct ContractWorkspace{T, VT <: AbstractVector, PT <: AbstractVector} # Macro-block offset buffers and their per-sliver descriptors. const m_buf_A::Vector{Int} const m_buf_C::Vector{Int} @@ -31,7 +32,7 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} const n_desc_B::Vector{BlockDescriptor} const n_desc_C::Vector{BlockDescriptor} - # Packed macro panels: the only allocator-routed buffers. + # Packed macro panels and the panel of C: the only allocator-routed buffers. const packed_a::VT const packed_b::VT @@ -52,10 +53,10 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} # The compute-type panel of C for an eltype of C narrower than `T` (see # `_PanelPath`); empty otherwise. - const c_panel::Vector{T} + const c_panel::PT # GUARDRAIL: the packed panels hold `real(T)`, never `T`. - function ContractWorkspace{T, VT}( + function ContractWorkspace{T, VT, PT}( m_buf_A::Vector{Int}, m_buf_C::Vector{Int}, n_buf_B::Vector{Int}, n_buf_C::Vector{Int}, k_buf_A::Vector{Int}, k_buf_B::Vector{Int}, @@ -65,15 +66,15 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector} tile_m_buf_A::Vector{Int}, tile_m_buf_C::Vector{Int}, tile_n_buf_B::Vector{Int}, tile_n_buf_C::Vector{Int}, tw_k_buf_A::Vector{Int}, tw_k_buf_B::Vector{Int}, - tw_packed_a::VT, tw_packed_b::VT, c_panel::Vector{T} - ) where {T, VT <: AbstractVector} + tw_packed_a::VT, tw_packed_b::VT, c_panel::PT + ) where {T, VT <: AbstractVector, PT <: AbstractVector} eltype(VT) === real(T) || throw( ArgumentError( "ContractWorkspace{$T,$VT}: the packed panels must hold $(real(T)) " * "(the real type of the compute type $T), got $(eltype(VT))" ) ) - return new{T, VT}( + return new{T, VT, PT}( m_buf_A, m_buf_C, n_buf_B, n_buf_C, k_buf_A, k_buf_B, m_desc_A, m_desc_C, n_desc_B, n_desc_C, packed_a, packed_b, @@ -118,9 +119,9 @@ end @inline function _build_workspace( ::Type{T}, s, ntw::Int, ints::F, - packed_a::VT, packed_b::VT, tw_packed_a::VT, tw_packed_b::VT, panel::Int - ) where {T, F, VT <: AbstractVector} - return ContractWorkspace{T, VT}( + packed_a::VT, packed_b::VT, tw_packed_a::VT, tw_packed_b::VT, c_panel::PT + ) where {T, F, VT <: AbstractVector, PT <: AbstractVector} + return ContractWorkspace{T, VT, PT}( ints(s.mc), ints(s.mc), ints(s.nc), ints(s.nc), ints(s.kc), ints(s.kc), @@ -130,7 +131,7 @@ end ints(s.mr), ints(s.mr), ints(s.nr), ints(s.nr), ints(ntw), ints(ntw), - tw_packed_a, tw_packed_b, Vector{T}(undef, panel), + tw_packed_a, tw_packed_b, c_panel, ) end @@ -142,8 +143,9 @@ end A workspace for compute type `T`, `kernel` and an effective `blocking`. `oracle = false` leaves the `execute_tilewise!` buffers empty. `panel` is the length of the compute-type panel of C (see [`plan_contract`](@ref)). Under a -non-default `allocator` the packed panels come from `tensoralloc`, are never -resized, and must be handed back with [`release!`](@ref). +non-default `allocator` the packed panels and the panel of C come from +`tensoralloc`, are never resized, and must be handed back with +[`release!`](@ref). """ function ContractWorkspace( ::Type{T}, kernel, blocking::Blocking, oracle::Bool, ::TO.DefaultAllocator, @@ -155,7 +157,7 @@ function ContractWorkspace( T, s, oracle ? s.kc : 0, n -> Vector{Int}(undef, n), Vector{R}(undef, s.packed_a), Vector{R}(undef, s.packed_b), Vector{R}(undef, oracle ? s.tw_packed_a : 0), - Vector{R}(undef, oracle ? s.tw_packed_b : 0), panel, + Vector{R}(undef, oracle ? s.tw_packed_b : 0), Vector{T}(undef, panel), ) end @@ -170,10 +172,11 @@ function ContractWorkspace( packed_b = _alloc_temp(R, s.packed_b, allocator) tw_packed_a = _alloc_temp(R, oracle ? s.tw_packed_a : 0, allocator) tw_packed_b = _alloc_temp(R, oracle ? s.tw_packed_b : 0, allocator) + c_panel = _alloc_temp(T, panel, allocator) return _build_workspace( T, s, oracle ? s.kc : 0, n -> _alloc_offsets(n, allocator), - packed_a, packed_b, tw_packed_a, tw_packed_b, panel + packed_a, packed_b, tw_packed_a, tw_packed_b, c_panel ) end @@ -188,8 +191,8 @@ end # for `kernel` at `blocking`. Only for the GC-owned path: an allocator's # temporaries must never be resized. function reserve!( - ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, oracle::Bool, - panel::Int = 0 + ws::ContractWorkspace{T, Vector{R}, Vector{T}}, kernel, blocking::Blocking, + oracle::Bool, panel::Int = 0 ) where {T, R} s = _workspace_sizes(kernel, blocking) @@ -232,11 +235,12 @@ end """ release!(ws::ContractWorkspace, allocator) -Hand `ws`'s packed panels back to `allocator` (`TensorOperations.tensorfree!`) -in reverse acquisition order; `ws` must not be used afterwards. Correct for -every allocator. +Hand `ws`'s packed panels and panel of C back to `allocator` +(`TensorOperations.tensorfree!`) in reverse acquisition order; `ws` must not be +used afterwards. Correct for every allocator. """ function release!(ws::ContractWorkspace, allocator) + TO.tensorfree!(ws.c_panel, allocator) TO.tensorfree!(ws.tw_packed_b, allocator) TO.tensorfree!(ws.tw_packed_a, allocator) TO.tensorfree!(ws.packed_b, allocator) @@ -272,7 +276,7 @@ end # The pool is keyed by the compute type alone, so also check the packed type # (folds at compile time). @inline function _reuse_workspace( - ::Type{T}, ws::ContractWorkspace{T, Vector{R}}, kernel, blocking::Blocking, + ::Type{T}, ws::ContractWorkspace{T, Vector{R}, Vector{T}}, kernel, blocking::Blocking, oracle::Bool, panel::Int ) where {T, R} R === realtype(kernel) || _throw_packed_eltype_mismatch(T, ws, kernel) @@ -296,7 +300,7 @@ end ArgumentError( "plan_contract: cannot reuse a $(typeof(ws)) for a contraction of compute " * "type $T on the default allocator; only a " * - "ContractWorkspace{$T,Vector{$(realtype(kernel))}} -- compute " * + "ContractWorkspace{$T,Vector{$(realtype(kernel))},Vector{$T}} -- compute " * "type $T, packed panels of $(realtype(kernel)) -- is `reserve!`-able" ) ) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index 4960b58..631611a 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -47,8 +47,8 @@ end ws = get(pool, T, nothing) # Both assertions are needed: without them the branches join to the abstract # `ContractWorkspace`, and the call into `_planned` boxes its arguments. - ws === nothing || return ws::ContractWorkspace{T, Vector{real(T)}} - return _qs_build_task_workspace!(pool, T)::ContractWorkspace{T, Vector{real(T)}} + ws === nothing || return ws::ContractWorkspace{T, Vector{real(T)}, Vector{T}} + return _qs_build_task_workspace!(pool, T)::ContractWorkspace{T, Vector{real(T)}, Vector{T}} end @noinline function _qs_build_task_workspace!(pool::Dict{DataType, ContractWorkspace}, ::Type{T}) where {T} diff --git a/src/planning/plan.jl b/src/planning/plan.jl index cd10d06..61d3e12 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -13,7 +13,7 @@ after an M/N orientation swap the `A*` fields describe the original `B`. """ struct ContractPlan{ T, Kern, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, - TA, TB, VT <: AbstractVector, + TA, TB, VT <: AbstractVector, PT <: AbstractVector, } kernel::Kern mgroup::GM @@ -32,7 +32,7 @@ struct ContractPlan{ atransform::TA btransform::TB - workspace::ContractWorkspace{T, VT} + workspace::ContractWorkspace{T, VT, PT} end """ diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index aa8a0ff..fbb0be3 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -234,7 +234,7 @@ end plan = _mm_plan(Cmat, Amat, Bmat; mc = 8, kc = 6, nc = 7) @test isconcretetype(typeof(plan)) @test all(isconcretetype, fieldtypes(typeof(plan.workspace))) - @test typeof(plan.workspace) === QuasiStrided.ContractWorkspace{Float64, Vector{Float64}} + @test typeof(plan.workspace) === QuasiStrided.ContractWorkspace{Float64, Vector{Float64}, Vector{Float64}} plan_argtypes = (typeof(Cv), typeof(Av), NTuple{2, Int}, typeof(Bv), NTuple{2, Int}, NTuple{2, Int}) @test isempty(_ws_nonconcrete_types(plan_contract, plan_argtypes)) for f in (execute!, execute_tilewise!) From e5854c351763cd3c692f07a76d1d979091ddd518 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:44:00 -0400 Subject: [PATCH 09/27] Reject a non-Float32/Float64 accumulator in _compute_type QuasiStridedBackend{Float16}() skipped the keyword constructor's check and failed later with a MethodError. _compute_type now takes only Float32 and Float64 accumulators and throws an ArgumentError otherwise. The backend calls it first instead of repeating its eltype and complex-into-real checks, so a rejected call still never touches the workspace pool. Co-Authored-By: Claude Opus 5.5 --- src/integrations/tensoroperations.jl | 27 +++++++--------------- src/planning/plan.jl | 7 +++++- test/integrations/test_tensoroperations.jl | 4 ++++ 3 files changed, 18 insertions(+), 20 deletions(-) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index 631611a..f668cb2 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -75,30 +75,20 @@ end @noinline _qs_throw(msg::AbstractString) = throw(ArgumentError(msg)) -@noinline function _qs_check_eligible(C, A, B) - f = TO.tensorcontract! - TA, TB, TC = eltype(A), eltype(B), eltype(C) - all(in(_QS_ELTYPES), (TA, TB, TC)) || _qs_throw( - "QuasiStridedBackend requires every tensor of $f to have an element type " * - "out of Float32, Float64, ComplexF32 and ComplexF64, got " * - join((TC, TA, TB), ", ") - ) - (TC <: Complex || (TA <: Real && TB <: Real)) || _qs_throw( - "QuasiStridedBackend: a complex input of $f needs a complex output, got " * - join((TC, TA, TB), ", ") - ) +@noinline function _qs_check_strided(C, A, B) all(isstrided, (A, B, C)) || _qs_throw( - "QuasiStridedBackend requires strided arrays for $f, got " * + "QuasiStridedBackend requires strided arrays for $(TO.tensorcontract!), got " * join(map(typeof, (C, A, B)), ", ") ) return nothing end -# Checks shared by both `tensorcontract!` methods. The conjugated-C check -# duplicates `plan_contract`'s so that a rejected call never acquires or grows -# a pooled workspace (the workspace is an argument to `plan_contract`). +# Checks shared by both `tensorcontract!` methods. The eltype and conjugated-C +# checks precede `plan_contract`'s so that a rejected call never acquires or +# grows a pooled workspace (the workspace is an argument to `plan_contract`). @inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β, accumulator) - _qs_check_eligible(C, A, B) + T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) + _qs_check_strided(C, A, B) TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) @@ -114,9 +104,8 @@ end "`StridedView.op` on store, so a conjugated `C` would be silently wrong" ) - # Dropping `Zero()`/`One()` is safe: the kernels branch on `iszero(alpha/beta)`. - T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) indA, indB, indC = _qs_labels(pA, pB, pAB) + # Dropping `Zero()`/`One()` is safe: the kernels branch on `iszero(alpha/beta)`. return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 61d3e12..f0014ee 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -230,9 +230,14 @@ const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) (TC <: Real && !(TA <: Real && TB <: Real)) && _throw_complex_into_real(TA, TB, TC) return promote_type(TA, TB, TC) end -@inline function _compute_type(::Type{TA}, ::Type{TB}, ::Type{TC}, ::Type{R}) where {TA, TB, TC, R} +@inline function _compute_type( + ::Type{TA}, ::Type{TB}, ::Type{TC}, ::Type{R} + ) where {TA, TB, TC, R <: Union{Float32, Float64}} return _compute_type(TA, TB, TC, nothing) <: Complex ? Complex{R} : R end +@noinline _compute_type(TA, TB, TC, accumulator) = throw( + ArgumentError("accumulator must be nothing, Float32 or Float64, got $accumulator") +) @noinline _throw_eltypes(TA, TB, TC) = throw( ArgumentError( diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 39f9703..647b5de 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -142,6 +142,10 @@ end @test_throws ArgumentError tensorcontract!(zeros(4, 5), D4, pA, false, randn(4, 5), pB, false, pAB, 1, 0, qsbackend) @test_throws ArgumentError tensorcontract!(zeros(4, 5), randn(4, 5), pA, false, D5, pB, false, pAB, 1, 0, qsbackend) @test_throws ArgumentError QuasiStrided.QuasiStridedBackend(accumulator = Float16) + @test_throws ArgumentError tensorcontract!( + zeros(4, 5), randn(4, 4), pA, false, randn(4, 5), pB, false, pAB, 1, 0, + QuasiStrided.QuasiStridedBackend{Float16}() + ) end @testset "@tensor with mixed eltypes and an accumulator" begin From 02b6cd41ee13e256f5fe69b185789473b704bce2 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:44:00 -0400 Subject: [PATCH 10/27] Recheck the panel decision and capacity when a hinted path is taken The TO route predicts the path at the default kc. _hint_holds now also requires that the plan's own blocking needs the panel for a hinted _PanelPath (with a workspace panel large enough) and does not need it for a hinted _NestPath; otherwise the plan's own path runs behind the barrier. _c_panel_length is shared with plan sizing. For a same-eltype plan the check folds away: code_llvm of _hint_holds and _execute_hinted! (F64, F32, C64, C32, permuted C, dot, outer) is identical to 565ad27 after normalising symbol numbering. Co-Authored-By: Claude Opus 5.5 --- src/execution/execute.jl | 27 +++++++++++++++++++++------ src/planning/plan.jl | 2 +- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index bf08a69..e34995f 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -94,6 +94,10 @@ end return Qk > (kc === nothing ? _resolved_defaults(T).real_row.kc : kc) end +# The panel length at the plan's effective `blocking`, 0 if none is needed. +@inline _c_panel_length(::Type{T}, Cstorage, Qm::Int, Qn::Int, Qk::Int, blocking::Blocking) where {T} = + _c_panel_needed(T, Cstorage, Qk, blocking.kc) ? Qm * min(blocking.nc, Qn) : 0 + # Run `path` behind a dynamic call. The plan crosses in the workspace slot # with its storages stripped; they cross as arguments. @inline function _execute_across_barrier!(plan::ContractPlan{T}, alphaT::T, betaT::T, path) where {T} @@ -160,14 +164,25 @@ end return nothing end -# The prediction differs from `_select_path(plan)` in two inputs only: the -# kernel's unpacked-B eligibility (predicted from its method; folds) and the -# dot path's workspace capacity (assumed). -@inline _hint_holds(plan::ContractPlan, ::_NestPath) = - _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) +# The prediction differs from `_select_path(plan)` in three inputs only: the +# kernel's unpacked-B eligibility (predicted from its method; folds), the +# dot path's workspace capacity (assumed) and the panel decision (made at the +# default `kc`; folds to `false` unless C is narrower than `T`). +@inline _hint_holds(plan::ContractPlan, ::_NestPath) = _panel_holds(plan, false) && _unpacked_b_holds(plan) +@inline _hint_holds(plan::ContractPlan, ::_PanelPath) = _panel_holds(plan, true) && _unpacked_b_holds(plan) @inline _hint_holds(plan::ContractPlan, ::_DotPath) = _dot_capacity_ok(plan) @inline _hint_holds(plan::ContractPlan, ::_OuterPath) = true -@inline _hint_holds(plan::ContractPlan, ::_PanelPath{P}) where {P} = _hint_holds(plan, P()) + +@inline _unpacked_b_holds(plan::ContractPlan) = + _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) + +@inline function _panel_holds(plan::ContractPlan{T}, panel::Bool) where {T} + n = _c_panel_length( + T, plan.Cstorage, axis_length(plan.mgroup), axis_length(plan.ngroup), + axis_length(plan.kgroup), plan.blocking + ) + return panel ? n > 0 && length(plan.workspace.c_panel) >= n : n == 0 +end _execute_path!(plan::ContractPlan, alphaT, betaT, ::_DotPath{MATB, W}) where {MATB, W} = (_execute_dot!(plan, alphaT, betaT, MATB, Val(W)); nothing) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index f0014ee..e75bc2d 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -337,7 +337,7 @@ function _plan_contract( nc_eff = Qn == 0 ? NRk : min(_roundup(requested.nc, NRk), _roundup(Qn, NRk)) kc_eff = Qk == 0 ? 1 : min(requested.kc, Qk) blocking = Blocking(mc_eff, kc_eff, nc_eff) - panel = _c_panel_needed(T, req.Cstorage, Qk, kc_eff) ? Qm * min(nc_eff, Qn) : 0 + panel = _c_panel_length(T, req.Cstorage, Qm, Qn, Qk, blocking) ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel) plan = ContractPlan( From 00b31f9af8bf25706e5b78882648048270c84676 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:46:28 -0400 Subject: [PATCH 11/27] Rewrap two long lines; drop a redundant _dense_lanes clause _lane_convertible(T, T) holds for every supported T, so the DenseVector{T} test added nothing. Same-eltype code_llvm of the nest (packing inlined) and the micro-tile specialisations for F64 and C64 is unchanged. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +++--- src/microkernels/fmaddsub.jl | 4 ++-- src/packing/pack_contiguous.jl | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 41ca46c..0f7611f 100644 --- a/README.md +++ b/README.md @@ -121,9 +121,9 @@ Implemented: needs a complex output). The compute type is the promoted eltype, or the precision chosen with `accumulator = Float32`/`Float64` (`plan_contract(...; accumulator)`, `QuasiStridedBackend(; accumulator)`); - the result is rounded to `eltype(C)` once. A real operand of a complex contraction is - promoted at pack time. `conjA`/`conjB` and each operand's `StridedView.op` - are applied during packing. + the result is rounded to `eltype(C)` once. A real operand of a complex + contraction is promoted at pack time. `conjA`/`conjB` and each operand's + `StridedView.op` are applied during packing. - Zero steady-state allocation on Julia >= 1.11 for `execute!` on a reused plan and for `tensorcontract!` through the backend (which pools workspaces per task). On Julia 1.10 `SIMDKernel`'s accumulator is not kept in diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 78ee65d..58f0d08 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -174,8 +174,8 @@ end end # One full `W÷2`-row block, already in `Complex`'s memory order, so no -# interleave shuffle. `C` converts from and to its storage type `RC`. Base's `Complex` expression trees in lanes, as planar's -# `_planar_store_block!`: +# interleave shuffle. `C` converts from and to its storage type `RC`. Base's +# `Complex` expression trees in lanes, as planar's `_planar_store_block!`: # beta == 0: addsub(ar*r, ai*swap(r)) # beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) # otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 6b60eb3..1caea2c 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -15,7 +15,7 @@ using SIMD: shufflevector # the vector stores. const _LaneFloat = Union{Float32, Float64} @inline _dense_lanes(storage::S, ::Type{T}) where {S, T} = - storage isa DenseVector{T} || (storage isa DenseVector && _lane_convertible(eltype(S), T)) + storage isa DenseVector && _lane_convertible(eltype(S), T) _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 From 51a0306681889eed3d50a7fa268f0eb513e598c8 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:46:28 -0400 Subject: [PATCH 12/27] Name RealComplexKernel's accumulator layout in _planar_acc_index Any other kernel type is now an ArgumentError at generator time instead of silently getting the column-pair layout. The PlanarKernel branch is unchanged, so the planar store is generated as before. Co-Authored-By: Claude Opus 5.5 --- src/microkernels/planar.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index a30b417..a3c6926 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -135,10 +135,12 @@ function Base.accumulate( end # Generator-time accumulator indices of the re/im vectors of block `v`, column -# `j`: planar's two planes, or `RealComplexKernel`'s column pairs. +# `j`: planar's two planes, or `RealComplexKernel`'s column pairs. A `<:` test +# rather than dispatch, as `RealComplexKernel` is defined after this file. function _planar_acc_index(kernel::Type, MV::Int, NR::Int, v::Int, j::Int) kernel <: PlanarKernel && return (v + MV * j + 1, MV * NR + v + MV * j + 1) - return (v + MV * 2j + 1, v + MV * (2j + 1) + 1) + kernel <: RealComplexKernel && return (v + MV * 2j + 1, v + MV * (2j + 1) + 1) + throw(ArgumentError("_planar_acc_index: no planar accumulator layout for $kernel")) end # Scalar fallback store, for every destination the vector store cannot take. From 179c545c9dcdbcd7621d2d2ba8870a2cbce5f673 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:50:50 -0400 Subject: [PATCH 13/27] test_mixed_kernel.jl: reuse one CR/RC shape pair per precision The pack and end-to-end testsets now use the mk_contract shapes, so the file builds 6 distinct mixed-domain kernels instead of 10. File time (after the scalar/simd/planar kernel files, workstation) 24.7 s -> 24.5 s: the six mk_contract shapes take about 15.6 s of it and stay. Co-Authored-By: Claude Opus 5.5 --- test/microkernels/test_mixed_kernel.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index cfdd951..06d604b 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -51,8 +51,8 @@ end @testset "pack -> accumulate -> store: $(nameof(K)) $T from $SA x $SB" for (K, T, SA, SB, MR, NR, W) in ( (ComplexRealKernel, ComplexF32, ComplexF64, Float64, 8, 5, 8), (RealComplexKernel, ComplexF32, Float64, ComplexF64, 16, 3, 8), - (ComplexRealKernel, ComplexF64, ComplexF64, Float32, 4, 3, 4), - (RealComplexKernel, ComplexF64, Float32, ComplexF32, 8, 2, 4), + (ComplexRealKernel, ComplexF64, ComplexF64, Float32, 12, 8, 8), + (RealComplexKernel, ComplexF64, Float32, ComplexF32, 24, 4, 8), ) k = K(Val(MR), Val(NR), T, Val(W)) R = real(T) @@ -86,9 +86,9 @@ end @testset "end to end: $(nameof(typeof(k))) $TA x $TB -> $TC, B $mode" for (k, TA, TB, TC, acc, mode) in ( (ComplexRealKernel(Val(12), Val(8), ComplexF64, Val(8)), ComplexF64, Float64, ComplexF64, nothing, :always), (ComplexRealKernel(Val(12), Val(8), ComplexF64, Val(8)), ComplexF64, Float64, ComplexF64, nothing, :never), - (ComplexRealKernel(Val(8), Val(5), ComplexF32, Val(16)), ComplexF32, Float64, ComplexF32, Float32, :auto), + (ComplexRealKernel(Val(8), Val(5), ComplexF32, Val(8)), ComplexF32, Float64, ComplexF32, Float32, :auto), (RealComplexKernel(Val(24), Val(4), ComplexF64, Val(8)), Float64, ComplexF64, ComplexF64, nothing, :auto), - (RealComplexKernel(Val(16), Val(3), ComplexF32, Val(16)), Float64, ComplexF32, ComplexF32, Float32, :auto), + (RealComplexKernel(Val(16), Val(3), ComplexF32, Val(8)), Float64, ComplexF32, ComplexF32, Float32, :auto), ) rng = MersenneTwister(4242) A, B, C0 = rand(rng, TA, 37, 41), rand(rng, TB, 41, 23), rand(rng, TC, 37, 23) From 1558b711c02b838f21dcd98630d88e4ac29a1045 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 11:56:43 -0400 Subject: [PATCH 14/27] test_mixed.jl: a deterministic round-once test for the panel of C Replaces the statistical comparison with the all-Float32 result. Along K the contributions 1, 2^-30 and -1 fall in separate kc blocks, so the exact 2^-30 (plus a small beta*C) survives only if nothing is rounded to eltype(C) between blocks. Six rows: several jc blocks under ManualAllocator (checking the panel is allocator-owned), packed B, a complex panel on the RC kernel with strided C, the TensorOperations hint route at the default kc, M == 1, and beta = 0 on a NaN-filled C. With _c_panel_needed forced to false every row fails. The RC row of the main table now uses scattered operands. File time (after test_execute.jl, workstation): 35.3 s -> 22.6 s. Co-Authored-By: Claude Opus 5.5 --- test/execution/test_mixed.jl | 56 ++++++++++++++++++++++-------------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index cc47506..7fb765a 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -25,7 +25,7 @@ const _PROMOTED = Union{QuasiStrided.PlanarKernel, QuasiStrided.FMAddSubKernel} @testset "mixed eltypes: $TA x $TB -> $TC, accumulator = $acc, $layout" for (TA, TB, TC, acc, layout, MKN, beta, conjA, conjB, path, K) in ( (Float32, Float64, Float64, nothing, :dense, (23, 37, 19), 0, false, false, QuasiStrided._NestPath, SIMDKernel), (Float64, Float32, Float32, nothing, :dense, (32, 40, 12), 1, false, false, QuasiStrided._NestPath, SIMDKernel), - (Float64, ComplexF64, ComplexF64, nothing, :dense, (21, 33, 17), 0.3 - 0.7im, false, true, QuasiStrided._NestPath, _RC), + (Float64, ComplexF64, ComplexF64, nothing, :scattered, (32, 32, 32), 0.3 - 0.7im, false, true, QuasiStrided._NestPath, _RC), (ComplexF32, Float64, ComplexF64, nothing, :scattered, (32, 32, 32), 0.5, true, false, QuasiStrided._NestPath, _CR), (ComplexF32, ComplexF64, ComplexF64, nothing, :dense, (19, 29, 23), 1, true, true, QuasiStrided._NestPath, _PROMOTED), (Float64, Float64, ComplexF64, nothing, :dense, (18, 25, 14), 0.25 + 0.5im, false, false, QuasiStrided._NestPath, _PROMOTED), @@ -60,28 +60,42 @@ end @test _steady_allocs!(execute!, _mm_plan(Cmat, Amat, Bmat), Cmat) == 0 skip = (VERSION < v"1.11") end -@testset "Float32 C, Float64 compute type, K over several kc blocks: within one ulp ($TA, $layout, beta = $beta)" for (TA, acc, layout) in ( - (Float32, Float64, :dense), (Float64, nothing, :dense), (Float32, Float64, :scattered), - ), beta in (0, 0.5) - fx = if layout === :scattered - Cv = scattered_fixture(Float32, 32, 512)[1] - copyto!(Cv, randn(MersenneTwister(4), Float32, size(Cv))) - (Cv, scattered_fixture(TA, 32, 512)[2:end]...) +# Along K, 1, 2^-30 and -1 fall in different kc blocks: the 2^-30 and the small +# beta*C survive only if the partial sums stay in the compute type between blocks. +@testset "one rounding to $TC across kc blocks: $TA x $TB, M = $M, $case, beta = $beta" for (TA, TB, TC, acc, M, N, K, case, beta) in ( + (Float64, Float64, Float32, nothing, 13, 40, 20, :manual_nc16, 0.5), + (Float32, Float32, Float32, Float64, 13, 9, 20, :packed_b, 1), + (Float64, ComplexF64, ComplexF32, nothing, 13, 9, 20, :strided_c, 0.5), + (Float32, Float32, Float32, Float64, 7, 5, 3000, :backend, 0), + (Float64, Float64, Float32, nothing, 1, 9, 20, :strided_c, 1), + (Float64, Float64, Float32, nothing, 13, 9, 20, :strided_c, 0), + ) + b = TB <: Complex ? TB(1, 1) : one(TB) + A = zeros(TA, M, K) + A[:, 1] .= 1 + A[:, K ÷ 2 + 1] .= 2.0^-30 + A[:, K] .= -1 + Av, Bv = StridedView(A), StridedView(fill(b, K, N)) + case === :packed_b && (Bv = permutedims(StridedView(permutedims(Array(Bv))), (2, 1))) + Cv = case === :strided_c ? StridedView(zeros(TC, 2M, 3N))[2:2:2M, 1:3:3N] : StridedView(zeros(TC, M, N)) + fill!(Cv, iszero(beta) ? NaN : 2.0^-28) + if case === :backend + @test K ÷ 2 > QuasiStrided._resolved_defaults(Float64).real_row.kc + backend = QuasiStrided.QuasiStridedBackend(accumulator = acc) + TO.tensorcontract!(Cv, Av, ((1,), (2,)), false, Bv, ((1,), (2,)), false, ((1, 2), ()), 1, beta, backend) else - rng = MersenneTwister(5) - A, B, C = randn(rng, TA, 64, 512), randn(rng, TA, 512, 48), randn(rng, Float32, 64, 48) - (StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3)) + allocator = case === :manual_nc16 ? TO.ManualAllocator() : TO.DefaultAllocator() + nc = case === :manual_nc16 ? 16 : nothing + plan = plan_contract(Cv, Av, (1, 2), Bv, (2, 3), (1, 3); accumulator = acc, kc = 4, nc, allocator) + @test _path_of(plan) isa QuasiStrided._PanelPath + @test axis_length(plan.mgroup) == M + case === :manual_nc16 && @test N > plan.blocking.nc && !(plan.workspace.c_panel isa Vector) + case === :packed_b && @test _path_of(plan) isa QuasiStrided._PanelPath{<:QuasiStrided._NestPath{false}} + TB <: Complex && @test plan.kernel isa _RC + execute!(plan, 1, beta) + QuasiStrided.release!(plan.workspace, allocator) end - ref = Float32.(_mixed_ref(fx, 1, beta)) - C32 = copy(fx[1]) - execute!(plan_contract(C32, map(x -> x isa StridedView ? StridedView(Float32.(x)) : x, fx[2:end])...), 1, beta) - plan = plan_contract(fx...; accumulator = acc) - @test _path_of(plan) isa QuasiStrided._PanelPath - @test axis_length(plan.kgroup) > plan.blocking.kc - execute!(plan, 1, beta) - err = abs.(Array(fx[1]) .- ref) - @test all(err .<= eps.(ref)) - @test maximum(err) < maximum(abs.(Array(C32) .- ref)) + @test all(==(TC(2.0^-30 * b + beta * 2.0^-28)), Array(Cv)) end @testset "Float32 C, Float64 compute type: steady-state execute! allocates nothing" begin From 8492cfd06211733b7fbf1cc9b988da4b03134f63 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:05:45 -0400 Subject: [PATCH 15/27] Drop the mixed-domain kernels' redundant constructor checks The descriptor already rejects a real T with an ArgumentError, and a RealComplexKernel's inner SIMDKernel already checks that W divides MR. The three-argument constructors forward to the four-argument ones. Co-Authored-By: Claude Opus 5.5 --- src/microkernels/mixed.jl | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index 01e95a7..c45c1be 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -46,7 +46,6 @@ struct RealComplexKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, descriptor::ComplexKernelDescriptor{MR, NR, T, RealFormat, InterleavedFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} - _check_vector_shape("RealComplexKernel", MR, W) KI === SIMDKernel{MR, 2 * NR, real(T), W} || throw( ArgumentError("RealComplexKernel's inner kernel must be SIMDKernel{$MR,$(2 * NR),$(real(T)),$W}, got $KI") ) @@ -57,31 +56,21 @@ end const _MixedKernel = Union{ComplexRealKernel, RealComplexKernel} function ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} - T <: Complex || - throw(ArgumentError("ComplexRealKernel requires a complex element type, got $T")) descriptor = ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), RealFormat()) inner = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) return ComplexRealKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) end function RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} - T <: Complex || - throw(ArgumentError("RealComplexKernel requires a complex element type, got $T")) descriptor = ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()) inner = SIMDKernel(Val(MR), Val(2 * NR), real(T), Val(W)) return RealComplexKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) end -function ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} - T <: Complex || - throw(ArgumentError("ComplexRealKernel requires a complex element type, got $T")) - return ComplexRealKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) -end -function RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} - T <: Complex || - throw(ArgumentError("RealComplexKernel requires a complex element type, got $T")) - return RealComplexKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) -end +ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = + ComplexRealKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) +RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = + RealComplexKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) complex_method(::ComplexRealKernel) = ComplexRealMethod() complex_method(::RealComplexKernel) = RealComplexMethod() From 3a6b53ff4b024878ad2eaf13edd23d12cdab7a12 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:05:46 -0400 Subject: [PATCH 16/27] Fold the panel hint check into _hint_holds; the panel helpers take the plan A plan's workspace always holds the panel its own blocking needs (_plan_contract sizes or grows it from the same inputs), so the hint only has to agree with _c_panel_needed at the plan's kc. _c_panel_length is inlined at its one remaining use, and the panel copy reads the workspace and panel from the target plan. Same-eltype code_llvm of execute!, _execute_path!, _execute_nest!, _execute_hinted! and _hint_holds (F64, F32, C64, C32, permuted C, dot, outer) is identical to 1558b71 after normalising symbol numbering. Co-Authored-By: Claude Opus 5.5 --- src/execution/execute.jl | 44 ++++++++++++++-------------------------- src/planning/plan.jl | 2 +- 2 files changed, 16 insertions(+), 30 deletions(-) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index e34995f..aecb10a 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -94,10 +94,6 @@ end return Qk > (kc === nothing ? _resolved_defaults(T).real_row.kc : kc) end -# The panel length at the plan's effective `blocking`, 0 if none is needed. -@inline _c_panel_length(::Type{T}, Cstorage, Qm::Int, Qn::Int, Qk::Int, blocking::Blocking) where {T} = - _c_panel_needed(T, Cstorage, Qk, blocking.kc) ? Qm * min(blocking.nc, Qn) : 0 - # Run `path` behind a dynamic call. The plan crosses in the workspace slot # with its storages stripped; they cross as arguments. @inline function _execute_across_barrier!(plan::ContractPlan{T}, alphaT::T, betaT::T, path) where {T} @@ -168,22 +164,12 @@ end # kernel's unpacked-B eligibility (predicted from its method; folds), the # dot path's workspace capacity (assumed) and the panel decision (made at the # default `kc`; folds to `false` unless C is narrower than `T`). -@inline _hint_holds(plan::ContractPlan, ::_NestPath) = _panel_holds(plan, false) && _unpacked_b_holds(plan) -@inline _hint_holds(plan::ContractPlan, ::_PanelPath) = _panel_holds(plan, true) && _unpacked_b_holds(plan) +@inline _hint_holds(plan::ContractPlan{T}, path::Union{_NestPath, _PanelPath}) where {T} = + _c_panel_needed(T, plan.Cstorage, axis_length(plan.kgroup), plan.blocking.kc) === (path isa _PanelPath) && + _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) @inline _hint_holds(plan::ContractPlan, ::_DotPath) = _dot_capacity_ok(plan) @inline _hint_holds(plan::ContractPlan, ::_OuterPath) = true -@inline _unpacked_b_holds(plan::ContractPlan) = - _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) - -@inline function _panel_holds(plan::ContractPlan{T}, panel::Bool) where {T} - n = _c_panel_length( - T, plan.Cstorage, axis_length(plan.mgroup), axis_length(plan.ngroup), - axis_length(plan.kgroup), plan.blocking - ) - return panel ? n > 0 && length(plan.workspace.c_panel) >= n : n == 0 -end - _execute_path!(plan::ContractPlan, alphaT, betaT, ::_DotPath{MATB, W}) where {MATB, W} = (_execute_dot!(plan, alphaT, betaT, MATB, Val(W)); nothing) @@ -240,26 +226,26 @@ end end # `target`: `nothing`, or the plan whose C the panel stands in for. -@inline _panel_enter!(::Nothing, plan, ws, jc, nblock, betaT) = plan -@inline _panel_exit!(::Nothing, plan, ws, jc, nblock) = nothing +@inline _panel_enter!(::Nothing, plan, jc, nblock, betaT) = plan +@inline _panel_exit!(::Nothing, jc, nblock) = nothing -function _panel_enter!(target::ContractPlan, plan::ContractPlan, ws, jc::Int, nblock::Int, betaT) - iszero(betaT) || _panel_copy!(target, plan.Cstorage, ws, jc, nblock, true) - Qm = axis_length(plan.mgroup) +function _panel_enter!(target::ContractPlan, plan::ContractPlan, jc::Int, nblock::Int, betaT) + iszero(betaT) || _panel_copy!(target, jc, nblock, true) return ContractPlan( plan.kernel, plan.mgroup, plan.ngroup, plan.kgroup, plan.blocking, - plan.Astorage, plan.Abase, plan.Bstorage, plan.Bbase, plan.Cstorage, -jc * Qm, - plan.atransform, plan.btransform, plan.workspace + plan.Astorage, plan.Abase, plan.Bstorage, plan.Bbase, plan.Cstorage, + -jc * axis_length(plan.mgroup), plan.atransform, plan.btransform, plan.workspace ) end -_panel_exit!(target::ContractPlan, plan::ContractPlan, ws, jc::Int, nblock::Int) = - _panel_copy!(target, plan.Cstorage, ws, jc, nblock, false) +_panel_exit!(target::ContractPlan, jc::Int, nblock::Int) = _panel_copy!(target, jc, nblock, false) # Columns `jc .+ (0:nblock-1)` of `target`'s C into (`load`) or out of the # panel, converting to the destination's eltype. Borrows the M/N offset # buffers, which the nest refills before reading them again. -function _panel_copy!(target::ContractPlan, panel::AbstractVector, ws, jc::Int, nblock::Int, load::Bool) +function _panel_copy!(target::ContractPlan, jc::Int, nblock::Int, load::Bool) + ws = target.workspace + panel = ws.c_panel C = target.Cstorage Qm = axis_length(target.mgroup) mc = target.blocking.mc @@ -310,7 +296,7 @@ function _execute_nest!( jc = 0 while jc < Qn nblock = min(nc_eff, Qn - jc) - cplan = _panel_enter!(target, plan, ws, jc, nblock, betaT) + cplan = _panel_enter!(target, plan, jc, nblock, betaT) n_slivers = cld(nblock, NRk) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( @@ -420,7 +406,7 @@ function _execute_nest!( pc += kblock end - _panel_exit!(target, plan, ws, jc, nblock) + _panel_exit!(target, jc, nblock) jc += nblock end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index e75bc2d..f0014ee 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -337,7 +337,7 @@ function _plan_contract( nc_eff = Qn == 0 ? NRk : min(_roundup(requested.nc, NRk), _roundup(Qn, NRk)) kc_eff = Qk == 0 ? 1 : min(requested.kc, Qk) blocking = Blocking(mc_eff, kc_eff, nc_eff) - panel = _c_panel_length(T, req.Cstorage, Qm, Qn, Qk, blocking) + panel = _c_panel_needed(T, req.Cstorage, Qk, kc_eff) ? Qm * min(nc_eff, Qn) : 0 ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel) plan = ContractPlan( From dfd9b081372217f9c06f31123e4a109a8b526062 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:05:46 -0400 Subject: [PATCH 17/27] Check a named mixed-domain kernel's real side by dispatch on its type _check_kernel_domain dispatches on ComplexRealKernel/RealComplexKernel instead of forwarding through complex_method, and the run_m condition is inlined. Co-Authored-By: Claude Opus 5.5 --- src/planning/plan.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index f0014ee..3e44578 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -172,8 +172,7 @@ function _planned( norder = _order_free_labels(nlabels, indC, C) # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. - runs = T <: Real || method isa _MixedMethod - run_m = runs ? _leading_unit_run(morder, indC, C) : 0 + run_m = T <: Real || method isa _MixedMethod ? _leading_unit_run(morder, indC, C) : 0 run_n = T <: Real ? _leading_unit_run(norder, indC, C) : 0 mgroup = _build_pair_group(morder, indA, A, indC, C) # maps: (A, C) @@ -211,12 +210,10 @@ end const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) # A named mixed-domain kernel's `RealFormat` side packs a real operand only. -@inline _check_kernel_domain(kernel, ::Type{TA}, ::Type{TB}) where {TA, TB} = - _check_kernel_domain(complex_method(kernel), kernel, TA, TB) -@inline _check_kernel_domain(::ComplexMethod, kernel, ::Type, ::Type) = nothing -@inline _check_kernel_domain(::ComplexRealMethod, kernel, ::Type, ::Type{TB}) where {TB} = +@inline _check_kernel_domain(kernel, ::Type, ::Type) = nothing +@inline _check_kernel_domain(kernel::ComplexRealKernel, ::Type, ::Type{TB}) where {TB} = TB <: Real || _throw_kernel_domain(kernel, "B", TB) -@inline _check_kernel_domain(::RealComplexMethod, kernel, ::Type{TA}, ::Type) where {TA} = +@inline _check_kernel_domain(kernel::RealComplexKernel, ::Type{TA}, ::Type) where {TA} = TA <: Real || _throw_kernel_domain(kernel, "A", TA) @noinline _throw_kernel_domain(kernel, side, T) = throw( From 892793ce1e33337c9895740da98381fc8da39099 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:05:46 -0400 Subject: [PATCH 18/27] tensoroperations.jl: keep _qs_check_eligible's name and layout Only the eltype check leaves it; the rename and the moved comment were unrelated churn. Co-Authored-By: Claude Opus 5.5 --- src/integrations/tensoroperations.jl | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index f668cb2..b50bb90 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -75,9 +75,10 @@ end @noinline _qs_throw(msg::AbstractString) = throw(ArgumentError(msg)) -@noinline function _qs_check_strided(C, A, B) +@noinline function _qs_check_eligible(C, A, B) + f = TO.tensorcontract! all(isstrided, (A, B, C)) || _qs_throw( - "QuasiStridedBackend requires strided arrays for $(TO.tensorcontract!), got " * + "QuasiStridedBackend requires strided arrays for $f, got " * join(map(typeof, (C, A, B)), ", ") ) return nothing @@ -88,7 +89,7 @@ end # grows a pooled workspace (the workspace is an argument to `plan_contract`). @inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β, accumulator) T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) - _qs_check_strided(C, A, B) + _qs_check_eligible(C, A, B) TO.argcheck_tensorcontract(C, A, pA, B, pB, pAB) TO.dimcheck_tensorcontract(C, A, pA, B, pB, pAB) @@ -104,8 +105,8 @@ end "`StridedView.op` on store, so a conjugated `C` would be silently wrong" ) - indA, indB, indC = _qs_labels(pA, pB, pAB) # Dropping `Zero()`/`One()` is safe: the kernels branch on `iszero(alpha/beta)`. + indA, indB, indC = _qs_labels(pA, pB, pAB) return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end From dc7621c04b129968f4ab936133c5d64e215a6f91 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:05:46 -0400 Subject: [PATCH 19/27] README: list the mixed-domain kernels; drop the stale promotion wording A real operand of a complex contraction is packed as real by ComplexRealKernel/RealComplexKernel, not promoted, so "promoted at pack time" and "Not implemented: dedicated real-times-complex kernels" were wrong. Co-Authored-By: Claude Opus 5.5 --- README.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 0f7611f..82c85a4 100644 --- a/README.md +++ b/README.md @@ -108,8 +108,10 @@ Implemented: kernel-specific panel formats and contiguous fast paths. - Microkernels: `ScalarKernel` (reference), `SIMDKernel` (default for real eltypes), `PlanarKernel` (split-complex, default for complex eltypes), - `OneMKernel` (the 1m method) and `FMAddSubKernel` (interleaved complex, used - for small-M complex contractions on AVX-512). Any of them can be named via + `OneMKernel` (the 1m method), `FMAddSubKernel` (interleaved complex, used + for small-M complex contractions on AVX-512) and `ComplexRealKernel`/ + `RealComplexKernel` (a complex times a real operand on the real kernel, + default for those). Any of them can be named via `plan_contract(...; kernel = ...)`. - Register shapes derived from the detected ISA, demoted when `M` is too short for a full tile or when a smaller tile keeps stores into `C` unit-stride. @@ -121,9 +123,9 @@ Implemented: needs a complex output). The compute type is the promoted eltype, or the precision chosen with `accumulator = Float32`/`Float64` (`plan_contract(...; accumulator)`, `QuasiStridedBackend(; accumulator)`); - the result is rounded to `eltype(C)` once. A real operand of a complex - contraction is promoted at pack time. `conjA`/`conjB` and each operand's - `StridedView.op` are applied during packing. + operands are converted during packing and the result is rounded to + `eltype(C)` once. `conjA`/`conjB` and each operand's `StridedView.op` are + applied during packing. - Zero steady-state allocation on Julia >= 1.11 for `execute!` on a reused plan and for `tensorcontract!` through the backend (which pools workspaces per task). On Julia 1.10 `SIMDKernel`'s accumulator is not kept in @@ -136,7 +138,6 @@ Not implemented: - Other eltypes, a complex input with a real output, non-strided operands, an output aliased with an input, and a conjugated output view: `tensorcontract!` throws an `ArgumentError` instead of falling back to another backend. -- Dedicated real-times-complex kernels. - Being the default backend: `StridedBLAS()` is faster on most shapes. - Batch axes, traces/diagonals, threading, GPU execution, the 3m complex method, autotuning. From 966d37cdfc09cd9a5222ae7f2da913ca5b404f9f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:06:25 -0400 Subject: [PATCH 20/27] test_kernel_selection.jl: drop the same-eltype golden shapes from the mixed testset The same-eltype defaults per ISA are covered by the real and complex shape-selection testsets; the mixed testset keeps the mapping to the real default of real(T), now also at a large extent, on the three synthetic targets. Co-Authored-By: Claude Opus 5.5 --- test/planning/test_kernel_selection.jl | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index 9a8eb8a..196ff96 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -232,7 +232,6 @@ end @testset "mixed-domain selection: the real default of real(T), mapped" begin CR, RC = QuasiStrided.ComplexRealMethod(), QuasiStrided.RealComplexMethod() mapped(m, (MR, NR, W)) = m === CR ? (MR ÷ 2, NR, W) : (MR, NR ÷ 2, W) - dshape = QuasiStrided._default_shape dmethod = QuasiStrided._default_method for T in (ComplexF64, ComplexF32) R = real(T) @@ -249,28 +248,15 @@ end @test dmethod(T, R, R) === dmethod(T, T, T) === dmethod(T, ComplexF32, T) === PlanarMethod() @test dmethod(R, Float32, Float64) === RealMethod() end - # (isa, vb, nregisters) => (F64, F32, C64, C32) default shapes at a large extent. - expected = ( - (:avx512, 64, 32) => ((32, 6, 8), (64, 6, 16), (24, 3, 8), (48, 3, 16)), - (:avx2, 32, 16) => ((8, 6, 4), (16, 6, 8), (4, 5, 4), (8, 5, 8)), - (:neon, 16, 32) => ((8, 6, 4), (8, 6, 8), (4, 6, 2), (8, 6, 4)), - ) + # The real extent, store and small-M demotions carry over. saved = QuasiStrided._TARGET[] try - for ((key, vb, nreg), shapes) in expected + for (key, vb, nreg) in ((:avx512, 64, 32), (:avx2, 32, 16), (:neon, 16, 32)) QuasiStrided._TARGET[] = synthetic(key, vb; nregisters = nreg) - Q = 4096 - for (T, shape) in zip((Float64, Float32, ComplexF64, ComplexF32), shapes) - @test dshape(T, Q, Q, Q) === (shape, QuasiStrided._default_method(T)) - @test dshape(T, QuasiStrided._default_method(T), Q, Q, Q) === dshape(T, Q, Q, Q) - end - for (T, rshape) in zip((ComplexF64, ComplexF32), shapes[1:2]), m in (CR, RC) - @test dshape(T, m, Q, Q, Q) === (mapped(m, rshape), m) - # The real extent, store and small-M demotions carry over. - for Qm in (1, 3, 17, 40, 100), run in (0, 1, 8, Qm) - real_problem = m === CR ? (2 * Qm, 50, 2 * run) : (Qm, 100, run) - @test dshape(T, m, Qm, 50, run) === (mapped(m, dshape(real(T), real_problem...)[1]), m) - end + for T in (ComplexF64, ComplexF32), m in (CR, RC), Qm in (1, 3, 17, 40, 4096), run in (0, 1, 8, Qm) + real_problem = m === CR ? (2 * Qm, 50, 2 * run) : (Qm, 100, run) + @test QuasiStrided._default_shape(T, m, Qm, 50, run) === + (mapped(m, QuasiStrided._default_shape(real(T), real_problem...)[1]), m) end end finally From 9cd431aff29c588698b78e6eabe41a7834025ce0 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:06:25 -0400 Subject: [PATCH 21/27] Mixed tests: merge the two allocation testsets; trim restated constants test_mixed.jl checks steady-state allocations for the mixed-domain nest and the panel path in one table. test_mixed_kernel.jl drops the checks of a_reals/b_reals/accumulator_planes/complex_method (restating their definitions; the selection tests check complex_method) and runs its named-kernel end-to-end rows unconjugated and with both operands conjugated; test_mixed.jl covers one-sided conjugation on the automatic CR and RC plans. Co-Authored-By: Claude Opus 5.5 --- test/execution/test_mixed.jl | 18 +++++++----------- test/microkernels/test_mixed_kernel.jl | 9 ++------- 2 files changed, 9 insertions(+), 18 deletions(-) diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 7fb765a..1f2677d 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -54,12 +54,6 @@ end end end -@testset "mixed eltypes: steady-state execute! allocates nothing" begin - rng = MersenneTwister(3) - Amat, Bmat, Cmat = randn(rng, Float32, 40, 30), randn(rng, ComplexF64, 30, 20), zeros(ComplexF64, 40, 20) - @test _steady_allocs!(execute!, _mm_plan(Cmat, Amat, Bmat), Cmat) == 0 skip = (VERSION < v"1.11") -end - # Along K, 1, 2^-30 and -1 fall in different kc blocks: the 2^-30 and the small # beta*C survive only if the partial sums stay in the compute type between blocks. @testset "one rounding to $TC across kc blocks: $TA x $TB, M = $M, $case, beta = $beta" for (TA, TB, TC, acc, M, N, K, case, beta) in ( @@ -98,10 +92,12 @@ end @test all(==(TC(2.0^-30 * b + beta * 2.0^-28)), Array(Cv)) end -@testset "Float32 C, Float64 compute type: steady-state execute! allocates nothing" begin - rng = MersenneTwister(6) - Amat, Bmat, Cmat = randn(rng, Float32, 64, 700), randn(rng, Float32, 700, 40), zeros(Float32, 64, 40) - plan = _mm_plan(Cmat, Amat, Bmat; accumulator = Float64) - @test _path_of(plan) isa QuasiStrided._PanelPath +@testset "steady-state execute! allocates nothing: $TA x $TB -> $TC, accumulator = $acc" for (TA, TB, TC, acc, K) in ( + (Float32, ComplexF64, ComplexF64, nothing, 30), (Float32, Float32, Float32, Float64, 700), + ) + rng = MersenneTwister(3) + Amat, Bmat, Cmat = randn(rng, TA, 64, K), randn(rng, TB, K, 40), zeros(TC, 64, 40) + plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc) + @test _path_of(plan) isa (TC === Float32 ? QuasiStrided._PanelPath : QuasiStrided._NestPath) @test _steady_allocs!(execute!, plan, Cmat) == 0 skip = (VERSION < v"1.11") end diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index 06d604b..c501be3 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -1,4 +1,4 @@ -using QuasiStrided: ComplexRealKernel, RealComplexKernel, ComplexRealMethod, RealComplexMethod, packed_panel +using QuasiStrided: ComplexRealKernel, RealComplexKernel, packed_panel mk_optypes(k::ComplexRealKernel) = (scalartype(k), real(scalartype(k))) mk_optypes(k::RealComplexKernel) = (real(scalartype(k)), scalartype(k)) @@ -33,11 +33,6 @@ end @testset "construction" begin cr = ComplexRealKernel(Val(8), Val(4), ComplexF32) rc = RealComplexKernel(Val(8), Val(4), ComplexF64) - @test complex_method(cr) === ComplexRealMethod() - @test complex_method(rc) === RealComplexMethod() - @test (a_reals(ComplexRealMethod()), b_reals(ComplexRealMethod())) == (2, 1) - @test (a_reals(RealComplexMethod()), b_reals(RealComplexMethod())) == (1, 2) - @test accumulator_planes(ComplexRealMethod()) == accumulator_planes(RealComplexMethod()) == 1 @test cr.inner isa SIMDKernel{16, 4, Float32, 8} @test rc.inner isa SIMDKernel{8, 8, Float64, 4} @test (packed_a_per_k(cr), packed_b_per_k(cr), packed_a_per_k(rc), packed_b_per_k(rc)) == (16, 4, 8, 8) @@ -96,7 +91,7 @@ end saved = QuasiStrided._UNPACKED_B_MODE[] QuasiStrided._UNPACKED_B_MODE[] = mode try - for (cA, cB) in ((false, false), (true, false), (false, true)) + for (cA, cB) in ((false, false), (true, true)) C = copy(C0) plan = plan_contract( StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3); From 785b264447cefa736fd973ad7eb04e7c016dca18 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:06:25 -0400 Subject: [PATCH 22/27] Trim the eltype rejection lists to one case per check Co-Authored-By: Claude Opus 5.5 --- test/integrations/test_tensoroperations.jl | 2 +- test/planning/test_plan_contract.jl | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 647b5de..01204af 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -133,7 +133,7 @@ end for (TA, TB, TC) in ( (Float16, Float16, Float16), (Complex{Float16}, Complex{Float16}, Complex{Float16}), (Complex{Int}, Complex{Int}, Complex{Int}), (Complex{BigFloat}, Complex{BigFloat}, Complex{BigFloat}), - (Float64, Float32, Float16), (ComplexF64, Float64, Float64), (Float32, ComplexF32, Float32), + (Float64, Float32, Float16), (ComplexF64, Float64, Float64), ) A, B, C = ones(TA, (3, 4)), ones(TB, (4, 5)), zeros(TC, (3, 5)) @test_throws ArgumentError tensorcontract!(C, A, pA, false, B, pB, false, pAB, 1, 0, qsbackend) diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 4a96f21..df71923 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -27,10 +27,7 @@ end @testset "driver: eltype validation errors" begin mk(TA, TB, TC) = (StridedView(zeros(TC, 3, 5)), StridedView(ones(TA, 3, 4)), (1, 2), StridedView(ones(TB, 4, 5)), (2, 3), (1, 3)) - for eltypes in ( - (Float16, Float16, Float16), (Int, Float64, Float64), (Float64, Float64, BigFloat), - (ComplexF64, Float64, Float64), (Float32, ComplexF32, Float32), - ) + for eltypes in ((Float16, Float16, Float16), (Float64, Float64, BigFloat), (ComplexF64, Float64, Float64)) @test_throws ArgumentError plan_contract(mk(eltypes...)...) end ws = plan_contract(mk(Float32, Float32, Float32)...).workspace From 27f9ab03fec01f6f59f1b2a9b2307539c805f36e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:06:25 -0400 Subject: [PATCH 23/27] bench_mixed.jl: shorten the header Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_mixed.jl | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index 5a981a6..355aafa 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -1,13 +1,11 @@ # Mixed-domain GEMM, single core: complex x real (CR) and real x complex (RC) -# through the named mixed-domain kernel, the default planar kernel (which -# promotes the real operand to complex), and the real GEMM of the same real FMA count -# (2M x N x K for CR, M x 2N x K for RC). +# through the mixed-domain kernel, the planar kernel (the real operand promoted) +# and the real GEMM of equal FMA count (2M x N x K for CR, M x 2N x K for RC). +# The inner real tile is the host's default real shape, so `real/mixed` is the +# fraction of real efficiency reached. # # julia --project=benchmark benchmark/bench_mixed.jl [--dtypes ComplexF64,ComplexF32] # [--sizes 512,2048] [--reps 5] [--outdir DIR] -# -# The mixed kernels' inner real tile is the host's default real shape, so -# `real/mixed` is the fraction of real efficiency reached. include(joinpath(@__DIR__, "harness.jl")) From 6ae02cc53447a0c7b045590a216e178c22a3374b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:08:26 -0400 Subject: [PATCH 24/27] Forward the mixed-domain kernels' accumulate without restating its signature Co-Authored-By: Claude Opus 5.5 --- src/microkernels/mixed.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index c45c1be..0add24d 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -78,9 +78,7 @@ lanewidth(kernel::_MixedKernel) = lanewidth(kernel.inner) zero_accumulator(kernel::_MixedKernel) = zero_accumulator(kernel.inner) -@inline Base.accumulate( - kernel::K, acc::NTuple{NV, Vec{W, R}}, packed_a::PA, packed_b::PB, kc::Int -) where {K <: _MixedKernel, NV, W, R, PA, PB} = +@inline Base.accumulate(kernel::_MixedKernel, acc::NTuple, packed_a, packed_b, kc::Int) = accumulate(kernel.inner, acc, packed_a, packed_b, kc) # Inlined or not as the fmaddsub and planar stores each reuses. From 44f6371a1ed19227ae9d08a0fce6ef93334f7f29 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:09:38 -0400 Subject: [PATCH 25/27] Validate the backend's accumulator once, in its inner constructor QuasiStridedBackend{Float16}() now throws the keyword constructor's ArgumentError itself, so _compute_type no longer needs a throwing fallback; plan_contract and contract! restrict the keyword's type. Co-Authored-By: Claude Opus 5.5 --- src/integrations/tensoroperations.jl | 15 ++++++++------- src/planning/plan.jl | 10 ++-------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index b50bb90..a10e588 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -22,15 +22,16 @@ falls back to another backend. `tensoradd!` and `tensortrace!` are forwarded to The backend is not registered with `TensorOperations.select_backend`. """ -struct QuasiStridedBackend{A} <: TO.AbstractBackend end - -function QuasiStridedBackend(; accumulator = nothing) - accumulator in (nothing, Float32, Float64) || throw( - ArgumentError("QuasiStridedBackend: accumulator must be nothing, Float32 or Float64, got $accumulator") - ) - return QuasiStridedBackend{accumulator}() +struct QuasiStridedBackend{A} <: TO.AbstractBackend + function QuasiStridedBackend{A}() where {A} + A in (nothing, Float32, Float64) || + throw(ArgumentError("QuasiStridedBackend: accumulator must be nothing, Float32 or Float64, got $A")) + return new{A}() + end end +QuasiStridedBackend(; accumulator = nothing) = QuasiStridedBackend{accumulator}() + # Task-local pool of workspaces, keyed by the compute type only: all complex # methods pack into `Vector{real(T)}` and `reserve!` is grow-only, so they can # share one. diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 3e44578..9a92051 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -227,14 +227,8 @@ const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) (TC <: Real && !(TA <: Real && TB <: Real)) && _throw_complex_into_real(TA, TB, TC) return promote_type(TA, TB, TC) end -@inline function _compute_type( - ::Type{TA}, ::Type{TB}, ::Type{TC}, ::Type{R} - ) where {TA, TB, TC, R <: Union{Float32, Float64}} - return _compute_type(TA, TB, TC, nothing) <: Complex ? Complex{R} : R -end -@noinline _compute_type(TA, TB, TC, accumulator) = throw( - ArgumentError("accumulator must be nothing, Float32 or Float64, got $accumulator") -) +@inline _compute_type(::Type{TA}, ::Type{TB}, ::Type{TC}, ::Type{R}) where {TA, TB, TC, R <: Union{Float32, Float64}} = + _compute_type(TA, TB, TC, nothing) <: Complex ? Complex{R} : R @noinline _throw_eltypes(TA, TB, TC) = throw( ArgumentError( From 7f681045e6a2484a8f4747f75534ea252ee40341 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 29 Sep 2026 13:10:25 -0400 Subject: [PATCH 26/27] Restore pre-existing kernel comments the lane conversion did not invalidate Co-Authored-By: Claude Opus 5.5 --- src/microkernels/fmaddsub.jl | 4 ++-- src/microkernels/interface.jl | 3 +-- src/microkernels/planar.jl | 8 +++----- src/microkernels/simd.jl | 8 +++----- src/packing/pack_contiguous.jl | 3 +-- 5 files changed, 10 insertions(+), 16 deletions(-) diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 58f0d08..48ec904 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -174,8 +174,8 @@ end end # One full `W÷2`-row block, already in `Complex`'s memory order, so no -# interleave shuffle. `C` converts from and to its storage type `RC`. Base's -# `Complex` expression trees in lanes, as planar's `_planar_store_block!`: +# interleave shuffle. Base's `Complex` expression trees in lanes, as planar's +# `_planar_store_block!`: # beta == 0: addsub(ar*r, ai*swap(r)) # beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) # otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 8597eae..f7ce9a7 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -113,8 +113,7 @@ end _default_lanewidth(::Type{Float64}) = 4 _default_lanewidth(::Type{Float32}) = 8 -# `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. These -# helpers load C as `T`, compute in `T` and round on store. +# `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. function scale_tile!(destination::QSTile, beta::T) where {T} m = nrows(destination) n = ncols(destination) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index a3c6926..401a4bb 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -182,9 +182,8 @@ end end # Vector store: unit-stride rows into rank-1 dense `Complex` storage (so the -# storage can be reinterpreted as `2W` consecutive reals per `W` rows) of `T` or -# another complex type the lanes convert to and from, on an ISA the complex -# fast paths ship for (shared with the complex pack fast path). +# storage can be reinterpreted as `2W` consecutive reals per `W` rows), on an +# ISA the complex fast paths ship for (shared with the complex pack fast path). @inline _complex_vector_eligible(tile::QSTile, ::Type{T}) where {T} = _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) && _complex_fastpath_isa_eligible() @@ -208,8 +207,7 @@ end return :(shufflevector(re, im, Val($idx))) end -# One full `W`-row block; `at` is the zero-based index of its first real. `C` -# is `RC` in memory and converts to and from `R` around the arithmetic. +# One full `W`-row block; `at` is the zero-based index of its first real. # The arithmetic transcribes Base's `Complex` expression trees (the ones # `_axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, # `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index 2584fe4..982bb08 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -89,10 +89,9 @@ end return acc end -# Vector store eligibility: unit-stride rows into rank-1 dense storage of `T` -# or of another real type the lanes convert to and from, exactly what SIMD.jl's -# array `vload`/`vstore` accept. Must admit `Memory`: that is the `parent` of an -# Array-backed `StridedView` on Julia >= 1.11. +# Vector store eligibility: unit-stride rows into rank-1 dense real storage, +# exactly what SIMD.jl's array `vload`/`vstore` accept. Must admit `Memory{T}`: +# that is the `parent` of an Array-backed `StridedView` on Julia >= 1.11. @inline _vector_store_eligible(tile::QSTile, ::Type{T}) where {T} = _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) @@ -129,7 +128,6 @@ end # Whole `W`-row blocks are one vector load/store; a block straddling `m` is # stored lane by lane, so nothing outside the valid rectangle is touched. # `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. -# `C` is loaded into and rounded back from `T` lanes. @generated function _store_tile_vector!( destination::QSTile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 1caea2c..bbe739f 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -11,8 +11,7 @@ using SIMD: shufflevector @inline _copies_unchanged(::Any, ::Type) = false # Dense rank-1 storage whose elements load as SIMD lanes and convert lane-wise -# to `T`: `T` itself, or another supported type of the same domain. Shared with -# the vector stores. +# to `T`: `T` itself, or another supported type of the same domain. const _LaneFloat = Union{Float32, Float64} @inline _dense_lanes(storage::S, ::Type{T}) where {S, T} = storage isa DenseVector && _lane_convertible(eltype(S), T) From e1fb81a6d068a5e787d3c343d33b3f763f60f36e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 30 Sep 2026 09:58:19 -0400 Subject: [PATCH 27/27] test_mixed.jl: pin kc in the panel allocation test K = 700 spans more than one kc block only when the default kc is below 700, which is false on targets with a larger cache model (the macOS aarch64 CI runners), so the plan took the plain nest. Co-Authored-By: Claude Opus 5.5 --- test/execution/test_mixed.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 8a24a84..1486a60 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -100,7 +100,7 @@ end ) rng = MersenneTwister(3) Amat, Bmat, Cmat = randn(rng, TA, 64, K), randn(rng, TB, K, 40), zeros(TC, 64, 40) - plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc) + plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc, kc = 256) @test _path_of(plan) isa (TC === Float32 ? QuasiStrided._PanelPath : QuasiStrided._NestPath) @test _steady_allocs!(execute!, plan, Cmat) == 0 skip = (VERSION < v"1.11") end