diff --git a/README.md b/README.md index af1eab4..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. @@ -117,8 +119,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)`); + 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 @@ -127,11 +134,10 @@ 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. - 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/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl new file mode 100644 index 0000000..355aafa --- /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 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] + +include(joinpath(@__DIR__, "harness.jl")) + +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")) +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) + planar = _kernel_from_shape(_derived_shape(target_profile(), T), T, PlanarMethod()) + 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; 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( + "%-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) diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 19df77d..5ff6419 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -33,6 +33,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") @@ -72,6 +73,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/barrier.jl b/src/execution/barrier.jl index f136952..c51460e 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -48,28 +48,32 @@ struct _NestPath{UNPACKED_B, AFF, SPLIT} end # N's B flag stays `false` when B is read in place: that path builds no B # sliver axis, so the flag would only split specialisations. A split group is -# enumerated in another order, so its ramp flags are `false`. +# enumerated in another order, so its ramp flags are `false`. `panel`: C is +# `_PanelPath`'s dense panel, whose M and N maps are ramps. @inline function _nest_path( unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, - split_a::Bool, split_b::Bool + split_a::Bool, split_b::Bool, panel::Bool = false ) split_b &= !unpacked_b return _nest_path_from_flags( - unpacked_b, - !split_a && _is_ramp_map(mgroup, 1), !split_a && _is_ramp_map(mgroup, 2), - !unpacked_b && !split_b && _is_ramp_map(ngroup, 1), !split_b && _is_ramp_map(ngroup, 2), + panel ? _PANEL_PATHS : _NEST_PATHS, unpacked_b, + !split_a && _is_ramp_map(mgroup, 1), !split_a && (panel || _is_ramp_map(mgroup, 2)), + !unpacked_b && !split_b && _is_ramp_map(ngroup, 1), !split_b && (panel || _is_ramp_map(ngroup, 2)), _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2), split_a, split_b ) end +# The nest run against a dense panel of C in the compute type. +struct _PanelPath{P <: _NestPath} end + # Table lookup rather than `_NestPath{u, aff, split}()` built at runtime (a # dynamic type application) or a 512-leaf branch tree (slow to infer). -@inline function _nest_path_from_flags(flags::Vararg{Bool, 9}) +@inline function _nest_path_from_flags(table::Vector{Any}, flags::Vararg{Bool, 9}) index = 1 for i in 1:9 index += flags[i] << (i - 1) end - return @inbounds _NEST_PATHS[index] + return @inbounds table[index] end const _NEST_PATHS = let @@ -81,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/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 bb07c9d..51ad95a 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -58,22 +58,24 @@ end @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, - _is_split(plan.mpack), _is_split(plan.npack) + plan.blocking.kc, _is_split(plan.mpack), _is_split(plan.npack) ) # 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. `split_a`/`split_b`: the plan packs A/B line by line. +# `nothing` to assume it does. `kc`: the requested K block, `nothing` for +# the default. `split_a`/`split_b`: the plan packs A/B line by line. @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, split_a::Bool = false, split_b::Bool = false ) 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) @@ -81,10 +83,18 @@ 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, split_a, split_b + unpack_ok && _unpacked_b_rule(mgroup, kgroup), mgroup, ngroup, kgroup, split_a, split_b, panel ) 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 # with its storages stripped; they cross as arguments. @inline function _execute_across_barrier!(plan::ContractPlan{T}, alphaT::T, betaT::T, path) where {T} @@ -129,7 +139,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} = @@ -154,15 +165,20 @@ end end @inline _splits(plan::ContractPlan, hint) = false -@inline _splits(plan::ContractPlan, ::_NestPath) = _is_split(plan.mpack) || _is_split(plan.npack) +@inline _splits(plan::ContractPlan, ::Union{_NestPath, _PanelPath}) = _is_split(plan.mpack) || _is_split(plan.npack) @inline _split_path(plan::ContractPlan, ::_NestPath{U}) where {U} = _nest_path( U, plan.mgroup, plan.ngroup, plan.kgroup, _is_split(plan.mpack), _is_split(plan.npack) ) +@inline _split_path(plan::ContractPlan, ::_PanelPath{<:_NestPath{U}}) where {U} = _nest_path( + U, plan.mgroup, plan.ngroup, plan.kgroup, _is_split(plan.mpack), _is_split(plan.npack), true +) -# 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) = +# 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{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 @@ -186,16 +202,90 @@ 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, plan.mpack, plan.npack + ) + 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, jc, nblock, betaT) = plan +@inline _panel_exit!(::Nothing, jc, nblock) = nothing + +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 * axis_length(plan.mgroup), plan.atransform, plan.btransform, plan.workspace, + plan.mpack, plan.npack + ) +end + +_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, 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 + 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, SPLIT} + alphaT::T, betaT::T, ::_NestPath{UNPACKED_B, AFF, SPLIT}, target ) where {T, K, UNPACKED_B, AFF, SPLIT} # GUARDRAIL: reals per sliver per K step address the packed panels; # `MRk`/`NRk` count register-tile rows. They differ for complex kernels. @@ -227,6 +317,7 @@ function _execute_nest!( jc = 0 while jc < Qn nblock = min(nc_eff, Qn - jc) + cplan = _panel_enter!(target, plan, jc, nblock, betaT) n_slivers = cld(nblock, NRk) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( @@ -311,7 +402,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) if split_a _pack_block_transposed!( @@ -335,12 +426,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 @@ -352,6 +443,7 @@ function _execute_nest!( pc += kblock end + _panel_exit!(target, jc, nblock) jc += nblock end @@ -388,25 +480,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/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/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..952a84b 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 @@ -39,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/execution/workspace.jl b/src/execution/workspace.jl index 5a4f49e..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 storage element type; the packed panels (`VT`) -hold `real(T)`. Field layout is an implementation detail. +[`release!`](@ref)d. `T` is the plan's compute type; the packed panels (`VT`) +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-one GC-tracked fields. -mutable struct ContractWorkspace{T, VT <: AbstractVector} +# pointer to it instead of copying twenty-two GC-tracked fields. +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 @@ -50,8 +51,12 @@ 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::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}, @@ -61,20 +66,20 @@ 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 - ) 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 storage element type $T), got $(eltype(VT))" + "(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, 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,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 - ) 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), @@ -126,21 +131,25 @@ 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, c_panel, ) end """ ContractWorkspace(T, kernel, blocking::Blocking; - oracle = true, allocator = TensorOperations.DefaultAllocator()) - -A workspace for storage 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). + 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. `panel` is the +length of the compute-type panel of C (see [`plan_contract`](@ref)). Under a +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 + ::Type{T}, kernel, blocking::Blocking, oracle::Bool, ::TO.DefaultAllocator, + panel::Int = 0 ) where {T} s = _workspace_sizes(kernel, blocking) R = realtype(kernel) @@ -148,12 +157,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), Vector{T}(undef, 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) @@ -163,25 +172,27 @@ 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 + packed_a, packed_b, tw_packed_a, tw_packed_b, c_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}, Vector{T}}, kernel, blocking::Blocking, + oracle::Bool, panel::Int = 0 ) where {T, R} s = _workspace_sizes(kernel, blocking) @@ -204,6 +215,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) @@ -223,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) @@ -238,16 +251,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 +270,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 `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, - oracle::Bool + ::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) - 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,13 +293,14 @@ 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( - "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))},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 ec30a40..a10e588 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -5,25 +5,36 @@ 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 + 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 -# 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. +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. const _QS_WORKSPACE_KEY = :quasistrided_contract_workspaces @inline function _qs_workspace_pool() @@ -37,8 +48,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} @@ -63,17 +74,10 @@ 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)), ", ") - ) all(isstrided, (A, B, C)) || _qs_throw( "QuasiStridedBackend requires strided arrays for $f, got " * join(map(typeof, (C, A, B)), ", ") @@ -81,10 +85,11 @@ const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) 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`). -@inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β) +# 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) + T = _compute_type(eltype(A), eltype(B), eltype(C), 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 +107,6 @@ end ) # Dropping `Zero()`/`One()` is safe: the kernels branch on `iszero(alpha/beta)`. - T = eltype(C) indA, indB, indC = _qs_labels(pA, pB, pAB) return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end @@ -116,14 +120,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 +141,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/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 534e769..48ec904 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -180,31 +180,32 @@ end # 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 # 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 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/interface.jl b/src/microkernels/interface.jl index 787acbb..f7ce9a7 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() @@ -118,24 +126,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/microkernels/mixed.jl b/src/microkernels/mixed.jl new file mode 100644 index 0000000..0add24d --- /dev/null +++ b/src/microkernels/mixed.jl @@ -0,0 +1,111 @@ +# 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} + 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} + 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} + 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 + +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() +lanewidth(kernel::_MixedKernel) = lanewidth(kernel.inner) + +zero_accumulator(kernel::_MixedKernel) = zero_accumulator(kernel.inner) + +@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. +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 2da99d2..401a4bb 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -134,10 +134,19 @@ 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. 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) + 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. @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 +155,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 @@ -176,7 +185,7 @@ end # 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) && 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 @@ -206,15 +215,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 +235,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 @@ -234,12 +243,13 @@ 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 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) @@ -248,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, @@ -291,7 +301,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..982bb08 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -89,11 +89,11 @@ end return acc end -# Vector store eligibility: unit-stride rows into rank-1 dense storage of `T`, +# 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) && 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}}, @@ -134,6 +134,8 @@ end 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 +147,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/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/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 27f4347..bbe739f 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -10,27 +10,36 @@ 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. +const _LaneFloat = Union{Float32, Float64} +@inline _dense_lanes(storage::S, ::Type{T}) where {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 + # 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 +60,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 +111,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 +137,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/src/packing/transposed.jl b/src/packing/transposed.jl index da0f5f1..a5f84e1 100644 --- a/src/packing/transposed.jl +++ b/src/packing/transposed.jl @@ -48,6 +48,6 @@ end ) where {PO, FMT, PK, DS, R} r, t = divrem(i, R) po = (plane, idx, pp) -> r * panel + plane_offset(descriptor, plane, idx, pp) - _pack_emit!(buffer, format, po, t, p, convert(scalartype(descriptor), x)) + _pack_emit!(buffer, format, po, t, p, convert(_element_type(format, scalartype(descriptor)), x)) return nothing end 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 0f94161..bff24b5 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} # Line-by-line packing of A (`mpack`) and B (`npack`); the groups above keep # their natural order, and the nest path enumerates the split ones. @@ -49,7 +49,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, @@ -57,12 +57,21 @@ 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)` 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 type, a [`PlanarKernel`](@ref) for a - complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512). - [`OneMKernel`](@ref) is used only when named. + [`SIMDKernel`](@ref) for a real `T`, a [`PlanarKernel`](@ref) for a + 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 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. @@ -73,10 +82,10 @@ invalid input. with the lines' other elements needed only after more lines than fit L2, is packed line by line (`PackSplit`); its `mc` (or `nc`) then becomes a whole number of line groups, up to `requested kc / kc` times the request. - * `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` @@ -94,17 +103,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, @@ -147,13 +157,12 @@ 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) + 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`. @@ -171,8 +180,8 @@ 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. + 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) @@ -186,7 +195,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 @@ -197,24 +206,57 @@ 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, ::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(kernel::RealComplexKernel, ::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) || + _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 _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( + "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. -@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 @@ -224,11 +266,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 @@ -273,7 +315,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) @@ -296,13 +338,28 @@ function _plan_contract( mc_eff = Qm == 0 ? MRk : min(mc_rounded, _roundup(Qm, MRk)) nc_eff = Qn == 0 ? NRk : min(nc_rounded, _roundup(Qn, NRk)) kc_eff = Qk == 0 ? 1 : min(requested.kc, Qk) + panel = _c_panel_needed(T, req.Cstorage, Qk, kc_eff) mpack = npack = _NO_SPLIT + # Cache lines hold each operand's storage eltype, and the block walk costs + # what its packed format's scatter does. B is not split under a panel of C: + # the panel holds `jc` blocks in C's own N order, which a split N group does + # not enumerate contiguously. if Qm > 0 && Qn > 0 && Qk > 0 - mc_eff, mpack = _pack_split(req.mgroup, req.kgroup, 1, MRk, sizeof(T), T <: Complex, kc_eff, mc_eff, mc_rounded, requested.kc) - nc_eff, npack = _pack_split(req.ngroup, req.kgroup, 2, NRk, sizeof(T), T <: Complex, kc_eff, nc_eff, nc_rounded, requested.kc) + mc_eff, mpack = _pack_split( + req.mgroup, req.kgroup, 1, MRk, sizeof(eltype(req.Astorage)), + !(a_format(kernel) isa RealFormat), kc_eff, mc_eff, mc_rounded, requested.kc + ) + if !panel + nc_eff, npack = _pack_split( + req.ngroup, req.kgroup, 2, NRk, sizeof(eltype(req.Bstorage)), + !(b_format(kernel) isa RealFormat), kc_eff, nc_eff, nc_rounded, requested.kc + ) + end end blocking = Blocking(mc_eff, kc_eff, nc_eff) - ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator) + ws = _resolve_workspace( + T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel ? Qm * min(nc_eff, Qn) : 0 + ) plan = ContractPlan( kernel, req.mgroup, req.ngroup, req.kgroup, blocking, diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index a60aecf..69f7389 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!) @@ -249,9 +249,9 @@ end const _SP_I7 = ((1, 2, 3, 4, 5), (4, 6), (5, 3, 2, 6, 1)) const _SP_BOTH = ((2, 3, 1), (5, 3, 4), (1, 2, 4, 5)) -function _sp_views(T, (iA, iB, iC), ext) - arr(I) = StridedView(randn(T, map(l -> ext[l], I))) - return (arr(iC), arr(iA), iA, arr(iB), iB, iC) +function _sp_views(T, (iA, iB, iC), ext, TA = T) + arr(S, I) = StridedView(randn(S, map(l -> ext[l], I))) + return (arr(T, iC), arr(TA, iA), iA, arr(T, iB), iB, iC) end # The plan with every structurally eligible group split, whatever this host's L2: diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl new file mode 100644 index 0000000..1486a60 --- /dev/null +++ b/test/execution/test_mixed.jl @@ -0,0 +1,106 @@ +# 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, `:split` is intensli_7 with +# M^4 rows. +function _mixed_fixture(TA, TB, TC, M, K, N, layout, seed) + layout === :split && return _sp_views(TC, _SP_I7, (M, M, M, K, M, N), TA) + 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 + +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, :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), + (Float64, ComplexF64, ComplexF64, nothing, :split, (8, 3, 5), 0.5, false, true, _NestPath{false, <:Any, (true, false)}, _RC), + (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 + ref = _mixed_ref(fx, alpha, beta; conjA, conjB) + plan = (layout === :split ? _sp_forced_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) + 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 + +# 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 + 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 + @test all(==(TC(2.0^-30 * b + beta * 2.0^-28)), Array(Cv)) +end + +@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, 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 diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 9053cfd..6076c8b 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -139,13 +139,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), ) 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) @@ -153,6 +151,28 @@ 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) + @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 + 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/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl new file mode 100644 index 0000000..c501be3 --- /dev/null +++ b/test/microkernels/test_mixed_kernel.jl @@ -0,0 +1,108 @@ +using QuasiStrided: ComplexRealKernel, RealComplexKernel, 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 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, 12, 8, 8), + (RealComplexKernel, ComplexF64, Float32, ComplexF32, 24, 4, 8), + ) + 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(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(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) + 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, 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_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_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/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_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/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 diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index b87b59f..196ff96 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -229,6 +229,41 @@ 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) + 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 + # The real extent, store and small-M demotions carry over. + saved = QuasiStrided._TARGET[] + try + for (key, vb, nreg) in ((:avx512, 64, 32), (:avx2, 32, 16), (:neon, 16, 32)) + QuasiStrided._TARGET[] = synthetic(key, vb; nregisters = nreg) + 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 + 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()), diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 9ab1342..df71923 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -25,6 +25,16 @@ @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), (Float64, Float64, BigFloat), (ComplexF64, Float64, Float64)) + @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 e2b1c94..cda2519 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") @@ -55,6 +56,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")