diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 453d153..19df77d 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -24,6 +24,7 @@ include("packing/format.jl") include("packing/panel.jl") include("packing/pack.jl") include("packing/pack_contiguous.jl") +include("packing/transposed.jl") # --- Microkernels: accumulate over one packed K panel, store into C --- include("microkernels/interface.jl") @@ -39,6 +40,7 @@ include("planning/conjugation.jl") include("planning/kernel_selection.jl") include("planning/blocking.jl") include("planning/defaults.jl") +include("planning/pack_split.jl") include("execution/workspace.jl") include("execution/barrier.jl") include("planning/plan.jl") diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index a6ecb7a..f136952 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -40,36 +40,43 @@ end # The five-loop nest. `UNPACKED_B`: B read in place instead of packed. `AFF`: # per map, in the order M-in-A, M-in-C, N-in-B, N-in-C, K-in-A, K-in-B, # whether it is an affine ramp, so that its sliver axes are statically -# `AffineAxis` and the `PtrScatterAxis` arm is never compiled. -struct _NestPath{UNPACKED_B, AFF} end +# `AffineAxis` and the `PtrScatterAxis` arm is never compiled. `SPLIT`: whether +# A and B are packed line by line (`plan.mpack`/`plan.npack`). +struct _NestPath{UNPACKED_B, AFF, SPLIT} end @inline _is_ramp_map(g::AxisGroup, p::Int) = _map_ramp_step(g, p) !== nothing # 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. -@inline _nest_path(unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup) = - _nest_path_from_flags( - unpacked_b, - _is_ramp_map(mgroup, 1), _is_ramp_map(mgroup, 2), - !unpacked_b && _is_ramp_map(ngroup, 1), _is_ramp_map(ngroup, 2), - _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2) -) +# sliver axis, so the flag would only split specialisations. A split group is +# enumerated in another order, so its ramp flags are `false`. +@inline function _nest_path( + unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, + split_a::Bool, split_b::Bool + ) + 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), + _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2), split_a, split_b + ) +end -# Table lookup rather than `_NestPath{u, aff}()` built at runtime (a dynamic -# type application) or a 128-leaf branch tree (slow to infer). -@inline function _nest_path_from_flags(flags::Vararg{Bool, 7}) +# 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}) index = 1 - for i in 1:7 + for i in 1:9 index += flags[i] << (i - 1) end return @inbounds _NEST_PATHS[index] end const _NEST_PATHS = let - table = Vector{Any}(undef, 128) - for index in 1:128 - bits = ntuple(i -> ((index - 1) >> (i - 1)) & 1 == 1, 7) - table[index] = _NestPath{bits[1], bits[2:end]}() + table = Vector{Any}(undef, 512) + for index in 1:512 + bits = ntuple(i -> ((index - 1) >> (i - 1)) & 1 == 1, 9) + table[index] = _NestPath{bits[1], bits[2:7], bits[8:9]}() end table end diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 6c4e0de..bb07c9d 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -57,16 +57,18 @@ end # outer-product path, else the nest (B packed or read in place). @inline _select_path(plan::ContractPlan{T}) where {T} = _select_path( T, _unpacked_b_kernel_eligible(plan.kernel), - plan.Astorage, plan.Bstorage, plan.Cstorage, plan.mgroup, plan.ngroup, plan.kgroup, plan + plan.Astorage, plan.Bstorage, plan.Cstorage, plan.mgroup, plan.ngroup, plan.kgroup, plan, + _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. +# `nothing` to assume it does. `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, + split_a::Bool = false, split_b::Bool = false ) where {T} Qm = axis_length(mgroup) Qn = axis_length(ngroup) @@ -78,7 +80,9 @@ end 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) + return _nest_path( + unpack_ok && _unpacked_b_rule(mgroup, kgroup), mgroup, ngroup, kgroup, split_a, split_b + ) end # Run `path` behind a dynamic call. The plan crosses in the workspace slot @@ -104,12 +108,12 @@ end @inline _strip_storage(p::ContractPlan) = ContractPlan( p.kernel, p.mgroup, p.ngroup, p.kgroup, p.blocking, nothing, p.Abase, nothing, p.Bbase, nothing, p.Cbase, - p.atransform, p.btransform, p.workspace + p.atransform, p.btransform, p.workspace, p.mpack, p.npack ) @inline _with_storage(p::ContractPlan, Astorage, Bstorage, Cstorage) = ContractPlan( p.kernel, p.mgroup, p.ngroup, p.kgroup, p.blocking, Astorage, p.Abase, Bstorage, p.Bbase, Cstorage, p.Cbase, - p.atransform, p.btransform, p.workspace + p.atransform, p.btransform, p.workspace, p.mpack, p.npack ) # The continuation the TensorOperations adapter hands `_planned`: `execute!` @@ -135,17 +139,26 @@ end # Runs the predicted path statically when it provably equals the plan's own # `_select_path` (always, for an automatically chosen kernel); otherwise -# falls back to `execute!`'s barrier. +# falls back to `execute!`'s barrier. The prediction assumes no split, so a +# split plan crosses the barrier to its own nest path. @inline function _execute_hinted!(plan::ContractPlan{T}, alphaT::T, betaT::T, hint) where {T} _execute_short_circuit!(plan, alphaT, betaT) && return nothing - if _hint_holds(plan, hint) - _execute_path!(plan, alphaT, betaT, hint) - else + if !_hint_holds(plan, hint) _execute_across_barrier!(plan, alphaT, betaT, _select_path(plan)) + elseif _splits(plan, hint) + _execute_across_barrier!(plan, alphaT, betaT, _split_path(plan, hint)) + else + _execute_path!(plan, alphaT, betaT, hint) end return nothing end +@inline _splits(plan::ContractPlan, hint) = false +@inline _splits(plan::ContractPlan, ::_NestPath) = _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) +) + # 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). @@ -182,8 +195,8 @@ 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} - ) where {T, K, UNPACKED_B, AFF} + alphaT::T, betaT::T, ::_NestPath{UNPACKED_B, AFF, SPLIT} + ) 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. MRp = packed_a_per_k(kernel) @@ -192,9 +205,16 @@ function _execute_nest!( atransform = plan.atransform btransform = plan.btransform - # Ramp composites get closed-form block descriptors, no offset buffers. - (m_ramp, m_step) = affine_ramp(plan.mgroup) - (n_ramp, n_step) = affine_ramp(plan.ngroup) + split_a, split_b = SPLIT + mgroup = split_a ? _split_group(plan.mgroup, plan.mpack) : plan.mgroup + ngroup = split_b ? _split_group(plan.ngroup, plan.npack) : plan.ngroup + + # Ramp composites get closed-form block descriptors, no offset buffers; the + # block pack reads the offsets. + (m_ramp, m_step) = affine_ramp(mgroup) + (n_ramp, n_step) = affine_ramp(ngroup) + m_ramp &= !split_a + n_ramp &= !split_b (k_ramp, k_step) = affine_ramp(plan.kgroup) lenA = length(plan.Astorage) @@ -214,7 +234,7 @@ function _execute_nest!( nblock, NRk, n_slivers ) else - fill_offsets!((ws.n_buf_B, ws.n_buf_C), plan.ngroup, jc, nblock) + fill_offsets!((ws.n_buf_B, ws.n_buf_C), ngroup, jc, nblock) _classify_slivers!( ws.n_desc_B, ws.n_desc_C, ws.n_buf_B, ws.n_buf_C, nblock, NRk, n_slivers @@ -253,7 +273,14 @@ function _execute_nest!( beta_eff = firstpanel ? betaT : one(T) - if !UNPACKED_B + if split_b + _pack_block_transposed!( + packed_b_plane_offset, b_format(kernel), + packed_panel(ws.packed_b, 1, NRp * kblock * n_slivers), kernel, Val(nr(kernel)), NRp, + plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, nblock, kblock, + plan.npack + ) + elseif !UNPACKED_B for s in 0:(n_slivers - 1) sfirst = s * NRk colsB = _axis_of(ws.n_desc_B[s + 1], ws.n_buf_B, sfirst, aff_nB) @@ -276,7 +303,7 @@ function _execute_nest!( mblock, MRk, m_slivers ) else - fill_offsets!((ws.m_buf_A, ws.m_buf_C), plan.mgroup, ic, mblock) + fill_offsets!((ws.m_buf_A, ws.m_buf_C), mgroup, ic, mblock) _classify_slivers!( ws.m_desc_A, ws.m_desc_C, ws.m_buf_A, ws.m_buf_C, mblock, MRk, m_slivers @@ -286,14 +313,23 @@ function _execute_nest!( checked_span_bounds(plan.Abase, rng_mA, rng_kA, lenA) checked_span_bounds(plan.Cbase, rng_mC, rng_nC, lenC) - for r in 0:(m_slivers - 1) - rfirst = r * MRk - rowsA = _axis_of(ws.m_desc_A[r + 1], ws.m_buf_A, rfirst, aff_mA) - apanel = _sliver_panel(ws.packed_a, MRp, kblock, r) - _pack_sliver!( - unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, - kernel, atransform + if split_a + _pack_block_transposed!( + packed_a_plane_offset, a_format(kernel), + packed_panel(ws.packed_a, 1, MRp * kblock * m_slivers), kernel, Val(mr(kernel)), MRp, + plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, mblock, kblock, + plan.mpack ) + else + for r in 0:(m_slivers - 1) + rfirst = r * MRk + rowsA = _axis_of(ws.m_desc_A[r + 1], ws.m_buf_A, rfirst, aff_mA) + apanel = _sliver_panel(ws.packed_a, MRp, kblock, r) + _pack_sliver!( + unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, + kernel, atransform + ) + end end # --- loop 2: jr over N-slivers; loop 1: ir over M-slivers --- diff --git a/src/packing/transposed.jl b/src/packing/transposed.jl new file mode 100644 index 0000000..da0f5f1 --- /dev/null +++ b/src/packing/transposed.jl @@ -0,0 +1,53 @@ +# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated +# by a split group, whole `E x L` groups) by `kc` K steps into the panels the +# per-sliver packer would write, reading each group line by line: coordinate +# `i0 + E * y2` is element `y2` of the line at `i0`. `@noinline`: one call per +# block, and inlined into the nest it slows the nest's micro-kernel loop. +@noinline function _pack_block_transposed!( + plane_offset::PO, format::FMT, buffer::PK, descriptor::DS, ::Val{R}, Rp::Int, + storage::ST, base::Int, fbuf::Vector{Int}, kaxis::KA, transform::F, + fcount::Int, kc::Int, split + ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: Axis, F} + E, L = Int(split.E), Int(split.L) + G = E * L + # The sliver and lane come from a per-element divrem by the constant `R`; + # hoisting them per line lets LLVM rewrite the loop so the line misses no + # longer overlap. + emit(kbase, i, p) = @inbounds _pack_line_element!( + plane_offset, format, buffer, descriptor, Val(R), Rp * kc, transform(storage[kbase + fbuf[i + 1]]), i, p + ) + if split.kinner + for g0 in 0:G:(fcount - 1), y1 in 0:(E - 1), p in 0:(kc - 1) + kbase = base + axis_offset(kaxis, p) + 1 + for y2 in 0:(L - 1) + emit(kbase, g0 + y1 + E * y2, p) + end + end + else + for g0 in 0:G:(fcount - 1), p in 0:(kc - 1) + kbase = base + axis_offset(kaxis, p) + 1 + for y1 in 0:(E - 1), y2 in 0:(L - 1) + emit(kbase, g0 + y1 + E * y2, p) + end + end + end + valid = fcount % R + if valid != 0 + rb = (fcount ÷ R) * Rp * kc + pad = (plane, idx, pp) -> rb + plane_offset(descriptor, plane, idx, pp) + for p in 0:(kc - 1), t in valid:(R - 1) + _pack_emit_zero!(buffer, format, pad, t, p, realtype(descriptor)) + end + end + return nothing +end + +# Element `x` of block coordinate `i` (sliver `i ÷ R`, lane `i % R`) at K step `p`. +@inline function _pack_line_element!( + plane_offset::PO, format::FMT, buffer::PK, descriptor::DS, ::Val{R}, panel::Int, x, i::Int, p::Int + ) 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)) + return nothing +end diff --git a/src/planning/pack_split.jl b/src/planning/pack_split.jl new file mode 100644 index 0000000..0aead6e --- /dev/null +++ b/src/planning/pack_split.jl @@ -0,0 +1,163 @@ +# Packing by cache lines. When a register sliver reads one element per line and +# the line's other elements belong to coordinates a whole sweep of the free +# composite later, lines rarely survive until reused; at power-of-two extents +# they alias in a few cache sets and the pack costs a memory latency per +# element. A split plan enumerates the free composite so that a macro block +# holds whole lines, and packs each block line by line +# (`_pack_block_transposed!`); the packed format and the K order are unchanged. + +""" + PackSplit(q, dj, L, kinner, Eq, E) + +How one operand's free group is enumerated and packed under a split: axis `q` +in chunks of `Eq`, the unit-stride axis `dj > q` in chunks of `L` (a line's +worth) placed right after the `Eq` chunk. A group is then `E x L` +coordinates, `E` = (lengths before `q`) * `Eq`, and the block walk reads its +lines K step by K step, or with K innermost when `kinner`. `L == 0`: no split. +""" +struct PackSplit + # Narrow fields: the plan carries two of these through every call. + q::UInt8 + dj::UInt8 + L::UInt8 + kinner::Bool + Eq::Int32 + E::Int32 +end + +const _NO_SPLIT = PackSplit(0, 0, 0, false, 0, 0) + +_is_split(s::PackSplit) = s.L != 0 + +# The split for the operand behind map 1 of `g` (A for M, B for N) and the block +# extent it needs, or `(eff, _NO_SPLIT)`. `eff` is the plan's extent, `rounded` +# the requested one rounded to `R`; a split block may take the budget the +# blocking reserved for the requested `kc_req` when the K extent clamps `kc`. +# `l2bytes = nothing`: `_split_capacity`, doubled for complex, whose block +# walk costs more per element. +@inline function _pack_split( + g::AxisGroup{D}, kg::AxisGroup, kmap::Int, R::Int, S::Int, complex::Bool, kc::Int, + eff::Int, rounded::Int, kc_req::Int, l2bytes::Union{Int, Nothing} = nothing + ) where {D} + D < 256 || return (eff, _NO_SPLIT) + line = _K_LINE_BYTES + d1 = 0 + for d in 1:D + g.lengths[d] > 1 && (d1 = d; break) + end + (d1 > 0 && abs(g.strides[1][d1]) * S >= line) || return (eff, _NO_SPLIT) + dj = 0 + for d in (d1 + 1):D + g.lengths[d] > 1 && abs(g.strides[1][d]) == 1 && (dj = d; break) + end + dj == 0 && return (eff, _NO_SPLIT) + return _pack_split_window(g, kg, kmap, d1, dj, R, S, complex, kc, eff, rounded, kc_req, l2bytes) +end + +# Out of line: shared by every eltype, and most plans never get here. +@noinline function _pack_split_window( + g::AxisGroup, kg::AxisGroup{DK}, kmap::Int, d1::Int, dj::Int, R::Int, S::Int, complex::Bool, + kc::Int, eff::Int, rounded::Int, kc_req::Int, l2bytes::Union{Int, Nothing} + ) where {DK} + line = _K_LINE_BYTES + L = _largest_divisor_upto(g.lengths[dj], max(1, line ÷ S)) + L >= 2 || return (eff, _NO_SPLIT) + + ks = 0 + for d in 1:DK + kg.lengths[d] > 1 && (ks = abs(kg.strides[kmap][d]); break) + end + # K steps within a page make each sliver element's lines a stream the + # prefetchers follow; a split then only saves refetching those lines, which + # does not pay for the block walk's costlier complex scatter. + kinner = ks * S < _K_WALK_FAR_BYTES + kinner && complex && return (eff, _NO_SPLIT) + + # A line of `dj` is reused `psi` coordinates later; while the lines touched + # meanwhile fit the cache, the per-sliver walk is already cache-friendly. + psi = 1 + for d in d1:(dj - 1) + psi *= g.lengths[d] + end + klines = ks * S >= line ? kc : cld(kc * ks * S, line) + aliased = abs(g.strides[1][d1]) * S % _K_WALK_FAR_BYTES == 0 + cap = something(l2bytes, _split_capacity(target_profile(), aliased) << complex) + widemul(psi, max(1, klines) * line) > cap || return (eff, _NO_SPLIT) + return _pack_split_dynamic(g.lengths, d1, dj, L, psi, R, eff, rounded, kc, kc_req, kinner) +end + +# The cache that can hold the per-sliver walk's reuse window: the core's L2 +# share, plus its L3 share unless the sliver steps a whole number of pages, +# which folds the lines it gathers into a few sets of every level. +function _split_capacity(profile::TargetProfile, aliased::Bool) + l2 = _l2_core_bytes(profile) + l3 = profile.l3 + (aliased || l3.bytes <= 0) && return l2 + return l2 + l3.bytes ÷ max(1, l3.sharing ÷ max(1, profile.l1d.sharing)) +end + +# Past the tests above the plan is for a large contraction; a dynamic call keeps +# the search from being compiled for every other plan. +@noinline _pack_split_dynamic(args...) = Base.inferencebarrier(_pack_split_blocks)(args...)::Tuple{Int, PackSplit} + +function _pack_split_blocks( + lengths::NTuple{D, Int}, d1::Int, dj::Int, L::Int, psi::Int, R::Int, + eff::Int, rounded::Int, kc::Int, kc_req::Int, kinner::Bool + ) where {D} + budget = max(rounded, Int(min(widemul(rounded, kc_req) ÷ kc, typemax(Int32))) ÷ R * R) + # Prefixes of `d1:dj-1`: whole axes, then a divisor chunk of the next one. + # `E` is also the run of consecutive C coordinates a block stores, so take + # the largest whose whole groups fit the budget, preferring `R | E`. A + # prefix shorter than a sliver would make slivers straddle chunks. + best = anyc = (0, 0, 0) + pre = 1 + for q in d1:(dj - 1) + len = lengths[q] + limit = budget ÷ (pre * L) + limit >= 1 || break + for Eq in min(len, limit):-1:1 + len % Eq == 0 || continue + E = pre * Eq + ((E >= R || E == psi) && lcm(R, E * L) <= budget) || continue + E > anyc[3] && (anyc = (q, Eq, E)) + E % R == 0 && E > best[3] && (best = (q, Eq, E)) + end + len <= limit || break + pre *= len + end + (q, Eq, E) = best[3] > 0 ? best : anyc + E == 0 && return (eff, _NO_SPLIT) + blk = lcm(R, E * L) + neweff = min(budget ÷ blk * blk, _roundup(_unchecked_axis_length(lengths), R)) + return (neweff, PackSplit(q, dj, L, kinner, Eq, E)) +end + +@inline function _largest_divisor_upto(n::Int, cap::Int) + for L in min(n, cap):-1:1 + n % L == 0 && return L + end + return 1 +end + +# `g` in the enumeration `s` describes: `q` and `dj` each replaced by an (inner, +# outer) pair of axes, the inner chunk of `dj` moved right after that of `q`. +function _split_group(g::AxisGroup{D, P}, s::PackSplit) where {D, P} + ax(k, p) = _split_axis(g, s, k, p) + return AxisGroup( + ntuple(k -> ax(k, 0), Val(D + 2)), ntuple(p -> ntuple(k -> ax(k, p), Val(D + 2)), Val(P)) + ) +end + +# Length (`p == 0`) or map-`p` stride of axis `k` of `_split_group(g, s)`. +@inline function _split_axis(g::AxisGroup, s::PackSplit, k::Int, p::Int) + len(d) = g.lengths[d] + str(d) = g.strides[p][d] + q, dj, Eq, L = Int(s.q), Int(s.dj), Int(s.Eq), Int(s.L) + k < q && return p == 0 ? len(k) : str(k) + k == q && return p == 0 ? Eq : str(q) + k == q + 1 && return p == 0 ? L : str(dj) + k == q + 2 && return p == 0 ? len(q) ÷ Eq : Eq * str(q) + k <= dj + 1 && return p == 0 ? len(k - 2) : str(k - 2) + k == dj + 2 && return p == 0 ? len(dj) ÷ L : L * str(dj) + return p == 0 ? len(k - 2) : str(k - 2) +end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index ce22dcc..0f94161 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -33,6 +33,11 @@ struct ContractPlan{ btransform::TB workspace::ContractWorkspace{T, VT} + + # 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. + mpack::PackSplit + npack::PackSplit end """ @@ -64,6 +69,10 @@ invalid input. * `mc`/`kc`/`nc` override the fields of `default_blocking(kernel)` (each `>= 1`); `mc`/`nc` are rounded up to `mr`/`nr` multiples and all three are clamped to the extents. + * An operand whose register slivers would read one element per cache line, + 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 @@ -282,16 +291,23 @@ function _plan_contract( NRk = nr(kernel) # Empty extents: the drivers never read these; the floors keep them valid. - mc_eff = Qm == 0 ? MRk : min(_roundup(requested.mc, MRk), _roundup(Qm, MRk)) - nc_eff = Qn == 0 ? NRk : min(_roundup(requested.nc, NRk), _roundup(Qn, NRk)) + mc_rounded = _roundup(requested.mc, MRk) + nc_rounded = _roundup(requested.nc, NRk) + 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) + mpack = npack = _NO_SPLIT + 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) + end blocking = Blocking(mc_eff, kc_eff, nc_eff) ws = _resolve_workspace(T, req.workspace, kernel, blocking, req.oracle, req.allocator) plan = ContractPlan( kernel, req.mgroup, req.ngroup, req.kgroup, blocking, req.Astorage, req.Abase, req.Bstorage, req.Bbase, req.Cstorage, req.Cbase, - atransform, btransform, ws + atransform, btransform, ws, mpack, npack ) return _continue(req.f, plan, hint) end diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index aa8a0ff..a60aecf 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -187,7 +187,7 @@ end ContractPlan( p.kernel, p.mgroup, p.ngroup, p.kgroup, p.blocking, p.Astorage, p.Abase, p.Bstorage, p.Bbase, p.Cstorage, p.Cbase, - conj, conj, p.workspace, + conj, conj, p.workspace, p.mpack, p.npack, ), ) @test parent(StridedView(Cmat)) isa DenseVector{Float64} @@ -242,3 +242,62 @@ end end @test isconcretetype(only(Base.return_types(execute!, (typeof(plan), Float64, Float64)))) end + +# Line-by-line packing. intensli_7, C[e,c,b,f,a] = A[a,b,c,d,e] * B[d,f], splits A +# (K steps of a page or more, every eltype); C[a1,au,f1,f2] = A[au,k,a1] * B[f2,k,f1] +# splits both operands (K steps within a page, real only). +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) +end + +# The plan with every structurally eligible group split, whatever this host's L2: +# the planner's decision at a zero L2 threshold, replanned at its block extents. +function _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; kw...) + p = plan_contract(Cv, Av, iA, Bv, iB, iC; kw...) + k, kc, d, T = p.kernel, p.blocking.kc, default_blocking(p.kernel), eltype(Cv) + split(g, map, R, eff, req) = _pack_split(g, p.kgroup, map, R, sizeof(T), T <: Complex, kc, eff, cld(req, R) * R, d.kc, 0) + (mc, ms) = split(p.mgroup, 1, mr(k), p.blocking.mc, something(get(kw, :mc, nothing), d.mc)) + (nc, ns) = split(p.ngroup, 2, nr(k), p.blocking.nc, d.nc) + q = plan_contract(Cv, Av, iA, Bv, iB, iC; kw..., kernel = k, mc, nc) + @assert q.mgroup == p.mgroup && q.ngroup == p.ngroup + return ContractPlan( + k, q.mgroup, q.ngroup, q.kgroup, q.blocking, q.Astorage, q.Abase, + q.Bstorage, q.Bbase, q.Cstorage, q.Cbase, q.atransform, q.btransform, q.workspace, ms, ns + ) +end + +@testset "line-by-line packing: the planner splits A of intensli_7 past the cache" begin + plan_of(d) = plan_contract(_sp_views(Float64, _SP_I7, ntuple(_ -> d, 6))...) + p = plan_of(16) + k, b = p.kernel, p.blocking + splits(cap) = QuasiStrided._is_split(_pack_split(p.mgroup, p.kgroup, 1, mr(k), 8, false, b.kc, b.mc, b.mc, default_blocking(k).kc, cap)[2]) + @test splits(2^20) && !splits(2^24) # a 4 MB reuse window + host = splits(QuasiStrided._split_capacity(target_profile(), true)) + @test QuasiStrided._is_split(p.mpack) == host + host && @test _path_of(p) isa _NestPath{false, <:Any, (true, false)} + @test _path_of(plan_of(4)) isa _NestPath{<:Any, <:Any, (false, false)} +end + +@testset "line-by-line packing: contractions ($T)" for T in (Float64, Float32, ComplexF64, ComplexF32) + kernels = T === ComplexF64 ? (nothing, OneMKernel(Val(4), Val(4), T), FMAddSubKernel(Val(8), Val(4), T)) : (nothing,) + for (ind, ext, mc) in ( + (_SP_I7, (8, 8, 8, 3, 6, 5), nothing), (_SP_I7, (6, 8, 11, 5, 10, 7), 48), + (_SP_BOTH, (20, 16, 9, 20, 8), nothing), (_SP_BOTH, (13, 12, 5, 11, 12), 64), + ), kernel in kernels, conjB in (T <: Complex ? (false, true) : (false,)), beta in (0, 0.7) + T <: Complex && ind === _SP_BOTH && continue # K steps within a page: complex never splits + Cv, Av, iA, Bv, iB, iC = _sp_views(T, ind, ext) + iszero(beta) && fill!(Cv, NaN) + Cref = _lo_reference(iszero(beta) ? zero(Array(Cv)) : Array(Cv), Av, iA, Bv, iB, iC; conjB, alpha = 1.3, beta) + plan = _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; conjB, mc, kernel) + @test _path_of(plan) isa _NestPath{false, <:Any, (true, ind === _SP_BOTH)} + execute!(plan, 1.3, beta) + @test Array(Cv) ≈ Cref rtol = 100 * eps(real(T)) + end + plan = _sp_forced_plan(_sp_views(T, T <: Complex ? _SP_I7 : _SP_BOTH, (20, 16, 9, 20, 8, 5))...) + execute!(plan, 1, 0) + @test (@allocated execute!(plan, 1, 0)) == 0 skip = (VERSION < v"1.11") +end diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 5d3cb9c..9053cfd 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -40,6 +40,16 @@ const _MATMUL_PAB = ((1,), (2,)), ((1,), (2,)), ((1, 2), ()) end end +# A plan that packs A line by line runs past the planner's predicted path. +@testset "tensorcontract! through a split plan" begin + A, B = randn(16, 16, 16, 16, 16), randn(16, 16) + Cn, Cq = zeros(16, 16, 16, 16, 16), zeros(16, 16, 16, 16, 16) + pA, pB, pAB = ((1, 2, 3, 5), (4,)), ((1,), (2,)), ((4, 3, 2, 5, 1), ()) + tensorcontract!(Cn, A, pA, false, B, pB, false, pAB, 1.0, 0.0, to_native) + tensorcontract!(Cq, A, pA, false, B, pB, false, pAB, 1.0, 0.0, qsbackend) + @test Cq ≈ Cn +end + # `conjA`/`conjB` and each operand's `StridedView.op` compose by xor. @testset "conjugation: flags x StridedView.op (eltype = $T)" for T in complex_eltypes Random.seed!(20260914) diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index d130210..7722866 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -214,7 +214,7 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() p = _pcf_Plan( base.kernel, base.mgroup, base.ngroup, base.kgroup, base.blocking, astore, 0, bstore, 0, cstore, 0, - base.atransform, base.btransform, base.workspace, + base.atransform, base.btransform, base.workspace, base.mpack, base.npack, ) astore.n = 0; bstore.n = 0; cstore.n = 0 _pcf_exec(p, 1.0, 0.0) @@ -249,7 +249,7 @@ end return _pcf_Plan( base.kernel, base.mgroup, base.ngroup, base.kgroup, base.blocking, Astorage, Abase, Bstorage, Bbase, Cstorage, Cbase, - base.atransform, base.btransform, base.workspace, + base.atransform, base.btransform, base.workspace, base.mpack, base.npack, ) end @@ -311,7 +311,7 @@ end pshort = _pcf_Plan( base.kernel, base.mgroup, base.ngroup, base.kgroup, base.blocking, base.Astorage, base.Abase, base.Bstorage, base.Bbase, short_C, base.Cbase, - base.atransform, base.btransform, base.workspace, + base.atransform, base.btransform, base.workspace, base.mpack, base.npack, ) @test_throws BoundsError _pcf_exec(pshort, 1.0, 0.0) @test all(iszero, short_C) # nothing written before the throw @@ -341,7 +341,7 @@ end base.kernel, base.mgroup, base.ngroup, base.kgroup, base.blocking, base.Astorage, base.Abase, base.Bstorage, base.Bbase, zeros(M * N1 * N2), base.Cbase - 1, - base.atransform, base.btransform, base.workspace, + base.atransform, base.btransform, base.workspace, base.mpack, base.npack, ) @test_throws BoundsError _pcf_exec(plow, 1.0, 0.0) end diff --git a/test/runtests.jl b/test/runtests.jl index 5c9d966..e2b1c94 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -15,7 +15,7 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip _rule_applies, _isa_nregisters, packed_a_per_k, packed_b_per_k, realtype, complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, a_reals, b_reals, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, - _planar_pressure + _planar_pressure, _pack_split, _NestPath # plan_contract, execute! and ContractPlan are bound in helpers.jl instead. # All files share one scope, so helper names must be unique across files.