From ff20abd3c1c8f58df2010dfcce23fde83d8c0ae6 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 30 Sep 2026 15:18:13 -0400 Subject: [PATCH 01/60] Add cleanup decision log Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 CLEANUP.md diff --git a/CLEANUP.md b/CLEANUP.md new file mode 100644 index 0000000..1d19efd --- /dev/null +++ b/CLEANUP.md @@ -0,0 +1,94 @@ +# Guided cleanup: decisions and status + +Working notes for the `cleanup` branch; delete before merging. + +## Process + +Per chunk: explain purpose, reading order and design decisions; review +together; apply the agreed changes as one commit; run that chunk's tests (and +TTFX when codegen is touched). + +## Chunks + +| # | Chunk | Files | Status | +|---|---|---|---| +| 0 | Tour | `QuasiStrided.jl`, `plan.jl`, `execute.jl` call path | done | +| 1 | Hardware | `hardware/target.jl` | | +| 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | | +| 3 | Tiles | `layout/tiles.jl` | | +| 4 | Packing formats | `packing/format.jl`, `panel.jl`, `transposed.jl` | | +| 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | | +| 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | | +| 7 | SIMD kernels | `simd.jl`, `planar.jl` | | +| 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | | +| 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | +| 10 | Kernel selection | `planning/kernel_selection.jl` | | +| 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl` | | +| 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | +| 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | +| 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | +| 15 | Alternative paths | `oracle.jl`, `unpackedb.jl` | | +| 16 | Degenerate paths | `dot.jl`, `outer.jl` | | +| 17 | TensorOperations backend | `integrations/tensoroperations.jl` | | +| 18 | Test infrastructure | `runtests.jl`, `helpers.jl`, `forced_isa_runner.jl`, `quality/` | | +| 19 | Benchmarks (optional) | `benchmark/` | | + +## Decisions + +### D1. Descriptive M/N/K names (applied up front, codebase-wide) + +Pattern `__`; M/N/K stay the organizing letters. + +| Old | New | +|---|---| +| `Qm, Qn, Qk` | `m_length, n_length, k_length` | +| `mc, kc, nc` (`Blocking` fields, `plan_contract` kwargs) | `m_block, k_block, n_block` | +| `ic, pc, jc` | `m_block_start, k_block_start, n_block_start` | +| `mblock, kblock, nblock`, `kc_len` | `m_block_length, k_block_length, n_block_length` | +| `mr(k), nr(k)` | `tile_size(k) -> (m_tile, n_tile)` | +| `MRk, NRk` | `m_tile, n_tile` | +| `packed_a_per_k(k), packed_b_per_k(k)` | `sliver_widths(k) -> (a_sliver_width, b_sliver_width)` | +| `MRp, NRp` | `a_sliver_width, b_sliver_width` | +| `m_slivers, n_slivers` | `m_tiles, n_tiles` | +| `r, s`, `rfirst, sfirst` | `m_tile_index, n_tile_index`, `m_tile_start, n_tile_start` | +| `loop 5: jc` comments | `loop over N blocks` etc. | + +Public names are renamed too (0.1, unregistered). Deferred: the `MR`/`NR` +type parameters and names inside the microkernels' generated code, reviewed +in chunks 6–8. + +### D2. No leading underscores (per chunk) + +Drop the `_` prefix from internal names as each chunk is reviewed; rename +where the underscore meant "internal helper of `foo`" (known clashes: +`_pack_a!`/`_pack_b!` vs `pack_a!`/`pack_b!`, `_plan_contract` vs +`plan_contract`). No submodules: nothing else clashes with Base, Core, +LinearAlgebra, SIMD or the package's own names. Explicit imports instead of +`using LinearAlgebra`, so later shadowing is an error (applied up front). + +### D3. Dispatch barriers (chunks 12–14) + +Keep barrier #1 (kernel choice). Remove barrier #2 (execution path) by +deciding the path at plan time and making it a `ContractPlan` type parameter; +this deletes `_path_hint`/`_continue`/`_execute_hinted!`/`_hint_holds`/ +`_splits`/`_split_path`. To verify first: the path is fully fixed at plan +time, including dot-path workspace capacity under a reused grow-only +workspace. Measured: one crossing ~13 ns; a 2×2 `tensorcontract!` ~320 ns. + +### D4. Precompile workload (later, separate) + +Add a PrecompileTools workload for common cases (ranks ≤ 3, `Vector` +storage, the four eltypes) on top of the barriers, not instead of them. +Measure precompile time and pkgimage size first. + +## Open items noted in the tour (to revisit in their chunk) + +- 16-argument `ContractPlan(...)` spelled out five times, `_PlanRequest` + twice more: one "copy with fields replaced" helper (chunk 13/14). +- `execute.jl` mixes public API, barrier/hint machinery, the C-panel path and + the nest; `_c_panel_needed` lives there but is used by `plan.jl` (chunk 14). +- `_select_path` has 12 positional arguments, two optional (chunk 14). +- `oracle = true` default makes every user plan allocate test-path buffers; + consider defaulting to `false` (chunk 13/15). +- `workspace.jl`/`barrier.jl` live in `execution/` but are included in the + planning section (chunk 12). From 564a7401ae22f0ce0c0f95fe440136cdd037608c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 30 Sep 2026 15:44:54 -0400 Subject: [PATCH 02/60] Descriptive M/N/K block and tile names Decision D1: names follow `__` codebase-wide. Public renames: - `Blocking` fields `mc, kc, nc` -> `m_block, k_block, n_block` - `plan_contract` keywords `mc/kc/nc` -> `m_block/k_block/n_block` - `mr(kernel)`, `nr(kernel)` -> `tile_size(kernel) -> (m_tile, n_tile)` - `packed_a_per_k`, `packed_b_per_k` -> `sliver_widths(kernel)` (`tile_size` and `sliver_widths` are now in the `public` list) `using LinearAlgebra` is dropped from src: it used no LinearAlgebra names (`adjoint`/`transpose` are Base's). Full test suite: 34542/34542 pass. TTFX unchanged (Float64 first call 3.3-3.5 s before, 3.2-3.7 s after). Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_blocking_model.jl | 34 +-- benchmark/bench_kernels.jl | 28 +- benchmark/bench_mixed.jl | 2 +- benchmark/harness.jl | 6 +- src/QuasiStrided.jl | 3 +- src/execution/barrier.jl | 3 +- src/execution/dot.jl | 79 +++--- src/execution/execute.jl | 261 ++++++++++--------- src/execution/macrokernel.jl | 60 ++--- src/execution/oracle.jl | 57 ++-- src/execution/outer.jl | 38 +-- src/execution/unpackedb.jl | 36 +-- src/execution/workspace.jl | 74 +++--- src/microkernels/fmaddsub.jl | 8 +- src/microkernels/interface.jl | 57 ++-- src/microkernels/mixed.jl | 13 +- src/microkernels/onem.jl | 21 +- src/microkernels/planar.jl | 8 +- src/microkernels/scalar.jl | 8 +- src/microkernels/simd.jl | 8 +- src/packing/format.jl | 19 +- src/packing/pack.jl | 58 ++--- src/packing/pack_contiguous.jl | 12 +- src/packing/transposed.jl | 23 +- src/planning/blocking.jl | 54 ++-- src/planning/defaults.jl | 34 +-- src/planning/kernel_selection.jl | 63 ++--- src/planning/labels.jl | 16 +- src/planning/pack_split.jl | 18 +- src/planning/plan.jl | 129 +++++---- test/execution/test_dot.jl | 6 +- test/execution/test_execute.jl | 26 +- test/execution/test_macro_blocking.jl | 24 +- test/execution/test_manual_pipeline.jl | 4 +- test/execution/test_mixed.jl | 22 +- test/execution/test_outer.jl | 4 +- test/execution/test_scalar_vs_simd.jl | 6 +- test/execution/test_unpackedb.jl | 8 +- test/execution/test_workspace.jl | 34 +-- test/microkernels/test_fmaddsub_kernel.jl | 20 +- test/microkernels/test_mixed_kernel.jl | 28 +- test/microkernels/test_onem_kernel.jl | 6 +- test/microkernels/test_scalar_kernel.jl | 70 ++--- test/packing/test_kernel_descriptor.jl | 12 +- test/packing/test_pack_complex.jl | 36 +-- test/packing/test_pack_complex_contiguous.jl | 10 +- test/packing/test_pack_real.jl | 54 ++-- test/planning/test_kernel_selection.jl | 100 +++---- test/planning/test_per_call_overhead.jl | 30 +-- test/planning/test_plan_contract.jl | 94 +++---- test/runtests.jl | 4 +- 51 files changed, 916 insertions(+), 912 deletions(-) diff --git a/benchmark/bench_blocking_model.jl b/benchmark/bench_blocking_model.jl index 2f3dd3b..1f238db 100644 --- a/benchmark/bench_blocking_model.jl +++ b/benchmark/bench_blocking_model.jl @@ -1,13 +1,13 @@ -# (mc, kc, nc) plateau sweep: the analytical blocking model and the fallback -# row against a log-spaced grid, at the engine's default kernel. +# (m_block, k_block, n_block) plateau sweep: the analytical blocking model and +# the fallback row against a log-spaced grid, at the engine's default kernel. # # julia --project=benchmark benchmark/bench_blocking_model.jl [--reps 21] [--smoke] [--outdir DIR] # # Per real dtype: the named rows (`model` = `_modelled_blocking`, `fallback` = # `_fallback_blocking`) and the `wide` grid on every harness shape, plus 1-D -# `nc` and `mc` slices through the model row on a 2048x256x2048 shape (the only -# one large enough for `nc` to bind). Per complex dtype: the named rows scaled -# by `_scale_blocking`. Every configuration of a shape is timed back to back. +# `n_block` and `m_block` slices through the model row on a 2048x256x2048 shape +# (the only one large enough for `n_block` to bind). Per complex dtype: the +# named rows scaled by `_scale_blocking`. Every configuration of a shape is timed back to back. # The score of a configuration is the geomean over shapes of its time over the # best time at that shape. `--smoke` runs 1 rep on a few points. @@ -34,27 +34,27 @@ const OUTDIR = outdir() const CSV_PATH = joinpath(OUTDIR, "blocking_model.csv") const SUMMARY_PATH = joinpath(OUTDIR, "summary_blocking_model.txt") -tup(b::Blocking) = (b.mc, b.kc, b.nc) +tup(b::Blocking) = (b.m_block, b.k_block, b.n_block) named_points(rows) = [(string(k), tup(v)) for (k, v) in pairs(rows) if v !== nothing] named_rows(::Type{T}) where {T <: Real} = (model = _modelled_blocking(PROFILE, T), fallback = _fallback_blocking(T)) csv = open(CSV_PATH, "w") -println(csv, "set,dtype,shape,Ma,Ka,Na,mc,kc,nc,mc_eff,kc_eff,nc_eff,reps,median_seconds,gflops") +println(csv, "set,dtype,shape,Ma,Ka,Na,m_block_requested,k_block_requested,n_block_requested,m_block,k_block,n_block,reps,median_seconds,gflops") function sweep_shape!(raw, kernel, ::Type{T}, spec, points) where {T} fx = build_plain(T, spec, MersenneTwister(0xB10C)) - for (set, (mc, kc, nc)) in points + for (set, (m_block, k_block, n_block)) in points plan = plan_contract( - fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; kernel = kernel, mc = mc, kc = kc, nc = nc + fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block ) t = median_time_s(() -> execute!(plan, one(T), zero(T)); reps = REPS) b = plan.blocking println( - csv, "$set,$T,$(spec.name),$(spec.Ma),$(spec.Ka),$(spec.Na),$mc,$kc,$nc,$(b.mc),$(b.kc),$(b.nc),$REPS,", + csv, "$set,$T,$(spec.name),$(spec.Ma),$(spec.Ka),$(spec.Na),$m_block,$k_block,$n_block,$(b.m_block),$(b.k_block),$(b.n_block),$REPS,", @sprintf("%.9f,%.4f", t, gflops(T, spec.Ma, spec.Ka, spec.Na, t)) ) flush(csv) - push!(raw, (; set, dtype = T, shape = spec.name, mc, kc, nc, t)) + push!(raw, (; set, dtype = T, shape = spec.name, m_block, k_block, n_block, t)) end return nothing end @@ -81,7 +81,7 @@ function summarize(io, raw, canaries) grid_names = [s.name for s in GRID_SHAPES] for T in DTYPES println(io, "\n## $T ", named_rows(T)) - w = score(raw, T, grid_names, r -> (r.mc, r.kc, r.nc); sets = ("wide",)) + w = score(raw, T, grid_names, r -> (r.m_block, r.k_block, r.n_block); sets = ("wide",)) @printf( io, "wide grid: best %.4f %s worst %.4f %s spread %.1f%%\n", w[1][2], w[1][1], w[end][2], w[end][1], 100 * (w[end][2] / w[1][2] - 1) @@ -92,7 +92,7 @@ function summarize(io, raw, canaries) for (set, g) in score(raw, T, grid_names, r -> r.set; sets = ("model", "fallback")) @printf(io, " %-10s %.4f\n", set, g) end - for (s, f) in (("nslice", r -> r.nc), ("mslice", r -> r.mc)) + for (s, f) in (("nslice", r -> r.n_block), ("mslice", r -> r.m_block)) pts = sort([r for r in raw if r.dtype == T && r.set == s]; by = f) isempty(pts) && continue tb = minimum(r.t for r in pts) @@ -116,7 +116,7 @@ canaries = [run_canary(crng, "start")] for T in DTYPES kernel = _default_kernel(T) rows = named_rows(T) - println("\n$T kernel $(mr(kernel))x$(nr(kernel))/W$(lanewidth(kernel)) ", rows) + println("\n$T kernel $(tile_size(kernel)[1])x$(tile_size(kernel)[2])/W$(lanewidth(kernel)) ", rows) rows.model === nothing && @warn "no cache geometry detected: model undefined" grid = [("wide", c) for c in thin(WIDE[T])] for spec in thin(GRID_SHAPES) @@ -124,8 +124,8 @@ for T in DTYPES end m = something(rows.model, rows.fallback) slices = vcat( - named_points(rows), [("nslice", (m.mc, m.kc, n)) for n in thin(NSLICE)], - [("mslice", (x, m.kc, m.nc)) for x in thin(MSLICE)], + named_points(rows), [("nslice", (m.m_block, m.k_block, n)) for n in thin(NSLICE)], + [("mslice", (x, m.k_block, m.n_block)) for x in thin(MSLICE)], ) sweep_shape!(raw, kernel, T, BIG_SHAPE, slices) push!(canaries, run_canary(crng, "after-$T")) @@ -133,7 +133,7 @@ end for T in CDTYPES kernel = _default_kernel(T) rows = map(v -> v === nothing ? nothing : _scale_blocking(v, complex_method(kernel)), named_rows(real(T))) - println("\n$T kernel $(mr(kernel))x$(nr(kernel))/W$(lanewidth(kernel)) ", rows) + println("\n$T kernel $(tile_size(kernel)[1])x$(tile_size(kernel)[2])/W$(lanewidth(kernel)) ", rows) for spec in thin(MAIN_SHAPES) sweep_shape!(raw, kernel, T, spec, named_points(rows)) end diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index 2a8b1b9..c6e13f1 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -6,7 +6,7 @@ # # Arms: # kernel (real dtypes) the ceiling of each menu tile: `unsafe_execute_tile!` -# on one L1-resident packed A and B sliver at the tile's default `kc`. +# on one L1-resident packed A and B sliver at the tile's default `k_block`. # engine per dtype and shape: OpenBLAS, the default plan, and every menu # kernel (of every complex method, for complex dtypes) at its own # default blocking; all configurations of a shape back to back. @@ -35,17 +35,17 @@ const OUTDIR = outdir() const CSV_PATH = joinpath(OUTDIR, "bench_kernels.csv") const SUMMARY_PATH = joinpath(OUTDIR, "summary_kernels.txt") -tag(k) = "$(mr(k))x$(nr(k))/W$(lanewidth(k))" +tag(k) = "$(tile_size(k)[1])x$(tile_size(k)[2])/W$(lanewidth(k))" menu_methods(::Type{T}) where {T} = T <: Complex ? (PlanarMethod(), OneMMethod(), FMAddSubMethod()) : (RealMethod(),) csv = open(CSV_PATH, "w") -println(csv, "arm,dtype,shape,M,K,N,config,kernel,mc,kc,nc,seconds,gflops,frac_openblas") +println(csv, "arm,dtype,shape,M,K,N,config,kernel,m_block,k_block,n_block,seconds,gflops,frac_openblas") rows = NamedTuple[] function record!(arm, T, spec, config, kernel, blk, t, t_blas) gf = gflops(T, spec.Ma, spec.Ka, spec.Na, t) println( csv, "$arm,$T,$(spec.name),$(spec.Ma),$(spec.Ka),$(spec.Na),$config,$kernel,", - blk === nothing ? ",," : "$(blk.mc),$(blk.kc),$(blk.nc)", ",", + blk === nothing ? ",," : "$(blk.m_block),$(blk.k_block),$(blk.n_block)", ",", @sprintf("%.9f,%.3f,%.4f", t, gf, t_blas / t) ) flush(csv) @@ -56,13 +56,13 @@ print_env_header(stdout, "bench_kernels.jl") println("target = ", target_profile()) canaries = [run_canary(MersenneTwister(0xCA), "start")] -function kernel_hot!(kernel, C, apack, bpack, kc, reps) +function kernel_hot!(kernel, C, apack, bpack, k_block, reps) GC.@preserve apack bpack begin ap = packed_panel(apack, 1, length(apack)) bp = packed_panel(bpack, 1, length(bpack)) - dest = DestinationTile(C, 0, AffineAxis(0, 1, mr(kernel)), AffineAxis(0, mr(kernel), nr(kernel))) + dest = DestinationTile(C, 0, AffineAxis(0, 1, tile_size(kernel)[1]), AffineAxis(0, tile_size(kernel)...)) for _ in 1:reps - unsafe_execute_tile!(kernel, dest, ap, bp, kc, one(eltype(C)), one(eltype(C))) + unsafe_execute_tile!(kernel, dest, ap, bp, k_block, one(eltype(C)), one(eltype(C))) end end return nothing @@ -71,12 +71,12 @@ end for T in filter(t -> t <: Real, RUN_DTYPES), sh in kernel_shapes(T) kernel = _kernel_from_shape(sh, T) MR, NR = sh - kc = default_blocking(kernel).kc - apack = rand(T, MR * kc); bpack = rand(T, NR * kc); C = zeros(T, MR * NR) - kernel_hot!(kernel, C, apack, bpack, kc, 10) - reps = max(100, round(Int, 0.05 / @elapsed(kernel_hot!(kernel, C, apack, bpack, kc, 100)) * 100)) - t = median_time_s(() -> kernel_hot!(kernel, C, apack, bpack, kc, reps); reps = REPS) / reps - record!("kernel", T, ShapeSpec("l1_tile", MR, kc, NR), "menu", tag(kernel), nothing, t, NaN) + k_block = default_blocking(kernel).k_block + apack = rand(T, MR * k_block); bpack = rand(T, NR * k_block); C = zeros(T, MR * NR) + kernel_hot!(kernel, C, apack, bpack, k_block, 10) + reps = max(100, round(Int, 0.05 / @elapsed(kernel_hot!(kernel, C, apack, bpack, k_block, 100)) * 100)) + t = median_time_s(() -> kernel_hot!(kernel, C, apack, bpack, k_block, reps); reps = REPS) / reps + record!("kernel", T, ShapeSpec("l1_tile", MR, k_block, NR), "menu", tag(kernel), nothing, t, NaN) end function time_plan(fx, ::Type{T}; kw...) where {T} @@ -94,7 +94,7 @@ for spec in SHAPES, T in RUN_DTYPES for m in menu_methods(T), sh in kernel_shapes(T, m) k = _kernel_from_shape(sh, T, m) b = default_blocking(k) - t, _ = time_plan(fx, T; kernel = k, mc = b.mc, kc = b.kc, nc = b.nc) + t, _ = time_plan(fx, T; kernel = k, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block) record!("engine", T, spec, lowercase(replace(string(nameof(typeof(m))), "Method" => "")), tag(k), b, t, tb) end println(" ", rpad(string(T), 11), rpad(spec.name, 22), @sprintf("default %7.2f GF/s %.2fx OpenBLAS", gflops(T, spec.Ma, spec.Ka, spec.Na, td), tb / td)) diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index 355aafa..6ff3e8c 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -46,7 +46,7 @@ function run(io) 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))" + tag = "$(tile_size(kernel)[1])x$(tile_size(kernel)[2])/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 diff --git a/benchmark/harness.jl b/benchmark/harness.jl index bbd4e81..d4cd9b4 100644 --- a/benchmark/harness.jl +++ b/benchmark/harness.jl @@ -3,7 +3,7 @@ # parsing and output paths. using QuasiStrided -using QuasiStrided: SIMDKernel, mr, nr, lanewidth, plan_contract, execute! +using QuasiStrided: SIMDKernel, tile_size, lanewidth, plan_contract, execute! using StridedViews: StridedView using LinearAlgebra using Statistics: median @@ -61,7 +61,7 @@ function build_plain(::Type{T}, spec::ShapeSpec, rng) where {T} ) end -full_grid(mcs, kcs, ncs) = [(mc, kc, nc) for kc in kcs for mc in mcs for nc in ncs] +full_grid(m_blocks, k_blocks, n_blocks) = [(m_block, k_block, n_block) for k_block in k_blocks for m_block in m_blocks for n_block in n_blocks] const DTYPES = (Float64, Float32) const CDTYPES = (ComplexF64, ComplexF32) @@ -105,7 +105,7 @@ function run_canary(rng, label::String) fx = build_plain(Float64, ShapeSpec("canary_64^3", 64, 64, 64), rng) plan = plan_contract( fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; - kernel = SIMDKernel(Val(8), Val(6), Float64), mc = 128, kc = 256, nc = 1536 + kernel = SIMDKernel(Val(8), Val(6), Float64), m_block = 128, k_block = 256, n_block = 1536 ) t = median_time_s(() -> execute!(plan, 1.0, 0.0); reps = 15) println("canary[$label] median = $(t) s") diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 5ff6419..66d075f 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -1,6 +1,5 @@ module QuasiStrided -using LinearAlgebra using StridedViews: StridedView, offset # TensorOperations names are always qualified: a bare `using` collides on `scalartype`. @@ -71,7 +70,7 @@ end eval( Expr( :public, :contract!, :plan_contract, :execute!, :ContractPlan, - :ContractWorkspace, :Blocking, :default_blocking, + :ContractWorkspace, :Blocking, :default_blocking, :tile_size, :sliver_widths, :ScalarKernel, :SIMDKernel, :PlanarKernel, :OneMKernel, :FMAddSubKernel, :ComplexRealKernel, :RealComplexKernel, :target_profile, :cache_topology, diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index c51460e..4f4ef27 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -87,7 +87,8 @@ 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`). +# The dot path at lane width `W`; `MATB`: the matrix operand is B +# (`m_length == 1`). struct _DotPath{MATB, W} end # The outer-product path at lane width `W`. diff --git a/src/execution/dot.jl b/src/execution/dot.jl index 3e9dd0b..eb440d8 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -1,6 +1,6 @@ # The dot-product ("gemv") path for `M == 1` or `N == 1` when the matrix # operand (`Qfree x K`) is stored with K unit-stride. The register microkernel -# vectorizes along M, which is 1/mr full here, and gathering the K-major +# vectorizes along M, which is 1/m_tile full here, and gathering the K-major # matrix into packed panels would dominate the call. This path vectorizes # along K instead: one W-wide FMA per W matrix elements, the matrix read once # and unpacked, one horizontal reduction per output. Only the K-vector is @@ -46,11 +46,11 @@ end # Whether the dot path applies (all but the workspace capacity): a degenerate # free extent, a matrix operand with unit-ramp K in dense storage (raw-pointer # loads), and at least one vector of K. -function _dot_applicable(::Type{T}, Astorage, Bstorage, kgroup::AxisGroup, Qm::Int, Qn::Int, Qk::Int) where {T} +function _dot_applicable(::Type{T}, Astorage, Bstorage, kgroup::AxisGroup, m_length::Int, n_length::Int, k_length::Int) where {T} _DOT_MODE[] === :never && return false - (Qm == 1 || Qn == 1) || return false - Qk >= _dot_lanewidth(T) || return false - if Qm == 1 + (m_length == 1 || n_length == 1) || return false + k_length >= _dot_lanewidth(T) || return false + if m_length == 1 Bstorage isa DenseVector{T} || return false _map_ramp_step(kgroup, 2) == 1 || return false else @@ -64,7 +64,7 @@ end # packed-A buffer. Always true for an automatically chosen kernel; `nothing` # is the prediction's "assume so". function _dot_capacity_ok(plan::ContractPlan{T}) where {T} - need = (T <: Complex ? 2 : 1) * plan.blocking.kc + need = (T <: Complex ? 2 : 1) * plan.blocking.k_block return (length(plan.workspace.packed_a) * sizeof(real(T))) ÷ sizeof(T) >= need end _dot_capacity_ok(::Nothing) = true @@ -75,14 +75,14 @@ function _execute_dot!(plan::ContractPlan{T}, alphaT::T, betaT::T, matB::Bool, : _dot_nest!( plan, ws, plan.Bstorage, plan.Bbase, plan.btransform, plan.Astorage, plan.Abase, plan.atransform, - plan.ngroup, ws.n_buf_B, ws.n_buf_C, plan.blocking.nc, ws.k_buf_A, + plan.ngroup, ws.n_buf_B, ws.n_buf_C, plan.blocking.n_block, ws.k_buf_A, alphaT, betaT, Val(W) ) else _dot_nest!( plan, ws, plan.Astorage, plan.Abase, plan.atransform, plan.Bstorage, plan.Bbase, plan.btransform, - plan.mgroup, ws.m_buf_A, ws.m_buf_C, plan.blocking.mc, ws.k_buf_B, + plan.mgroup, ws.m_buf_A, ws.m_buf_C, plan.blocking.m_block, ws.k_buf_B, alphaT, betaT, Val(W) ) end @@ -111,8 +111,8 @@ function _dot_nest!( blocklen::Int, vkbuf::Vector{Int}, alphaT::T, betaT::T, ::Val{W} ) where {T, SM, FM, SV, FV, G <: AxisGroup, W} Qf = axis_length(g) - Qk = axis_length(plan.kgroup) - kc_eff = plan.blocking.kc + k_length = axis_length(plan.kgroup) + k_block = plan.blocking.k_block Cstorage = plan.Cstorage cbase = plan.Cbase lenm = length(mstorage) @@ -125,35 +125,36 @@ function _dot_nest!( GC.@preserve ws mstorage begin gptr = reinterpret(Ptr{T}, pointer(ws.packed_a)) mptr = pointer(mstorage) - pc = 0 + k_block_start = 0 firstblock = true - while pc < Qk - kblock = min(kc_eff, Qk - pc) + while k_block_start < k_length + k_block_length = min(k_block, k_length - k_block_start) beta_eff = firstblock ? betaT : one(T) - fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, pc, kblock) - checked_span_bounds(vbase, _buffer_range(vkbuf, kblock), (0, 0), lenv) + fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, k_block_start, k_block_length) + checked_span_bounds(vbase, _buffer_range(vkbuf, k_block_length), (0, 0), lenv) if vflip - _dot_gather!(gptr, vstorage, vbase, vkbuf, kblock, conj) + _dot_gather!(gptr, vstorage, vbase, vkbuf, k_block_length, conj) else - _dot_gather!(gptr, vstorage, vbase, vkbuf, kblock, identity) + _dot_gather!(gptr, vstorage, vbase, vkbuf, k_block_length, identity) end q0 = 0 while q0 < Qf qcount = min(blocklen, Qf - q0) fill_offsets!((bufM, bufC), g, q0, qcount) - # The matrix's K offsets are `pc .. pc+kblock-1` (unit ramp). - checked_span_bounds(mbase, _buffer_range(bufM, qcount), (pc, pc + kblock - 1), lenm) + # The matrix's K offsets are + # `k_block_start .. k_block_start+k_block_length-1` (unit ramp). + checked_span_bounds(mbase, _buffer_range(bufM, qcount), (k_block_start, k_block_start + k_block_length - 1), lenm) checked_span_bounds(cbase, _buffer_range(bufC, qcount), (0, 0), lenc) if mconj _dot_block!( - Cstorage, cbase, bufC, mptr, mbase + pc, bufM, qcount, gptr, kblock, + Cstorage, cbase, bufC, mptr, mbase + k_block_start, bufM, qcount, gptr, k_block_length, alphaT, beta_eff, Val(W), Val(NB), Val(true) ) else _dot_block!( - Cstorage, cbase, bufC, mptr, mbase + pc, bufM, qcount, gptr, kblock, + Cstorage, cbase, bufC, mptr, mbase + k_block_start, bufM, qcount, gptr, k_block_length, alphaT, beta_eff, Val(W), Val(NB), Val(false) ) end @@ -161,23 +162,23 @@ function _dot_nest!( end firstblock = false - pc += kblock + k_block_start += k_block_length end end return nothing end # Gather the transformed vector into `gptr`; for a complex `T` also its -# pair-swapped copy at `gptr + kblock`, which gives the imaginary part without -# a per-step shuffle. +# pair-swapped copy at `gptr + k_block_length`, which gives the imaginary part +# without a per-step shuffle. @inline function _dot_gather!( - gptr::Ptr{T}, vstorage::SV, vbase::Int, vkbuf::Vector{Int}, kblock::Int, transform::F + gptr::Ptr{T}, vstorage::SV, vbase::Int, vkbuf::Vector{Int}, k_block_length::Int, transform::F ) where {T, SV, F} - @inbounds for t in 0:(kblock - 1) + @inbounds for t in 0:(k_block_length - 1) 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))) + unsafe_store!(gptr + sizeof(T) * (k_block_length + t), Complex(imag(z), real(z))) end end return nothing @@ -187,14 +188,14 @@ end # the last valid one (computed, never stored). @inline function _dot_block!( Cstorage::SC, cbase::Int, bufC::Vector{Int}, mptr::Ptr{T}, mbase_k::Int, - bufM::Vector{Int}, qcount::Int, gptr::Ptr{T}, kblock::Int, + bufM::Vector{Int}, qcount::Int, gptr::Ptr{T}, k_block_length::Int, alpha::T, beta::T, ::Val{W}, ::Val{NB}, ::Val{MCONJ} ) where {SC, T, W, NB, MCONJ} q = 0 while q < qcount nvalid = min(NB, qcount - q) mbases = _dot_bases(bufM, q, nvalid, mbase_k, Val(NB)) - sums = _dot_group(mptr, gptr, mbases, kblock, Val(W), Val(NB)) + sums = _dot_group(mptr, gptr, mbases, k_block_length, Val(W), Val(NB)) @inbounds for j in 1:nvalid s = MCONJ ? conj(sums[j]) : sums[j] _axpby_at!(Cstorage, cbase + bufC[q + j] + 1, alpha, s, beta) @@ -218,18 +219,18 @@ end # bodies are `@generated` so every tuple index is a literal. @inline function _dot_group( - mptr::Ptr{T}, gptr::Ptr{T}, mbases::NTuple{NB, Int}, kblock::Int, ::Val{W}, ::Val{NB} + mptr::Ptr{T}, gptr::Ptr{T}, mbases::NTuple{NB, Int}, k_block_length::Int, ::Val{W}, ::Val{NB} ) where {T <: Real, NB, W} accs = ntuple(_ -> zero(Vec{W, T}), Val(NB)) - kmain = (kblock ÷ W) * W + kmain = (k_block_length ÷ W) * W t = 0 while t < kmain accs = _dot_step_real(accs, mptr, gptr, mbases, t, nothing) t += W end # Masked-off lanes are never read, so nothing past a row's end is touched. - if t < kblock - accs = _dot_step_real(accs, mptr, gptr, mbases, t, _dot_tailmask(Val(W), kblock - t)) + if t < k_block_length + accs = _dot_step_real(accs, mptr, gptr, mbases, t, _dot_tailmask(Val(W), k_block_length - t)) end return _dot_reduce_real(accs) end @@ -270,21 +271,21 @@ end # so two FMAs per matrix vector and no shuffle. `accs[NB+j]` is `P2_j`. @inline function _dot_group( - mptr::Ptr{T}, gptr::Ptr{T}, mbases::NTuple{NB, Int}, kblock::Int, ::Val{W}, ::Val{NB} + mptr::Ptr{T}, gptr::Ptr{T}, mbases::NTuple{NB, Int}, k_block_length::Int, ::Val{W}, ::Val{NB} ) where {T <: Complex, NB, W} R = real(T) WC = W ÷ 2 accs = ntuple(_ -> zero(Vec{W, R}), Val(2 * NB)) - kmain = (kblock ÷ WC) * WC + kmain = (k_block_length ÷ WC) * WC t = 0 while t < kmain - accs = _dot_step_cplx(accs, mptr, gptr, gptr + sizeof(T) * kblock, mbases, t, nothing) + accs = _dot_step_cplx(accs, mptr, gptr, gptr + sizeof(T) * k_block_length, mbases, t, nothing) t += WC end - if t < kblock + if t < k_block_length accs = _dot_step_cplx( - accs, mptr, gptr, gptr + sizeof(T) * kblock, mbases, t, - _dot_tailmask(Val(W), 2 * (kblock - t)) + accs, mptr, gptr, gptr + sizeof(T) * k_block_length, mbases, t, + _dot_tailmask(Val(W), 2 * (k_block_length - t)) ) end return _dot_reduce_cplx(accs, T) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 51ad95a..f93d1b0 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -1,24 +1,25 @@ -# `C *= beta` tile by tile, never reading A or B: the `Qk == 0 || alpha == 0` -# pass of `execute!` and `execute_tilewise!`. Uses the MR/NR-sized `tile_*` -# buffers, which is why those exist even under `oracle = false`. -function _scale_all_of_C!(plan, betaT::T, MRk::Int, NRk::Int, Qm::Int, Qn::Int) where {T} +# `C *= beta` tile by tile, never reading A or B: the +# `k_length == 0 || alpha == 0` pass of `execute!` and `execute_tilewise!`. Uses +# the tile-sized `tile_*` buffers, which is why those exist even under +# `oracle = false`. +function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::Int, n_length::Int) where {T} ws = plan.workspace m_bufs = (ws.tile_m_buf_A, ws.tile_m_buf_C) n_bufs = (ws.tile_n_buf_B, ws.tile_n_buf_C) - mfirst = 0 - while mfirst < Qm - mcount = min(MRk, Qm - mfirst) - (_, dM_C) = block_descriptors!(m_bufs, plan.mgroup, mfirst, mcount) + m_tile_start = 0 + while m_tile_start < m_length + m_tile_length = min(m_tile, m_length - m_tile_start) + (_, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) - nfirst = 0 - while nfirst < Qn - ncount = min(NRk, Qn - nfirst) - (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, nfirst, ncount) + n_tile_start = 0 + while n_tile_start < n_length + n_tile_length = min(n_tile, n_length - n_tile_start) + (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) - nfirst += ncount + n_tile_start += n_tile_length end - mfirst += mcount + m_tile_start += m_tile_length end return nothing end @@ -43,11 +44,11 @@ end # The short-circuits every path shares; `true` when the call is finished. @inline function _execute_short_circuit!(plan::ContractPlan{T}, alphaT::T, betaT::T) where {T} - Qm = axis_length(plan.mgroup) - Qn = axis_length(plan.ngroup) - (Qm == 0 || Qn == 0) && return true + m_length = axis_length(plan.mgroup) + n_length = axis_length(plan.ngroup) + (m_length == 0 || n_length == 0) && return true if axis_length(plan.kgroup) == 0 || iszero(alphaT) - _scale_all_of_C!(plan, betaT, mr(plan.kernel), nr(plan.kernel), Qm, Qn) + _scale_all_of_C!(plan, betaT, tile_size(plan.kernel)..., m_length, n_length) return true end return false @@ -58,29 +59,29 @@ 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, - plan.blocking.kc, _is_split(plan.mpack), _is_split(plan.npack) + plan.blocking.k_block, _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. `kc`: the requested K block, `nothing` for +# `nothing` to assume it does. `k_block`: 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, kc, + mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, capacity, k_block, split_a::Bool = false, split_b::Bool = false ) where {T} - Qm = axis_length(mgroup) - Qn = axis_length(ngroup) - Qk = axis_length(kgroup) - panel = _c_panel_needed(T, Cstorage, Qk, kc) - if (Qm == 1 || Qn == 1) && !panel && _dot_applicable(T, Astorage, Bstorage, kgroup, Qm, Qn, Qk) && + m_length = axis_length(mgroup) + n_length = axis_length(ngroup) + k_length = axis_length(kgroup) + panel = _c_panel_needed(T, Cstorage, k_length, k_block) + if (m_length == 1 || n_length == 1) && !panel && _dot_applicable(T, Astorage, Bstorage, kgroup, m_length, n_length, k_length) && _dot_capacity_ok(capacity) W = _dot_lanewidth(T) - return Qm == 1 ? _lane_path(_DotPath{true}, W) : _lane_path(_DotPath{false}, W) + return m_length == 1 ? _lane_path(_DotPath{true}, W) : _lane_path(_DotPath{false}, W) end - Qk == 1 && _outer_applicable(T, Astorage, Cstorage, mgroup, Qm) && + k_length == 1 && _outer_applicable(T, Astorage, Cstorage, mgroup, m_length) && return _lane_path(_OuterPath, _dot_lanewidth(T)) return _nest_path( unpack_ok && _unpacked_b_rule(mgroup, kgroup), mgroup, ngroup, kgroup, split_a, split_b, panel @@ -90,9 +91,9 @@ 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} +@inline function _c_panel_needed(::Type{T}, Cstorage, k_length::Int, k_block) where {T} sizeof(real(eltype(Cstorage))) < sizeof(real(T)) || return false - return Qk > (kc === nothing ? _resolved_defaults(T).real_row.kc : kc) + return k_length > (k_block === nothing ? _resolved_defaults(T).real_row.k_block : k_block) end # Run `path` behind a dynamic call. The plan crosses in the workspace slot @@ -140,7 +141,7 @@ end @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.kc + req.k_block ) @inline _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} = @@ -176,9 +177,9 @@ end # 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`). +# default `k_block`; 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) && + _c_panel_needed(T, plan.Cstorage, axis_length(plan.kgroup), plan.blocking.k_block) === (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 @@ -200,33 +201,33 @@ function _execute_path!( # Panels and `PtrScatterAxis`es borrow pointers into `ws`. GC.@preserve ws begin _execute_nest!( - plan, ws, kernel, mr(kernel), nr(kernel), + plan, ws, kernel, tile_size(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, nothing + plan.blocking.m_block, plan.blocking.k_block, plan.blocking.n_block, 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 +# dense M/N maps, one N 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) + m_length = axis_length(plan.mgroup) pplan = ContractPlan( - kernel, _dense_second_map(plan.mgroup, 1), _dense_second_map(plan.ngroup, Qm), + kernel, _dense_second_map(plan.mgroup, 1), _dense_second_map(plan.ngroup, m_length), 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 + pplan, ws, kernel, tile_size(kernel)..., + m_length, axis_length(plan.ngroup), axis_length(plan.kgroup), + plan.blocking.m_block, plan.blocking.k_block, plan.blocking.n_block, alphaT, betaT, P(), plan ) end return nothing @@ -239,58 +240,58 @@ end 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 +@inline _panel_enter!(::Nothing, plan, n_block_start, n_block_length, betaT) = plan +@inline _panel_exit!(::Nothing, n_block_start, n_block_length) = nothing -function _panel_enter!(target::ContractPlan, plan::ContractPlan, jc::Int, nblock::Int, betaT) - iszero(betaT) || _panel_copy!(target, jc, nblock, true) +function _panel_enter!(target::ContractPlan, plan::ContractPlan, n_block_start::Int, n_block_length::Int, betaT) + iszero(betaT) || _panel_copy!(target, n_block_start, n_block_length, 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, + -n_block_start * 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) +_panel_exit!(target::ContractPlan, n_block_start::Int, n_block_length::Int) = _panel_copy!(target, n_block_start, n_block_length, 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) +# Columns `n_block_start .+ (0:n_block_length-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, n_block_start::Int, n_block_length::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 + m_length = axis_length(target.mgroup) + m_block = target.blocking.m_block + fill_offsets!((ws.n_buf_B, ws.n_buf_C), target.ngroup, n_block_start, n_block_length) + m_block_start = 0 + while m_block_start < m_length + m_block_length = min(m_block, m_length - m_block_start) + fill_offsets!((ws.m_buf_A, ws.m_buf_C), target.mgroup, m_block_start, m_block_length) + for j in 1:n_block_length, i in 1:m_block_length c = target.Cbase + ws.m_buf_C[i] + ws.n_buf_C[j] + 1 - q = ic + i + (j - 1) * Qm + q = m_block_start + i + (j - 1) * m_length if load panel[q] = C[c] else C[c] = panel[q] end end - ic += mblock + m_block_start += m_block_length 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, + plan::ContractPlan{T}, ws, kernel::K, m_tile::Int, n_tile::Int, + m_length::Int, n_length::Int, k_length::Int, m_block::Int, k_block::Int, n_block::Int, 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. - MRp = packed_a_per_k(kernel) - NRp = packed_b_per_k(kernel) + # `m_tile`/`n_tile` count register-tile rows. They differ for complex + # kernels. + a_sliver_width, b_sliver_width = sliver_widths(kernel) atransform = plan.atransform btransform = plan.btransform @@ -313,41 +314,41 @@ function _execute_nest!( aff_mA, aff_mC, aff_nB, aff_nC, aff_kA, aff_kB = map(Val, AFF) - # --- loop 5: jc over N in steps of nc_eff --- - jc = 0 - while jc < Qn - nblock = min(nc_eff, Qn - jc) - cplan = _panel_enter!(target, plan, jc, nblock, betaT) - n_slivers = cld(nblock, NRk) + # --- loop over N blocks --- + n_block_start = 0 + while n_block_start < n_length + n_block_length = min(n_block, n_length - n_block_start) + cplan = _panel_enter!(target, plan, n_block_start, n_block_length, betaT) + n_tiles = cld(n_block_length, n_tile) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( - ws.n_desc_B, ws.n_desc_C, n_step[1], n_step[2], jc, - nblock, NRk, n_slivers + ws.n_desc_B, ws.n_desc_C, n_step[1], n_step[2], n_block_start, + n_block_length, n_tile, n_tiles ) else - fill_offsets!((ws.n_buf_B, ws.n_buf_C), ngroup, jc, nblock) + fill_offsets!((ws.n_buf_B, ws.n_buf_C), ngroup, n_block_start, n_block_length) _classify_slivers!( ws.n_desc_B, ws.n_desc_C, ws.n_buf_B, ws.n_buf_C, - nblock, NRk, n_slivers + n_block_length, n_tile, n_tiles ) end - # --- loop 4: pc over K in steps of kc_eff --- - pc = 0 + # --- loop over K blocks --- + k_block_start = 0 firstpanel = true - while pc < Qk - kblock = min(kc_eff, Qk - pc) + while k_block_start < k_length + k_block_length = min(k_block, k_length - k_block_start) dK_A, dK_B, rng_kA, rng_kB = if k_ramp ( - _ramp_descriptor(k_step[1], pc, kblock), - _ramp_descriptor(k_step[2], pc, kblock), - _ramp_offset_range(k_step[1], pc, kblock), - _ramp_offset_range(k_step[2], pc, kblock), + _ramp_descriptor(k_step[1], k_block_start, k_block_length), + _ramp_descriptor(k_step[2], k_block_start, k_block_length), + _ramp_offset_range(k_step[1], k_block_start, k_block_length), + _ramp_offset_range(k_step[2], k_block_start, k_block_length), ) else - fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, pc, kblock) - dA = describe_block(ws.k_buf_A, 0, kblock) - dB = describe_block(ws.k_buf_B, 0, kblock) + fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, k_block_start, k_block_length) + dA = describe_block(ws.k_buf_A, 0, k_block_length) + dB = describe_block(ws.k_buf_B, 0, k_block_length) ( dA, dB, descriptor_offset_range(dA, ws.k_buf_A, 0), @@ -357,7 +358,7 @@ function _execute_nest!( colsA_k = _axis_of(dK_A, ws.k_buf_A, 0, aff_kA) rowsB_k = _axis_of(dK_B, ws.k_buf_B, 0, aff_kB) - # Hoisted bounds checks (B here, A and C per ic block): each + # Hoisted bounds checks (B here, A and C per M block): each # rectangle is exactly the union of the per-sliver/per-tile # regions, so the `unsafe_*` calls below read nothing unchecked. checked_span_bounds(plan.Bbase, rng_kB, rng_nB, lenB) @@ -367,15 +368,15 @@ function _execute_nest!( 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, + packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel)[2]), b_sliver_width, + plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, 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) - bpanel = _sliver_panel(ws.packed_b, NRp, kblock, s) + for n_tile_index in 0:(n_tiles - 1) + n_tile_start = n_tile_index * n_tile + colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) + bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) _pack_sliver!( unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, kernel, btransform @@ -383,21 +384,21 @@ function _execute_nest!( end end - # --- loop 3: ic over M in steps of mc_eff --- - ic = 0 - while ic < Qm - mblock = min(mc_eff, Qm - ic) - m_slivers = cld(mblock, MRk) + # --- loop over M blocks --- + m_block_start = 0 + while m_block_start < m_length + m_block_length = min(m_block, m_length - m_block_start) + m_tiles = cld(m_block_length, m_tile) (rng_mA, rng_mC) = if m_ramp _ramp_slivers!( - ws.m_desc_A, ws.m_desc_C, m_step[1], m_step[2], ic, - mblock, MRk, m_slivers + ws.m_desc_A, ws.m_desc_C, m_step[1], m_step[2], m_block_start, + m_block_length, m_tile, m_tiles ) else - fill_offsets!((ws.m_buf_A, ws.m_buf_C), mgroup, ic, mblock) + fill_offsets!((ws.m_buf_A, ws.m_buf_C), mgroup, m_block_start, m_block_length) _classify_slivers!( ws.m_desc_A, ws.m_desc_C, ws.m_buf_A, ws.m_buf_C, - mblock, MRk, m_slivers + m_block_length, m_tile, m_tiles ) end @@ -407,15 +408,15 @@ function _execute_nest!( 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, + packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel)[1]), a_sliver_width, + plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, 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) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) _pack_sliver!( unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, kernel, atransform @@ -423,52 +424,52 @@ function _execute_nest!( end end - # --- loop 2: jr over N-slivers; loop 1: ir over M-slivers --- + # --- loops over N tiles and M tiles --- if UNPACKED_B _micro_tiles_unpacked_b!( - kernel, cplan, ws, rowsB_k, m_slivers, n_slivers, - MRk, NRk, MRp, kblock, alphaT, beta_eff, aff_mC, aff_nC + kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC ) else _micro_tiles_packed_b!( - kernel, cplan, ws, m_slivers, n_slivers, - MRk, NRk, MRp, NRp, kblock, alphaT, beta_eff, aff_mC, aff_nC + kernel, cplan, ws, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, b_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC ) end - ic += mblock + m_block_start += m_block_length end firstpanel = false - pc += kblock + k_block_start += k_block_length end - _panel_exit!(target, jc, nblock) - jc += nblock + _panel_exit!(target, n_block_start, n_block_length) + n_block_start += n_block_length end return plan.Cstorage end -# Loops 2/1 over one (jc, pc, ic) block. `@noinline`: one call per block, and +# The tile loops over one (N, K, M) block. `@noinline`: one call per block, and # the inlined microkernel is most of a nest's size; compile cost grows faster # than linearly with function size. @noinline function _micro_tiles_packed_b!( - kernel::K, plan::ContractPlan, ws, m_slivers::Int, n_slivers::Int, - MRk::Int, NRk::Int, MRp::Int, NRp::Int, kblock::Int, alphaT, beta_eff, + kernel::K, plan::ContractPlan, ws, m_tiles::Int, n_tiles::Int, + m_tile::Int, n_tile::Int, a_sliver_width::Int, b_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, aff_mC::Val{MC}, aff_nC::Val{NC} ) where {K, MC, NC} - for s in 0:(n_slivers - 1) - sfirst = s * NRk - colsC = _axis_of(ws.n_desc_C[s + 1], ws.n_buf_C, sfirst, aff_nC) - bpanel = _sliver_panel(ws.packed_b, NRp, kblock, s) - for r in 0:(m_slivers - 1) - rfirst = r * MRk - rowsC = _axis_of(ws.m_desc_C[r + 1], ws.m_buf_C, rfirst, aff_mC) - apanel = _sliver_panel(ws.packed_a, MRp, kblock, r) + for n_tile_index in 0:(n_tiles - 1) + n_tile_start = n_tile_index * n_tile + colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) + bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) unsafe_execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bpanel, kblock, alphaT, beta_eff + apanel, bpanel, k_block_length, alphaT, beta_eff ) end end diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index 8fc0f87..b4893bb 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -34,20 +34,20 @@ end @inline function _execute_micro_tile!( kernel, storage::S, base::Int, rows::R, cols::C, - packed_a::PA, packed_b::PB, kc_len::Int, alpha, beta + packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta ) where {PA, PB, S, R <: Axis, C <: Axis} destination = DestinationTile(storage, base, rows, cols) - execute_tile!(kernel, destination, packed_a, packed_b, kc_len, alpha, beta) + execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) return nothing end # The caller has bounds-checked the whole macro block this tile belongs to. @inline function unsafe_execute_micro_tile!( kernel, storage::S, base::Int, rows::R, cols::C, - packed_a::PA, packed_b::PB, kc_len::Int, alpha, beta + packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta ) where {PA, PB, S, R <: Axis, C <: Axis} destination = DestinationTile(storage, base, rows, cols) - unsafe_execute_tile!(kernel, destination, packed_a, packed_b, kc_len, alpha, beta) + unsafe_execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) return nothing end @@ -59,13 +59,13 @@ end return nothing end -# Sliver `s` of a packed panel at the current block's depth `kc_len`, shared by -# the packing and the consuming step so the two cannot disagree. GUARDRAIL: -# `reg_tile` counts reals per K step (`packed_a_per_k`), not `mr`; they differ -# for complex kernels. -@inline function _sliver_panel(buffer, reg_tile::Int, kc_len::Int, s::Int) - stride = reg_tile * kc_len - return packed_panel(buffer, s * stride + 1, stride) +# Sliver `tile_index` of a packed panel at the current block's depth +# `k_block_length`, shared by the packing and the consuming step so the two +# cannot disagree. GUARDRAIL: `sliver_width` counts reals per K step +# (`sliver_widths`), not the tile size; they differ for complex kernels. +@inline function _sliver_panel(buffer, sliver_width::Int, k_block_length::Int, tile_index::Int) + stride = sliver_width * k_block_length + return packed_panel(buffer, tile_index * stride + 1, stride) end # Classify each register sliver of a just-filled macro block, and return the @@ -74,22 +74,22 @@ end @inline function _classify_slivers!( desc1::Vector{BlockDescriptor}, desc2::Vector{BlockDescriptor}, buf1::Vector{Int}, buf2::Vector{Int}, - blocklen::Int, reg_tile::Int, nslivers::Int + block_length::Int, tile::Int, tile_count::Int ) lo1 = typemax(Int); hi1 = typemin(Int) lo2 = typemax(Int); hi2 = typemin(Int) - for s in 0:(nslivers - 1) - sfirst = s * reg_tile - scount = min(reg_tile, blocklen - sfirst) - d1 = describe_block(buf1, sfirst, scount) - d2 = describe_block(buf2, sfirst, scount) - desc1[s + 1] = d1 - desc2[s + 1] = d2 - (l1, h1) = descriptor_offset_range(d1, buf1, sfirst) + for tile_index in 0:(tile_count - 1) + tile_start = tile_index * tile + tile_length = min(tile, block_length - tile_start) + d1 = describe_block(buf1, tile_start, tile_length) + d2 = describe_block(buf2, tile_start, tile_length) + desc1[tile_index + 1] = d1 + desc2[tile_index + 1] = d2 + (l1, h1) = descriptor_offset_range(d1, buf1, tile_start) if h1 >= l1 lo1 = min(lo1, l1); hi1 = max(hi1, h1) end - (l2, h2) = descriptor_offset_range(d2, buf2, sfirst) + (l2, h2) = descriptor_offset_range(d2, buf2, tile_start) if h2 >= l2 lo2 = min(lo2, l2); hi2 = max(hi2, h2) end @@ -116,17 +116,17 @@ end @inline function _ramp_slivers!( desc1::Vector{BlockDescriptor}, desc2::Vector{BlockDescriptor}, step1::Int, step2::Int, first::Int, - blocklen::Int, reg_tile::Int, nslivers::Int + block_length::Int, tile::Int, tile_count::Int ) - for s in 0:(nslivers - 1) - sfirst = s * reg_tile - scount = min(reg_tile, blocklen - sfirst) - q0 = first + sfirst - desc1[s + 1] = _ramp_descriptor(step1, q0, scount) - desc2[s + 1] = _ramp_descriptor(step2, q0, scount) + for tile_index in 0:(tile_count - 1) + tile_start = tile_index * tile + tile_length = min(tile, block_length - tile_start) + q0 = first + tile_start + desc1[tile_index + 1] = _ramp_descriptor(step1, q0, tile_length) + desc2[tile_index + 1] = _ramp_descriptor(step2, q0, tile_length) end return ( - _ramp_offset_range(step1, first, blocklen), - _ramp_offset_range(step2, first, blocklen), + _ramp_offset_range(step1, first, block_length), + _ramp_offset_range(step2, first, block_length), ) end diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl index 2ddf32f..533588f 100644 --- a/src/execution/oracle.jl +++ b/src/execution/oracle.jl @@ -1,8 +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`. Unlike `execute!`, rounds a C narrower -# than the compute type once per `kc` block. +# `k_block` 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`. Unlike `execute!`, rounds +# a C narrower than the compute type once per K 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`. @@ -17,44 +17,43 @@ function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) w betaT = convert(T, beta) kernel = plan.kernel - MRk = mr(kernel) - NRk = nr(kernel) + m_tile, n_tile = tile_size(kernel) - Qm = axis_length(plan.mgroup) - Qn = axis_length(plan.ngroup) - Qk = axis_length(plan.kgroup) + m_length = axis_length(plan.mgroup) + n_length = axis_length(plan.ngroup) + k_length = axis_length(plan.kgroup) - (Qm == 0 || Qn == 0) && return plan.Cstorage + (m_length == 0 || n_length == 0) && return plan.Cstorage - if Qk == 0 || iszero(alphaT) - _scale_all_of_C!(plan, betaT, MRk, NRk, Qm, Qn) + if k_length == 0 || iszero(alphaT) + _scale_all_of_C!(plan, betaT, m_tile, n_tile, m_length, n_length) return plan.Cstorage end m_bufs = (ws.tile_m_buf_A, ws.tile_m_buf_C) n_bufs = (ws.tile_n_buf_B, ws.tile_n_buf_C) k_bufs = (ws.tw_k_buf_A, ws.tw_k_buf_B) - kc_panel = plan.blocking.kc + k_block = plan.blocking.k_block - mfirst = 0 - while mfirst < Qm - mcount = min(MRk, Qm - mfirst) - (dM_A, dM_C) = block_descriptors!(m_bufs, plan.mgroup, mfirst, mcount) + m_tile_start = 0 + while m_tile_start < m_length + m_tile_length = min(m_tile, m_length - m_tile_start) + (dM_A, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) rowsA = _axis_of(dM_A, ws.tile_m_buf_A, 0) rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) - nfirst = 0 - while nfirst < Qn - ncount = min(NRk, Qn - nfirst) - (dN_B, dN_C) = block_descriptors!(n_bufs, plan.ngroup, nfirst, ncount) + n_tile_start = 0 + while n_tile_start < n_length + n_tile_length = min(n_tile, n_length - n_tile_start) + (dN_B, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) colsB = _axis_of(dN_B, ws.tile_n_buf_B, 0) colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) - kfirst = 0 + k_block_start = 0 firstpanel = true - while kfirst < Qk - kcount = min(kc_panel, Qk - kfirst) - (dK_A, dK_B) = block_descriptors!(k_bufs, plan.kgroup, kfirst, kcount) + while k_block_start < k_length + k_block_length = min(k_block, k_length - k_block_start) + (dK_A, dK_B) = block_descriptors!(k_bufs, plan.kgroup, k_block_start, k_block_length) colsK_A = _axis_of(dK_A, ws.tw_k_buf_A, 0) rowsK_B = _axis_of(dK_B, ws.tw_k_buf_B, 0) @@ -71,16 +70,16 @@ function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) w beta_eff = firstpanel ? betaT : one(T) _execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - ws.tw_packed_a, ws.tw_packed_b, kcount, alphaT, beta_eff + ws.tw_packed_a, ws.tw_packed_b, k_block_length, alphaT, beta_eff ) firstpanel = false - kfirst += kcount + k_block_start += k_block_length end - nfirst += ncount + n_tile_start += n_tile_length end - mfirst += mcount + m_tile_start += m_tile_length end return plan.Cstorage diff --git a/src/execution/outer.jl b/src/execution/outer.jl index 7df91f4..80f7e91 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -10,17 +10,17 @@ const _OUTER_MODE = Ref{Symbol}(:auto) # N and B may have any layout: B is read one scalar per column. -function _outer_applicable(::Type{T}, Astorage, Cstorage, mgroup::AxisGroup, Qm::Int) where {T} +function _outer_applicable(::Type{T}, Astorage, Cstorage, mgroup::AxisGroup, m_length::Int) where {T} _OUTER_MODE[] === :never && return false T <: Real || return false (Astorage isa DenseVector{T} && Cstorage isa DenseVector{T}) || return false - Qm >= _dot_lanewidth(T) || return false + m_length >= _dot_lanewidth(T) || return false _map_ramp_step(mgroup, 1) == 1 || return false return _map_ramp_step(mgroup, 2) == 1 end function _execute_outer!( - plan::ContractPlan{T}, alphaT::T, betaT::T, Qm::Int, Qn::Int, ::Val{W} + plan::ContractPlan{T}, alphaT::T, betaT::T, m_length::Int, n_length::Int, ::Val{W} ) where {T, W} ws = plan.workspace Astorage = plan.Astorage @@ -29,26 +29,26 @@ function _execute_outer!( Bbase = plan.Bbase bufB = ws.n_buf_B bufC = ws.n_buf_C - nc_eff = plan.blocking.nc + n_block = plan.blocking.n_block lenB = length(Bstorage) lenC = length(Cstorage) - checked_span_bounds(plan.Abase, (0, Qm - 1), (0, 0), length(Astorage)) + checked_span_bounds(plan.Abase, (0, m_length - 1), (0, 0), length(Astorage)) GC.@preserve ws Astorage Cstorage begin ap = pointer(Astorage) + sizeof(T) * plan.Abase cp = pointer(Cstorage) + sizeof(T) * plan.Cbase - jc = 0 - while jc < Qn - nblock = min(nc_eff, Qn - jc) - fill_offsets!((bufB, bufC), plan.ngroup, jc, nblock) - checked_span_bounds(Bbase, _buffer_range(bufB, nblock), (0, 0), lenB) - checked_span_bounds(plan.Cbase, (0, Qm - 1), _buffer_range(bufC, nblock), lenC) + n_block_start = 0 + while n_block_start < n_length + n_block_length = min(n_block, n_length - n_block_start) + fill_offsets!((bufB, bufC), plan.ngroup, n_block_start, n_block_length) + checked_span_bounds(Bbase, _buffer_range(bufB, n_block_length), (0, 0), lenB) + checked_span_bounds(plan.Cbase, (0, m_length - 1), _buffer_range(bufC, n_block_length), lenC) _outer_block!( - cp, bufC, ap, Bstorage, Bbase, bufB, nblock, Qm, + cp, bufC, ap, Bstorage, Bbase, bufB, n_block_length, m_length, alphaT, betaT, plan.btransform, Val(W) ) - jc += nblock + n_block_start += n_block_length end end return plan.Cstorage @@ -58,11 +58,11 @@ end # chosen once per block. @inline function _outer_block!( cp::Ptr{T}, bufC::Vector{Int}, ap::Ptr{T}, Bstorage::SB, Bbase::Int, - bufB::Vector{Int}, nblock::Int, Qm::Int, alpha::T, beta::T, btransform::F, ::Val{W} + bufB::Vector{Int}, n_block_length::Int, m_length::Int, alpha::T, beta::T, btransform::F, ::Val{W} ) where {T, SB, F, W} sz = sizeof(T) - mmain = (Qm ÷ W) * W - @inbounds for n in 1:nblock + mmain = (m_length ÷ W) * W + @inbounds for n in 1:n_block_length b = convert(T, btransform(Bstorage[Bbase + bufB[n] + 1])) bv = Vec{W, T}(b) cpn = cp + sz * bufC[n] @@ -73,7 +73,7 @@ end vstore(alpha * r, cpn + sz * m) m += W end - while m < Qm + while m < m_length r = unsafe_load(ap + sz * m) * b unsafe_store!(cpn + sz * m, alpha * r) m += 1 @@ -84,7 +84,7 @@ end vstore(muladd(alpha, r, vload(Vec{W, T}, cpn + sz * m)), cpn + sz * m) m += W end - while m < Qm + while m < m_length r = unsafe_load(ap + sz * m) * b unsafe_store!(cpn + sz * m, muladd(alpha, r, unsafe_load(cpn + sz * m))) m += 1 @@ -96,7 +96,7 @@ end vstore(muladd(alpha, r, betav * vload(Vec{W, T}, cpn + sz * m)), cpn + sz * m) m += W end - while m < Qm + while m < m_length r = unsafe_load(ap + sz * m) * b unsafe_store!(cpn + sz * m, muladd(alpha, r, beta * unsafe_load(cpn + sz * m))) m += 1 diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 952a84b..119bc3f 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -2,7 +2,7 @@ # panel. The kernels consume B one (complex) scalar at a time and broadcast # it, so any K axis, column offsets and `AbstractVector` storage work. # Packing B buys locality and padding columns at an O(N*K) gather per -# (jc, pc) block, which at small M costs as much as the kernel itself. +# (N, K) block, which at small M costs as much as the kernel itself. # Stands in for a packed B micro-panel of `NR` columns: column `j`, K step # `p` is storage address `colbase[j+1] + axis_offset(ksteps, p)` (base @@ -52,7 +52,7 @@ const _UNPACKED_B_MODE = Ref{Symbol}(:auto) # Whether B is read in place (given an eligible kernel): small M, and every B # column contiguous along K. Packing pays off only through reuse across M -# slivers, while `nr` contiguous columns read in place cost the kernel +# slivers, while `n_tile` contiguous columns read in place cost the kernel # nothing. A large K stride in B (each step its own cache line, and for a # power of two only a few L1 sets) is what packing exists for. @inline function _unpacked_b_rule(mgroup::AxisGroup, kgroup::AxisGroup) @@ -67,46 +67,46 @@ end # Beyond this M the packed B's reuse wins. const _UNPACKED_B_MMAX = 256 -# The view for one N sliver; `buf`/`sfirst` locate an irregular sliver's +# The view for one N sliver; `buf`/`n_tile_start` locate an irregular sliver's # offsets in `ws.n_buf_B`. @inline function _unpacked_b_view( kernel::DescriptorKernel{MR, NR, T}, storage::S, base::Int, - d::BlockDescriptor, buf::Vector{Int}, sfirst::Int, ksteps::K, transform::F + d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F ) where {MR, NR, T, S, K <: Axis, F} last = d.count - 1 colbase = ntuple(Val(NR)) do j1 jj = min(j1 - 1, last) - off = d.regular ? d.base + jj * d.stride : (@inbounds buf[sfirst + jj + 1]) + off = d.regular ? d.base + jj * d.stride : (@inbounds buf[n_tile_start + jj + 1]) base + off end - return UnpackedBView(storage, colbase, ksteps, packed_b_per_k(kernel), transform) + return UnpackedBView(storage, colbase, ksteps, sliver_widths(kernel)[2], transform) end -# Loops 2/1 of `_execute_nest!` with B read in place. A barrier over the K +# The tile loops of `_execute_nest!` with B read in place. A barrier over the K # axis type, so each view is concretely typed. `@noinline` so that the # packed-B nest does not grow by this body. The view's addresses are within # the caller's B check. @noinline function _micro_tiles_unpacked_b!( - kernel::K, plan::ContractPlan, ws, rowsB_k::KA, m_slivers::Int, n_slivers::Int, - MRk::Int, NRk::Int, MRp::Int, kblock::Int, alphaT, beta_eff, + kernel::K, plan::ContractPlan, ws, rowsB_k::KA, m_tiles::Int, n_tiles::Int, + m_tile::Int, n_tile::Int, a_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, aff_mC::Val{MC}, aff_nC::Val{NC} ) where {K, KA <: Axis, MC, NC} Bstorage = plan.Bstorage Bbase = plan.Bbase btransform = plan.btransform - for s in 0:(n_slivers - 1) - sfirst = s * NRk - colsC = _axis_of(ws.n_desc_C[s + 1], ws.n_buf_C, sfirst, aff_nC) + for n_tile_index in 0:(n_tiles - 1) + n_tile_start = n_tile_index * n_tile + colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) bview = _unpacked_b_view( - kernel, Bstorage, Bbase, ws.n_desc_B[s + 1], ws.n_buf_B, sfirst, rowsB_k, btransform + kernel, Bstorage, Bbase, ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, rowsB_k, btransform ) - for r in 0:(m_slivers - 1) - rfirst = r * MRk - rowsC = _axis_of(ws.m_desc_C[r + 1], ws.m_buf_C, rfirst, aff_mC) - apanel = _sliver_panel(ws.packed_a, MRp, kblock, r) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) unsafe_execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bview, kblock, alphaT, beta_eff + apanel, bview, k_block_length, alphaT, beta_eff ) end end diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 9080ab6..698571a 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -86,20 +86,20 @@ end # Buffer lengths for `kernel` at the effective `blocking`, shared by the # constructors and `reserve!`. GUARDRAIL: `packed_a_length` already counts -# reals, the sliver counts are in logical `mr`/`nr`; do not rescale either. +# reals, the sliver counts are in logical `m_tile`/`n_tile`; do not rescale +# either. @inline function _workspace_sizes(kernel, blocking::Blocking) - MRk = mr(kernel) - NRk = nr(kernel) - kc = blocking.kc - pa = packed_a_length(kernel, kc) - pb = packed_b_length(kernel, kc) - m_slivers = cld(blocking.mc, MRk) - n_slivers = cld(blocking.nc, NRk) + m_tile, n_tile = tile_size(kernel) + k_block = blocking.k_block + pa = packed_a_length(kernel, k_block) + pb = packed_b_length(kernel, k_block) + m_tiles = cld(blocking.m_block, m_tile) + n_tiles = cld(blocking.n_block, n_tile) return ( - mc = blocking.mc, nc = blocking.nc, kc = kc, - mr = MRk, nr = NRk, - m_slivers = m_slivers, n_slivers = n_slivers, - packed_a = m_slivers * pa, packed_b = n_slivers * pb, + m_block = blocking.m_block, n_block = blocking.n_block, k_block = k_block, + m_tile = m_tile, n_tile = n_tile, + m_tiles = m_tiles, n_tiles = n_tiles, + packed_a = m_tiles * pa, packed_b = n_tiles * pb, tw_packed_a = pa, tw_packed_b = pb, ) end @@ -122,14 +122,14 @@ end 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), - _alloc_descriptors(s.m_slivers), _alloc_descriptors(s.m_slivers), - _alloc_descriptors(s.n_slivers), _alloc_descriptors(s.n_slivers), + ints(s.m_block), ints(s.m_block), + ints(s.n_block), ints(s.n_block), + ints(s.k_block), ints(s.k_block), + _alloc_descriptors(s.m_tiles), _alloc_descriptors(s.m_tiles), + _alloc_descriptors(s.n_tiles), _alloc_descriptors(s.n_tiles), packed_a, packed_b, - ints(s.mr), ints(s.mr), - ints(s.nr), ints(s.nr), + ints(s.m_tile), ints(s.m_tile), + ints(s.n_tile), ints(s.n_tile), ints(ntw), ints(ntw), tw_packed_a, tw_packed_b, c_panel, ) @@ -154,7 +154,7 @@ function ContractWorkspace( s = _workspace_sizes(kernel, blocking) R = realtype(kernel) return _build_workspace( - T, s, oracle ? s.kc : 0, n -> Vector{Int}(undef, n), + T, s, oracle ? s.k_block : 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{T}(undef, panel), @@ -175,7 +175,7 @@ function ContractWorkspace( c_panel = _alloc_temp(T, panel, allocator) return _build_workspace( - T, s, oracle ? s.kc : 0, n -> _alloc_offsets(n, allocator), + T, s, oracle ? s.k_block : 0, n -> _alloc_offsets(n, allocator), packed_a, packed_b, tw_packed_a, tw_packed_b, c_panel ) end @@ -196,30 +196,30 @@ function reserve!( ) where {T, R} s = _workspace_sizes(kernel, blocking) - _grow!(ws.m_buf_A, s.mc) - _grow!(ws.m_buf_C, s.mc) - _grow!(ws.n_buf_B, s.nc) - _grow!(ws.n_buf_C, s.nc) - _grow!(ws.k_buf_A, s.kc) - _grow!(ws.k_buf_B, s.kc) + _grow!(ws.m_buf_A, s.m_block) + _grow!(ws.m_buf_C, s.m_block) + _grow!(ws.n_buf_B, s.n_block) + _grow!(ws.n_buf_C, s.n_block) + _grow!(ws.k_buf_A, s.k_block) + _grow!(ws.k_buf_B, s.k_block) - _grow!(ws.m_desc_A, s.m_slivers) - _grow!(ws.m_desc_C, s.m_slivers) - _grow!(ws.n_desc_B, s.n_slivers) - _grow!(ws.n_desc_C, s.n_slivers) + _grow!(ws.m_desc_A, s.m_tiles) + _grow!(ws.m_desc_C, s.m_tiles) + _grow!(ws.n_desc_B, s.n_tiles) + _grow!(ws.n_desc_C, s.n_tiles) _grow!(ws.packed_a, s.packed_a) _grow!(ws.packed_b, s.packed_b) - _grow!(ws.tile_m_buf_A, s.mr) - _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.tile_m_buf_A, s.m_tile) + _grow!(ws.tile_m_buf_C, s.m_tile) + _grow!(ws.tile_n_buf_B, s.n_tile) + _grow!(ws.tile_n_buf_C, s.n_tile) _grow!(ws.c_panel, panel) if oracle - _grow!(ws.tw_k_buf_A, s.kc) - _grow!(ws.tw_k_buf_B, s.kc) + _grow!(ws.tw_k_buf_A, s.k_block) + _grow!(ws.tw_k_buf_B, s.k_block) _grow!(ws.tw_packed_a, s.tw_packed_a) _grow!(ws.tw_packed_b, s.tw_packed_b) end diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 48ec904..19f7b26 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -155,11 +155,11 @@ end function Base.accumulate( kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, kc::Int + packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NA, PA, PB} - kc == 0 && return acc - kc > 0 || _throw_negative_kc(:accumulate, kc) - @inbounds for p in 0:(kc - 1) + k_block_length == 0 && return acc + k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + @inbounds for p in 0:(k_block_length - 1) acc = _accumulate_step_fmaddsub(kernel, acc, packed_a, packed_b, p) end return acc diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index f7ce9a7..3a4c5c7 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -1,7 +1,8 @@ # What every microkernel shares. Each kernel file implements only # `zero_accumulator`, `accumulate` and `store_tile!`, with one contract: -# * `accumulate(kernel, acc, packed_a, packed_b, kc)` adds `kc` K steps to -# `acc`; `kc == 0` returns `acc` without reading the panels. +# * `accumulate(kernel, acc, packed_a, packed_b, k_block_length)` adds +# `k_block_length` K steps to `acc`; `k_block_length == 0` returns `acc` +# without reading the panels. # * `store_tile!(dest, acc, alpha, beta, kernel)` writes `alpha*acc + beta*C` # over the valid rectangle only: `alpha == 0` never reads `acc`, `beta == 0` # never reads old `C`, padding lanes are never read. @@ -18,13 +19,13 @@ abstract type DescriptorKernel{MR, NR, T} end abstract type ComplexMethod end 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 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. +# Reals per packed A / B element. Blocking divides the real `m_block`/`n_block` +# by these, so every method gets the same packed byte budget. a_reals(::RealMethod) = 1 b_reals(::RealMethod) = 1 a_reals(::PlanarMethod) = 2 @@ -49,17 +50,15 @@ accumulator_planes(::RealComplexMethod) = 1 complex_method(::Any) = RealMethod() realtype(k::DescriptorKernel) = realtype(k.descriptor) -packed_a_per_k(k::DescriptorKernel) = packed_a_per_k(k.descriptor) -packed_b_per_k(k::DescriptorKernel) = packed_b_per_k(k.descriptor) +sliver_widths(k::DescriptorKernel) = sliver_widths(k.descriptor) a_format(k::DescriptorKernel) = a_format(k.descriptor) b_format(k::DescriptorKernel) = b_format(k.descriptor) -mr(k::DescriptorKernel) = mr(k.descriptor) -nr(k::DescriptorKernel) = nr(k.descriptor) +tile_size(k::DescriptorKernel) = tile_size(k.descriptor) scalartype(k::DescriptorKernel) = scalartype(k.descriptor) packed_a_offset(k::DescriptorKernel, i::Int, p::Int) = packed_a_offset(k.descriptor, i, p) packed_b_offset(k::DescriptorKernel, j::Int, p::Int) = packed_b_offset(k.descriptor, j, p) -packed_a_length(k::DescriptorKernel, kc::Int) = packed_a_length(k.descriptor, kc) -packed_b_length(k::DescriptorKernel, kc::Int) = packed_b_length(k.descriptor, kc) +packed_a_length(k::DescriptorKernel, k_block_length::Int) = packed_a_length(k.descriptor, k_block_length) +packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.descriptor, k_block_length) # Only these resolve for a complex descriptor; the single-plane offsets above # are a MethodError there, by design. @inline packed_a_plane_offset(k::DescriptorKernel, plane::Int, i::Int, p::Int) = @@ -162,19 +161,19 @@ end # GUARDRAIL: every throw reachable from the per-tile prologue sits behind a # `@noinline` helper. An inline interpolated message pulls `print_to_string`, a # GC frame and ~1 KB of stack into the hot tile function. -@noinline _throw_packed_short(which::Symbol, got::Int, need::Int, kc::Int) = throw( +@noinline _throw_packed_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( DimensionMismatch( "execute_tile!: packed_$which has length $got, " * - "need at least packed_$(which)_length(kernel, kc=$kc) = $need" + "need at least packed_$(which)_length(kernel, k_block_length=$k_block_length) = $need" ) ) @noinline _throw_tile_extent(which::Symbol, got::Int, limit::Int) = throw( ArgumentError( - "destination valid $which extent $got exceeds $(which === :row ? "mr" : "nr")(kernel) = $limit" + "destination valid $which extent $got exceeds tile_size(kernel)[$(which === :row ? 1 : 2)] = $limit" ) ) -@noinline _throw_negative_kc(where::Symbol, kc::Int) = - throw(ArgumentError("$where requires kc >= 0, got kc = $kc")) +@noinline _throw_negative_k_block_length(where::Symbol, k_block_length::Int) = + throw(ArgumentError("$where requires k_block_length >= 0, got k_block_length = $k_block_length")) # `execute_tile!`'s validation, in order. Returns `(run, alphaT, betaT)`; # `run == false` means the call is finished. `Val(false)` drops the storage @@ -182,13 +181,13 @@ end # GUARDRAIL: `@inline` and one bound type parameter per argument (hot path). @inline function _execute_tile_prologue!( kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, - kc::Int, alpha, beta, ::Val{BOUNDS} + k_block_length::Int, alpha, beta, ::Val{BOUNDS} ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB, BOUNDS} m = nrows(destination) n = ncols(destination) m <= MR || _throw_tile_extent(:row, m, MR) n <= NR || _throw_tile_extent(:column, n, NR) - kc >= 0 || _throw_negative_kc(:execute_tile!, kc) + k_block_length >= 0 || _throw_negative_k_block_length(:execute_tile!, k_block_length) alphaT = convert(T, alpha) betaT = convert(T, beta) @@ -197,15 +196,15 @@ end BOUNDS && checked_tile_storage_bounds(destination) - if kc == 0 || iszero(alphaT) + if k_block_length == 0 || iszero(alphaT) scale_tile!(destination, betaT) return (false, alphaT, betaT) end - need_a = packed_a_length(kernel, kc) - length(packed_a) >= need_a || _throw_packed_short(:a, length(packed_a), need_a, kc) - need_b = packed_b_length(kernel, kc) - length(packed_b) >= need_b || _throw_packed_short(:b, length(packed_b), need_b, kc) + need_a = packed_a_length(kernel, k_block_length) + length(packed_a) >= need_a || _throw_packed_short(:a, length(packed_a), need_a, k_block_length) + need_b = packed_b_length(kernel, k_block_length) + length(packed_b) >= need_b || _throw_packed_short(:b, length(packed_b), need_b, k_block_length) return (true, alphaT, betaT) end @@ -213,13 +212,13 @@ end # One checked K panel: `zero_accumulator`, `accumulate`, `store_tile!`. function execute_tile!( kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, - kc::Int, alpha, beta + k_block_length::Int, alpha, beta ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB} run, alphaT, betaT = _execute_tile_prologue!( - kernel, destination, packed_a, packed_b, kc, alpha, beta, Val(true) + kernel, destination, packed_a, packed_b, k_block_length, alpha, beta, Val(true) ) run || return destination - acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, kc) + acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, k_block_length) return store_tile!(destination, acc, alphaT, betaT, kernel) end @@ -230,12 +229,12 @@ end # range extremes. @inline function unsafe_execute_tile!( kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, - kc::Int, alpha, beta + k_block_length::Int, alpha, beta ) where {K, PA, PB} run, alphaT, betaT = _execute_tile_prologue!( - kernel, destination, packed_a, packed_b, kc, alpha, beta, Val(false) + kernel, destination, packed_a, packed_b, k_block_length, alpha, beta, Val(false) ) run || return destination - acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, kc) + acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, k_block_length) return store_tile!(destination, acc, alphaT, betaT, kernel) end diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index 0add24d..3a85395 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -1,10 +1,11 @@ # 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`. +# * complex A x real B: interleaved A is a real `2MR x k_block_length` +# panel; the accumulator is 1m/fmaddsub's interleaved layout. +# * real A x complex B: interleaved B is a real `k_block_length x 2NR` panel; +# accumulator column `2j` holds the real and `2j+1` the imaginary part of +# column `j`. using SIMD: Vec @@ -78,8 +79,8 @@ 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) +@inline Base.accumulate(kernel::_MixedKernel, acc::NTuple, packed_a, packed_b, k_block_length::Int) = + accumulate(kernel.inner, acc, packed_a, packed_b, k_block_length) # Inlined or not as the fmaddsub and planar stores each reuses. function store_tile!( diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index a1c8029..76a35db 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -1,10 +1,10 @@ # 1m (Van Zee's induced method): the real `SIMDKernel` of `2MR x NR`, run over -# `2*kc` real K steps against 1e-packed A and planar B. At logical K step `p`, -# `OneEFormat` A holds `(re_0, im_0, re_1, im_1, ...)` then -# `(-im_0, re_0, -im_1, re_1, ...)`, and planar B `re_0..` then `im_0..`, so -# accumulator real row `2i` is the real and `2i+1` the imaginary part of -# complex row `i`. There is deliberately no FMA loop here: 1m reuses the real -# kernel body verbatim. +# `2*k_block_length` real K steps against 1e-packed A and planar B. At logical K +# step `p`, `OneEFormat` A holds `(re_0, im_0, re_1, im_1, ...)` then `(-im_0, +# re_0, -im_1, re_1, ...)`, and planar B `re_0..` then `im_0..`, so accumulator +# real row `2i` is the real and `2i+1` the imaginary part of complex row `i`. +# There is deliberately no FMA loop here: 1m reuses the real kernel body +# verbatim. using SIMD: Vec @@ -53,11 +53,12 @@ zero_accumulator(kernel::OneMKernel) = zero_accumulator(kernel.inner) function Base.accumulate( kernel::OneMKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, R}}, - packed_a::PA, packed_b::PB, kc::Int + packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NV, PA, PB} - # Checked here so the message reports the logical kc, not the doubled one. - kc >= 0 || throw(ArgumentError("accumulate requires kc >= 0, got kc = $kc")) - return accumulate(kernel.inner, acc, packed_a, packed_b, 2 * kc) + # Checked here so the message reports the logical k_block_length, not the + # doubled one. + k_block_length >= 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) + return accumulate(kernel.inner, acc, packed_a, packed_b, 2 * k_block_length) end # Scalar store for an interleaved accumulator (1m and fmaddsub): with `W` even diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 401a4bb..d23793c 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -124,11 +124,11 @@ end function Base.accumulate( kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, kc::Int + packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NA, PA, PB} - kc == 0 && return acc - kc > 0 || _throw_negative_kc(:accumulate, kc) - @inbounds for p in 0:(kc - 1) + k_block_length == 0 && return acc + k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + @inbounds for p in 0:(k_block_length - 1) acc = _accumulate_step_planar(kernel, acc, packed_a, packed_b, p) end return acc diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl index f5b4bc2..b09dbae 100644 --- a/src/microkernels/scalar.jl +++ b/src/microkernels/scalar.jl @@ -16,11 +16,11 @@ zero_accumulator(kernel::ScalarKernel{MR, NR, T}) where {MR, NR, T} = zeros(T, M # Extends `Base.accumulate` so `using QuasiStrided` does not clash with it. function Base.accumulate( kernel::ScalarKernel{MR, NR, T}, acc::AbstractMatrix{T}, - packed_a::PA, packed_b::PB, kc::Int + packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, PA, PB} - kc == 0 && return acc - kc > 0 || throw(ArgumentError("accumulate requires kc >= 0, got kc = $kc")) - @inbounds for p in 0:(kc - 1) + k_block_length == 0 && return acc + k_block_length > 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) + @inbounds for p in 0:(k_block_length - 1) for j in 0:(NR - 1) bj = panel_load(packed_b, packed_b_offset(kernel, j, p)) for i in 0:(MR - 1) diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index 982bb08..f780226 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -79,11 +79,11 @@ end @inline function Base.accumulate( kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, - packed_a::PA, packed_b::PB, kc::Int + packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, NV, PA, PB} - kc == 0 && return acc - kc > 0 || _throw_negative_kc(:accumulate, kc) - @inbounds for p in 0:(kc - 1) + k_block_length == 0 && return acc + k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + @inbounds for p in 0:(k_block_length - 1) acc = _accumulate_step(kernel, acc, packed_a, packed_b, p) end return acc diff --git a/src/packing/format.jl b/src/packing/format.jl index 07d1a9f..bf8b1ba 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -2,7 +2,7 @@ # and the descriptor fixing a kernel's register tile, element type and panel # formats. Complex panels are "N planes of real(T)"; every format is addressed by # -# p * reg_tile * reals_per_element + plane * reg_tile + i +# p * MR * reals_per_element + plane * MR + i # # which at RealFormat (one real per element, plane 0) reduces to `i + MR*p`. @@ -25,7 +25,7 @@ reals_per_element(::InterleavedFormat) = 2 reals_per_element(::OneEFormat) = 4 # `MR`/`NR` are logical extents (complex elements for a complex `T`); every -# length and offset below takes the logical `kc` and counts reals. +# length and offset below takes the logical `k_block_length` and counts reals. struct Descriptor{MR, NR, T, FA <: PackFormat, FB <: PackFormat} function Descriptor{MR, NR, T, FA, FB}() where {MR, NR, T, FA, FB} _check_descriptor(MR, NR, T, FA, FB) @@ -71,8 +71,7 @@ function _check_descriptor(MR, NR, T, FA, FB) return nothing end -mr(::Descriptor{MR}) where {MR} = MR -nr(::Descriptor{MR, NR}) where {MR, NR} = NR +tile_size(::Descriptor{MR, NR}) where {MR, NR} = (MR, NR) scalartype(::Descriptor{MR, NR, T}) where {MR, NR, T} = T a_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FA() b_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FB() @@ -80,18 +79,18 @@ b_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FB() realtype(::Descriptor{MR, NR, T}) where {MR, NR, T} = real(T) # Reals per A / B sliver per logical K step. -packed_a_per_k(d::Descriptor{MR}) where {MR} = MR * reals_per_element(a_format(d)) -packed_b_per_k(d::Descriptor{MR, NR}) where {MR, NR} = NR * reals_per_element(b_format(d)) +sliver_widths(d::Descriptor{MR, NR}) where {MR, NR} = + (MR * reals_per_element(a_format(d)), NR * reals_per_element(b_format(d))) -packed_a_length(d::Descriptor, kc::Int) = packed_a_per_k(d) * kc -packed_b_length(d::Descriptor, kc::Int) = packed_b_per_k(d) * kc +packed_a_length(d::Descriptor, k_block_length::Int) = sliver_widths(d)[1] * k_block_length +packed_b_length(d::Descriptor, k_block_length::Int) = sliver_widths(d)[2] * k_block_length # Zero-based offsets, in reals. For 1e and interleaved, `i`/`j` runs over reals # (0:2MR-1) and 1e's second region is `plane == 2`. @inline packed_a_plane_offset(d::Descriptor{MR}, plane::Int, i::Int, p::Int) where {MR} = - p * packed_a_per_k(d) + plane * MR + i + p * sliver_widths(d)[1] + plane * MR + i @inline packed_b_plane_offset(d::Descriptor{MR, NR}, plane::Int, j::Int, p::Int) where {MR, NR} = - p * packed_b_per_k(d) + plane * NR + j + p * sliver_widths(d)[2] + plane * NR + j # Real descriptors only: a complex element has no single offset. packed_a_offset(kernel::KernelDescriptor{MR}, i::Int, p::Int) where {MR} = i + MR * p diff --git a/src/packing/pack.jl b/src/packing/pack.jl index dcfde8e..ac6ef23 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -15,13 +15,13 @@ end # Out of line so the packers carry no string formatting (and no GC frame). @noinline _throw_pack_extent(which::Symbol, name::Symbol, got::Int, limit::Int) = throw( ArgumentError( - "$(which)!: source $(name) count $got must satisfy 0 <= $got <= $(which === :pack_a ? "mr" : "nr")(kernel)=$limit" + "$(which)!: source $(name) count $got must satisfy 0 <= $got <= tile_size(kernel)[$(which === :pack_a ? 1 : 2)]=$limit" ) ) -@noinline _throw_pack_short(which::Symbol, got::Int, need::Int, kc::Int) = throw( +@noinline _throw_pack_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( DimensionMismatch( "$(which)!: packed buffer has length $got, " * - "need at least packed_$(which === :pack_a ? "a" : "b")_length(kernel, kc=$kc) = $need" + "need at least packed_$(which === :pack_a ? "a" : "b")_length(kernel, k_block_length=$k_block_length) = $need" ) ) @@ -65,15 +65,15 @@ end ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} _check_packed_eltype(packed, kernel) m = nrows(source) - kc = ncols(source) + k_block_length = ncols(source) (0 <= m <= MR) || _throw_pack_extent(:pack_a, :row, m, MR) - needed = packed_a_length(kernel, kc) - length(packed) >= needed || _throw_pack_short(:pack_a, length(packed), needed, kc) + needed = packed_a_length(kernel, k_block_length) + length(packed) >= needed || _throw_pack_short(:pack_a, length(packed), needed, k_block_length) # `<= 0`, not `== 0` (counts are never negative): LLVM may then assume - # `kc >= 1` in the loops below. - kc <= 0 && return packed + # `k_block_length >= 1` in the loops below. + k_block_length <= 0 && return packed BOUNDS && checked_tile_storage_bounds(source) - return _pack_a_sliver!(FA(), packed, source, kernel, transform, m, kc) + return _pack_a_sliver!(FA(), packed, source, kernel, transform, m, k_block_length) end @inline function _pack_b!( @@ -81,55 +81,55 @@ end transform::F, ::Val{BOUNDS} ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} _check_packed_eltype(packed, kernel) - kc = nrows(source) + k_block_length = nrows(source) n = ncols(source) (0 <= n <= NR) || _throw_pack_extent(:pack_b, :column, n, NR) - needed = packed_b_length(kernel, kc) - length(packed) >= needed || _throw_pack_short(:pack_b, length(packed), needed, kc) - kc <= 0 && return packed + needed = packed_b_length(kernel, k_block_length) + length(packed) >= needed || _throw_pack_short(:pack_b, length(packed), needed, k_block_length) + k_block_length <= 0 && return packed BOUNDS && checked_tile_storage_bounds(source) - return _pack_b_sliver!(FB(), packed, source, kernel, transform, n, kc) + return _pack_b_sliver!(FB(), packed, source, kernel, transform, n, k_block_length) end # A's packed index runs along `source.rows`, B's along `source.cols`; each # contiguous fast path needs that lane axis to be unit-stride. @inline function _pack_a_sliver!( format::FMT, packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2}, - transform::F, m::Int, kc::Int + transform::F, m::Int, k_block_length::Int ) where {FMT, V, MR, NR, T2, F} if format isa RealFormat 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) + return _pack_a_contiguous!(packed, source.storage, rowbase, source.cols, Val(MR), k_block_length) end elseif _pack_complex_contiguous_eligible( packed, source.storage, source.rows, transform, format, m, Val(MR), T2 ) elembase = source.base + source.rows.base return _pack_complex_contiguous!( - format, packed, source.storage, elembase, source.cols, Val(MR), kc, transform + format, packed, source.storage, elembase, source.cols, Val(MR), k_block_length, transform ) 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, _element_type(format, T2), format, Val(MR), kc, m, transform, load, plane_offset) + return _pack_panel!(packed, _element_type(format, T2), format, Val(MR), k_block_length, m, transform, load, plane_offset) end @inline function _pack_b_sliver!( format::FMT, packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2}, - transform::F, n::Int, kc::Int + transform::F, n::Int, k_block_length::Int ) where {FMT, V, MR, NR, T2, F} if _pack_complex_contiguous_eligible( packed, source.storage, source.cols, transform, format, n, Val(NR), T2 ) elembase = source.base + source.cols.base return _pack_complex_contiguous!( - format, packed, source.storage, elembase, source.rows, Val(NR), kc, transform + format, packed, source.storage, elembase, source.rows, Val(NR), k_block_length, transform ) 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, _element_type(format, T2), format, Val(NR), kc, n, transform, load, plane_offset) + return _pack_panel!(packed, _element_type(format, T2), format, Val(NR), k_block_length, n, transform, load, plane_offset) end # What a packed element converts to: the real operand of a mixed-domain kernel @@ -137,24 +137,24 @@ end @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 -# `t < valid ? load : zero` would compile to a branch around the load and keep -# the loop scalar. +# One sliver over `k_block_length` 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 `t < valid ? load : zero` would compile to a branch around the load +# and keep the loop scalar. @inline function _pack_panel!( - packed::V, ::Type{T}, format::FMT, ::Val{PD}, kc::Int, valid::Int, + packed::V, ::Type{T}, format::FMT, ::Val{PD}, k_block_length::Int, valid::Int, transform::F, load::L, plane_offset::P ) where {V, T, FMT <: PackFormat, PD, F, L, P} if valid == PD - @inbounds for p in 0:(kc - 1) + @inbounds for p in 0:(k_block_length - 1) for t in 0:(PD - 1) z = convert(T, transform(load(t, p)))::T _pack_emit!(packed, format, plane_offset, t, p, z) end end else - @inbounds for p in 0:(kc - 1) + @inbounds for p in 0:(k_block_length - 1) for t in 0:(valid - 1) z = convert(T, transform(load(t, p)))::T _pack_emit!(packed, format, plane_offset, t, p, z) diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index bbe739f..834ba40 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -32,12 +32,12 @@ end # Real A: one `Vec{MR}` load/convert/store per K step. @inline function _pack_a_contiguous!( packed::PackedPanel{T}, storage::DenseVector{S}, rowbase::Int, cols::C, - ::Val{MR}, kc::Int + ::Val{MR}, k_block_length::Int ) where {T, S, C, MR} GC.@preserve storage begin sp = pointer(storage) dp = packed.ptr - for p in 0:(kc - 1) + for p in 0:(k_block_length - 1) 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 @@ -115,14 +115,14 @@ _single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a # lanes convert to `R` before the shuffle. @inline function _pack_complex_contiguous!( format::Union{PlanarFormat, InterleavedFormat}, packed::PackedPanel{R}, - storage::DenseVector{Complex{RS}}, elembase::Int, steps::C, ::Val{PD}, kc::Int, + storage::DenseVector{Complex{RS}}, elembase::Int, steps::C, ::Val{PD}, k_block_length::Int, transform::F ) where {R, RS, C, PD, F} shuffle = _single_region_shuffle(format) GC.@preserve storage begin sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr - for p in 0:(kc - 1) + for p in 0:(k_block_length - 1) src = convert( Vec{2 * PD, R}, vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) @@ -138,12 +138,12 @@ end @inline function _pack_complex_contiguous!( ::OneEFormat, packed::PackedPanel{R}, storage::DenseVector{Complex{RS}}, - elembase::Int, steps::C, ::Val{PD}, kc::Int, transform::F + elembase::Int, steps::C, ::Val{PD}, k_block_length::Int, transform::F ) where {R, RS, C, PD, F} GC.@preserve storage begin sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr - for p in 0:(kc - 1) + for p in 0:(k_block_length - 1) src = convert( Vec{2 * PD, R}, vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) diff --git a/src/packing/transposed.jl b/src/packing/transposed.jl index a5f84e1..c5575f6 100644 --- a/src/packing/transposed.jl +++ b/src/packing/transposed.jl @@ -1,12 +1,13 @@ -# 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. +# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated by +# a split group, whole `E x L` groups) by `k_block_length` 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 + fcount::Int, k_block_length::Int, split ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: Axis, F} E, L = Int(split.E), Int(split.L) G = E * L @@ -14,17 +15,17 @@ # 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 + plane_offset, format, buffer, descriptor, Val(R), Rp * k_block_length, 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) + for g0 in 0:G:(fcount - 1), y1 in 0:(E - 1), p in 0:(k_block_length - 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) + for g0 in 0:G:(fcount - 1), p in 0:(k_block_length - 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) @@ -33,9 +34,9 @@ end valid = fcount % R if valid != 0 - rb = (fcount ÷ R) * Rp * kc + rb = (fcount ÷ R) * Rp * k_block_length pad = (plane, idx, pp) -> rb + plane_offset(descriptor, plane, idx, pp) - for p in 0:(kc - 1), t in valid:(R - 1) + for p in 0:(k_block_length - 1), t in valid:(R - 1) _pack_emit_zero!(buffer, format, pad, t, p, realtype(descriptor)) end end diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index dcfe971..6210af8 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -1,29 +1,29 @@ """ - Blocking(mc::Int, kc::Int, nc::Int) + Blocking(m_block::Int, k_block::Int, n_block::Int) -Cache-blocking factors for the macro-blocking driver: `mc` (M block extent, -loop 3), `kc` (K block/panel depth, loop 4), `nc` (N block extent, loop 5). +Cache-blocking factors for the macro-blocking driver: `m_block` (M block +extent), `k_block` (K block/panel depth) and `n_block` (N block extent). All three fields must be `>= 1`; the constructor throws `ArgumentError` otherwise. """ struct Blocking - mc::Int - kc::Int - nc::Int + m_block::Int + k_block::Int + n_block::Int - function Blocking(mc::Int, kc::Int, nc::Int) - mc >= 1 || throw(ArgumentError("Blocking requires mc >= 1, got mc = $mc")) - kc >= 1 || throw(ArgumentError("Blocking requires kc >= 1, got kc = $kc")) - nc >= 1 || throw(ArgumentError("Blocking requires nc >= 1, got nc = $nc")) - return new(mc, kc, nc) + function Blocking(m_block::Int, k_block::Int, n_block::Int) + m_block >= 1 || throw(ArgumentError("Blocking requires m_block >= 1, got m_block = $m_block")) + k_block >= 1 || throw(ArgumentError("Blocking requires k_block >= 1, got k_block = $k_block")) + n_block >= 1 || throw(ArgumentError("Blocking requires n_block >= 1, got n_block = $n_block")) + return new(m_block, k_block, n_block) end end # Complex blocking is the real row divided by the packed reals per element of -# each operand, so every method gets the same packed BYTE budget (1m's `mc` is -# half planar's). +# each operand, so every method gets the same packed BYTE budget (1m's +# `m_block` is half planar's). @inline _scale_blocking(base::Blocking, m::ComplexMethod) = - Blocking(max(1, base.mc ÷ a_reals(m)), base.kc, max(1, base.nc ÷ b_reals(m))) + Blocking(max(1, base.m_block ÷ a_reals(m)), base.k_block, max(1, base.n_block ÷ b_reals(m))) # For a host whose L1d or L2 size is undetected. _fallback_blocking(::Type{Float64}) = Blocking(128, 256, 768) @@ -32,17 +32,19 @@ _fallback_blocking(::Type{Float32}) = Blocking(96, 768, 1152) _real_blocking_row(profile::TargetProfile, ::Type{T}) where {T <: Union{Float32, Float64}} = something(_modelled_blocking(profile, T), _fallback_blocking(T)) -# Analytical real row from the cache geometry. Per (jc, pc) the packed B panel -# is `nc*kc` elements, per (jc, pc, ic) the A block `mc*kc`, and one `NR x kc` -# B sliver is reused by every A sliver of the block: +# Analytical real row from the cache geometry. Per (N, K) block the packed B +# panel is `n_block*k_block` elements, per (N, K, M) block the A block +# `m_block*k_block`, and one `NR x k_block` B sliver is reused by every A +# sliver of the block: # -# NR*kc*S <= L1/2 the reused B sliver -# mc*kc*S <= L2core/2 the A block -# nc*kc*S <= L2core + L3core the B panel +# NR*k_block*S <= L1/2 the reused B sliver +# m_block*k_block*S <= L2core/2 the A block +# n_block*k_block*S <= L2core + L3core the B panel # # rounded down to MR/NR multiples, shared levels divided by the cores sharing -# them. `nc` takes the whole per-core capacity, not half the L3: an oversized B -# panel only re-streams B, an undersized one repacks A once per `jc` block. +# them. `n_block` takes the whole per-core capacity, not half the L3: an +# oversized B panel only re-streams B, an undersized one repacks A once per N +# block. # `nothing` when L1d or L2 is undetected. function _modelled_blocking(profile::TargetProfile, ::Type{T}, MR::Int, NR::Int) where {T} l1, l2, l3 = profile.l1d, profile.l2, profile.l3 @@ -52,10 +54,10 @@ function _modelled_blocking(profile::TargetProfile, ::Type{T}, MR::Int, NR::Int) l2core = core_share(l2) l3core = l3.bytes > 0 ? core_share(l3) : 0 S = sizeof(T) - kc = max(1, (l1.bytes ÷ 2) ÷ (NR * S)) - mc = max(MR, ((l2core ÷ 2) ÷ (kc * S)) ÷ MR * MR) - nc = max(NR, ((l2core + l3core) ÷ (kc * S)) ÷ NR * NR) - return Blocking(mc, kc, nc) + k_block = max(1, (l1.bytes ÷ 2) ÷ (NR * S)) + m_block = max(MR, ((l2core ÷ 2) ÷ (k_block * S)) ÷ MR * MR) + n_block = max(NR, ((l2core + l3core) ÷ (k_block * S)) ÷ NR * NR) + return Blocking(m_block, k_block, n_block) end function _modelled_blocking(profile::TargetProfile, ::Type{T}) where {T <: Real} diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index 8ac74f0..32450b0 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -59,37 +59,37 @@ end _default_kernel(::Type{T}) where {T} = _kernel_from_shape(_resolved_defaults(T).shape, T, _default_method(T)) -# The automatic `(shape, method)` for extents `Qm`/`Qn` and C's unit-stride run -# along M (`run = Qm`: no layout known). Returned as plain values so -# `plan_contract` never holds a menu-wide kernel Union. The extent and store -# step-downs apply first; then an `Qm` that cannot fill one tile demotes to the -# fitted shape, or on AVX-512 complex to FMAddSub. -@inline function _default_shape(::Type{T}, Qm::Int, Qn::Int, run::Int = Qm) where {T} +# The automatic `(shape, method)` for extents `m_length`/`n_length` and C's +# unit-stride run along M (`run = m_length`: no layout known). Returned as plain +# values so `plan_contract` never holds a menu-wide kernel Union. The extent and +# store step-downs apply first; then an `m_length` that cannot fill one tile +# demotes to the fitted shape, or on AVX-512 complex to FMAddSub. +@inline function _default_shape(::Type{T}, m_length::Int, n_length::Int, run::Int = m_length) where {T} d = _resolved_defaults(T) method = _default_method(T) - shape = _store_shape(_extent_shape(d.shape, T, method, Qm), T, method, Qm, run) - (Qm > 0 && Qm < shape[1]) || return (shape, method) + shape = _store_shape(_extent_shape(d.shape, T, method, m_length), T, method, m_length, run) + (m_length > 0 && m_length < shape[1]) || return (shape, method) # Static, so a real `T`'s method stays a concrete `RealMethod`. T <: Complex || return (d.fitted, method) - small = _small_m_shape(d.small_m, Qm) + small = _small_m_shape(d.small_m, m_length) small === nothing || return (small, FMAddSubMethod()) 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) +@inline _default_shape(::Type{T}, ::ComplexMethod, m_length::Int, n_length::Int, run::Int) where {T} = + _default_shape(T, m_length, n_length, 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)...) +@inline function _default_shape(::Type{T}, method::_MixedMethod, m_length::Int, n_length::Int, run::Int) where {T} + shape, _ = _default_shape(real(T), _real_problem(method, m_length, n_length, 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) +@noinline function _default_kernel(::Type{T}, m_length::Int, n_length::Int) where {T} + shape, method = _default_shape(T, m_length, n_length) return _kernel_from_shape(shape, T, method) end @@ -99,8 +99,8 @@ end Cache-blocking factors for `kernel` on this host: an analytical model of the detected cache geometry ([`target_profile`](@ref)), or fixed constants where L1d or L2 is undetected, scaled for complex methods by packed reals per -element. `plan_contract` rounds `mc`/`nc` to `mr`/`nr` multiples and clamps -all three to the contraction's extents. +element. `plan_contract` rounds `m_block`/`n_block` to multiples of the tile +size and clamps all three to the contraction's extents. """ function default_blocking(kernel) T = scalartype(kernel) diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index b1bfd9f..69ba47a 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -71,8 +71,8 @@ _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) +_real_problem(::ComplexRealMethod, m_length::Int, n_length::Int, run::Int) = (2 * m_length, n_length, 2 * run) +_real_problem(::RealComplexMethod, m_length::Int, n_length::Int, run::Int) = (m_length, 2 * n_length, run) # The kernel type implementing `method` for `T`, or `nothing` (exactly the # pairs with a menu). @@ -228,50 +228,50 @@ end @noinline _kernel_for(profile::TargetProfile, ::Type{T}) where {T} = _kernel_from_shape(_derived_shape(profile, T), T, _default_method(T)) -# The real MV = 4 shape steps down to its MV = 2 sibling when `Qm` is below one -# tall tile, or below two and the half tile pads to fewer rows. Above `2MR` the -# tall tile's padding excess is under 1.2x, its speed advantage. -@inline _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method, Qm::Int) where {T} = shape +# The real MV = 4 shape steps down to its MV = 2 sibling when `m_length` is +# below one tall tile, or below two and the half tile pads to fewer rows. Above +# `2MR` the tall tile's padding excess is under 1.2x, its speed advantage. +@inline _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int) where {T} = shape -@inline function _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, Qm::Int) where {T} +@inline function _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int) where {T} MR, NR, W = shape - (Qm > 0 && MR == 4 * W && Qm < 2 * MR) || return shape + (m_length > 0 && MR == 4 * W && m_length < 2 * MR) || return shape half = (2 * W, NR, W) half in kernel_shapes(T, method) || return shape # Below one tall tile, always: the tall shape would fall to the fitted # shape (two steps down) in `_default_shape`. - Qm < MR && return half - return cld(Qm, half[1]) * half[1] < cld(Qm, MR) * MR ? half : shape + m_length < MR && return half + return cld(m_length, half[1]) * half[1] < cld(m_length, MR) * MR ? half : shape end # The real MV = 4 shape steps down to MV = 2 when C's unit-stride run along M # (`run`) can't fill the tall tile's slivers but fills a half tile: a sliver # that is not unit-stride in C takes the scattered store, which costs more than # the tall tile gains. Below one half tile both scatter, and the tall one stays. -@inline _store_shape(shape::NTuple{3, Int}, ::Type{T}, method, Qm::Int, run::Int) where {T} = shape +@inline _store_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int, run::Int) where {T} = shape -@inline function _store_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, Qm::Int, run::Int) where {T} +@inline function _store_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int, run::Int) where {T} MR, NR, W = shape - (MR == 4 * W && Qm != run && run % MR != 0 && run >= 2 * W) || return shape + (MR == 4 * W && m_length != run && run % MR != 0 && run >= 2 * W) || return shape half = (2 * W, NR, W) return half in kernel_shapes(T, method) ? half : shape end # Small-M demotion for complex `T` on AVX-512, where the planar fitted shape is -# the spilling `MR = 2W` tile: the native-width FMAddSub shape that pads `Qm` -# least, ties by the larger tile. `_small_m_candidates` is the cached, -# host-dependent half (empty where the rule does not apply). +# the spilling `MR = 2W` tile: the native-width FMAddSub shape that pads +# `m_length` least, ties by the larger tile. `_small_m_candidates` is the +# cached, host-dependent half (empty where the rule does not apply). _small_m_candidates(::Val, ::TargetProfile, ::Type) = NTuple{3, Int}[] function _small_m_candidates(::Val{:avx512}, profile::TargetProfile, ::Type{T}) where {T <: Complex} lanes = profile.vector_bytes ÷ sizeof(real(T)) return [shape for shape in kernel_shapes(T, FMAddSubMethod()) if shape[3] == lanes] end -function _small_m_shape(candidates::Vector{NTuple{3, Int}}, Qm::Int) +function _small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) best = nothing for shape in candidates MR, NR, _ = shape - key = (-(cld(Qm, MR) * MR), MR * NR) + key = (-(cld(m_length, MR) * MR), MR * NR) if best === nothing || key > best[1] best = (key, shape) end @@ -280,30 +280,31 @@ function _small_m_shape(candidates::Vector{NTuple{3, Int}}, Qm::Int) end # Run-length demotion (real `T` only): the vectorized store needs every -# register sliver unit-stride in C, i.e. `Qm == run || run % mr == 0` (not -# `mr <= run`). Otherwise demote to the largest menu shape whose `mr` divides -# `run`, unless `Qk` is deep enough that the smaller kernel's cost dominates. +# register sliver unit-stride in C, i.e. `m_length == run || run % m_tile == 0` +# (not `m_tile <= run`). Otherwise demote to the largest menu shape whose +# `m_tile` divides `run`, unless `k_length` is deep enough that the smaller +# kernel's cost dominates. @inline function _demote_shape_for_run( - ::Type{T}, shape::NTuple{3, Int}, method, run::Int, Qm::Int, Qk::Int + ::Type{T}, shape::NTuple{3, Int}, method, run::Int, m_length::Int, k_length::Int ) where {T} - target = _run_demotion_target(T, shape[1], method, run, Qm, Qk) + target = _run_demotion_target(T, shape[1], method, run, m_length, k_length) 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 + ::Type{T}, shape::NTuple{3, Int}, method::_MixedMethod, run::Int, m_length::Int, k_length::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) + m, _, r = _real_problem(method, m_length, 0, run) + real_shape = _demote_shape_for_run(real(T), _real_shape(method, shape), RealMethod(), r, m, k_length) return _mixed_shape(method, real_shape) end -function _run_demotion_target(::Type{T}, mrk::Int, method, run::Int, Qm::Int, Qk::Int) where {T} +function _run_demotion_target(::Type{T}, m_tile::Int, method, run::Int, m_length::Int, k_length::Int) where {T} T <: Real || return nothing kmax = T === Float64 ? _RUN_DEMOTE_KMAX_F64 : _RUN_DEMOTE_KMAX_F32 - Qk > kmax && return nothing - Qm == run && return nothing - run % mrk == 0 && return nothing + k_length > kmax && return nothing + m_length == run && return nothing + run % m_tile == 0 && return nothing best = nothing for shape in kernel_shapes(T, method) m = shape[1] @@ -314,6 +315,6 @@ function _run_demotion_target(::Type{T}, mrk::Int, method, run::Int, Qm::Int, Qk return best end -# Deepest `Qk` at which run-length demotion still wins. +# Deepest `k_length` at which run-length demotion still wins. const _RUN_DEMOTE_KMAX_F64 = 32 const _RUN_DEMOTE_KMAX_F32 = 64 diff --git a/src/planning/labels.jl b/src/planning/labels.jl index 6dccc4e..db8927f 100644 --- a/src/planning/labels.jl +++ b/src/planning/labels.jl @@ -250,10 +250,10 @@ end klabels::NTuple{DK, Int}, indA::NTuple{NA, Int}, A::StridedView, morder::NTuple{DM, Int}, indB::NTuple{NB, Int}, B::StridedView, norder::NTuple{DN, Int}, - Qm::Int, Qn::Int + m_length::Int, n_length::Int ) where {DK, NA, NB, DM, DN} DK <= 1 && return klabels - return _choose_k_order(klabels, indA, A, morder, indB, B, norder, Qm, Qn, nothing) + return _choose_k_order(klabels, indA, A, morder, indB, B, norder, m_length, n_length, nothing) end # `l2bytes = nothing`: the core's L2 share, looked up only once a cost is needed. @@ -261,7 +261,7 @@ function _choose_k_order( klabels::NTuple{DK, Int}, indA::NTuple{NA, Int}, A::StridedView, morder::NTuple{DM, Int}, indB::NTuple{NB, Int}, B::StridedView, norder::NTuple{DN, Int}, - Qm::Int, Qn::Int, l2bytes::Union{Int, Nothing} + m_length::Int, n_length::Int, l2bytes::Union{Int, Nothing} ) where {DK, NA, NB, DM, DN} DK <= 1 && return klabels opA = _k_operand(klabels, indA, A, morder) @@ -271,7 +271,7 @@ function _choose_k_order( permB = _sort_perm(id, opB.st) permA == id && permB == id && return klabels l2 = l2bytes === nothing ? _l2_core_bytes(eltype(A)) : l2bytes - cost(perm) = _k_order_cost(perm, opA, Qm, l2) + _k_order_cost(perm, opB, Qn, l2) + cost(perm) = _k_order_cost(perm, opA, m_length, l2) + _k_order_cost(perm, opB, n_length, l2) best = id bestcost = cost(id) for cand in (permA, permB) @@ -307,7 +307,7 @@ end # Swap the operand roles (B feeds M) when the as-is M run cannot fill a register # sliver of the kernel it would run and the swapped one can: the vectorized -# store needs `mr` consecutive M coordinates unit-stride in C, so a shorter run -# buys nothing. -_prefer_swap(run_m::Int, run_n::Int, mr_asis::Int, mr_swapped::Int = mr_asis) = - run_m < mr_asis && run_n >= mr_swapped +# store needs `m_tile` consecutive M coordinates unit-stride in C, so a shorter +# run buys nothing. +_prefer_swap(run_m::Int, run_n::Int, m_tile_asis::Int, m_tile_swapped::Int = m_tile_asis) = + run_m < m_tile_asis && run_n >= m_tile_swapped diff --git a/src/planning/pack_split.jl b/src/planning/pack_split.jl index 0aead6e..ca54993 100644 --- a/src/planning/pack_split.jl +++ b/src/planning/pack_split.jl @@ -32,12 +32,12 @@ _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`. +# blocking reserved for `k_block_requested` when the K extent clamps `k_block`. # `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 + g::AxisGroup{D}, kg::AxisGroup, kmap::Int, R::Int, S::Int, complex::Bool, k_block::Int, + eff::Int, rounded::Int, k_block_requested::Int, l2bytes::Union{Int, Nothing} = nothing ) where {D} D < 256 || return (eff, _NO_SPLIT) line = _K_LINE_BYTES @@ -51,13 +51,13 @@ _is_split(s::PackSplit) = s.L != 0 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) + return _pack_split_window(g, kg, kmap, d1, dj, R, S, complex, k_block, eff, rounded, k_block_requested, 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} + k_block::Int, eff::Int, rounded::Int, k_block_requested::Int, l2bytes::Union{Int, Nothing} ) where {DK} line = _K_LINE_BYTES L = _largest_divisor_upto(g.lengths[dj], max(1, line ÷ S)) @@ -79,11 +79,11 @@ end for d in d1:(dj - 1) psi *= g.lengths[d] end - klines = ks * S >= line ? kc : cld(kc * ks * S, line) + klines = ks * S >= line ? k_block : cld(k_block * 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) + return _pack_split_dynamic(g.lengths, d1, dj, L, psi, R, eff, rounded, k_block, k_block_requested, kinner) end # The cache that can hold the per-sliver walk's reuse window: the core's L2 @@ -102,9 +102,9 @@ end 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 + eff::Int, rounded::Int, k_block::Int, k_block_requested::Int, kinner::Bool ) where {D} - budget = max(rounded, Int(min(widemul(rounded, kc_req) ÷ kc, typemax(Int32))) ÷ R * R) + budget = max(rounded, Int(min(widemul(rounded, k_block_requested) ÷ k_block, 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 diff --git a/src/planning/plan.jl b/src/planning/plan.jl index bff24b5..93bf642 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -46,7 +46,7 @@ end indC::NTuple{NC,Int}; kernel = nothing, conjA = false, conjB = false, - mc = nothing, kc = nothing, nc = nothing, + m_block = nothing, k_block = nothing, n_block = nothing, workspace = nothing, allocator = TensorOperations.DefaultAllocator(), oracle = true, accumulator = nothing) -> ContractPlan @@ -63,8 +63,8 @@ invalid input. `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. + lower precision than `T`, a K longer than `k_block` accumulates in a panel + of `T` in the workspace, of `M * min(N, n_block)` elements. * `kernel = nothing` picks one from the hardware profile and the extents: a [`SIMDKernel`](@ref) for a real `T`, a [`PlanarKernel`](@ref) for a complex one (an [`FMAddSubKernel`](@ref) for a short M on AVX-512), a @@ -75,13 +75,14 @@ invalid input. operand roles are swapped (B feeds M) when only the N side gives `C` a unit-stride run long enough for the kernel's register tile; the result is the same either way. - * `mc`/`kc`/`nc` override the fields of `default_blocking(kernel)` (each - `>= 1`); `mc`/`nc` are rounded up to `mr`/`nr` multiples and all three are - clamped to the extents. + * `m_block`/`k_block`/`n_block` override the fields of + `default_blocking(kernel)` (each `>= 1`); `m_block`/`n_block` are rounded + up to multiples of the tile size 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. + is packed line by line (`PackSplit`); its `m_block` (or `n_block`) then + becomes a whole number of line groups, up to `requested k_block / k_block` + times the request. * `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`, @@ -98,9 +99,9 @@ function plan_contract( kernel = nothing, conjA::Bool = false, conjB::Bool = false, - mc::Union{Int, Nothing} = nothing, - kc::Union{Int, Nothing} = nothing, - nc::Union{Int, Nothing} = nothing, + m_block::Union{Int, Nothing} = nothing, + k_block::Union{Int, Nothing} = nothing, + n_block::Union{Int, Nothing} = nothing, workspace::Union{Nothing, ContractWorkspace} = nothing, allocator = TO.DefaultAllocator(), oracle::Bool = true, @@ -108,7 +109,7 @@ function plan_contract( ) where {NA, NB, NC} return _planned( identity, C, A, indA, B, indB, indC, - kernel, conjA, conjB, mc, kc, nc, workspace, allocator, oracle, accumulator + kernel, conjA, conjB, m_block, k_block, n_block, workspace, allocator, oracle, accumulator ) end @@ -130,9 +131,9 @@ struct _PlanRequest{ Cstorage::SC Cbase::Int run::Int - mc::Union{Int, Nothing} - kc::Union{Int, Nothing} - nc::Union{Int, Nothing} + m_block::Union{Int, Nothing} + k_block::Union{Int, Nothing} + n_block::Union{Int, Nothing} workspace::WS allocator::AL oracle::Bool @@ -141,11 +142,11 @@ end @inline function _plan_request( ::Type{T}, f::F, mgroup::GM, ngroup::GN, kgroup::GK, Astorage::SA, Abase::Int, Bstorage::SB, Bbase::Int, Cstorage::SC, Cbase::Int, - run::Int, mc, kc, nc, workspace::WS, allocator::AL, oracle::Bool + run::Int, m_block, k_block, n_block, workspace::WS, allocator::AL, oracle::Bool ) where {T, F, GM, GN, GK, SA, SB, SC, WS, AL} return _PlanRequest{T, F, GM, GN, GK, SA, SB, SC, WS, AL}( f, mgroup, ngroup, kgroup, Astorage, Abase, Bstorage, Bbase, Cstorage, Cbase, - run, mc, kc, nc, workspace, allocator, oracle + run, m_block, k_block, n_block, workspace, allocator, oracle ) end @@ -156,7 +157,7 @@ function _planned( f::F, C::StridedView, A::StridedView, indA::NTuple{NA, Int}, 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}, + m_block::Union{Int, Nothing}, k_block::Union{Int, Nothing}, n_block::Union{Int, Nothing}, workspace::Union{Nothing, ContractWorkspace}, allocator, oracle::Bool, accumulator::AC ) where {F, NA, NB, NC, AC} @@ -187,31 +188,31 @@ function _planned( mgroup = _build_pair_group(morder, indA, A, indC, C) # maps: (A, C) ngroup = _build_pair_group(norder, indB, B, indC, C) # maps: (B, C) - Qm = axis_length(mgroup) - Qn = axis_length(ngroup) + m_length = axis_length(mgroup) + n_length = axis_length(ngroup) - korder = _order_contract_labels(klabels, indA, A, morder, indB, B, norder, Qm, Qn) + korder = _order_contract_labels(klabels, indA, A, morder, indB, B, norder, m_length, n_length) kgroup = _build_pair_group(korder, indA, A, indB, B) # maps: (A, B) - Qk = axis_length(kgroup) + k_length = axis_length(kgroup) - mr_asis, mr_swapped = _candidate_mrs(T, method, kernel, Qm, Qn, run_m, run_n) + m_tile_asis, m_tile_swapped = _candidate_m_tiles(T, method, kernel, m_length, n_length, run_m, run_n) - if T <: Real && _prefer_swap(run_m, run_n, mr_asis, mr_swapped) + if T <: Real && _prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) # B takes the M role: groups, K maps, storages, run and transforms move # together; the sum is unchanged. kgroup_swapped = _build_pair_group(korder, indB, B, indA, A) # maps: (B, A) req_swapped = _plan_request( T, f, ngroup, mgroup, kgroup_swapped, parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), - run_n, mc, kc, nc, workspace, allocator, oracle + run_n, m_block, k_block, n_block, workspace, allocator, oracle ) 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 + run_m, m_block, k_block, n_block, workspace, allocator, oracle ) return _plan_with_kernel(kernel, method, atransform, btransform, req) end @@ -249,15 +250,15 @@ end 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}, method, kernel, Qm::Int, Qn::Int, run_m::Int, run_n::Int) where {T} = - (mr(kernel), mr(kernel)) -@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 +# `m_tile` 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_m_tiles(::Type{T}, method, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = + (tile_size(kernel)[1], tile_size(kernel)[1]) +@inline function _candidate_m_tiles(::Type{T}, method, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} + m_tile_asis = _default_shape(T, method, m_length, n_length, run_m)[1][1] + m_tile_swapped = T <: Real ? _default_shape(T, method, n_length, m_length, run_n)[1][1] : m_tile_asis + return m_tile_asis, m_tile_swapped end # Kernel resolution. A named kernel goes straight through, never demoted. An @@ -269,9 +270,9 @@ end @inline _plan_with_kernel(kernel, method, atransform, btransform, req::_PlanRequest) = _plan_contract(kernel, atransform, btransform, req, nothing) @inline function _plan_with_kernel(::Nothing, method, atransform, btransform, req::_PlanRequest{T}) where {T} - Qm = axis_length(req.mgroup) - 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)) + m_length = axis_length(req.mgroup) + shape, method = _default_shape(T, method, m_length, axis_length(req.ngroup), req.run) + shape = _demote_shape_for_run(T, shape, method, req.run, m_length, axis_length(req.kgroup)) vshape = _menu_val(shape, T, method) # The execution path, predicted so the callee is specialised on it. Only a # `Bool` crosses: a call union-split on `method` is emitted out of line and @@ -298,12 +299,12 @@ end @inline _strip_storage(req::_PlanRequest{T}) where {T} = _plan_request( T, req.f, req.mgroup, req.ngroup, req.kgroup, nothing, req.Abase, nothing, req.Bbase, nothing, req.Cbase, - req.run, req.mc, req.kc, req.nc, req.workspace, req.allocator, req.oracle + req.run, req.m_block, req.k_block, req.n_block, req.workspace, req.allocator, req.oracle ) @inline _with_storage(req::_PlanRequest{T}, Astorage, Bstorage, Cstorage) where {T} = _plan_request( T, req.f, req.mgroup, req.ngroup, req.kgroup, Astorage, req.Abase, Bstorage, req.Bbase, Cstorage, req.Cbase, - req.run, req.mc, req.kc, req.nc, req.workspace, req.allocator, req.oracle + req.run, req.m_block, req.k_block, req.n_block, req.workspace, req.allocator, req.oracle ) # `nothing` for a continuation that does not execute (see src/execution/execute.jl). @@ -317,48 +318,46 @@ function _plan_contract( scalartype(kernel) === 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) - Qk = axis_length(req.kgroup) + m_length = axis_length(req.mgroup) + n_length = axis_length(req.ngroup) + k_length = axis_length(req.kgroup) defaults = default_blocking(kernel) - mc, kc, nc = req.mc, req.kc, req.nc requested = Blocking( - mc === nothing ? defaults.mc : mc, - kc === nothing ? defaults.kc : kc, - nc === nothing ? defaults.nc : nc + req.m_block === nothing ? defaults.m_block : req.m_block, + req.k_block === nothing ? defaults.k_block : req.k_block, + req.n_block === nothing ? defaults.n_block : req.n_block ) - MRk = mr(kernel) - NRk = nr(kernel) + m_tile, n_tile = tile_size(kernel) # Empty extents: the drivers never read these; the floors keep them valid. - 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) - panel = _c_panel_needed(T, req.Cstorage, Qk, kc_eff) + m_block_rounded = _roundup(requested.m_block, m_tile) + n_block_rounded = _roundup(requested.n_block, n_tile) + m_block = m_length == 0 ? m_tile : min(m_block_rounded, _roundup(m_length, m_tile)) + n_block = n_length == 0 ? n_tile : min(n_block_rounded, _roundup(n_length, n_tile)) + k_block = k_length == 0 ? 1 : min(requested.k_block, k_length) + panel = _c_panel_needed(T, req.Cstorage, k_length, k_block) 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 + # the panel holds N 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(eltype(req.Astorage)), - !(a_format(kernel) isa RealFormat), kc_eff, mc_eff, mc_rounded, requested.kc + if m_length > 0 && n_length > 0 && k_length > 0 + m_block, mpack = _pack_split( + req.mgroup, req.kgroup, 1, m_tile, sizeof(eltype(req.Astorage)), + !(a_format(kernel) isa RealFormat), k_block, m_block, m_block_rounded, requested.k_block ) 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 + n_block, npack = _pack_split( + req.ngroup, req.kgroup, 2, n_tile, sizeof(eltype(req.Bstorage)), + !(b_format(kernel) isa RealFormat), k_block, n_block, n_block_rounded, requested.k_block ) end end - blocking = Blocking(mc_eff, kc_eff, nc_eff) + blocking = Blocking(m_block, k_block, n_block) ws = _resolve_workspace( - T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel ? Qm * min(nc_eff, Qn) : 0 + T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel ? m_length * min(n_block, n_length) : 0 ) plan = ContractPlan( diff --git a/test/execution/test_dot.jl b/test/execution/test_dot.jl index 0b9d421..f510557 100644 --- a/test/execution/test_dot.jl +++ b/test/execution/test_dot.jl @@ -70,10 +70,10 @@ end @test !_dot_takes(plan_contract(mk()...)) end end - # N = 1: K tails, several K blocks (`kc = 64`) and output blocks (`mc = 7`). + # N = 1: K tails, several K blocks (`k_block = 64`) and output blocks (`m_block = 7`). W = QuasiStrided._dot_lanewidth(T) for (d, K) in ((3, W), (5, 3W - 1), (2, 1000)), (alpha, beta) in ab - _dot_check(_dot_n1_maker(T, d, K, 3K + d), alpha, beta; kc = 64, mc = 7) + _dot_check(_dot_n1_maker(T, d, K, 3K + d), alpha, beta; k_block = 64, m_block = 7) end # A scalar output, C[] = sum_k a[k] b[k]. K = 5W + 3 @@ -111,7 +111,7 @@ end end @testset "dot path: allocation-free, and through the backend ($T)" for T in _DOT_TYPES - for (mk, kw) in ((_dot_gemv_maker(T, 6, 4), (;)), (_dot_n1_maker(T, 5, 300, 4), (kc = 64, mc = 8))) + for (mk, kw) in ((_dot_gemv_maker(T, 6, 4), (;)), (_dot_n1_maker(T, 5, 300, 4), (k_block = 64, m_block = 8))) plan = plan_contract(mk()...; oracle = false, kw...) @test _dot_takes(plan) execute!(plan, 1.0, 0.0) diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 69f7389..542ff7d 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -134,7 +134,7 @@ end @test C == beta .* Array(mk()[1]) end # beta = 0 never reads a NaN C, across several blocks. - C, _ = _run_fresh(execute!, _mm_maker(Float64, 11, 10, 9, 2; Cfill = NaN), 2.5, 0.0; mc = 4, kc = 4, nc = 4) + C, _ = _run_fresh(execute!, _mm_maker(Float64, 11, 10, 9, 2; Cfill = NaN), 2.5, 0.0; m_block = 4, k_block = 4, n_block = 4) @test C ≈ _ref_of(_mm_maker(Float64, 11, 10, 9, 2), 2.5, 0.0) # Empty M or N: a no-op. for (M, N) in ((3, 0), (0, 3)) @@ -151,7 +151,7 @@ end @test C ≈ _ref_of(mk, 1.5, 0.5) end # Block sizes beyond the extents are clamped, not an error. - C, _ = _run_fresh(execute!, _mm_maker(Float64, 9, 7, 6, 4), 1.0, 0.5; kernel = ScalarKernel(Val(4), Val(3), Float64), mc = 10_000, kc = 10_000, nc = 10_000) + C, _ = _run_fresh(execute!, _mm_maker(Float64, 9, 7, 6, 4), 1.0, 0.5; kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 10_000, k_block = 10_000, n_block = 10_000) @test C ≈ _ref_of(_mm_maker(Float64, 9, 7, 6, 4), 1.0, 0.5) end @@ -183,7 +183,7 @@ end p = _mm_plan(Cmat, Amat, Bmat) plans = ( p, _mm_plan(Cmat, Amat, Bmat; workspace = p.workspace, oracle = false), - _mm_plan(Cmat, Amat, Bmat; kernel = kernel, mc = 16, kc = 5, nc = 8), + _mm_plan(Cmat, Amat, Bmat; kernel = kernel, m_block = 16, k_block = 5, n_block = 8), ContractPlan( p.kernel, p.mgroup, p.ngroup, p.kgroup, p.blocking, p.Astorage, p.Abase, p.Bstorage, p.Bbase, p.Cstorage, p.Cbase, @@ -200,7 +200,7 @@ end @test Cmat ≈ Amat * Bmat @test allocs_tw == 0 skip = (VERSION < v"1.11") # ScalarKernel's accumulator is a Matrix: bounded, one per tile call. - ps = _mm_plan(Cmat, Amat, Bmat; kernel = ScalarKernel(Val(4), Val(3), Float64), mc = 8, kc = 6, nc = 7) + ps = _mm_plan(Cmat, Amat, Bmat; kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 6, n_block = 7) @test _steady_allocs!(execute!, ps, Cmat) <= 176 * cld(Ma, 4) * cld(Na, 3) * cld(Ka, 6) + 1 end @@ -231,7 +231,7 @@ end @testset "plan_contract/execute!: no union-typed or partially-applied tile/workspace types" begin Amat, Bmat, Cmat = randn(19, 23), randn(23, 17), zeros(19, 17) Av, Bv, Cv = StridedView(Amat), StridedView(Bmat), StridedView(Cmat) - plan = _mm_plan(Cmat, Amat, Bmat; mc = 8, kc = 6, nc = 7) + plan = _mm_plan(Cmat, Amat, Bmat; m_block = 8, k_block = 6, n_block = 7) @test isconcretetype(typeof(plan)) @test all(isconcretetype, fieldtypes(typeof(plan.workspace))) @test typeof(plan.workspace) === QuasiStrided.ContractWorkspace{Float64, Vector{Float64}, Vector{Float64}} @@ -258,11 +258,11 @@ end # 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) + k, k_block, d, T = p.kernel, p.blocking.k_block, default_blocking(p.kernel), eltype(Cv) + split(g, map, R, eff, req) = _pack_split(g, p.kgroup, map, R, sizeof(T), T <: Complex, k_block, eff, cld(req, R) * R, d.k_block, 0) + (m_block, ms) = split(p.mgroup, 1, tile_size(k)[1], p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) + (n_block, ns) = split(p.ngroup, 2, tile_size(k)[2], p.blocking.n_block, d.n_block) + q = plan_contract(Cv, Av, iA, Bv, iB, iC; kw..., kernel = k, m_block, n_block) @assert q.mgroup == p.mgroup && q.ngroup == p.ngroup return ContractPlan( k, q.mgroup, q.ngroup, q.kgroup, q.blocking, q.Astorage, q.Abase, @@ -274,7 +274,7 @@ end 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]) + splits(cap) = QuasiStrided._is_split(_pack_split(p.mgroup, p.kgroup, 1, tile_size(k)[1], 8, false, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block, 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 @@ -284,7 +284,7 @@ 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 ( + for (ind, ext, m_block) 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) @@ -292,7 +292,7 @@ end 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) + plan = _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; conjB, m_block, 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)) diff --git a/test/execution/test_macro_blocking.jl b/test/execution/test_macro_blocking.jl index 391e2bc..71d7b9b 100644 --- a/test/execution/test_macro_blocking.jl +++ b/test/execution/test_macro_blocking.jl @@ -10,10 +10,10 @@ _macro_rtol(::Type{T}, Ka::Integer) where {T} = 50 * max(Ka, 1) * eps(real(T)) # Plan for C[m,n] = A[m,k] * B[k,n], with operands as `op`-carrying views. The # 5-argument constructor keeps `op` on the view instead of materializing. _macro_op_view(M::AbstractMatrix, op) = StridedView(M, size(M), strides(M), 0, op) -function _macro_plan(Cmat, Amat, Bmat, kernel, mc, kc, nc; conjA = false, conjB = false, opA = identity, opB = identity) +function _macro_plan(Cmat, Amat, Bmat, kernel, m_block, k_block, n_block; conjA = false, conjB = false, opA = identity, opB = identity) return QuasiStrided.plan_contract( StridedView(Cmat), _macro_op_view(Amat, opA), (1, 2), _macro_op_view(Bmat, opB), (2, 3), (1, 3); - kernel = kernel, conjA = conjA, conjB = conjB, mc = mc, kc = kc, nc = nc + kernel = kernel, conjA = conjA, conjB = conjB, m_block = m_block, k_block = k_block, n_block = n_block ) end @@ -57,7 +57,7 @@ end @test !isempty(kernels) for kernel in kernels, i in 1:(T <: Complex ? 12 : 5) Ma, Ka, Na = rand(rng, 1:37, 3) - mc, kc, nc = rand(rng, 1:13, 3) + m_block, k_block, n_block = rand(rng, 1:13, 3) conjA, conjB = rand(rng, Bool, 2) opA, opB = rand(rng, _MACRO_OPS, 2) alpha, beta = _macro_scalar(rng, T), _macro_scalar(rng, T) @@ -71,11 +71,11 @@ end expected = iszero(beta) ? alpha .* (Aeff * Beff) : alpha .* (Aeff * Beff) .+ beta .* Cstart C = copy(Cstart) - QuasiStrided.execute!(_macro_plan(C, Amat, Bmat, kernel, mc, kc, nc; kw...), alpha, beta) + QuasiStrided.execute!(_macro_plan(C, Amat, Bmat, kernel, m_block, k_block, n_block; kw...), alpha, beta) Ctw = copy(Cstart) - QuasiStrided.execute_tilewise!(_macro_plan(Ctw, Amat, Bmat, kernel, mc, kc, nc; kw...), alpha, beta) + QuasiStrided.execute_tilewise!(_macro_plan(Ctw, Amat, Bmat, kernel, m_block, k_block, n_block; kw...), alpha, beta) ok = isapprox(C, expected; rtol = _macro_rtol(T, Ka)) && isapprox(C, Ctw; rtol = _macro_rtol(T, Ka)) - ok || @error "macro driver case $i failed" kernel Ma Ka Na mc kc nc kw alpha beta + ok || @error "macro driver case $i failed" kernel Ma Ka Na m_block k_block n_block kw alpha beta @test ok end end @@ -92,7 +92,7 @@ end Cref = _brute_ref(Array(Aperm), (2, 3, 1), Array(Bneg), (2, 4), Cstart, (1, 4, 3), 1.3, 0.6) plan = QuasiStrided.plan_contract( StridedView(Csub), Aperm, (2, 3, 1), Bneg, (2, 4), (1, 4, 3); - kernel = ScalarKernel(Val(4), Val(3), Float64), mc = 3, kc = 4, nc = 3 + kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 3, k_block = 4, n_block = 3 ) QuasiStrided.execute!(plan, 1.3, 0.6) @test isapprox(Array(Csub), Cref; rtol = _macro_rtol(Float64, k_n)) @@ -111,7 +111,7 @@ end Csub .= Cstart plan = QuasiStrided.plan_contract( StridedView(Csub), StridedView(Amat), (1, 2, 3), StridedView(Bmat), (3, 4), (1, 2, 4); - kernel = ScalarKernel(Val(4), Val(3), Float64), mc = 8, kc = 5, nc = 3 + kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 5, n_block = 3 ) run!(plan, 1.7, -0.4) @test isapprox(Array(Csub), Cref; rtol = _macro_rtol(Float64, k_n)) @@ -119,16 +119,16 @@ end end @testset "macro driver: several blocks at the default blocking ($(nameof(typeof(k))){$T})" for T in (ComplexF64, ComplexF32), k in _macro_kernels(T) - QuasiStrided.mr(k) == 8 || continue + QuasiStrided.tile_size(k)[1] == 8 || continue b = QuasiStrided.default_blocking(k) - Ma, Ka, Na = b.mc + QuasiStrided.mr(k), b.kc + 1, b.nc + QuasiStrided.nr(k) + Ma, Ka, Na = b.m_block + QuasiStrided.tile_size(k)[1], b.k_block + 1, b.n_block + QuasiStrided.tile_size(k)[2] Amat, Bmat = randn(MersenneTwister(1), T, Ma, Ka), randn(MersenneTwister(2), T, Ka, Na) Cstart = randn(MersenneTwister(3), T, Ma, Na) C = copy(Cstart) alpha, beta = T(1.5, -0.25), T(-0.5, 0.75) # Conjugation from both sources: the flag on A, a conjugating op on B. plan = _macro_plan(C, Amat, Bmat, k, nothing, nothing, nothing; conjA = true, opB = adjoint) - @test plan.blocking.mc < Ma && plan.blocking.kc < Ka && plan.blocking.nc < Na + @test plan.blocking.m_block < Ma && plan.blocking.k_block < Ka && plan.blocking.n_block < Na QuasiStrided.execute!(plan, alpha, beta) @test isapprox(C, alpha .* (conj.(Amat) * conj.(Bmat)) .+ beta .* Cstart; rtol = _macro_rtol(T, Ka)) end @@ -158,7 +158,7 @@ end C = zeros(T, 9, 6) mkplan() = QuasiStrided.plan_contract( _macro_op_view(C, op), _macro_op_view(Amat, op), (1, 2), _macro_op_view(Bmat, op), (2, 3), (1, 3); - kernel = kernel, conjA = true, conjB = true, mc = 4, kc = 3, nc = 4 + kernel = kernel, conjA = true, conjB = true, m_block = 4, k_block = 3, n_block = 4 ) if T <: Complex && _macro_op_conjugates(op) @test_throws ArgumentError mkplan() diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index 93a5ef1..947af10 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -20,7 +20,7 @@ Cstorage = zeros(length(Cref)) destination = DestinationTile(Cstorage, 0, row_C, col_C) @test (nrows(destination), ncols(destination)) == (6, 4) - packed_a, packed_b = zeros(mr(kernel) * 5), zeros(nr(kernel) * 5) + packed_a, packed_b = zeros(tile_size(kernel)[1] * 5), zeros(tile_size(kernel)[2] * 5) Cstart = rand(MersenneTwister(1234), size(Cref)...) # One panel covering K, then panels of 2, 2, 1 with nontrivial alpha/beta. @@ -51,7 +51,7 @@ end @test_throws BoundsError execute_tile!(k, dst, zeros(4), zeros(4), 0, 1.0, 2.0) @test all(==(999.0), canary) - # Undersized packed buffers are rejected, except by the kc = 0 / alpha = 0 + # Undersized packed buffers are rejected, except by the k_block_length = 0 / alpha = 0 # short-circuits, which never read them. dst = DestinationTile(zeros(4), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) @test_throws DimensionMismatch execute_tile!(k, dst, zeros(4), zeros(4), 50, 1.0, 0.0) diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 1486a60..ea354e0 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -45,8 +45,8 @@ const _PROMOTED = Union{QuasiStrided.PlanarKernel, QuasiStrided.FMAddSubKernel} @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))) + k_length = layout === :scattered ? 32 : MKN[2] + rtol = 10 * sqrt(k_length) * max(eps(real(T)), eps(real(TC))) @test norm(Array(fx[1]) - ref) <= rtol * norm(ref) end @@ -57,10 +57,10 @@ end end end -# Along K, 1, 2^-30 and -1 fall in different kc blocks: the 2^-30 and the small +# Along K, 1, 2^-30 and -1 fall in different K 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), +@testset "one rounding to $TC across K 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_n_block_16, 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), @@ -77,16 +77,16 @@ end 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 + @test K ÷ 2 > QuasiStrided._resolved_defaults(Float64).real_row.k_block 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) + allocator = case === :manual_n_block_16 ? TO.ManualAllocator() : TO.DefaultAllocator() + n_block = case === :manual_n_block_16 ? 16 : nothing + plan = plan_contract(Cv, Av, (1, 2), Bv, (2, 3), (1, 3); accumulator = acc, k_block = 4, n_block, 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 === :manual_n_block_16 && @test N > plan.blocking.n_block && !(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) @@ -100,7 +100,7 @@ end ) rng = MersenneTwister(3) Amat, Bmat, Cmat = randn(rng, TA, 64, K), randn(rng, TB, K, 40), zeros(TC, 64, 40) - plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc, kc = 256) + plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc, k_block = 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/execution/test_outer.jl b/test/execution/test_outer.jl index f827987..ee33559 100644 --- a/test/execution/test_outer.jl +++ b/test/execution/test_outer.jl @@ -62,7 +62,7 @@ end end # Composite M, several N blocks; and a singleton M beside a long N, which # swaps roles and runs on the 9-long N with a scalar tail. - for (mk, kw) in ((_outer_maker(T, 32, 40, 50; variant = :multiM), (nc = 6,)), (_outer_maker(T, 1, 9, 77), (;))) + for (mk, kw) in ((_outer_maker(T, 32, 40, 50; variant = :multiM), (n_block = 6,)), (_outer_maker(T, 1, 9, 77), (;))) C_out, plan = _run_fresh(execute!, mk, 1.5, 0.5; kw...) @test _outer_takes(plan) == (axis_length(plan.mgroup) >= W) @test C_out ≈ _ref_of(mk, 1.5, 0.5) @@ -86,7 +86,7 @@ end @test_throws BoundsError execute!(plan, 1.0, 0.0) @test all(iszero, plan.Cstorage) end - plan = plan_contract(_outer_maker(T, 64, 50, 6)()...; oracle = false, nc = 12) + plan = plan_contract(_outer_maker(T, 64, 50, 6)()...; oracle = false, n_block = 12) @test _outer_takes(plan) execute!(plan, 1.0, 0.0) @test (@allocated execute!(plan, 1.0, 0.0)) == 0 skip = (VERSION < v"1.11") diff --git a/test/execution/test_scalar_vs_simd.jl b/test/execution/test_scalar_vs_simd.jl index 4c1079d..cc9282a 100644 --- a/test/execution/test_scalar_vs_simd.jl +++ b/test/execution/test_scalar_vs_simd.jl @@ -1,10 +1,10 @@ @testset "driver: ScalarKernel and SIMDKernel agree through execute!" begin rng = MersenneTwister(20260908) Amat, Bmat, Cstart = rand(rng, 11, 13), rand(rng, 13, 9), rand(rng, 11, 9) - for (m, n, kc) in ((Val(8), Val(6), 5), (Val(4), Val(3), 4)) + for (m, n, k_block) in ((Val(8), Val(6), 5), (Val(4), Val(3), 4)) Cs, Cv = copy(Cstart), copy(Cstart) - execute!(_mm_plan(Cs, Amat, Bmat; kernel = ScalarKernel(m, n, Float64), kc = kc), 2.5, 0.75) - execute!(_mm_plan(Cv, Amat, Bmat; kernel = SIMDKernel(m, n, Float64), kc = kc), 2.5, 0.75) + execute!(_mm_plan(Cs, Amat, Bmat; kernel = ScalarKernel(m, n, Float64), k_block = k_block), 2.5, 0.75) + execute!(_mm_plan(Cv, Amat, Bmat; kernel = SIMDKernel(m, n, Float64), k_block = k_block), 2.5, 0.75) @test Cs ≈ 2.5 .* (Amat * Bmat) .+ 0.75 .* Cstart @test isapprox(Cs, Cv; atol = 1.0e-10) end diff --git a/test/execution/test_unpackedb.jl b/test/execution/test_unpackedb.jl index a20f38a..9d0c75f 100644 --- a/test/execution/test_unpackedb.jl +++ b/test/execution/test_unpackedb.jl @@ -51,13 +51,13 @@ end (alpha, beta) in ((1.0, 0.0), (2.5, -0.75), (1.0, 1.0)) _ub_check(_mm_maker(T, M, K, N, 10 + idx; B), alpha, beta) end - # Several jc/pc/ic blocks with tails, explicit kernel shapes. + # Several N/K/M blocks with tails, explicit kernel shapes. k0 = plan_contract(_mm_maker(T, 64, 8, 8, 2)()...).kernel W = lanewidth(k0) kernels = T <: Real ? (k0, SIMDKernel(Val(8), Val(6), T)) : (k0, QuasiStrided.PlanarKernel(Val(W), Val(5), T, Val(W)), QuasiStrided.FMAddSubKernel(Val(W), Val(5), T, Val(W))) - for k in kernels, (mc, kc, nc) in ((8, 3, 6), (40, 1, 100)) - _ub_check(_mm_maker(T, 2 * mr(k) + 1, 13, 2 * nr(k) + 1, 3), 1.5, 0.5; kernel = k, mc, kc, nc) + for k in kernels, (m_block, k_block, n_block) in ((8, 3, 6), (40, 1, 100)) + _ub_check(_mm_maker(T, 2 * tile_size(k)[1] + 1, 13, 2 * tile_size(k)[2] + 1, 3), 1.5, 0.5; kernel = k, m_block, k_block, n_block) end # beta = 0 never reads a NaN C. _ub_check(_mm_maker(T, 10, 6, 7, 3; Cfill = NaN), 2.0, 0.0) @@ -93,7 +93,7 @@ end @testset "unpacked B: allocation-free, and through the backend ($T)" for T in _UB_TYPES Amat, Bmat, Cmat = randn(T, 40, 50), randn(T, 50, 37), zeros(T, 40, 37) - plan = _mm_plan(Cmat, Amat, Bmat; kc = 16, nc = 12) + plan = _mm_plan(Cmat, Amat, Bmat; k_block = 16, n_block = 12) @test _ub_takes(plan) allocs = _steady_allocs!(execute!, plan, Cmat) @test allocs == 0 skip = (VERSION < v"1.11") diff --git a/test/execution/test_workspace.jl b/test/execution/test_workspace.jl index 73c1b68..6fcf798 100644 --- a/test/execution/test_workspace.jl +++ b/test/execution/test_workspace.jl @@ -30,7 +30,7 @@ end @testset "ContractWorkspace: reuse across shapes is bitwise identical to fresh plans" begin kernel = ScalarKernel(Val(4), Val(3), Float64) - mc, kc, nc = 8, 6, 7 + m_block, k_block, n_block = 8, 6, 7 alpha, beta = 1.75, -0.5 # Not monotone: sized by the first, grown by the second, then oversized. @@ -44,12 +44,12 @@ end Cstart = randn(Ma, Na) Cfresh = copy(Cstart) - execute!(_mm_plan(Cfresh, Amat, Bmat; kernel = kernel, mc = mc, kc = kc, nc = nc), alpha, beta) + execute!(_mm_plan(Cfresh, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block), alpha, beta) Creuse = copy(Cstart) plan = _mm_plan( Creuse, Amat, Bmat; - kernel = kernel, mc = mc, kc = kc, nc = nc, workspace = ws + kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, workspace = ws ) execute!(plan, alpha, beta) @@ -72,13 +72,13 @@ end @testset "ContractWorkspace: an oversized reused buffer is not read beyond its live region" begin kernel = ScalarKernel(Val(4), Val(3), Float64) - mc, kc, nc = 8, 6, 7 + m_block, k_block, n_block = 8, 6, 7 alpha, beta = 2.5, -0.75 Random.seed!(606) Abig, Bbig = randn(23, 19), randn(19, 17) Cbig = zeros(23, 17) - big = _mm_plan(Cbig, Abig, Bbig; kernel = kernel, mc = mc, kc = kc, nc = nc) + big = _mm_plan(Cbig, Abig, Bbig; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block) execute!(big, 1.0, 0.0) ws = big.workspace @@ -87,7 +87,7 @@ end Cstart = randn(Ma, Na) Cref = copy(Cstart) - execute!(_mm_plan(Cref, Amat, Bmat; kernel = kernel, mc = mc, kc = kc, nc = nc), alpha, beta) + execute!(_mm_plan(Cref, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block), alpha, beta) # Poison every reused buffer: a stale packed slot yields NaN, a stale # offset or descriptor addresses far outside the operand. @@ -105,11 +105,11 @@ end Cpoisoned = copy(Cstart) plan = _mm_plan( Cpoisoned, Amat, Bmat; - kernel = kernel, mc = mc, kc = kc, nc = nc, workspace = ws + kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, workspace = ws ) @test _ws_lengths(ws) == lengths_before - @test length(ws.packed_a) > cld(Ma, 4) * 4 * min(kc, Ka) - @test length(ws.packed_b) > cld(Na, 3) * 3 * min(kc, Ka) + @test length(ws.packed_a) > cld(Ma, 4) * 4 * min(k_block, Ka) + @test length(ws.packed_b) > cld(Na, 3) * 3 * min(k_block, Ka) execute!(plan, alpha, beta) @test all(isfinite, Cpoisoned) @@ -123,7 +123,7 @@ end Amat, Bmat = randn(Ma, Ka), randn(Ka, Na) Cmat = zeros(Ma, Na) - plan = _mm_plan(Cmat, Amat, Bmat; kernel = kernel, kc = 4, oracle = false) + plan = _mm_plan(Cmat, Amat, Bmat; kernel = kernel, k_block = 4, oracle = false) ws = plan.workspace @test isempty(ws.tw_packed_a) @@ -145,7 +145,7 @@ end @test_throws ArgumentError execute_tilewise!(plan, 1.0, 0.0) Ctw = zeros(Ma, Na) - plan_tw = _mm_plan(Ctw, Amat, Bmat; kernel = kernel, kc = 4, workspace = ws, oracle = true) + plan_tw = _mm_plan(Ctw, Amat, Bmat; kernel = kernel, k_block = 4, workspace = ws, oracle = true) @test plan_tw.workspace === ws @test !isempty(ws.tw_packed_a) execute_tilewise!(plan_tw, 1.0, 0.0) @@ -156,11 +156,11 @@ end Random.seed!(1717) kernel = SIMDKernel(Val(4), Val(3), Float64) Ma, Ka, Na = 19, 23, 17 - mc, kc, nc = 8, 6, 7 + m_block, k_block, n_block = 8, 6, 7 Amat, Bmat = randn(Ma, Ka), randn(Ka, Na) Cdefault = zeros(Ma, Na) - default_plan = _mm_plan(Cdefault, Amat, Bmat; kernel = kernel, mc = mc, kc = kc, nc = nc) + default_plan = _mm_plan(Cdefault, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block) execute!(default_plan, 1.5, 0.0) @test default_plan.workspace.packed_a isa Vector{Float64} @@ -170,15 +170,15 @@ end Cmat = zeros(Ma, Na) plan = _mm_plan( Cmat, Amat, Bmat; - kernel = kernel, mc = mc, kc = kc, nc = nc, + kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, allocator = allocator, oracle = false ) ws = plan.workspace @test isconcretetype(typeof(ws)) @test all(isconcretetype, fieldtypes(typeof(ws))) - @test length(ws.packed_a) == cld(plan.blocking.mc, 4) * 4 * plan.blocking.kc - @test length(ws.packed_b) == cld(plan.blocking.nc, 3) * 3 * plan.blocking.kc + @test length(ws.packed_a) == cld(plan.blocking.m_block, 4) * 4 * plan.blocking.k_block + @test length(ws.packed_b) == cld(plan.blocking.n_block, 3) * 3 * plan.blocking.k_block @test isempty(ws.tw_packed_a) @test ws.m_buf_A isa Vector{Int} @test ws.k_buf_B isa Vector{Int} @@ -195,7 +195,7 @@ end Cmanual = zeros(Ma, Na) manual_plan = _mm_plan( Cmanual, Amat, Bmat; - kernel = kernel, mc = mc, kc = kc, nc = nc, allocator = manual, oracle = false + kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, allocator = manual, oracle = false ) @test_throws MethodError QuasiStrided.reserve!( manual_plan.workspace, kernel, manual_plan.blocking, false diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index e4fa535..fa8dbf3 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -113,24 +113,24 @@ const _QSF = QuasiStrided for T in (ComplexF64, ComplexF32), (MR, NR, W) in (kernel_shapes(T, FMAddSubMethod())..., (3, 2, 2)) R = real(T) k = FMAddSubKernel(Val(MR), Val(NR), T, Val(W)) - kc, lda, base = 7, MR + 3, 2 - vals = [T(10i + 1, -(10i + 2)) for i in 1:(lda * kc + 4)] + k_block_length, lda, base = 7, MR + 3, 2 + vals = [T(10i + 1, -(10i + 2)) for i in 1:(lda * k_block_length + 4)] vals[5] = T(0, 0) vals[9] = T(R(-0.0), R(0)) for f in (identity, conj), m in (MR, max(1, MR - 1)) - src = SourceTile(vals, base, AffineAxis(0, 1, m), AffineAxis(0, lda, kc)) - A = zeros(T, MR, kc) - A[1:m, :] = f.(reshape(vals[(base + 1):(base + lda * kc)], lda, kc)[1:m, :]) + src = SourceTile(vals, base, AffineAxis(0, 1, m), AffineAxis(0, lda, k_block_length)) + A = zeros(T, MR, k_block_length) + A[1:m, :] = f.(reshape(vals[(base + 1):(base + lda * k_block_length)], lda, k_block_length)[1:m, :]) want = mk_pack_a(k, A) - bufv = fill(R(-777), packed_a_length(k, kc)) # a Vector: the scalar loop + bufv = fill(R(-777), packed_a_length(k, k_block_length)) # a Vector: the scalar loop pack_a!(bufv, src, k, f) @test isequal(bufv, want) - bufp = fill(R(-777), packed_a_length(k, kc)) # a PackedPanel: the fast path where gated in + bufp = fill(R(-777), packed_a_length(k, k_block_length)) # a PackedPanel: the fast path where gated in GC.@preserve bufp pack_a!(packed_panel(bufp, 1, length(bufp)), src, k, f) @test isequal(bufp, want) end - src = SourceTile(vals, base, AffineAxis(0, 1, MR), AffineAxis(0, lda, kc)) - bufp = zeros(R, packed_a_length(k, kc)) + src = SourceTile(vals, base, AffineAxis(0, 1, MR), AffineAxis(0, lda, k_block_length)) + bufp = zeros(R, packed_a_length(k, k_block_length)) GC.@preserve bufp begin pp = packed_panel(bufp, 1, length(bufp)) @test _QSF._pack_complex_contiguous_eligible( @@ -144,7 +144,7 @@ const _QSF = QuasiStrided for T in (ComplexF64, ComplexF32) bf = default_blocking(FMAddSubKernel(Val(8), Val(8), T, Val(8))) bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) - @test (bf.mc, bf.kc, bf.nc) == (bp.mc, bp.kc, bp.nc) # planar's packed reals + @test (bf.m_block, bf.k_block, bf.n_block) == (bp.m_block, bp.k_block, bp.n_block) # planar's packed reals for s in kernel_shapes(T, FMAddSubMethod()) @test _kernel_from_shape(s, T, FMAddSubMethod()) isa FMAddSubKernel{s[1], s[2], T, s[3]} end diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index c501be3..f4c1140 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -9,12 +9,12 @@ mk_pack_b(::RealComplexKernel, B) = mk_cols(c -> mk_ilv(real(c), imag(c)), permu 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] + vec = acc[v + (2 * tile_size(k)[1] ÷ 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 + MV = tile_size(k)[1] ÷ W v, u = divrem(i, W) return Complex(acc[v + MV * 2j + 1][u + 1], acc[v + MV * (2j + 1) + 1][u + 1]) end @@ -35,7 +35,7 @@ end 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 (sliver_widths(cr)..., sliver_widths(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) @@ -52,17 +52,17 @@ end 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) + k_block_length, lda = 6, MR + 2 + Av, Bv = rand(rng, SA, lda * k_block_length + 1), rand(rng, SB, k_block_length * 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]) + srcA = SourceTile(Av, 1, AffineAxis(0, 1, m), AffineAxis(0, lda, k_block_length)) + srcB = SourceTile(Bv, 0, AffineAxis(0, NR, k_block_length), AffineAxis(0, 1, n)) + A = zeros(T, MR, k_block_length) + B = zeros(T, k_block_length, NR) + A[1:m, :] = f.(reshape(Av[2:(1 + lda * k_block_length)], lda, k_block_length)[1:m, :]) + B[:, 1:n] = f.(permutedims(reshape(Bv, NR, k_block_length))[:, 1:n]) TA, TB = mk_optypes(k) - pa, pb = zeros(R, packed_a_length(k, kc)), zeros(R, packed_b_length(k, kc)) + pa, pb = zeros(R, packed_a_length(k, k_block_length)), zeros(R, packed_b_length(k, k_block_length)) GC.@preserve pa pb begin dpa = panel ? packed_panel(pa, 1, length(pa)) : pa dpb = panel ? packed_panel(pb, 1, length(pb)) : pb @@ -73,7 +73,7 @@ end 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) + execute_tile!(k, DestinationTile(got, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)), pa, pb, k_block_length, alpha, beta) @test mk_close(got, alpha .* vec((A * B)[1:m, 1:n]) .+ beta .* C0, T) end end @@ -95,7 +95,7 @@ end 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 + kernel = k, k_block = 16, conjA = cA, conjB = cB, accumulator = acc ) execute!(plan, alpha, beta) want = alpha .* ((cA ? conj.(A) : A) * (cB ? conj.(B) : B)) .+ beta .* C0 diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 9b8c7d6..8f40fd8 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -26,20 +26,20 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEForm wrong = SIMDKernel(Val(12), Val(8), Float64, Val(4)) d = ComplexKernelDescriptor(Val(12), Val(8), ComplexF64, OneEFormat(), PlanarFormat()) @test_throws ArgumentError OneMKernel{12, 8, ComplexF64, 8, typeof(wrong)}(d, wrong) - # The error reports the logical kc, not the doubled real one. + # The error reports the logical k_block_length, not the doubled real one. err = try accumulate(k, zero_accumulator(k), Float64[], Float64[], -3) catch e e end - @test err isa ArgumentError && occursin("kc = -3", err.msg) + @test err isa ArgumentError && occursin("k_block_length = -3", err.msg) end @testset "blocking and selection" begin for T in (ComplexF64, ComplexF32) bm = default_blocking(OneMKernel(Val(8), Val(8), T, Val(8))) bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) - @test (bm.mc, bm.kc, bm.nc) == (bp.mc ÷ 2, bp.kc, bp.nc) # twice the packed A reals + @test (bm.m_block, bm.k_block, bm.n_block) == (bp.m_block ÷ 2, bp.k_block, bp.n_block) # twice the packed A reals for s in kernel_shapes(T, OneMMethod()) @test _kernel_from_shape(s, T, OneMMethod()) isa OneMKernel{s[1], s[2], T, s[3]} end diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index 5c47933..3099368 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -14,7 +14,7 @@ using StridedViews: StridedView mk_ilv(x, y) = vec(permutedims(hcat(x, y))) mk_cols(f, M) = reduce(vcat, [f(c) for c in eachcol(M)]) -# Packed A (MR x kc) and B (kc x NR), one K step after another. +# Packed A (MR x k_block_length) and B (k_block_length x NR), one K step after another. mk_pack_a(::Union{ScalarKernel, SIMDKernel}, A) = vec(A) mk_pack_a(::PlanarKernel, A) = mk_cols(c -> [real(c); imag(c)], A) mk_pack_a(::OneMKernel, A) = @@ -28,17 +28,17 @@ mk_pack(k, A, B) = (mk_pack_a(k, A), mk_pack_b(k, B)) mk_read(::ScalarKernel, acc, i, j) = acc[i + 1, j + 1] function mk_read(k::SIMDKernel, acc, i, j) W = lanewidth(k) - return acc[i ÷ W + (mr(k) ÷ W) * j + 1][i % W + 1] + return acc[i ÷ W + (tile_size(k)[1] ÷ W) * j + 1][i % W + 1] end function mk_read(k::PlanarKernel, acc, i, j) W = lanewidth(k) - idx = i ÷ W + (mr(k) ÷ W) * j + 1 + idx = i ÷ W + (tile_size(k)[1] ÷ W) * j + 1 return Complex(acc[idx][i % W + 1], acc[length(acc) ÷ 2 + idx][i % W + 1]) end function mk_read(k::Union{OneMKernel, FMAddSubKernel}, acc, i, j) W = lanewidth(k) v, u = divrem(i, W ÷ 2) - vec = acc[v + (2 * mr(k) ÷ W) * j + 1] + vec = acc[v + (2 * tile_size(k)[1] ÷ W) * j + 1] return Complex(vec[2u + 1], vec[2u + 2]) end @@ -57,13 +57,13 @@ mk_dense(v::AbstractVector{T}) where {T} = copyto!(parent(StridedView(zeros(T, l # 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) +mk_run_acc(k, pa, pb, k_block_length) = accumulate(k, zero_accumulator(k), pa, pb, k_block_length) +mk_run_exec(k, dst, pa, pb, k_block_length, alpha, beta) = (execute_tile!(k, dst, pa, pb, k_block_length, alpha, beta); nothing) # `full = false` checks only accumulate against the reference and allocations. # Separate functions, so the light check does not compile the full one. function mk_contract(k; full::Bool = true) - return @testset "$(nameof(typeof(k))){$(mr(k)),$(nr(k)),$(scalartype(k))}" begin + return @testset "$(nameof(typeof(k))){$(tile_size(k)[1]),$(tile_size(k)[2]),$(scalartype(k))}" begin mk_contract_light(k) full && mk_contract_full(k) end @@ -72,14 +72,14 @@ end function mk_contract_light(k) T = scalartype(k) TA, TB = mk_optypes(k) - MR, NR = mr(k), nr(k) + MR, NR = tile_size(k) rng = MersenneTwister(100MR + NR) - kc = 5 - A = rand(rng, TA, MR, kc) - B = rand(rng, TB, kc, NR) + k_block_length = 5 + A = rand(rng, TA, MR, k_block_length) + B = rand(rng, TB, k_block_length, 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) + @test (length(pa), length(pb)) == (packed_a_length(k, k_block_length), packed_b_length(k, k_block_length)) + acc = mk_run_acc(k, pa, pb, k_block_length) @test mk_close([mk_read(k, acc, i, j) for i in 0:(MR - 1), j in 0:(NR - 1)], A * B, T) # Allocation-free, including a scattered destination and a dense one @@ -91,11 +91,11 @@ function mk_contract_light(k) ab = (T(2), T(0.5)) scat = DestinationTile(zeros(T, MR * NR), 0, ScatterAxis(collect(0:(MR - 1)), MR), AffineAxis(0, MR, NR)) part = DestinationTile(mk_dense(zeros(T, m * NR)), 0, AffineAxis(0, 1, m), AffineAxis(0, m, NR)) - mk_run_acc(k, pa, pb, kc) - @test (@allocated mk_run_acc(k, pa, pb, kc)) == 0 skip = skip + mk_run_acc(k, pa, pb, k_block_length) + @test (@allocated mk_run_acc(k, pa, pb, k_block_length)) == 0 skip = skip for dst in (scat, part) - mk_run_exec(k, dst, pa, pb, kc, ab...) - @test (@allocated mk_run_exec(k, dst, pa, pb, kc, ab...)) == 0 skip = skip + mk_run_exec(k, dst, pa, pb, k_block_length, ab...) + @test (@allocated mk_run_exec(k, dst, pa, pb, k_block_length, ab...)) == 0 skip = skip end end return nothing @@ -105,23 +105,23 @@ function mk_contract_full(k) T = scalartype(k) TA, TB = mk_optypes(k) R = real(T) - MR, NR = mr(k), nr(k) + MR, NR = tile_size(k) rng = MersenneTwister(100MR + NR + 1) m = max(1, MR - 1) - kc = 5 - A = rand(rng, TA, MR, kc) - B = rand(rng, TB, kc, NR) + k_block_length = 5 + A = rand(rng, TA, MR, k_block_length) + B = rand(rng, TB, k_block_length, NR) pa, pb = mk_pack(k, A, B) - acc = mk_run_acc(k, pa, pb, kc) + acc = mk_run_acc(k, pa, pb, k_block_length) acc0 = zero_accumulator(k) @test all(iszero(mk_read(k, acc0, i, j)) for i in 0:(MR - 1), j in 0:(NR - 1)) @test accumulate(k, acc0, R[], R[], 0) === acc0 @test_throws ArgumentError accumulate(k, acc0, R[], R[], -1) - # Splitting kc across calls composes. - la, lb = length(pa) ÷ kc, length(pb) ÷ kc + # Splitting k_block_length across calls composes. + la, lb = length(pa) ÷ k_block_length, length(pb) ÷ k_block_length accs = zero_accumulator(k) - for p in 0:(kc - 1) + for p in 0:(k_block_length - 1) accs = accumulate(k, accs, view(pa, (p * la + 1):((p + 1) * la)), view(pb, (p * lb + 1):((p + 1) * lb)), 1) end @test all(mk_read(k, accs, i, j) ≈ mk_read(k, acc, i, j) for i in 0:(MR - 1), j in 0:(NR - 1)) @@ -130,15 +130,15 @@ function mk_contract_full(k) # alpha/beta branch, partial tiles; old C is NaN whenever beta == 0. 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, TA, MR, kc) - B = rand(rng, TB, kc, NR) - pa, pb = kc == 0 ? (R[], R[]) : mk_pack(k, A, B) + for k_block_length in (0, 1, 5), (alpha, beta) in mk_alphabeta(T), (m, n) in extents, scattered in (false, true) + A = rand(rng, TA, MR, k_block_length) + B = rand(rng, TB, k_block_length, NR) + pa, pb = k_block_length == 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)] scattered || (storage = mk_dense(storage)) rows = scattered ? ScatterAxis(collect(0:(m - 1)), m) : AffineAxis(0, 1, m) - execute_tile!(k, DestinationTile(storage, pad, rows, AffineAxis(0, m, n)), pa, pb, kc, alpha, beta) + execute_tile!(k, DestinationTile(storage, pad, rows, AffineAxis(0, m, n)), pa, pb, k_block_length, alpha, beta) AB = (A * B)[1:m, 1:n] want = iszero(beta) ? alpha .* vec(AB) : alpha .* vec(AB) .+ beta .* Cold @test mk_close(storage[(pad + 1):(pad + m * n)], want, T) @@ -159,7 +159,7 @@ function mk_contract_full(k) @test mk_close(storage, vec(A[1:m, :] * B[:, 1:n]), T) end - # alpha == 0 never reads acc or the panels; kc == 0 never reads the panels. + # alpha == 0 never reads acc or the panels; k_block_length == 0 never reads the panels. fulltile(s) = DestinationTile(s, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, NR)) storage = mk_dense(fill(T(2), MR * NR)) store_tile!(fulltile(storage), mk_fill_acc(k, R(NaN)), zero(T), T(3), k) @@ -179,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, TA, MR, kc) - B = rand(rng, TB, kc, NR) + A = rand(rng, TA, MR, k_block_length) + B = rand(rng, TB, k_block_length, NR) pa, pb = mk_pack(k, A, B) AB = A * B ld = m + 3 @@ -197,7 +197,7 @@ function mk_contract_full(k) (view(nanbuf(ld * n + 4), 3:(ld * n + 2)), 0, AffineAxis(0, 1, m), AffineAxis(0, ld, n), (i, j) -> i + ld * j), ) for (storage, base, rows, cols, addr) in layouts - execute_tile!(k, DestinationTile(storage, base, rows, cols), pa, pb, kc, one(T), zero(T)) + execute_tile!(k, DestinationTile(storage, base, rows, cols), pa, pb, k_block_length, one(T), zero(T)) got = [storage[addr(i, j) + 1] for i in 0:(m - 1), j in 0:(n - 1)] @test mk_close(got, AB[1:m, 1:n], T) @test count(!isnan, storage) == m * n @@ -220,7 +220,7 @@ function mk_e2e(T, kernel) C = copy(Cinit) plan = plan_contract( StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3); - kernel, kc = 16, conjA = cA, conjB = cB + kernel, k_block = 16, conjA = cA, conjB = cB ) execute!(plan, alpha, beta) want = alpha .* ((cA ? conj.(A) : A) * (cB ? conj.(B) : B)) .+ beta .* Cinit diff --git a/test/packing/test_kernel_descriptor.jl b/test/packing/test_kernel_descriptor.jl index ee171bb..83acf2c 100644 --- a/test/packing/test_kernel_descriptor.jl +++ b/test/packing/test_kernel_descriptor.jl @@ -1,9 +1,9 @@ using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEFormat, - InterleavedFormat, realtype, packed_a_per_k, packed_b_per_k + InterleavedFormat, realtype, sliver_widths @testset "KernelDescriptor" begin k = KernelDescriptor(Val(8), Val(6), Float64) - @test (mr(k), nr(k), scalartype(k), realtype(k)) == (8, 6, Float64, Float64) + @test (tile_size(k)..., scalartype(k), realtype(k)) == (8, 6, Float64, Float64) @test packed_a_offset(k, 3, 2) == 3 + 8 * 2 @test packed_b_offset(k, 5, 2) == 5 + 6 * 2 @test packed_a_length(k, 0) == 0 @@ -15,7 +15,7 @@ using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEForma @test_throws ArgumentError KernelDescriptor(Val(8), Val(6), Int) end -@testset "ComplexKernelDescriptor: lengths count reals at a logical kc ($T)" for +@testset "ComplexKernelDescriptor: lengths count reals at a logical k_block_length ($T)" for T in (ComplexF64, ComplexF32) MR, NR = 8, 6 @@ -23,14 +23,14 @@ end onem = ComplexKernelDescriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) fmas = ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) @test realtype(planar) === real(T) - @test (packed_a_per_k(planar), packed_a_per_k(onem), packed_a_per_k(fmas)) == + @test (sliver_widths(planar)[1], sliver_widths(onem)[1], sliver_widths(fmas)[1]) == (2MR, 4MR, 2MR) - @test packed_b_per_k(planar) == packed_b_per_k(onem) == packed_b_per_k(fmas) == 2NR + @test sliver_widths(planar)[2] == sliver_widths(onem)[2] == sliver_widths(fmas)[2] == 2NR @test packed_a_length(onem, 7) == 4MR * 7 @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 sliver_widths(ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()))[2] == 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.jl b/test/packing/test_pack_complex.jl index 2b3bf0a..643e181 100644 --- a/test/packing/test_pack_complex.jl +++ b/test/packing/test_pack_complex.jl @@ -28,29 +28,29 @@ const COMPLEX_FORMATS = ( storage = complex_storage(T, 4000) calls = Ref(0) counting = z -> (calls[] += 1; 2z + one(T)) # nonzero at zero - for (side, pack!, PD, fmt) in ((:a, pack_a!, MR, fa), (:b, pack_b!, NR, fb)), kc in (1, 5) + for (side, pack!, PD, fmt) in ((:a, pack_a!, MR, fa), (:b, pack_b!, NR, fb)), k_block_length in (1, 5) fixtures = ( - (AffineAxis(0, 1, PD), AffineAxis(0, 97, kc)), - (AffineAxis(40, -1, PD), AffineAxis(3000, -97, kc)), - (ScatterAxis(SCATTER_LANES, PD), ScatterAxis(SCATTER_STEPS, kc)), + (AffineAxis(0, 1, PD), AffineAxis(0, 97, k_block_length)), + (AffineAxis(40, -1, PD), AffineAxis(3000, -97, k_block_length)), + (ScatterAxis(SCATTER_LANES, PD), ScatterAxis(SCATTER_STEPS, k_block_length)), ) - len = _ref_rpe(fmt) * PD * kc + len = _ref_rpe(fmt) * PD * k_block_length for (lane, step) in fixtures, valid in (PD, 1, 0) src, g = pack_fixture(side, storage, 500, resized(lane, valid), step) for f in (identity, conj) got, _ = pack_into(pack!, :vector, R, len, src, kernel, f) - @test all(isequal.(got, ref_pack(fmt, T, PD, kc, valid, g, f))) + @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, f))) end calls[] = 0 got, canaries = pack_into(pack!, :view, R, len, src, kernel, counting) - @test all(isequal.(got, ref_pack(fmt, T, PD, kc, valid, g, z -> 2z + one(T)))) - @test calls[] == valid * kc # once per element, not per real half + @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, z -> 2z + one(T)))) + @test calls[] == valid * k_block_length # once per element, not per real half @test canaries end end end -@testset "complex packing: validation before any write, kc == 0 is a no-op" begin +@testset "complex packing: validation before any write, k_block_length == 0 is a no-op" begin kernel = ComplexKernelDescriptor(Val(4), Val(3), ComplexF64, PlanarFormat(), PlanarFormat()) storage = fill(ComplexF64(3, 4), 100) src = SourceTile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 8, 2)) @@ -79,32 +79,32 @@ end T in (ComplexF64, ComplexF32), (fa, fb) in COMPLEX_FORMATS[1:2] function run(::Type{T}, fa, fb) where {T} - MR, NR, kc = 8, 6, 4 + MR, NR, k_block_length = 8, 6, 4 kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, fb) R = real(T) storage = rand(T, 4000) - bufa = zeros(R, packed_a_length(kernel, kc)) - bufb = zeros(R, packed_b_length(kernel, kc)) + bufa = zeros(R, packed_a_length(kernel, k_block_length)) + bufb = zeros(R, packed_b_length(kernel, k_block_length)) ro, co, ro_b = collect(0:(MR - 1)) .* 3, [0, 100, 250, 400], collect(0:(NR - 1)) .* 5 fast = Int[] slow = Int[] GC.@preserve bufa bufb begin pa = packed_panel(bufa, 1, length(bufa)) pb = packed_panel(bufb, 1, length(bufb)) - full_a = SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, kc)) - full_b = SourceTile(storage, 0, AffineAxis(0, 1, kc), AffineAxis(0, kc, NR)) + full_a = SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)) + full_b = SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)) for f in (identity, conj) # the contiguous fast path, where the ISA has it push!(fast, steady_pack_allocs(pack_a!, pa, full_a, kernel, f)) push!(fast, steady_pack_allocs(pack_b!, pb, full_b, kernel, f)) end srcs_a = ( - full_a, SourceTile(storage, 0, ScatterAxis(ro, MR), ScatterAxis(co, kc)), - SourceTile(storage, 0, AffineAxis(0, 1, MR - 1), AffineAxis(0, MR, kc)), + full_a, SourceTile(storage, 0, ScatterAxis(ro, MR), ScatterAxis(co, k_block_length)), + SourceTile(storage, 0, AffineAxis(0, 1, MR - 1), AffineAxis(0, MR, k_block_length)), SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), ) srcs_b = ( - full_b, SourceTile(storage, 0, ScatterAxis(co, kc), ScatterAxis(ro_b, NR)), - SourceTile(storage, 0, AffineAxis(0, 1, kc), AffineAxis(0, kc, NR - 1)), + full_b, SourceTile(storage, 0, ScatterAxis(co, k_block_length), ScatterAxis(ro_b, NR)), + SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR - 1)), ) for s in srcs_a, f in (identity, conj) push!(slow, steady_pack_allocs(pack_a!, bufa, s, kernel, f)) diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index c6b2575..90819cb 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -56,16 +56,16 @@ end end @testset "complex pack fast path: scattered K steps, 1e B, ineligible shapes ($T)" for T in (ComplexF64, ComplexF32) - MR, NR, kc = 8, 6, 6 - koffs = [((p * 5) % 7) * 1013 + 3p for p in 0:(kc - 1)] + MR, NR, k_block_length = 8, 6, 6 + koffs = [((p * 5) % 7) * 1013 + 3p for p in 0:(k_block_length - 1)] for fa in (PlanarFormat(), OneEFormat(), InterleavedFormat()), f in (identity, conj) kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, OneEFormat()) # The lane axis must be unit-stride; the step axis may scatter. - check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), ScatterAxis(koffs, kc), f) - check_panel(:b, kernel, T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, kc), f) + check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), ScatterAxis(koffs, k_block_length), f) + check_panel(:b, kernel, T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, k_block_length), f) # Ineligible shapes fall back to the scalar loop, into the same panel. for lane in (AffineAxis(20, -1, MR), AffineAxis(0, 3, MR), AffineAxis(0, 1, MR - 3)) - check_panel(:a, kernel, T, MR, fa, lane, AffineAxis(3000, -97, kc), f) + check_panel(:a, kernel, T, MR, fa, lane, AffineAxis(3000, -97, k_block_length), f) end end end diff --git a/test/packing/test_pack_real.jl b/test/packing/test_pack_real.jl index 0d315a2..9253790 100644 --- a/test/packing/test_pack_real.jl +++ b/test/packing/test_pack_real.jl @@ -26,9 +26,9 @@ end _ref_rpe(::RealFormat) = 1 _ref_rpe(::Union{PlanarFormat, InterleavedFormat}) = 2 _ref_rpe(::OneEFormat) = 4 -function ref_pack(fmt, ::Type{T}, vr, kc, valid, g, f) where {T} - out = zeros(real(T), _ref_rpe(fmt) * vr * kc) - for p in 0:(kc - 1), t in 0:(valid - 1) +function ref_pack(fmt, ::Type{T}, vr, k_block_length, valid, g, f) where {T} + out = zeros(real(T), _ref_rpe(fmt) * vr * k_block_length) + for p in 0:(k_block_length - 1), t in 0:(valid - 1) _ref_emit!(out, fmt, vr, p, t, convert(T, f(g(t, p)))) end return out @@ -71,7 +71,7 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] @testset "pack_a!/pack_b! ($T): every stride kind and tail width vs direct indexing" for T in (Float64, Float32) - MR, NR, kc = 8, 6, 5 + MR, NR, k_block_length = 8, 6, 5 kernel = KernelDescriptor(Val(MR), Val(NR), T) storage = T.(collect(1.0:2000.0)) calls = Ref(0) @@ -80,14 +80,14 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] AffineAxis(0, 1, 8), AffineAxis(3, 7, 8), AffineAxis(40, -1, 8), AffineAxis(20, 0, 8), ScatterAxis(SCATTER_LANES, 8), ] - steps = Any[AffineAxis(0, 64, kc), AffineAxis(700, -97, kc), ScatterAxis(SCATTER_STEPS, kc)] + steps = Any[AffineAxis(0, 64, k_block_length), AffineAxis(700, -97, k_block_length), ScatterAxis(SCATTER_STEPS, k_block_length)] for (side, pack!, PD) in ((:a, pack_a!, MR), (:b, pack_b!, NR)), lane in lanes, step in steps, valid in (PD, 1, 0), dst in (:vector, :view) src, g = pack_fixture(side, storage, 11, resized(lane, valid), step) calls[] = 0 - got, canaries = pack_into(pack!, dst, T, PD * kc, src, kernel, counting) - @test got == ref_pack(RealFormat(), T, PD, kc, valid, g, x -> 3x + 1000) - @test calls[] == valid * kc + got, canaries = pack_into(pack!, dst, T, PD * k_block_length, src, kernel, counting) + @test got == ref_pack(RealFormat(), T, PD, k_block_length, valid, g, x -> 3x + 1000) + @test calls[] == valid * k_block_length @test canaries end # Every tail width, at the minimal and a longer K depth. @@ -117,13 +117,13 @@ end ] for (lane, eligible) in lanes, step in steps, f in (identity, conj, x -> -x) src, g = pack_fixture(:a, storage, 11, lane, step) - kc = axis_length(step) + k_block_length = axis_length(step) pp = packed_panel(zeros(T, 1), 1, 1) @test _pack_a_contiguous_eligible(pp, src, f, nrows(src), Val(MR), T) == (eligible && (f === identity || f === conj)) @test !_pack_a_contiguous_eligible(zeros(T, 1), src, f, nrows(src), Val(MR), T) - got, canaries = pack_into(pack_a!, :panel, T, MR * kc, src, kernel, f) - @test got == ref_pack(RealFormat(), T, MR, kc, nrows(src), g, f) + got, canaries = pack_into(pack_a!, :panel, T, MR * k_block_length, src, kernel, f) + @test got == ref_pack(RealFormat(), T, MR, k_block_length, nrows(src), g, f) @test canaries end end @@ -131,7 +131,7 @@ end @test !_copies_unchanged(conj, complex(T)) # conj is not the identity on complex end -@testset "pack_a!/pack_b!: kc == 0 reads and writes nothing" begin +@testset "pack_a!/pack_b!: k_block_length == 0 reads and writes nothing" begin kernel = KernelDescriptor(Val(4), Val(3), Float64) storage = fill(3.0, 10) read = Ref(false) @@ -161,33 +161,33 @@ end @testset "pack_a!/pack_b!: zero steady-state allocation" begin nontrivial(x) = 2x + 1 # Every destination kind, affine/scattered/pointer-scattered axes, tails and - # kc == 0; the ScalarKernel/SIMDKernel forwarding methods on a subset. + # k_block_length == 0; the ScalarKernel/SIMDKernel forwarding methods on a subset. function run(ctor, ::Type{T}, MR, NR, full) where {T} kernel = ctor(Val(MR), Val(NR), T) - kc = 7 + k_block_length = 7 storage = rand(T, 4000) koffs = [0, MR, 3MR, 2MR, 5MR, 4MR, 6MR] lanes = collect(0:(max(MR, NR) - 1)) .* 3 - bufa = zeros(T, MR * kc + 8) - bufb = zeros(T, NR * kc + 8) + bufa = zeros(T, MR * k_block_length + 8) + bufb = zeros(T, NR * k_block_length + 8) bytes = Int[] GC.@preserve bufa bufb koffs begin - pa = packed_panel(bufa, 1, MR * kc) - dsts_a = (bufa, view(bufa, 3:(2 + MR * kc)), pa) - dsts_b = (bufb, view(bufb, 3:(2 + NR * kc)), packed_panel(bufb, 1, NR * kc)) - pk = PtrScatterAxis(pointer(koffs), kc) + pa = packed_panel(bufa, 1, MR * k_block_length) + dsts_a = (bufa, view(bufa, 3:(2 + MR * k_block_length)), pa) + dsts_b = (bufb, view(bufb, 3:(2 + NR * k_block_length)), packed_panel(bufb, 1, NR * k_block_length)) + pk = PtrScatterAxis(pointer(koffs), k_block_length) srcs_a = ( - SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, kc)), + SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)), SourceTile(storage, 0, AffineAxis(0, 1, MR), pk), - SourceTile(storage, 0, AffineAxis(0, 2, MR - 1), AffineAxis(0, 2MR, kc)), - SourceTile(storage, 0, ScatterAxis(lanes, MR), ScatterAxis(koffs, kc)), + SourceTile(storage, 0, AffineAxis(0, 2, MR - 1), AffineAxis(0, 2MR, k_block_length)), + SourceTile(storage, 0, ScatterAxis(lanes, MR), ScatterAxis(koffs, k_block_length)), SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), ) srcs_b = ( - SourceTile(storage, 0, AffineAxis(0, 1, kc), AffineAxis(0, kc, NR)), - SourceTile(storage, 0, pk, AffineAxis(0, kc, NR - 1)), - SourceTile(storage, 0, ScatterAxis(koffs, kc), ScatterAxis(lanes, NR)), - SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, kc, NR)), + SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)), + SourceTile(storage, 0, pk, AffineAxis(0, k_block_length, NR - 1)), + SourceTile(storage, 0, ScatterAxis(koffs, k_block_length), ScatterAxis(lanes, NR)), + SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, k_block_length, NR)), ) if !full srcs_a, srcs_b, dsts_a, dsts_b = srcs_a[1:1], srcs_b[1:1], dsts_a[1:2], dsts_b[1:2] diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index 196ff96..051abcf 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -9,7 +9,7 @@ using StridedViews: StridedView @test _derived_shape(u, T) === _fallback_shape(T) k = _kernel_for(u, T) @test k isa SIMDKernel - @test (mr(k), nr(k), lanewidth(k)) === _fallback_shape(T) + @test (tile_size(k)..., lanewidth(k)) === _fallback_shape(T) @test _real_blocking_row(u, T) === _fallback_blocking(T) end @test _fallback_blocking(Float64) === Blocking(128, 256, 768) @@ -53,9 +53,9 @@ using StridedViews: StridedView for T in (Float64, Float32), (MR, NR, W) in kernel_shapes(T) @test MR % W == 0 k = SIMDKernel(Val(MR), Val(NR), T, Val(W)) - @test (mr(k), nr(k), lanewidth(k)) == (MR, NR, W) + @test (tile_size(k)..., lanewidth(k)) == (MR, NR, W) @test (MR ÷ W) * NR + (MR ÷ W) <= 32 - @test packed_a_per_k(k) === MR && packed_b_per_k(k) === NR + @test sliver_widths(k) === (MR, NR) @test realtype(k) === T === scalartype(k) @test complex_method(k) === RealMethod() @test packed_a_length(k, 7) === MR * 7 && packed_b_length(k, 7) === NR * 7 @@ -67,7 +67,7 @@ using StridedViews: StridedView @test kernel_shapes(T, RealMethod()) === kernel_shapes(T) k = _default_kernel(T) @test k isa SIMDKernel && scalartype(k) === T - @test (mr(k), nr(k), lanewidth(k)) in kernel_shapes(T) + @test (tile_size(k)..., lanewidth(k)) in kernel_shapes(T) end end @@ -82,26 +82,26 @@ using StridedViews: StridedView @testset "extent demotion when M cannot fill a register tile" begin for T in (Float64, Float32) - MR = mr(_default_kernel(T)) - @test mr(_default_kernel(T, 4 * MR, 256)) == MR + MR = tile_size(_default_kernel(T))[1] + @test tile_size(_default_kernel(T, 4 * MR, 256))[1] == MR small = _default_kernel(T, 1, 256) - @test (mr(small), nr(small), lanewidth(small)) === _fallback_shape(T) - @test mr(_default_kernel(T, 0, 256)) == MR # empty must not demote + @test (tile_size(small)..., lanewidth(small)) === _fallback_shape(T) + @test tile_size(_default_kernel(T, 0, 256))[1] == MR # empty must not demote end end @testset "_extent_shape: MV = 4 steps down to MV = 2 where it pads less" begin - ext(shape, T, Qm) = QuasiStrided._extent_shape(shape, T, RealMethod(), Qm) + ext(shape, T, m_length) = QuasiStrided._extent_shape(shape, T, RealMethod(), m_length) for T in (Float64, Float32) tall = _derived_shape(synthetic(:avx512, 64), T) MR, NR, W = tall half = (2 * W, NR, W) @test ext(tall, T, 0) === tall - for Qm in (1, 2 * W - 1, 2 * W, MR - 1, MR + 1, MR + W, MR + 2 * W) - @test ext(tall, T, Qm) === half + for m_length in (1, 2 * W - 1, 2 * W, MR - 1, MR + 1, MR + W, MR + 2 * W) + @test ext(tall, T, m_length) === half end - for Qm in (MR, MR + 2 * W + 1, 2 * MR - 1, 2 * MR, 2 * MR + 1, 3 * MR + W, 100 * MR) - @test ext(tall, T, Qm) === tall + for m_length in (MR, MR + 2 * W + 1, 2 * MR - 1, 2 * MR, 2 * MR + 1, 3 * MR + W, 100 * MR) + @test ext(tall, T, m_length) === tall end # Never below MV = 2, and complex methods are untouched. avx2 = _derived_shape(synthetic(:avx2, 32), T) @@ -110,10 +110,10 @@ using StridedViews: StridedView @test QuasiStrided._extent_shape(cs, ComplexF64, PlanarMethod(), 3) === cs # Through `_default_kernel`, where the host's shape is the tall one. if _derived_shape(target_profile(), T) === tall - @test (mr(_default_kernel(T, MR + W, 256)), nr(_default_kernel(T, MR + W, 256))) == (2 * W, NR) - @test mr(_default_kernel(T, 2 * W, 256)) == 2 * W - @test mr(_default_kernel(T, 2 * W - 1, 256)) == _fallback_shape(T)[1] - @test mr(_default_kernel(T, 2 * MR, 256)) == MR + @test tile_size(_default_kernel(T, MR + W, 256)) == (2 * W, NR) + @test tile_size(_default_kernel(T, 2 * W, 256))[1] == 2 * W + @test tile_size(_default_kernel(T, 2 * W - 1, 256))[1] == _fallback_shape(T)[1] + @test tile_size(_default_kernel(T, 2 * MR, 256))[1] == MR end end end @@ -124,32 +124,32 @@ using StridedViews: StridedView tall = _derived_shape(synthetic(:avx512, 64), T) MR, NR, W = tall half = (2 * W, NR, W) - Qm = 64 * MR - for run in (Qm, MR, 2 * MR, 5 * MR) # every tall sliver whole - @test st(tall, T, RealMethod(), Qm, run) === tall + m_length = 64 * MR + for run in (m_length, MR, 2 * MR, 5 * MR) # every tall sliver whole + @test st(tall, T, RealMethod(), m_length, run) === tall end @test st(tall, T, RealMethod(), 3 * W, 3 * W) === tall - for run in (2 * W, 3 * W, 6 * W, MR + 2 * W, Qm - 1) - @test st(tall, T, RealMethod(), Qm, run) === half + for run in (2 * W, 3 * W, 6 * W, MR + 2 * W, m_length - 1) + @test st(tall, T, RealMethod(), m_length, run) === half end for run in (1, W, 2 * W - 1) # below one half tile - @test st(tall, T, RealMethod(), Qm, run) === tall + @test st(tall, T, RealMethod(), m_length, run) === tall end - @test st(half, T, RealMethod(), Qm, 3 * W) === half - @test st(_fallback_shape(T), T, RealMethod(), Qm, 3 * W) === _fallback_shape(T) + @test st(half, T, RealMethod(), m_length, 3 * W) === half + @test st(_fallback_shape(T), T, RealMethod(), m_length, 3 * W) === _fallback_shape(T) cs = _derived_shape(synthetic(:avx512, 64), ComplexF64) - @test st(cs, ComplexF64, PlanarMethod(), Qm, 3 * W) === cs + @test st(cs, ComplexF64, PlanarMethod(), m_length, 3 * W) === cs if _derived_shape(target_profile(), T) === tall - @test QuasiStrided._default_shape(T, Qm, 256, 2 * W)[1] === half - @test QuasiStrided._default_shape(T, Qm, 256, MR)[1] === tall - @test QuasiStrided._default_shape(T, Qm, 256)[1] === tall # no run given + @test QuasiStrided._default_shape(T, m_length, 256, 2 * W)[1] === half + @test QuasiStrided._default_shape(T, m_length, 256, MR)[1] === tall + @test QuasiStrided._default_shape(T, m_length, 256)[1] === tall # no run given end end end @testset "_store_shape feeds the swap: ccsd_t_2 / ao2mo_2 at dim 16" begin # Regression: C's only unit run is 16 long, on the N side. Judged - # against mr = 32 the swap was declined and every tile stored scattered. + # against m_tile = 32 the swap was declined and every tile stored scattered. T = Float64 if _derived_shape(target_profile(), T)[1] == 32 d = 16 @@ -159,7 +159,7 @@ using StridedViews: StridedView StridedView(zeros(T, d, d, d, d, d, d)), StridedView(randn(T, d, d, d, d)), (i, j, m, b), Bv, (m, k, a, c), (a, b, c, i, j, k); oracle = false ) - @test mr(plan.kernel) == 16 + @test tile_size(plan.kernel)[1] == 16 @test plan.Astorage === parent(Bv) q, r, s = -1, 3, 4 B2 = StridedView(randn(T, d, d, d, d)) @@ -167,7 +167,7 @@ using StridedViews: StridedView StridedView(zeros(T, d, d, d, d)), StridedView(randn(T, d, d)), (q, b), B2, (a, q, r, s), (a, b, r, s); oracle = false ) - @test mr(plan2.kernel) == 16 + @test tile_size(plan2.kernel)[1] == 16 @test plan2.Astorage === parent(B2) @test axis_length(plan2.mgroup) == d^3 end @@ -240,7 +240,7 @@ end @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 + @test complex_method(k) === m && (tile_size(k)..., lanewidth(k)) === shape end end @test dmethod(T, T, R) === dmethod(T, ComplexF32, Float64) === CR @@ -253,9 +253,9 @@ end 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) === + for T in (ComplexF64, ComplexF32), m in (CR, RC), m_length in (1, 3, 17, 40, 4096), run in (0, 1, 8, m_length) + real_problem = m === CR ? (2 * m_length, 50, 2 * run) : (m_length, 100, run) + @test QuasiStrided._default_shape(T, m, m_length, 50, run) === (mapped(m, QuasiStrided._default_shape(real(T), real_problem...)[1]), m) end end @@ -274,7 +274,7 @@ end shape = _derived_shape(synthetic(isakey, vb; nregisters = nreg), T, m) @test shape in kernel_shapes(T, m) k = QuasiStrided._kernel_from_shape(shape, T, m) - @test (mr(k), nr(k), lanewidth(k)) === shape + @test (tile_size(k)..., lanewidth(k)) === shape end end @@ -296,8 +296,8 @@ end for p in (nosmt, smt2), T in (Float64, Float32) b = _modelled_blocking(p, T) MR, NR, _ = _derived_shape(p, T) - @test b.mc % MR == 0 && b.nc % NR == 0 - @test NR * b.kc * sizeof(T) <= p.l1d.bytes ÷ 2 + @test b.m_block % MR == 0 && b.n_block % NR == 0 + @test NR * b.k_block * sizeof(T) <= p.l1d.bytes ÷ 2 @test _real_blocking_row(p, T) === b end # No L3: the B panel is budgeted from the L2 alone. @@ -305,7 +305,7 @@ end :neon, :aarch64, "", 16, 32, CacheLevel(64KiB, 4, 64, 1), CacheLevel(4MiB, 16, 64, 4), CacheLevel() ) - @test _modelled_blocking(nol3, Float64, 4, 6).nc == (1MiB ÷ (682 * 8)) ÷ 6 * 6 + @test _modelled_blocking(nol3, Float64, 4, 6).n_block == (1MiB ÷ (682 * 8)) ÷ 6 * 6 tiny = TargetProfile( :avx2, :x86_64, "", 32, 16, CacheLevel(64, 1, 64, 1), CacheLevel(64, 1, 64, 1), CacheLevel() ) @@ -321,9 +321,9 @@ end for base in (_fallback_blocking(Float64), _fallback_blocking(Float32), Blocking(1, 5, 1)) for m in (PlanarMethod(), OneMMethod(), FMAddSubMethod()) b = _scale_blocking(base, m) - @test b.kc === base.kc - @test b.mc === max(1, base.mc ÷ a_reals(m)) - @test b.nc === max(1, base.nc ÷ b_reals(m)) + @test b.k_block === base.k_block + @test b.m_block === max(1, base.m_block ÷ a_reals(m)) + @test b.n_block === max(1, base.n_block ÷ b_reals(m)) end end @test _scale_blocking(_fallback_blocking(Float64), PlanarMethod()) === Blocking(64, 256, 384) @@ -343,7 +343,7 @@ end @test kernel isa QuasiStrided.PlanarKernel @test scalartype(kernel) === T && realtype(kernel) === real(T) @test complex_method(kernel) === PlanarMethod() === QuasiStrided._default_method(T) - shape = (mr(kernel), nr(kernel), lanewidth(kernel)) + shape = (tile_size(kernel)..., lanewidth(kernel)) profile = target_profile() @test _planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarMethod())) @@ -355,7 +355,7 @@ end shape1m = first(kernel_shapes(T, OneMMethod())) k1m = QuasiStrided._kernel_from_shape(shape1m, T, OneMMethod()) @test k1m isa QuasiStrided.OneMKernel && complex_method(k1m) === OneMMethod() - @test (mr(k1m), nr(k1m), lanewidth(k1m)) === shape1m + @test (tile_size(k1m)..., lanewidth(k1m)) === shape1m # A method with no kernel for `T` throws, naming itself. err = try QuasiStrided._kernel_from_shape(shape1m, T, RealMethod()) @@ -375,16 +375,16 @@ end end @testset "_small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin - sms(T, Qm, p = synthetic(:avx512, 64)) = - QuasiStrided._small_m_shape(QuasiStrided._small_m_candidates(Val(p.isa), p, T), Qm) + sms(T, m_length, p = synthetic(:avx512, 64)) = + QuasiStrided._small_m_shape(QuasiStrided._small_m_candidates(Val(p.isa), p, T), m_length) # Least padded rows, then the larger tile. @test sms(ComplexF64, 12) === (12, 8, 8) @test sms(ComplexF64, 16) === (8, 8, 8) @test sms(ComplexF64, 20) === (12, 8, 8) @test sms(ComplexF32, 12) === (16, 8, 16) @test sms(ComplexF32, 16) === (16, 8, 16) - for T in (ComplexF64, ComplexF32), Qm in 1:47 - shape = sms(T, Qm) + for T in (ComplexF64, ComplexF32), m_length in 1:47 + shape = sms(T, m_length) @test shape in kernel_shapes(T, FMAddSubMethod()) @test shape[3] == 64 ÷ sizeof(real(T)) end @@ -394,7 +394,7 @@ end end end -@testset "_demote_shape_for_run: largest menu shape whose mr divides the run" begin +@testset "_demote_shape_for_run: largest menu shape whose m_tile divides the run" begin demote = QuasiStrided._demote_shape_for_run for T in (Float64, Float32), shape in kernel_shapes(T) MR = shape[1] diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 7722866..950bd05 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -201,11 +201,11 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() (1, 3); kernel = kernel ) - # One (jc, pc, ic) block, so the hoisted checks run once. - @test base.blocking.mc >= Ma && base.blocking.nc >= Na && base.blocking.kc >= Ka - m_slivers = cld(Ma, mr(kernel)) - n_slivers = cld(Na, nr(kernel)) - ntiles = m_slivers * n_slivers + # One (N, K, M) block, so the hoisted checks run once. + @test base.blocking.m_block >= Ma && base.blocking.n_block >= Na && base.blocking.k_block >= Ka + m_tiles = cld(Ma, tile_size(kernel)[1]) + n_tiles = cld(Na, tile_size(kernel)[2]) + ntiles = m_tiles * n_tiles @test ntiles >= 20 # a per-tile check would be at least this many calls cstore = CountingStorage(zeros(Ma * Na)) @@ -286,18 +286,18 @@ end base = _pcf_plan(Cv, Av, (1, 2), Bv, (2, 3, 4), (1, 3, 4); kernel = kernel) @test base.Astorage === parent(Av) # no M/N swap: M's run is 20 >= 8 @test !first(QS.affine_ramp(base.ngroup)) # the buffer path, not the ramp path - Qn = axis_length(base.ngroup) - @test Qn == N1 * N2 - @test base.blocking.nc >= Qn # one jc block, so 7 slivers - nsliv = cld(Qn, nr(kernel)) + n_length = axis_length(base.ngroup) + @test n_length == N1 * N2 + @test base.blocking.n_block >= n_length # one N block, so 7 slivers + nsliv = cld(n_length, tile_size(kernel)[2]) @test nsliv == 7 - noffs = [offsets(base.ngroup, q)[2] for q in 0:(Qn - 1)] + noffs = [offsets(base.ngroup, q)[2] for q in 0:(n_length - 1)] binding = argmax(noffs) - 1 # zero-based logical coordinate @test binding == N1 - 1 - @test 0 < binding ÷ nr(kernel) < nsliv - 1 # a strictly interior sliver - @test maximum(noffs[1:nr(kernel)]) < noffs[binding + 1] # not the first - @test maximum(noffs[(1 + (nsliv - 1) * nr(kernel)):end]) < noffs[binding + 1] # not the last + @test 0 < binding ÷ tile_size(kernel)[2] < nsliv - 1 # a strictly interior sliver + @test maximum(noffs[1:tile_size(kernel)[2]]) < noffs[binding + 1] # not the first + @test maximum(noffs[(1 + (nsliv - 1) * tile_size(kernel)[2]):end]) < noffs[binding + 1] # not the last _pcf_exec(base, 1.0, 0.0) ref = zeros(M, N1, N2) @@ -319,7 +319,7 @@ end # A first- or last-sliver-only range would accept it. moffs = [offsets(base.mgroup, q)[2] for q in 0:(axis_length(base.mgroup) - 1)] mrange = (minimum(moffs), maximum(moffs)) - NR = nr(kernel) + NR = tile_size(kernel)[2] truerange = (minimum(noffs), maximum(noffs)) firstonly = (minimum(noffs[1:NR]), maximum(noffs[1:NR])) lastonly = let tail = noffs[(1 + (nsliv - 1) * NR):end] @@ -509,7 +509,7 @@ end # Several macro blocks, so ramp descriptors start at a nonzero `first`. check( (40, 30), (40, 21), (21, 30), (1, 2), (2, 3), (1, 3); - kernel = kernel, mc = 8, kc = 5, nc = 6 + kernel = kernel, m_block = 8, k_block = 5, n_block = 6 ) # Non-ramp composites on both sides (rank-4 operands, non-folding order). d = 5 diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index df71923..9810c1e 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -38,7 +38,7 @@ end @testset "Blocking: field validation" begin b = Blocking(4, 8, 16) - @test b.mc == 4 && b.kc == 8 && b.nc == 16 + @test b.m_block == 4 && b.k_block == 8 && b.n_block == 16 @test_throws ArgumentError Blocking(0, 8, 16) @test_throws ArgumentError Blocking(4, 0, 16) @@ -51,28 +51,28 @@ end bf32 = default_blocking(ScalarKernel(Val(8), Val(6), Float32)) @test bf64 isa Blocking @test bf32 isa Blocking - @test bf64.mc >= 1 && bf64.kc >= 1 && bf64.nc >= 1 - @test bf32.mc >= 1 && bf32.kc >= 1 && bf32.nc >= 1 + @test bf64.m_block >= 1 && bf64.k_block >= 1 && bf64.n_block >= 1 + @test bf32.m_block >= 1 && bf32.k_block >= 1 && bf32.n_block >= 1 # Same kernel shape, different scalar type, through SIMDKernel too. @test default_blocking(SIMDKernel(Val(8), Val(6), Float64)) == bf64 end -@testset "plan_contract: mc/kc/nc keywords are validated and rounded" begin +@testset "plan_contract: m_block/k_block/n_block keywords are validated and rounded" begin kernel = ScalarKernel(Val(4), Val(3), Float64) Ma, Ka, Na = 9, 10, 8 Amat, Bmat, Cmat = randn(Ma, Ka), randn(Ka, Na), zeros(Ma, Na) Av, Bv, Cv = StridedView(Amat), StridedView(Bmat), StridedView(Cmat) indA, indB, indC = (1, 2), (2, 3), (1, 3) - @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, mc = 0) - @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, kc = 0) - @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, nc = -3) + @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, m_block = 0) + @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, k_block = 0) + @test_throws ArgumentError plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel, n_block = -3) - # mc=5 with MR=4 rounds up to 8, then clamps to roundup(Ma=9,4)=12 -> 8. - plan = _mm_plan(Cmat, Amat, Bmat; kernel = kernel, mc = 5, kc = 100, nc = 100) - @test plan.blocking.mc == 8 - @test plan.blocking.kc == 10 # clamped to Qk - @test plan.blocking.nc == 9 # NR=3: roundup(8,3)=9, requested 100 clamped down to that + # m_block=5 with MR=4 rounds up to 8, then clamps to roundup(Ma=9,4)=12 -> 8. + plan = _mm_plan(Cmat, Amat, Bmat; kernel = kernel, m_block = 5, k_block = 100, n_block = 100) + @test plan.blocking.m_block == 8 + @test plan.blocking.k_block == 10 # clamped to k_length + @test plan.blocking.n_block == 9 # NR=3: roundup(8,3)=9, requested 100 clamped down to that end @@ -167,8 +167,8 @@ end # Label order within M/N and the M/N orientation swap. const _lo_order = QuasiStrided._order_free_labels const _lo_run = QuasiStrided._leading_unit_run -_lo_swap(morder, norder, indC, C, mr_asis, mr_swapped = mr_asis) = QuasiStrided._prefer_swap( - _lo_run(morder, indC, C), _lo_run(norder, indC, C), mr_asis, mr_swapped +_lo_swap(morder, norder, indC, C, m_tile_asis, m_tile_swapped = m_tile_asis) = QuasiStrided._prefer_swap( + _lo_run(morder, indC, C), _lo_run(norder, indC, C), m_tile_asis, m_tile_swapped ) # A view with these strides over dummy data (only strides/size are read). @@ -279,7 +279,7 @@ end @test _lo_run((10, 20), (10, 20), _lo_view((6, 1), (1, 17))) == 6 @test _lo_run((10, 20), (10, 20), _lo_view((0, 6), (1, 1))) == 0 - # The swap rule on the ccsd_t shapes at dim 4, with `mr` played by hand. + # The swap rule on the ccsd_t shapes at dim 4, with `m_tile` played by hand. # ccsd_t_2: sorted M = (b,i,j) (run 1), sorted N = (a,c,k) (run 4). @test _lo_swap((b, i, j), (a, c, k), indC, C6, 4) # 4-wide kernel: swap @test !_lo_swap((b, i, j), (a, c, k), indC, C6, 8) # 8-wide: 4 < 8, do not swap @@ -287,10 +287,10 @@ end @test _lo_swap((c, i, j), (a, b, k), indC, C6, 8) @test _lo_swap((c, i, j), (a, b, k), indC, C6, 16) @test !_lo_swap((c, i, j), (a, b, k), indC, C6, 32) - # ccsd_t_1: M already has the run; never swap, whatever mr says. + # ccsd_t_1: M already has the run; never swap, whatever m_tile says. @test !_lo_swap((a, i, j), (b, c, k), indC, C6, 4) @test !_lo_swap((a, i, j), (b, c, k), indC, C6, 8) - # Two-mr form: each orientation is judged against the kernel it would run. + # Two-m_tile form: each orientation is judged against the kernel it would run. @test _lo_swap((b, i, j), (a, c, k), indC, C6, 8, 4) @test !_lo_swap((b, i, j), (a, c, k), indC, C6, 8, 8) @test !_lo_swap((a, i, j), (b, c, k), indC, C6, 4, 4) @@ -312,8 +312,8 @@ end Cr = _lo_view(lens, strides) nlabels = rand(rng, 1:D) order = Tuple(Random.shuffle(rng, collect(1:D))[1:nlabels]) - Qm = prod(lens[l] for l in order) - Qm == 1 && continue + m_length = prod(lens[l] for l in order) + m_length == 1 && continue ntested += 1 run = QuasiStrided._leading_unit_run(order, indCr, Cr) @@ -323,7 +323,7 @@ end cgroup = QuasiStrided.AxisGroup(clens, (cstrides,)) (isramp, steps) = QuasiStrided.affine_ramp(cgroup) - @test (run == Qm) == (isramp && steps[1] == 1) + @test (run == m_length) == (isramp && steps[1] == 1) end end @@ -352,7 +352,7 @@ end @test (name == "ccsd_t_1") == (mrun == d) @test nrun == (name == "ccsd_t_1" ? 1 : name == "ccsd_t_3" ? d^2 : d) - # Named kernels, host-independent: mr = 4 swaps 2/3/4, mr = 8 only 3. + # Named kernels, host-independent: m_tile = 4 swaps 2/3/4, m_tile = 8 only 3. for (kernel, expect_swap) in ( (ScalarKernel(Val(4), Val(3), Float64), name != "ccsd_t_1"), (SIMDKernel(Val(8), Val(6), Float64), name == "ccsd_t_3"), @@ -360,7 +360,7 @@ end plan = plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel) swapped = plan.Astorage === parent(Bv) @test swapped == expect_swap - @test swapped == _lo_swap(msorted, nsorted, indC, Cv, mr(kernel)) + @test swapped == _lo_swap(msorted, nsorted, indC, Cv, tile_size(kernel)[1]) if swapped # B feeds M: mgroup's maps are (B, C) over the sorted N labels. @test plan.mgroup.strides[2] == Tuple(cstride(l) for l in nsorted) @@ -381,16 +381,16 @@ end @test plan.ngroup.lengths == ntuple(_ -> d, 3) end - # Default kernel: the same rule at this machine's `mr`. + # Default kernel: the same rule at this machine's `m_tile`. plan = plan_contract(Cv, Av, indA, Bv, indB, indC) - MRk = mr(plan.kernel) - @test (plan.Astorage === parent(Bv)) == (mrun < MRk && nrun >= MRk) + m_tile = tile_size(plan.kernel)[1] + @test (plan.Astorage === parent(Bv)) == (mrun < m_tile && nrun >= m_tile) end end @testset "run-length-aware kernel demotion" begin - # ccsd_t_1 with C's leading run exactly `d` and Qm != run, so only - # `run % mr == 0` avoids demotion. Expectations are derived from the + # ccsd_t_1 with C's leading run exactly `d` and m_length != run, so only + # `run % m_tile == 0` avoids demotion. Expectations are derived from the # menu, so this holds on every ISA. d = 16 extra = 4 @@ -409,23 +409,23 @@ end msorted = _lo_order(mlab, indC, Cv) run = _lo_run(msorted, indC, Cv) cpos(l) = findfirst(==(l), indC)::Int - Qm = prod(size(Cv, cpos(l)) for l in mlab) - Qn = prod(size(Cv, cpos(l)) for l in nlab) + m_length = prod(size(Cv, cpos(l)) for l in mlab) + n_length = prod(size(Cv, cpos(l)) for l in nlab) plan = plan_contract(Cv, Av, indA, Bv, indB, indC) @test plan.Astorage === parent(Av) - default_kernel = QuasiStrided._default_kernel(T, Qm, Qn) - default_mr = mr(default_kernel) - if Qm == run || run % default_mr == 0 + default_kernel = QuasiStrided._default_kernel(T, m_length, n_length) + default_m_tile = tile_size(default_kernel)[1] + if m_length == run || run % default_m_tile == 0 @test plan.kernel === default_kernel else candidates = [sh[1] for sh in QuasiStrided.kernel_shapes(T) if run % sh[1] == 0] if isempty(candidates) @test plan.kernel === default_kernel else - @test run % mr(plan.kernel) == 0 - @test mr(plan.kernel) == maximum(candidates) + @test run % tile_size(plan.kernel)[1] == 0 + @test tile_size(plan.kernel)[1] == maximum(candidates) end end @@ -443,7 +443,7 @@ end @test Cex ≈ Cref end - # Plain GEMM: `Qm == run`, so no demotion. + # Plain GEMM: `m_length == run`, so no demotion. for T in (Float64, Float32) Ma, Ka, Na = 37, 11, 23 Amat = randn(T, Ma, Ka) @@ -460,7 +460,7 @@ end end @testset "run-length demotion K-depth guard: deep-K does not demote, shallow-K does" begin - # C[a,b,c,i,j,k] = A[i,j,m,a] * B[m,k,b,c], i=j=k=b=c=6: Qk = m is swept. + # C[a,b,c,i,j,k] = A[i,j,m,a] * B[m,k,b,c], i=j=k=b=c=6: k_length = m is swept. IA = (:i, :j, :m, :a) IB = (:m, :k, :b, :c) IC = (:a, :b, :c, :i, :j, :k) @@ -478,9 +478,9 @@ end kmax_of(::Type{Float32}) = QuasiStrided._RUN_DEMOTE_KMAX_F32 for (T, a) in ((Float64, 8), (Float32, 16)) - Qm = a * 36 - Qn = 216 - default_kernel = QuasiStrided._default_kernel(T, Qm, Qn) + m_length = a * 36 + n_length = 216 + default_kernel = QuasiStrided._default_kernel(T, m_length, n_length) # Deep K never demotes; the fixture never swaps. m_deep = 2 * kmax_of(T) @@ -489,7 +489,7 @@ end @test plan_deep.Astorage === parent(Av) @test plan_deep.kernel === default_kernel - # Shallow K demotes where the default's `mr` breaks the run (ISA-dependent). + # Shallow K demotes where the default's `m_tile` breaks the run (ISA-dependent). m_shallow = 8 Cv2, Av2, Bv2 = _run_demote_fixture(T, a, m_shallow) plan_shallow = plan_contract(Cv2, Av2, indA, Bv2, indB, indC) @@ -497,19 +497,19 @@ end mlab, = QuasiStrided._classify_labels(indA, indB, indC) msorted = _lo_order(mlab, indC, Cv2) run = _lo_run(msorted, indC, Cv2) - if Qm == run || run % mr(default_kernel) == 0 + if m_length == run || run % tile_size(default_kernel)[1] == 0 @test plan_shallow.kernel === default_kernel else @test plan_shallow.kernel !== default_kernel @test typeof(plan_shallow.kernel) !== typeof(default_kernel) - @test run % mr(plan_shallow.kernel) == 0 + @test run % tile_size(plan_shallow.kernel)[1] == 0 end - # The guard is `Qk > kmax`: `Qk == kmax` may still demote. + # The guard is `k_length > kmax`: `k_length == kmax` may still demote. Cv_b, Av_b, Bv_b = _run_demote_fixture(T, a, kmax_of(T)) plan_b = plan_contract(Cv_b, Av_b, indA, Bv_b, indB, indC) run_b = _lo_run(msorted, indC, Cv_b) # same M order at every m in this fixture - if Qm == run_b || run_b % mr(default_kernel) == 0 + if m_length == run_b || run_b % tile_size(default_kernel)[1] == 0 @test plan_b.kernel === default_kernel else @test plan_b.kernel !== default_kernel @@ -537,7 +537,7 @@ end execute!(p, 1.0, 0.0) @test Cmat ≈ Amat * Bmat - # Into C's transpose: N's 12-wide run >= mr = 8, so B feeds M. + # Into C's transpose: N's 12-wide run >= m_tile = 8, so B feeds M. Ct = zeros(Na, Ma) pt = plan_contract(StridedView(Ct), Av, (1, 2), Bv, (2, 3), (3, 1); kernel = kernel) @test pt.Astorage === parent(Bv) @@ -592,13 +592,13 @@ end end @testset "label order: the swap never fires for complex kernels (guarded by T <: Real, deliberately deferred)" begin - # ccsd_t_3 would swap for a real dtype at this mr; the transforms follow + # ccsd_t_3 would swap for a real dtype at this m_tile; the transforms follow # the flag XOR A's `conj` op. d = 4 for T in (ComplexF64, ComplexF32) W = QuasiStrided._default_lanewidth(real(T)) kernel = QuasiStrided.PlanarKernel(Val(W), Val(8), T, Val(W)) - @test mr(kernel) <= 16 + @test tile_size(kernel)[1] <= 16 (indA, indB, indC), _ = _lo_labels(_LO_CASES[3][2], _LO_CASES[3][3]) A = randn(T, d, d, d, d) B = randn(T, d, d, d, d) diff --git a/test/runtests.jl b/test/runtests.jl index cda2519..f2ef324 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,7 +3,7 @@ using Random using QuasiStrided # Internal names the test files use unqualified. using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescriptor, - describe_block, block_descriptors!, normalize_group, KernelDescriptor, mr, nr, + describe_block, block_descriptors!, normalize_group, KernelDescriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, AffineAxis, ScatterAxis, SourceTile, DestinationTile, axis_from_descriptor, nrows, ncols, axis_offset_range, checked_tile_storage_bounds, pack_a!, pack_b!, @@ -12,7 +12,7 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip target_profile, cache_topology, unknown_target, _detect_isa, _detect_target, _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, _fallback_blocking, kernel_shapes, _parse_size, _count_cpu_list, NR_DEFAULT, - _rule_applies, _isa_nregisters, packed_a_per_k, packed_b_per_k, realtype, + _rule_applies, _isa_nregisters, sliver_widths, realtype, complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, a_reals, b_reals, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, _planar_pressure, _pack_split, _NestPath From 58b98e85837fc41a8fa752638389ecd40e41cac7 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 30 Sep 2026 16:19:28 -0400 Subject: [PATCH 03/60] tile_size(k, i) and sliver_width(k, i), like size(A, d) `sliver_widths` becomes `sliver_width`, matching `size`; both accessors gain an index method, replacing the `tile_size(k)[i]` spelling. Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 4 ++-- benchmark/bench_blocking_model.jl | 4 ++-- benchmark/bench_kernels.jl | 4 ++-- benchmark/bench_mixed.jl | 2 +- src/QuasiStrided.jl | 2 +- src/execution/execute.jl | 6 +++--- src/execution/macrokernel.jl | 8 ++++---- src/execution/unpackedb.jl | 2 +- src/microkernels/interface.jl | 4 ++-- src/packing/format.jl | 12 +++++++----- src/packing/pack.jl | 2 +- src/planning/plan.jl | 2 +- test/execution/test_execute.jl | 6 +++--- test/execution/test_macro_blocking.jl | 4 ++-- test/execution/test_manual_pipeline.jl | 2 +- test/execution/test_unpackedb.jl | 2 +- test/microkernels/test_mixed_kernel.jl | 6 +++--- test/microkernels/test_scalar_kernel.jl | 8 ++++---- test/packing/test_kernel_descriptor.jl | 8 ++++---- test/planning/test_kernel_selection.jl | 18 +++++++++--------- test/planning/test_per_call_overhead.jl | 14 +++++++------- test/planning/test_plan_contract.jl | 18 +++++++++--------- test/runtests.jl | 2 +- 23 files changed, 71 insertions(+), 69 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 1d19efd..086dab8 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -45,9 +45,9 @@ Pattern `__`; M/N/K stay the organizing letters. | `mc, kc, nc` (`Blocking` fields, `plan_contract` kwargs) | `m_block, k_block, n_block` | | `ic, pc, jc` | `m_block_start, k_block_start, n_block_start` | | `mblock, kblock, nblock`, `kc_len` | `m_block_length, k_block_length, n_block_length` | -| `mr(k), nr(k)` | `tile_size(k) -> (m_tile, n_tile)` | +| `mr(k), nr(k)` | `tile_size(k) -> (m_tile, n_tile)`, `tile_size(k, i)` (like `size`) | | `MRk, NRk` | `m_tile, n_tile` | -| `packed_a_per_k(k), packed_b_per_k(k)` | `sliver_widths(k) -> (a_sliver_width, b_sliver_width)` | +| `packed_a_per_k(k), packed_b_per_k(k)` | `sliver_width(k) -> (a_sliver_width, b_sliver_width)`, `sliver_width(k, i)` | | `MRp, NRp` | `a_sliver_width, b_sliver_width` | | `m_slivers, n_slivers` | `m_tiles, n_tiles` | | `r, s`, `rfirst, sfirst` | `m_tile_index, n_tile_index`, `m_tile_start, n_tile_start` | diff --git a/benchmark/bench_blocking_model.jl b/benchmark/bench_blocking_model.jl index 1f238db..9dd1656 100644 --- a/benchmark/bench_blocking_model.jl +++ b/benchmark/bench_blocking_model.jl @@ -116,7 +116,7 @@ canaries = [run_canary(crng, "start")] for T in DTYPES kernel = _default_kernel(T) rows = named_rows(T) - println("\n$T kernel $(tile_size(kernel)[1])x$(tile_size(kernel)[2])/W$(lanewidth(kernel)) ", rows) + println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) rows.model === nothing && @warn "no cache geometry detected: model undefined" grid = [("wide", c) for c in thin(WIDE[T])] for spec in thin(GRID_SHAPES) @@ -133,7 +133,7 @@ end for T in CDTYPES kernel = _default_kernel(T) rows = map(v -> v === nothing ? nothing : _scale_blocking(v, complex_method(kernel)), named_rows(real(T))) - println("\n$T kernel $(tile_size(kernel)[1])x$(tile_size(kernel)[2])/W$(lanewidth(kernel)) ", rows) + println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) for spec in thin(MAIN_SHAPES) sweep_shape!(raw, kernel, T, spec, named_points(rows)) end diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index c6e13f1..e396431 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -35,7 +35,7 @@ const OUTDIR = outdir() const CSV_PATH = joinpath(OUTDIR, "bench_kernels.csv") const SUMMARY_PATH = joinpath(OUTDIR, "summary_kernels.txt") -tag(k) = "$(tile_size(k)[1])x$(tile_size(k)[2])/W$(lanewidth(k))" +tag(k) = "$(tile_size(k, 1))x$(tile_size(k, 2))/W$(lanewidth(k))" menu_methods(::Type{T}) where {T} = T <: Complex ? (PlanarMethod(), OneMMethod(), FMAddSubMethod()) : (RealMethod(),) csv = open(CSV_PATH, "w") @@ -60,7 +60,7 @@ function kernel_hot!(kernel, C, apack, bpack, k_block, reps) GC.@preserve apack bpack begin ap = packed_panel(apack, 1, length(apack)) bp = packed_panel(bpack, 1, length(bpack)) - dest = DestinationTile(C, 0, AffineAxis(0, 1, tile_size(kernel)[1]), AffineAxis(0, tile_size(kernel)...)) + dest = DestinationTile(C, 0, AffineAxis(0, 1, tile_size(kernel, 1)), AffineAxis(0, tile_size(kernel)...)) for _ in 1:reps unsafe_execute_tile!(kernel, dest, ap, bp, k_block, one(eltype(C)), one(eltype(C))) end diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index 6ff3e8c..9e3a604 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -46,7 +46,7 @@ function run(io) 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 = "$(tile_size(kernel)[1])x$(tile_size(kernel)[2])/W$(lanewidth(kernel))" + tag = "$(tile_size(kernel, 1))x$(tile_size(kernel, 2))/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 diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 66d075f..c3ee8d4 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -70,7 +70,7 @@ end eval( Expr( :public, :contract!, :plan_contract, :execute!, :ContractPlan, - :ContractWorkspace, :Blocking, :default_blocking, :tile_size, :sliver_widths, + :ContractWorkspace, :Blocking, :default_blocking, :tile_size, :sliver_width, :ScalarKernel, :SIMDKernel, :PlanarKernel, :OneMKernel, :FMAddSubKernel, :ComplexRealKernel, :RealComplexKernel, :target_profile, :cache_topology, diff --git a/src/execution/execute.jl b/src/execution/execute.jl index f93d1b0..42235e6 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -291,7 +291,7 @@ function _execute_nest!( # GUARDRAIL: reals per sliver per K step address the packed panels; # `m_tile`/`n_tile` count register-tile rows. They differ for complex # kernels. - a_sliver_width, b_sliver_width = sliver_widths(kernel) + a_sliver_width, b_sliver_width = sliver_width(kernel) atransform = plan.atransform btransform = plan.btransform @@ -368,7 +368,7 @@ function _execute_nest!( if split_b _pack_block_transposed!( packed_b_plane_offset, b_format(kernel), - packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel)[2]), b_sliver_width, + packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, plan.npack ) @@ -408,7 +408,7 @@ function _execute_nest!( if split_a _pack_block_transposed!( packed_a_plane_offset, a_format(kernel), - packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel)[1]), a_sliver_width, + packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, plan.mpack ) diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index b4893bb..ab6f81a 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -61,10 +61,10 @@ end # Sliver `tile_index` of a packed panel at the current block's depth # `k_block_length`, shared by the packing and the consuming step so the two -# cannot disagree. GUARDRAIL: `sliver_width` counts reals per K step -# (`sliver_widths`), not the tile size; they differ for complex kernels. -@inline function _sliver_panel(buffer, sliver_width::Int, k_block_length::Int, tile_index::Int) - stride = sliver_width * k_block_length +# cannot disagree. GUARDRAIL: `width` counts reals per K step +# (`sliver_width`), not the tile size; they differ for complex kernels. +@inline function _sliver_panel(buffer, width::Int, k_block_length::Int, tile_index::Int) + stride = width * k_block_length return packed_panel(buffer, tile_index * stride + 1, stride) end diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 119bc3f..2baa77c 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -79,7 +79,7 @@ const _UNPACKED_B_MMAX = 256 off = d.regular ? d.base + jj * d.stride : (@inbounds buf[n_tile_start + jj + 1]) base + off end - return UnpackedBView(storage, colbase, ksteps, sliver_widths(kernel)[2], transform) + return UnpackedBView(storage, colbase, ksteps, sliver_width(kernel, 2), transform) end # The tile loops of `_execute_nest!` with B read in place. A barrier over the K diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 3a4c5c7..af6ec15 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -50,7 +50,7 @@ accumulator_planes(::RealComplexMethod) = 1 complex_method(::Any) = RealMethod() realtype(k::DescriptorKernel) = realtype(k.descriptor) -sliver_widths(k::DescriptorKernel) = sliver_widths(k.descriptor) +sliver_width(k::DescriptorKernel) = sliver_width(k.descriptor) a_format(k::DescriptorKernel) = a_format(k.descriptor) b_format(k::DescriptorKernel) = b_format(k.descriptor) tile_size(k::DescriptorKernel) = tile_size(k.descriptor) @@ -169,7 +169,7 @@ end ) @noinline _throw_tile_extent(which::Symbol, got::Int, limit::Int) = throw( ArgumentError( - "destination valid $which extent $got exceeds tile_size(kernel)[$(which === :row ? 1 : 2)] = $limit" + "destination valid $which extent $got exceeds tile_size(kernel, $(which === :row ? 1 : 2)) = $limit" ) ) @noinline _throw_negative_k_block_length(where::Symbol, k_block_length::Int) = diff --git a/src/packing/format.jl b/src/packing/format.jl index bf8b1ba..eff728f 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -72,6 +72,7 @@ function _check_descriptor(MR, NR, T, FA, FB) end tile_size(::Descriptor{MR, NR}) where {MR, NR} = (MR, NR) +tile_size(k, i::Int) = tile_size(k)[i] scalartype(::Descriptor{MR, NR, T}) where {MR, NR, T} = T a_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FA() b_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FB() @@ -79,18 +80,19 @@ b_format(::Descriptor{MR, NR, T, FA, FB}) where {MR, NR, T, FA, FB} = FB() realtype(::Descriptor{MR, NR, T}) where {MR, NR, T} = real(T) # Reals per A / B sliver per logical K step. -sliver_widths(d::Descriptor{MR, NR}) where {MR, NR} = +sliver_width(d::Descriptor{MR, NR}) where {MR, NR} = (MR * reals_per_element(a_format(d)), NR * reals_per_element(b_format(d))) +sliver_width(k, i::Int) = sliver_width(k)[i] -packed_a_length(d::Descriptor, k_block_length::Int) = sliver_widths(d)[1] * k_block_length -packed_b_length(d::Descriptor, k_block_length::Int) = sliver_widths(d)[2] * k_block_length +packed_a_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 1) * k_block_length +packed_b_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 2) * k_block_length # Zero-based offsets, in reals. For 1e and interleaved, `i`/`j` runs over reals # (0:2MR-1) and 1e's second region is `plane == 2`. @inline packed_a_plane_offset(d::Descriptor{MR}, plane::Int, i::Int, p::Int) where {MR} = - p * sliver_widths(d)[1] + plane * MR + i + p * sliver_width(d, 1) + plane * MR + i @inline packed_b_plane_offset(d::Descriptor{MR, NR}, plane::Int, j::Int, p::Int) where {MR, NR} = - p * sliver_widths(d)[2] + plane * NR + j + p * sliver_width(d, 2) + plane * NR + j # Real descriptors only: a complex element has no single offset. packed_a_offset(kernel::KernelDescriptor{MR}, i::Int, p::Int) where {MR} = i + MR * p diff --git a/src/packing/pack.jl b/src/packing/pack.jl index ac6ef23..8c35057 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -15,7 +15,7 @@ end # Out of line so the packers carry no string formatting (and no GC frame). @noinline _throw_pack_extent(which::Symbol, name::Symbol, got::Int, limit::Int) = throw( ArgumentError( - "$(which)!: source $(name) count $got must satisfy 0 <= $got <= tile_size(kernel)[$(which === :pack_a ? 1 : 2)]=$limit" + "$(which)!: source $(name) count $got must satisfy 0 <= $got <= tile_size(kernel, $(which === :pack_a ? 1 : 2))=$limit" ) ) @noinline _throw_pack_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 93bf642..0ef9074 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -254,7 +254,7 @@ end # named kernel either way, else `_default_shape`'s pick at that orientation's # extents and C run. @inline _candidate_m_tiles(::Type{T}, method, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = - (tile_size(kernel)[1], tile_size(kernel)[1]) + (tile_size(kernel, 1), tile_size(kernel, 1)) @inline function _candidate_m_tiles(::Type{T}, method, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} m_tile_asis = _default_shape(T, method, m_length, n_length, run_m)[1][1] m_tile_swapped = T <: Real ? _default_shape(T, method, n_length, m_length, run_n)[1][1] : m_tile_asis diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 542ff7d..262e166 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -260,8 +260,8 @@ function _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; kw...) p = plan_contract(Cv, Av, iA, Bv, iB, iC; kw...) k, k_block, d, T = p.kernel, p.blocking.k_block, default_blocking(p.kernel), eltype(Cv) split(g, map, R, eff, req) = _pack_split(g, p.kgroup, map, R, sizeof(T), T <: Complex, k_block, eff, cld(req, R) * R, d.k_block, 0) - (m_block, ms) = split(p.mgroup, 1, tile_size(k)[1], p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) - (n_block, ns) = split(p.ngroup, 2, tile_size(k)[2], p.blocking.n_block, d.n_block) + (m_block, ms) = split(p.mgroup, 1, tile_size(k, 1), p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) + (n_block, ns) = split(p.ngroup, 2, tile_size(k, 2), p.blocking.n_block, d.n_block) q = plan_contract(Cv, Av, iA, Bv, iB, iC; kw..., kernel = k, m_block, n_block) @assert q.mgroup == p.mgroup && q.ngroup == p.ngroup return ContractPlan( @@ -274,7 +274,7 @@ end 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, tile_size(k)[1], 8, false, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block, cap)[2]) + splits(cap) = QuasiStrided._is_split(_pack_split(p.mgroup, p.kgroup, 1, tile_size(k, 1), 8, false, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block, 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 diff --git a/test/execution/test_macro_blocking.jl b/test/execution/test_macro_blocking.jl index 71d7b9b..616d8a2 100644 --- a/test/execution/test_macro_blocking.jl +++ b/test/execution/test_macro_blocking.jl @@ -119,9 +119,9 @@ end end @testset "macro driver: several blocks at the default blocking ($(nameof(typeof(k))){$T})" for T in (ComplexF64, ComplexF32), k in _macro_kernels(T) - QuasiStrided.tile_size(k)[1] == 8 || continue + QuasiStrided.tile_size(k, 1) == 8 || continue b = QuasiStrided.default_blocking(k) - Ma, Ka, Na = b.m_block + QuasiStrided.tile_size(k)[1], b.k_block + 1, b.n_block + QuasiStrided.tile_size(k)[2] + Ma, Ka, Na = b.m_block + QuasiStrided.tile_size(k, 1), b.k_block + 1, b.n_block + QuasiStrided.tile_size(k, 2) Amat, Bmat = randn(MersenneTwister(1), T, Ma, Ka), randn(MersenneTwister(2), T, Ka, Na) Cstart = randn(MersenneTwister(3), T, Ma, Na) C = copy(Cstart) diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index 947af10..631674b 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -20,7 +20,7 @@ Cstorage = zeros(length(Cref)) destination = DestinationTile(Cstorage, 0, row_C, col_C) @test (nrows(destination), ncols(destination)) == (6, 4) - packed_a, packed_b = zeros(tile_size(kernel)[1] * 5), zeros(tile_size(kernel)[2] * 5) + packed_a, packed_b = zeros(tile_size(kernel, 1) * 5), zeros(tile_size(kernel, 2) * 5) Cstart = rand(MersenneTwister(1234), size(Cref)...) # One panel covering K, then panels of 2, 2, 1 with nontrivial alpha/beta. diff --git a/test/execution/test_unpackedb.jl b/test/execution/test_unpackedb.jl index 9d0c75f..ec9c9df 100644 --- a/test/execution/test_unpackedb.jl +++ b/test/execution/test_unpackedb.jl @@ -57,7 +57,7 @@ end kernels = T <: Real ? (k0, SIMDKernel(Val(8), Val(6), T)) : (k0, QuasiStrided.PlanarKernel(Val(W), Val(5), T, Val(W)), QuasiStrided.FMAddSubKernel(Val(W), Val(5), T, Val(W))) for k in kernels, (m_block, k_block, n_block) in ((8, 3, 6), (40, 1, 100)) - _ub_check(_mm_maker(T, 2 * tile_size(k)[1] + 1, 13, 2 * tile_size(k)[2] + 1, 3), 1.5, 0.5; kernel = k, m_block, k_block, n_block) + _ub_check(_mm_maker(T, 2 * tile_size(k, 1) + 1, 13, 2 * tile_size(k, 2) + 1, 3), 1.5, 0.5; kernel = k, m_block, k_block, n_block) end # beta = 0 never reads a NaN C. _ub_check(_mm_maker(T, 10, 6, 7, 3; Cfill = NaN), 2.0, 0.0) diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index f4c1140..158d05d 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -9,12 +9,12 @@ mk_pack_b(::RealComplexKernel, B) = mk_cols(c -> mk_ilv(real(c), imag(c)), permu function mk_read(k::ComplexRealKernel, acc, i, j) W = lanewidth(k) v, u = divrem(i, W ÷ 2) - vec = acc[v + (2 * tile_size(k)[1] ÷ W) * j + 1] + vec = acc[v + (2 * tile_size(k, 1) ÷ W) * j + 1] return Complex(vec[2u + 1], vec[2u + 2]) end function mk_read(k::RealComplexKernel, acc, i, j) W = lanewidth(k) - MV = tile_size(k)[1] ÷ W + MV = tile_size(k, 1) ÷ W v, u = divrem(i, W) return Complex(acc[v + MV * 2j + 1][u + 1], acc[v + MV * (2j + 1) + 1][u + 1]) end @@ -35,7 +35,7 @@ end 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 (sliver_widths(cr)..., sliver_widths(rc)...) == (16, 4, 8, 8) + @test (sliver_width(cr)..., sliver_width(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) diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index 3099368..fa92e8f 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -28,17 +28,17 @@ mk_pack(k, A, B) = (mk_pack_a(k, A), mk_pack_b(k, B)) mk_read(::ScalarKernel, acc, i, j) = acc[i + 1, j + 1] function mk_read(k::SIMDKernel, acc, i, j) W = lanewidth(k) - return acc[i ÷ W + (tile_size(k)[1] ÷ W) * j + 1][i % W + 1] + return acc[i ÷ W + (tile_size(k, 1) ÷ W) * j + 1][i % W + 1] end function mk_read(k::PlanarKernel, acc, i, j) W = lanewidth(k) - idx = i ÷ W + (tile_size(k)[1] ÷ W) * j + 1 + idx = i ÷ W + (tile_size(k, 1) ÷ W) * j + 1 return Complex(acc[idx][i % W + 1], acc[length(acc) ÷ 2 + idx][i % W + 1]) end function mk_read(k::Union{OneMKernel, FMAddSubKernel}, acc, i, j) W = lanewidth(k) v, u = divrem(i, W ÷ 2) - vec = acc[v + (2 * tile_size(k)[1] ÷ W) * j + 1] + vec = acc[v + (2 * tile_size(k, 1) ÷ W) * j + 1] return Complex(vec[2u + 1], vec[2u + 2]) end @@ -63,7 +63,7 @@ mk_run_exec(k, dst, pa, pb, k_block_length, alpha, beta) = (execute_tile!(k, dst # `full = false` checks only accumulate against the reference and allocations. # Separate functions, so the light check does not compile the full one. function mk_contract(k; full::Bool = true) - return @testset "$(nameof(typeof(k))){$(tile_size(k)[1]),$(tile_size(k)[2]),$(scalartype(k))}" begin + return @testset "$(nameof(typeof(k))){$(tile_size(k, 1)),$(tile_size(k, 2)),$(scalartype(k))}" begin mk_contract_light(k) full && mk_contract_full(k) end diff --git a/test/packing/test_kernel_descriptor.jl b/test/packing/test_kernel_descriptor.jl index 83acf2c..0e8f06f 100644 --- a/test/packing/test_kernel_descriptor.jl +++ b/test/packing/test_kernel_descriptor.jl @@ -1,5 +1,5 @@ using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEFormat, - InterleavedFormat, realtype, sliver_widths + InterleavedFormat, realtype, sliver_width @testset "KernelDescriptor" begin k = KernelDescriptor(Val(8), Val(6), Float64) @@ -23,14 +23,14 @@ end onem = ComplexKernelDescriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) fmas = ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) @test realtype(planar) === real(T) - @test (sliver_widths(planar)[1], sliver_widths(onem)[1], sliver_widths(fmas)[1]) == + @test (sliver_width(planar, 1), sliver_width(onem, 1), sliver_width(fmas, 1)) == (2MR, 4MR, 2MR) - @test sliver_widths(planar)[2] == sliver_widths(onem)[2] == sliver_widths(fmas)[2] == 2NR + @test sliver_width(planar, 2) == sliver_width(onem, 2) == sliver_width(fmas, 2) == 2NR @test packed_a_length(onem, 7) == 4MR * 7 @test packed_b_length(planar, 7) == 2NR * 7 @test_throws ArgumentError ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), PlanarFormat()) - @test sliver_widths(ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()))[2] == 2NR + @test sliver_width(ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()))[2] == 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/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index 051abcf..a607c0f 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -55,7 +55,7 @@ using StridedViews: StridedView k = SIMDKernel(Val(MR), Val(NR), T, Val(W)) @test (tile_size(k)..., lanewidth(k)) == (MR, NR, W) @test (MR ÷ W) * NR + (MR ÷ W) <= 32 - @test sliver_widths(k) === (MR, NR) + @test sliver_width(k) === (MR, NR) @test realtype(k) === T === scalartype(k) @test complex_method(k) === RealMethod() @test packed_a_length(k, 7) === MR * 7 && packed_b_length(k, 7) === NR * 7 @@ -82,11 +82,11 @@ using StridedViews: StridedView @testset "extent demotion when M cannot fill a register tile" begin for T in (Float64, Float32) - MR = tile_size(_default_kernel(T))[1] - @test tile_size(_default_kernel(T, 4 * MR, 256))[1] == MR + MR = tile_size(_default_kernel(T), 1) + @test tile_size(_default_kernel(T, 4 * MR, 256), 1) == MR small = _default_kernel(T, 1, 256) @test (tile_size(small)..., lanewidth(small)) === _fallback_shape(T) - @test tile_size(_default_kernel(T, 0, 256))[1] == MR # empty must not demote + @test tile_size(_default_kernel(T, 0, 256), 1) == MR # empty must not demote end end @@ -111,9 +111,9 @@ using StridedViews: StridedView # Through `_default_kernel`, where the host's shape is the tall one. if _derived_shape(target_profile(), T) === tall @test tile_size(_default_kernel(T, MR + W, 256)) == (2 * W, NR) - @test tile_size(_default_kernel(T, 2 * W, 256))[1] == 2 * W - @test tile_size(_default_kernel(T, 2 * W - 1, 256))[1] == _fallback_shape(T)[1] - @test tile_size(_default_kernel(T, 2 * MR, 256))[1] == MR + @test tile_size(_default_kernel(T, 2 * W, 256), 1) == 2 * W + @test tile_size(_default_kernel(T, 2 * W - 1, 256), 1) == _fallback_shape(T)[1] + @test tile_size(_default_kernel(T, 2 * MR, 256), 1) == MR end end end @@ -159,7 +159,7 @@ using StridedViews: StridedView StridedView(zeros(T, d, d, d, d, d, d)), StridedView(randn(T, d, d, d, d)), (i, j, m, b), Bv, (m, k, a, c), (a, b, c, i, j, k); oracle = false ) - @test tile_size(plan.kernel)[1] == 16 + @test tile_size(plan.kernel, 1) == 16 @test plan.Astorage === parent(Bv) q, r, s = -1, 3, 4 B2 = StridedView(randn(T, d, d, d, d)) @@ -167,7 +167,7 @@ using StridedViews: StridedView StridedView(zeros(T, d, d, d, d)), StridedView(randn(T, d, d)), (q, b), B2, (a, q, r, s), (a, b, r, s); oracle = false ) - @test tile_size(plan2.kernel)[1] == 16 + @test tile_size(plan2.kernel, 1) == 16 @test plan2.Astorage === parent(B2) @test axis_length(plan2.mgroup) == d^3 end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 950bd05..c928f65 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -203,8 +203,8 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() # One (N, K, M) block, so the hoisted checks run once. @test base.blocking.m_block >= Ma && base.blocking.n_block >= Na && base.blocking.k_block >= Ka - m_tiles = cld(Ma, tile_size(kernel)[1]) - n_tiles = cld(Na, tile_size(kernel)[2]) + m_tiles = cld(Ma, tile_size(kernel, 1)) + n_tiles = cld(Na, tile_size(kernel, 2)) ntiles = m_tiles * n_tiles @test ntiles >= 20 # a per-tile check would be at least this many calls @@ -289,15 +289,15 @@ end n_length = axis_length(base.ngroup) @test n_length == N1 * N2 @test base.blocking.n_block >= n_length # one N block, so 7 slivers - nsliv = cld(n_length, tile_size(kernel)[2]) + nsliv = cld(n_length, tile_size(kernel, 2)) @test nsliv == 7 noffs = [offsets(base.ngroup, q)[2] for q in 0:(n_length - 1)] binding = argmax(noffs) - 1 # zero-based logical coordinate @test binding == N1 - 1 - @test 0 < binding ÷ tile_size(kernel)[2] < nsliv - 1 # a strictly interior sliver - @test maximum(noffs[1:tile_size(kernel)[2]]) < noffs[binding + 1] # not the first - @test maximum(noffs[(1 + (nsliv - 1) * tile_size(kernel)[2]):end]) < noffs[binding + 1] # not the last + @test 0 < binding ÷ tile_size(kernel, 2) < nsliv - 1 # a strictly interior sliver + @test maximum(noffs[1:tile_size(kernel, 2)]) < noffs[binding + 1] # not the first + @test maximum(noffs[(1 + (nsliv - 1) * tile_size(kernel, 2)):end]) < noffs[binding + 1] # not the last _pcf_exec(base, 1.0, 0.0) ref = zeros(M, N1, N2) @@ -319,7 +319,7 @@ end # A first- or last-sliver-only range would accept it. moffs = [offsets(base.mgroup, q)[2] for q in 0:(axis_length(base.mgroup) - 1)] mrange = (minimum(moffs), maximum(moffs)) - NR = tile_size(kernel)[2] + NR = tile_size(kernel, 2) truerange = (minimum(noffs), maximum(noffs)) firstonly = (minimum(noffs[1:NR]), maximum(noffs[1:NR])) lastonly = let tail = noffs[(1 + (nsliv - 1) * NR):end] diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 9810c1e..78caa65 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -360,7 +360,7 @@ end plan = plan_contract(Cv, Av, indA, Bv, indB, indC; kernel = kernel) swapped = plan.Astorage === parent(Bv) @test swapped == expect_swap - @test swapped == _lo_swap(msorted, nsorted, indC, Cv, tile_size(kernel)[1]) + @test swapped == _lo_swap(msorted, nsorted, indC, Cv, tile_size(kernel, 1)) if swapped # B feeds M: mgroup's maps are (B, C) over the sorted N labels. @test plan.mgroup.strides[2] == Tuple(cstride(l) for l in nsorted) @@ -383,7 +383,7 @@ end # Default kernel: the same rule at this machine's `m_tile`. plan = plan_contract(Cv, Av, indA, Bv, indB, indC) - m_tile = tile_size(plan.kernel)[1] + m_tile = tile_size(plan.kernel, 1) @test (plan.Astorage === parent(Bv)) == (mrun < m_tile && nrun >= m_tile) end end @@ -416,7 +416,7 @@ end @test plan.Astorage === parent(Av) default_kernel = QuasiStrided._default_kernel(T, m_length, n_length) - default_m_tile = tile_size(default_kernel)[1] + default_m_tile = tile_size(default_kernel, 1) if m_length == run || run % default_m_tile == 0 @test plan.kernel === default_kernel else @@ -424,8 +424,8 @@ end if isempty(candidates) @test plan.kernel === default_kernel else - @test run % tile_size(plan.kernel)[1] == 0 - @test tile_size(plan.kernel)[1] == maximum(candidates) + @test run % tile_size(plan.kernel, 1) == 0 + @test tile_size(plan.kernel, 1) == maximum(candidates) end end @@ -497,19 +497,19 @@ end mlab, = QuasiStrided._classify_labels(indA, indB, indC) msorted = _lo_order(mlab, indC, Cv2) run = _lo_run(msorted, indC, Cv2) - if m_length == run || run % tile_size(default_kernel)[1] == 0 + if m_length == run || run % tile_size(default_kernel, 1) == 0 @test plan_shallow.kernel === default_kernel else @test plan_shallow.kernel !== default_kernel @test typeof(plan_shallow.kernel) !== typeof(default_kernel) - @test run % tile_size(plan_shallow.kernel)[1] == 0 + @test run % tile_size(plan_shallow.kernel, 1) == 0 end # The guard is `k_length > kmax`: `k_length == kmax` may still demote. Cv_b, Av_b, Bv_b = _run_demote_fixture(T, a, kmax_of(T)) plan_b = plan_contract(Cv_b, Av_b, indA, Bv_b, indB, indC) run_b = _lo_run(msorted, indC, Cv_b) # same M order at every m in this fixture - if m_length == run_b || run_b % tile_size(default_kernel)[1] == 0 + if m_length == run_b || run_b % tile_size(default_kernel, 1) == 0 @test plan_b.kernel === default_kernel else @test plan_b.kernel !== default_kernel @@ -598,7 +598,7 @@ end for T in (ComplexF64, ComplexF32) W = QuasiStrided._default_lanewidth(real(T)) kernel = QuasiStrided.PlanarKernel(Val(W), Val(8), T, Val(W)) - @test tile_size(kernel)[1] <= 16 + @test tile_size(kernel, 1) <= 16 (indA, indB, indC), _ = _lo_labels(_LO_CASES[3][2], _LO_CASES[3][3]) A = randn(T, d, d, d, d) B = randn(T, d, d, d, d) diff --git a/test/runtests.jl b/test/runtests.jl index f2ef324..c007900 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,7 +12,7 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip target_profile, cache_topology, unknown_target, _detect_isa, _detect_target, _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, _fallback_blocking, kernel_shapes, _parse_size, _count_cpu_list, NR_DEFAULT, - _rule_applies, _isa_nregisters, sliver_widths, realtype, + _rule_applies, _isa_nregisters, sliver_width, realtype, complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, a_reals, b_reals, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, _planar_pressure, _pack_split, _NestPath From 4a10fe14ccba33f1684294fe165ed2cba4271871 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 30 Sep 2026 16:35:33 -0400 Subject: [PATCH 04/60] Chunk 1: simplify hardware detection - The ISA is the CPUID probe's; the LLVM CPU-name table is gone (it could only narrow the probe, e.g. capping AVX-512 Zen 4/5 at AVX2 on Julia 1.10). - `CacheLevel` drops the unused `ways`; `TargetProfile` drops `arch` and derives `vector_bytes`/`nregisters` from `isa`, rejecting unknown ISAs. - `core_bytes(profile, level)` replaces three copies of the per-core cache share in blocking.jl, pack_split.jl and labels.jl. - No leading underscores in target.jl; `_TARGET` is `TARGET`. - Tests: synthetic profiles take only the ISA, so the sweeps over impossible (isa, vector_bytes, nregisters) combinations go (34542 -> 33444 tests); the forced-ISA runner takes QS_FAKE_ISA without QS_FAKE_VB/QS_FAKE_NREG. Full suite passes (15 min). Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 20 ++- src/QuasiStrided.jl | 2 +- src/hardware/target.jl | 143 +++++++----------- src/microkernels/planar.jl | 2 +- src/packing/pack_contiguous.jl | 2 +- src/planning/blocking.jl | 6 +- src/planning/labels.jl | 6 +- src/planning/pack_split.jl | 2 +- test/forced_isa_runner.jl | 23 ++- test/hardware/test_target.jl | 40 ++--- test/helpers.jl | 5 +- test/microkernels/test_fmaddsub_kernel.jl | 2 +- .../test_planar_store_fastpath.jl | 12 +- test/packing/test_pack_complex_contiguous.jl | 12 +- test/planning/test_kernel_selection.jl | 75 +++++---- test/runtests.jl | 6 +- 16 files changed, 151 insertions(+), 207 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 086dab8..233e1ab 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -13,7 +13,7 @@ TTFX when codegen is touched). | # | Chunk | Files | Status | |---|---|---|---| | 0 | Tour | `QuasiStrided.jl`, `plan.jl`, `execute.jl` call path | done | -| 1 | Hardware | `hardware/target.jl` | | +| 1 | Hardware | `hardware/target.jl` | done | | 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | | | 3 | Tiles | `layout/tiles.jl` | | | 4 | Packing formats | `packing/format.jl`, `panel.jl`, `transposed.jl` | | @@ -35,7 +35,7 @@ TTFX when codegen is touched). ## Decisions -### D1. Descriptive M/N/K names (applied up front, codebase-wide) +### D1. Descriptive M/N/K names (applied: 564a740, 58b98e8) Pattern `__`; M/N/K stay the organizing letters. @@ -81,8 +81,22 @@ Add a PrecompileTools workload for common cases (ranks ≤ 3, `Vector` storage, the four eltypes) on top of the barriers, not instead of them. Measure precompile time and pkgimage size first. -## Open items noted in the tour (to revisit in their chunk) +### D5. Hardware detection (chunk 1, applied) +Own implementation kept: CPUSummary/HostCPUFeatures detect at precompile +time, CpuId is x86-only without cache sharing, Hwloc needs a C library and +~1 s per process. The ISA is the CPUID probe's (name table dropped); +`CacheLevel` lost `ways`, `TargetProfile` lost `arch` and derives +`vector_bytes`/`nregisters` from `isa`; `core_bytes(profile, level)` replaces +three copies of the per-core share. + +## Open items (to revisit in their chunk) + +- Detected `l1d.line` is unused: `_K_LINE_BYTES = 64` in `labels.jl` (K-order + cost model) and `pack_split.jl`; thread the detected line size through + planning (chunks 9/11). +- `complex_fastpath_isa_eligible` is policy, not detection: move it next to + its consumers (chunks 5/7). - 16-argument `ContractPlan(...)` spelled out five times, `_PlanRequest` twice more: one "copy with fields replaced" helper (chunk 13/14). - `execute.jl` mixes public API, barrier/hint machinery, the C-panel path and diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index c3ee8d4..da4564e 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -62,7 +62,7 @@ export QuasiStridedBackend # Detect the hardware per process, not at precompile time: a cached .ji may be # loaded on a different CPU. function __init__() - _init_target!() + init_target!() return nothing end diff --git a/src/hardware/target.jl b/src/hardware/target.jl index 7f59555..928c5c8 100644 --- a/src/hardware/target.jl +++ b/src/hardware/target.jl @@ -3,59 +3,48 @@ # the fixed fallback shapes. """ - CacheLevel(bytes, ways, line, sharing) + CacheLevel(bytes, line, sharing) One detected cache level; any field may be `0` for "not detected". `sharing` is the number of logical CPUs sharing this level. """ struct CacheLevel bytes::Int - ways::Int line::Int sharing::Int end -CacheLevel() = CacheLevel(0, 0, 0, 0) +CacheLevel() = CacheLevel(0, 0, 0) """ - TargetProfile + TargetProfile(isa, cpu_name, l1d, l2, l3) What was detected about the host CPU. `isa` is one of `:avx512`, `:avx2`, -`:neon` or `:unknown`. +`:neon` or `:unknown`, and fixes `vector_bytes` and `nregisters`. """ struct TargetProfile isa::Symbol - arch::Symbol cpu_name::String vector_bytes::Int nregisters::Int l1d::CacheLevel l2::CacheLevel l3::CacheLevel -end -unknown_target() = TargetProfile( - :unknown, Sys.ARCH, "", 0, 0, CacheLevel(), CacheLevel(), CacheLevel() -) - -# LLVM CPU name -> vector ISA. Unlisted names fall through to the CPUID probe. -const _UARCH_ISA = Dict{String, Symbol}( - n => :avx512 for n in ( - "skylake-avx512", "cascadelake", "cooperlake", "cannonlake", - "icelake-client", "icelake-server", "tigerlake", "rocketlake", - "sapphirerapids", "emeraldrapids", "graniterapids", "knl", "knm", - "znver4", "znver5", - ) -) -for n in ( - "haswell", "broadwell", "skylake", "alderlake", "raptorlake", - "meteorlake", "sierraforest", "grandridge", "tremont", "goldmont", - "goldmont-plus", "znver1", "znver2", "znver3", "bdver4", - ) - _UARCH_ISA[n] = :avx2 + function TargetProfile(isa::Symbol, cpu_name::AbstractString, l1d::CacheLevel, l2::CacheLevel, l3::CacheLevel) + isa in (:avx512, :avx2, :neon, :unknown) || + throw(ArgumentError("TargetProfile: unknown ISA $(repr(isa))")) + return new(isa, cpu_name, isa_vector_bytes(isa), isa_nregisters(isa), l1d, l2, l3) + end end +unknown_target() = TargetProfile(:unknown, "", CacheLevel(), CacheLevel(), CacheLevel()) + +# The register count is the ISA's, not the lane width's: AVX512VL gives 32 ymm registers. +isa_vector_bytes(isa::Symbol) = isa === :avx512 ? 64 : isa === :avx2 ? 32 : isa === :neon ? 16 : 0 +isa_nregisters(isa::Symbol) = isa === :avx512 ? 32 : isa === :avx2 ? 16 : isa === :neon ? 32 : 0 + # `Base.BinaryPlatforms.CPUID` is undocumented, so a failure degrades to `:unknown`. -function _isa_from_cpuid() +function isa_from_cpuid() return try C = Base.BinaryPlatforms.CPUID C.test_cpu_feature(C.JL_X86_avx512f) ? :avx512 : @@ -65,39 +54,25 @@ function _isa_from_cpuid() end end -_isa_rank(k::Symbol) = k === :avx512 ? 2 : k === :avx2 ? 1 : 0 - -function _detect_isa() - if Sys.ARCH === :x86_64 || Sys.ARCH === :i686 - key = get(_UARCH_ISA, Sys.CPU_NAME, :miss) - key === :miss && return _isa_from_cpuid() - # A hypervisor can mask CPUID features without changing the CPU name; - # AVX-512 code would then SIGILL, so a narrower live probe wins. - probe = _isa_from_cpuid() - probe === :unknown && return key - return _isa_rank(probe) < _isa_rank(key) ? probe : key - elseif Sys.ARCH === :aarch64 - # SVE is not detected: its runtime vector length cannot be a fixed kernel shape. - return :neon - end +function detect_isa() + (Sys.ARCH === :x86_64 || Sys.ARCH === :i686) && return isa_from_cpuid() + # SVE is not detected: its runtime vector length cannot be a fixed kernel shape. + Sys.ARCH === :aarch64 && return :neon return :unknown end -# The register count is the ISA's, not the lane width's: AVX512VL gives 32 ymm registers. -_isa_vector_bytes(::Val{K}) where {K} = K === :avx512 ? 64 : K === :avx2 ? 32 : K === :neon ? 16 : 0 -_isa_nregisters(::Val{K}) where {K} = K === :avx512 ? 32 : K === :avx2 ? 16 : K === :neon ? 32 : 0 - # "32K" / "1024K" / "2M" as Linux sysfs writes them. -function _parse_size(s::AbstractString) +function parse_size(s::AbstractString) s = strip(s) isempty(s) && return 0 - mult = get(Dict('K' => 1024, 'M' => 1024^2, 'G' => 1024^3), uppercase(s[end]), 1) + unit = uppercase(s[end]) + mult = unit == 'K' ? 1024 : unit == 'M' ? 1024^2 : unit == 'G' ? 1024^3 : 1 mult == 1 || (s = s[1:(end - 1)]) return something(tryparse(Int, strip(s)), 0) * mult end # "0,16" -> 2; "0-7,16-23" -> 16; "" -> 0. -function _count_cpu_list(s::AbstractString) +function count_cpu_list(s::AbstractString) n = 0 for part in split(strip(s), ',') isempty(part) && continue @@ -110,31 +85,30 @@ function _count_cpu_list(s::AbstractString) return n end -_read_or(path, default = "") = try +read_or(path, default = "") = try isfile(path) ? chomp(read(path, String)) : default catch default end -_int_or(path) = something(tryparse(Int, _read_or(path)), 0) +int_or(path) = something(tryparse(Int, read_or(path)), 0) -function _cache_topology_linux() +function cache_topology_linux() base = "/sys/devices/system/cpu/cpu0/cache" isdir(base) || return nothing levels = Dict{Symbol, CacheLevel}() for entry in readdir(base) startswith(entry, "index") || continue d = joinpath(base, entry) - level = tryparse(Int, _read_or(joinpath(d, "level"))) + level = tryparse(Int, read_or(joinpath(d, "level"))) level === nothing && continue - kind = _read_or(joinpath(d, "type")) + kind = read_or(joinpath(d, "type")) key = level == 1 ? (kind == "Data" ? :l1d : :skip) : level == 2 ? :l2 : level == 3 ? :l3 : :skip key === :skip && continue levels[key] = CacheLevel( - _parse_size(_read_or(joinpath(d, "size"))), - _int_or(joinpath(d, "ways_of_associativity")), - _int_or(joinpath(d, "coherency_line_size")), - _count_cpu_list(_read_or(joinpath(d, "shared_cpu_list"))), + parse_size(read_or(joinpath(d, "size"))), + int_or(joinpath(d, "coherency_line_size")), + count_cpu_list(read_or(joinpath(d, "shared_cpu_list"))), ) end return ( @@ -143,25 +117,25 @@ function _cache_topology_linux() ) end -_sysctl_int(name) = try +sysctl_int(name) = try something(tryparse(Int, chomp(read(`sysctl -n $name`, String))), 0) catch 0 end # macOS exposes no associativity; `cpusperl2` gives the L2 sharing. -function _cache_topology_darwin() - line = _sysctl_int("hw.cachelinesize") - pick(a, b) = (v = _sysctl_int(a); v == 0 ? _sysctl_int(b) : v) +function cache_topology_darwin() + line = sysctl_int("hw.cachelinesize") + pick(a, b) = (v = sysctl_int(a); v == 0 ? sysctl_int(b) : v) return ( - l1d = CacheLevel(pick("hw.perflevel0.l1dcachesize", "hw.l1dcachesize"), 0, line, 1), + l1d = CacheLevel(pick("hw.perflevel0.l1dcachesize", "hw.l1dcachesize"), line, 1), l2 = CacheLevel( - pick("hw.perflevel0.l2cachesize", "hw.l2cachesize"), 0, line, - _sysctl_int("hw.perflevel0.cpusperl2") + pick("hw.perflevel0.l2cachesize", "hw.l2cachesize"), line, + sysctl_int("hw.perflevel0.cpusperl2") ), l3 = CacheLevel( - _sysctl_int("hw.l3cachesize"), 0, line, - _sysctl_int("hw.perflevel0.logicalcpu") + sysctl_int("hw.l3cachesize"), line, + sysctl_int("hw.perflevel0.logicalcpu") ), ) end @@ -174,25 +148,21 @@ Linux sysfs or macOS `sysctl`, or `nothing` if unavailable. """ function cache_topology() return try - Sys.islinux() ? _cache_topology_linux() : - Sys.isapple() ? _cache_topology_darwin() : nothing + Sys.islinux() ? cache_topology_linux() : + Sys.isapple() ? cache_topology_darwin() : nothing catch nothing end end -function _detect_target() - key = _detect_isa() +function detect_target() topo = cache_topology() l1d, l2, l3 = topo === nothing ? (CacheLevel(), CacheLevel(), CacheLevel()) : (topo.l1d, topo.l2, topo.l3) - return TargetProfile( - key, Sys.ARCH, Sys.CPU_NAME, - _isa_vector_bytes(Val(key)), _isa_nregisters(Val(key)), l1d, l2, l3 - ) + return TargetProfile(detect_isa(), Sys.CPU_NAME, l1d, l2, l3) end -const _TARGET = Ref{TargetProfile}(unknown_target()) +const TARGET = Ref{TargetProfile}(unknown_target()) """ target_profile() -> TargetProfile @@ -200,19 +170,16 @@ const _TARGET = Ref{TargetProfile}(unknown_target()) The [`TargetProfile`](@ref) detected for this process (all `:unknown` if detection failed). """ -target_profile() = _TARGET[] +target_profile() = TARGET[] -function _init_target!() - _TARGET[] = try - _detect_target() - catch - unknown_target() - end - return nothing -end +init_target!() = (TARGET[] = detect_target(); nothing) + +# This core's share of `level`: the CPUs sharing it over the SMT threads per +# core, which are the CPUs sharing L1d. +core_bytes(profile::TargetProfile, level::CacheLevel) = + level.bytes ÷ max(1, level.sharing ÷ max(1, profile.l1d.sharing)) # Shared by the deinterleaving complex packer and the planar store: both pay # off only with 512-bit vector registers. -@inline _complex_fastpath_isa_eligible(profile::TargetProfile) = - profile.vector_bytes == _isa_vector_bytes(Val(:avx512)) -@inline _complex_fastpath_isa_eligible() = _complex_fastpath_isa_eligible(target_profile()) +@inline complex_fastpath_isa_eligible(profile::TargetProfile) = profile.isa === :avx512 +@inline complex_fastpath_isa_eligible() = complex_fastpath_isa_eligible(target_profile()) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index d23793c..b9655bf 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -186,7 +186,7 @@ end # ISA the complex fast paths ship for (shared with the complex pack fast path). @inline _complex_vector_eligible(tile::QSTile, ::Type{T}) where {T} = _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) && - _complex_fastpath_isa_eligible() + complex_fastpath_isa_eligible() # Shuffle patterns built from `W` at specialization time, never hardcoded to # one ISA. `v` holds `W` complex values `[re_0, im_0, ..., re_{W-1}, im_{W-1}]`. diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 834ba40..d4e6a47 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -63,7 +63,7 @@ end 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() + complex_fastpath_isa_eligible() end # The shuffles read `re` lanes from `src` and `im` lanes from `alt`, so the diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index 6210af8..884b822 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -49,10 +49,8 @@ _real_blocking_row(profile::TargetProfile, ::Type{T}) where {T <: Union{Float32, function _modelled_blocking(profile::TargetProfile, ::Type{T}, MR::Int, NR::Int) where {T} l1, l2, l3 = profile.l1d, profile.l2, profile.l3 (l1.bytes > 0 && l2.bytes > 0) || return nothing - smt = max(1, l1.sharing) - core_share(c) = c.bytes ÷ max(1, c.sharing ÷ smt) - l2core = core_share(l2) - l3core = l3.bytes > 0 ? core_share(l3) : 0 + l2core = core_bytes(profile, l2) + l3core = core_bytes(profile, l3) S = sizeof(T) k_block = max(1, (l1.bytes ÷ 2) ÷ (NR * S)) m_block = max(MR, ((l2core ÷ 2) ÷ (k_block * S)) ÷ MR * MR) diff --git a/src/planning/labels.jl b/src/planning/labels.jl index db8927f..4f6713d 100644 --- a/src/planning/labels.jl +++ b/src/planning/labels.jl @@ -158,10 +158,8 @@ const _K_LINE_BYTES = 64 # The core's private L2 share, or 1 MB when undetected. function _l2_core_bytes(profile::TargetProfile) - l2 = profile.l2 - l2.bytes > 0 || return 1 << 20 - smt = max(1, profile.l1d.sharing) - return l2.bytes ÷ max(1, l2.sharing ÷ smt) + profile.l2.bytes > 0 || return 1 << 20 + return core_bytes(profile, profile.l2) end # Through the per-eltype cache; an eltype without a slot fails later in planning. diff --git a/src/planning/pack_split.jl b/src/planning/pack_split.jl index ca54993..7e7c7a6 100644 --- a/src/planning/pack_split.jl +++ b/src/planning/pack_split.jl @@ -93,7 +93,7 @@ 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)) + return l2 + core_bytes(profile, l3) end # Past the tests above the plan is for a large contraction; a dynamic call keeps diff --git a/test/forced_isa_runner.jl b/test/forced_isa_runner.jl index 6e90022..6efb852 100644 --- a/test/forced_isa_runner.jl +++ b/test/forced_isa_runner.jl @@ -1,15 +1,14 @@ # Run the whole test suite as if the host had a different vector ISA, to catch # tests that assert something true only on the host they were written on: # -# QS_FAKE_ISA=avx2 QS_FAKE_VB=32 QS_FAKE_NREG=16 \ -# julia --project=. test/forced_isa_runner.jl +# QS_FAKE_ISA=avx2 julia --project=. test/forced_isa_runner.jl # -# Useful triples: avx512/64/32, avx2/32/16, neon/16/32, unknown/0/0. The cache -# levels default to "not detected"; QS_FAKE_L1D, QS_FAKE_L2 (bytes) and -# QS_FAKE_L2_SHARING (CPUs) model a specific host, e.g. a 3-CPU M1: +# QS_FAKE_ISA is one of avx512, avx2, neon, unknown. The cache levels default to +# "not detected"; QS_FAKE_L1D, QS_FAKE_L2 (bytes) and QS_FAKE_L2_SHARING (CPUs) +# model a specific host, e.g. a 3-CPU M1: # -# QS_FAKE_ISA=neon QS_FAKE_VB=16 QS_FAKE_NREG=32 QS_FAKE_L1D=131072 \ -# QS_FAKE_L2=12582912 QS_FAKE_L2_SHARING=3 julia --project=. test/forced_isa_runner.jl +# QS_FAKE_ISA=neon QS_FAKE_L1D=131072 QS_FAKE_L2=12582912 QS_FAKE_L2_SHARING=3 \ +# julia --project=. test/forced_isa_runner.jl # # "runs on this host without throwing" in hardware/test_target.jl fails by # construction here. Run from the test environment for the Bumper testsets. @@ -17,15 +16,13 @@ using QuasiStrided const _QS = QuasiStrided -_QS._TARGET[] = _QS.TargetProfile( - Symbol(get(ENV, "QS_FAKE_ISA", "avx512")), Sys.ARCH, "forced-isa-runner", - parse(Int, get(ENV, "QS_FAKE_VB", "64")), - parse(Int, get(ENV, "QS_FAKE_NREG", "32")), +_QS.TARGET[] = _QS.TargetProfile( + Symbol(get(ENV, "QS_FAKE_ISA", "avx512")), "forced-isa-runner", let l1d = parse(Int, get(ENV, "QS_FAKE_L1D", "0")) - _QS.CacheLevel(l1d, 0, 0, l1d > 0 ? 1 : 0) + _QS.CacheLevel(l1d, 0, l1d > 0 ? 1 : 0) end, _QS.CacheLevel( - parse(Int, get(ENV, "QS_FAKE_L2", "0")), 0, 0, + parse(Int, get(ENV, "QS_FAKE_L2", "0")), 0, parse(Int, get(ENV, "QS_FAKE_L2_SHARING", "0")) ), _QS.CacheLevel() diff --git a/test/hardware/test_target.jl b/test/hardware/test_target.jl index 871c69c..43edfa4 100644 --- a/test/hardware/test_target.jl +++ b/test/hardware/test_target.jl @@ -6,35 +6,17 @@ using StridedViews: StridedView @testset "runs on this host without throwing" begin p = target_profile() # Logged so that a host-dependent CI failure can be read off the log. - @info "target profile" p Sys.CPU_NAME QuasiStrided._isa_from_cpuid() VERSION + @info "target profile" p Sys.CPU_NAME VERSION @test p.isa in VALID_ISAS - @test p.arch === Sys.ARCH - @test _detect_target().isa === p.isa - if p.isa === :unknown - @test (p.vector_bytes, p.nregisters) == (0, 0) - else - @test p.nregisters > 0 - @test p.vector_bytes > 0 && ispow2(p.vector_bytes) - end + @test detect_target().isa === p.isa + @test_throws ArgumentError synthetic(:sse) end - @testset "CPU name table and CPUID probe agree" begin - probe = QuasiStrided._isa_from_cpuid() - @test probe in (:avx512, :avx2, :unknown) - if Sys.ARCH === :x86_64 - rank = QuasiStrided._isa_rank - table = get(QuasiStrided._UARCH_ISA, Sys.CPU_NAME, :miss) - # A probe narrower than the table is a masked hypervisor, not a table - # bug. Julia 1.10 (LLVM 15) names AVX-512 Zen 4/5 "znver3", so there a - # wider probe is not a table bug either. - stale = VERSION < v"1.11" ? ("znver3",) : () - (table === :miss || probe === :unknown || Sys.CPU_NAME in stale) || - @test rank(probe) <= rank(table) - if table !== :miss && probe !== :unknown - @test _detect_isa() === (rank(probe) < rank(table) ? probe : table) - end - # If Base's undocumented CPUID names move, every unlisted CPU would - # silently lose its derived shape. + if Sys.ARCH === :x86_64 + @testset "the ISA is the CPUID probe's" begin + @test detect_isa() === QuasiStrided.isa_from_cpuid() + # If Base's undocumented CPUID names move, every x86 host would + # silently fall back to `:unknown`. C = Base.BinaryPlatforms.CPUID @test isdefined(C, :JL_X86_avx512f) @test isdefined(C, :JL_X86_avx2) @@ -47,7 +29,7 @@ using StridedViews: StridedView @test topo === nothing || topo isa NamedTuple topo === nothing && continue for lvl in (topo.l1d, topo.l2, topo.l3) - @test all(>=(0), (lvl.bytes, lvl.ways, lvl.line, lvl.sharing)) + @test all(>=(0), (lvl.bytes, lvl.line, lvl.sharing)) end end end @@ -57,10 +39,10 @@ using StridedViews: StridedView "32K" => 32 * 1024, "25344K" => 25344 * 1024, "2M" => 2 * 1024 * 1024, "512" => 512, "" => 0, "garbage" => 0, ) - @test _parse_size(str) == want + @test parse_size(str) == want end for (str, want) in ("0,16" => 2, "0-7,16-23" => 16, "3" => 1, "" => 0) - @test _count_cpu_list(str) == want + @test count_cpu_list(str) == want end end end diff --git a/test/helpers.jl b/test/helpers.jl index 40966d2..fa61167 100644 --- a/test/helpers.jl +++ b/test/helpers.jl @@ -45,10 +45,7 @@ const _INDC = (1, 4, 3) using QuasiStrided: TargetProfile, CacheLevel const VALID_ISAS = (:avx512, :avx2, :neon, :unknown) -synthetic(isakey, vb; nregisters::Int = 32) = TargetProfile( - isakey, Sys.ARCH, "synthetic", vb, nregisters, - CacheLevel(), CacheLevel(), CacheLevel() -) +synthetic(isakey) = TargetProfile(isakey, "synthetic", CacheLevel(), CacheLevel(), CacheLevel()) # Permuted A (with a zero-stride axis), negative-stride B, offset sliced C. function scattered_fixture(::Type{T}, a_n = 32, k_n = 32, b_n = 8, n_n = 32) where {T} diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index fa8dbf3..3c9db1a 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -135,7 +135,7 @@ const _QSF = QuasiStrided pp = packed_panel(bufp, 1, length(bufp)) @test _QSF._pack_complex_contiguous_eligible( pp, vals, src.rows, identity, InterleavedFormat(), MR, Val(MR), T - ) == _QSF._complex_fastpath_isa_eligible() + ) == _QSF.complex_fastpath_isa_eligible() end end end diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index e9a4441..97a1b89 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -14,7 +14,7 @@ using QuasiStrided: PlanarKernel, FMAddSubKernel, KERNEL_SHAPES_C64_FMADDSUB, KERNEL_SHAPES_C32_FMADDSUB, PtrScatterAxis, TargetProfile, CacheLevel, target_profile, unknown_target, KERNEL_SHAPES_C64_PLANAR, KERNEL_SHAPES_C32_PLANAR -const STORE_FASTPATH_ON = QuasiStrided._complex_fastpath_isa_eligible() +const STORE_FASTPATH_ON = QuasiStrided.complex_fastpath_isa_eligible() # `@noinline` identity: keeps LLVM from contracting separately rounded products. @noinline barrier(x) = x @@ -106,13 +106,11 @@ end end @testset "ISA gate: AVX-512 only" begin - profile(key, vb, nreg) = TargetProfile(key, Sys.ARCH, "t", vb, nreg, CacheLevel(), CacheLevel(), CacheLevel()) - @test QuasiStrided._complex_fastpath_isa_eligible(profile(:avx512, 64, 32)) - for (key, vb, nreg) in ((:avx2, 32, 16), (:neon, 16, 32), (:unknown, 0, 0)) - @test !QuasiStrided._complex_fastpath_isa_eligible(profile(key, vb, nreg)) + @test QuasiStrided.complex_fastpath_isa_eligible(synthetic(:avx512)) + for key in (:avx2, :neon, :unknown) + @test !QuasiStrided.complex_fastpath_isa_eligible(synthetic(key)) end - @test !QuasiStrided._complex_fastpath_isa_eligible(unknown_target()) - @test STORE_FASTPATH_ON == (target_profile().vector_bytes == QuasiStrided._isa_vector_bytes(Val(:avx512))) + @test STORE_FASTPATH_ON == (target_profile().isa === :avx512) end end diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index 90819cb..2424868 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -7,7 +7,7 @@ using QuasiStrided: target_profile, unknown_target, TargetProfile, CacheLevel, kernel_shapes, PlanarMethod, OneMMethod, FMAddSubMethod const QS = QuasiStrided -const FASTPATH_ON = QS._complex_fastpath_isa_eligible() +const FASTPATH_ON = QS.complex_fastpath_isa_eligible() # The extremes of each method's menu plus extents that are not multiples of any # lane width (interleaved shares 1e's shuffle); B is planar under every method. @@ -104,11 +104,9 @@ end end @testset "complex pack fast path: ISA gate is a register-width question" begin - profile(key, vb, nreg) = TargetProfile(key, Sys.ARCH, "t", vb, nreg, CacheLevel(), CacheLevel(), CacheLevel()) - @test QS._complex_fastpath_isa_eligible(profile(:avx512, 64, 32)) - for (key, vb, nreg) in ((:avx2, 32, 16), (:neon, 16, 32), (:unknown, 0, 0)) - @test !QS._complex_fastpath_isa_eligible(profile(key, vb, nreg)) + @test QS.complex_fastpath_isa_eligible(synthetic(:avx512)) + for key in (:avx2, :neon, :unknown) + @test !QS.complex_fastpath_isa_eligible(synthetic(key)) end - @test !QS._complex_fastpath_isa_eligible(unknown_target()) - @test FASTPATH_ON == (target_profile().vector_bytes == QS._isa_vector_bytes(Val(:avx512))) + @test FASTPATH_ON == (target_profile().isa === :avx512) end diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index a607c0f..7775764 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -22,7 +22,7 @@ using StridedViews: StridedView @testset "rule: MR = MV*W, NR = 6; MV = 4 on Intel :avx512, 2 on :avx2 and AMD" begin for (isakey, vb, MV) in ((:avx512, 64, 4), (:avx2, 32, 2)), T in (Float64, Float32) W = vb ÷ sizeof(T) - @test _derived_shape(synthetic(isakey, vb), T) === (MV * W, NR_DEFAULT, W) + @test _derived_shape(synthetic(isakey), T) === (MV * W, NR_DEFAULT, W) @test QuasiStrided._rule_mv(Val(isakey), RealMethod()) == MV @test QuasiStrided._rule_shape(vb, T, MV) === (MV * W, NR_DEFAULT, W) end @@ -31,16 +31,16 @@ using StridedViews: StridedView end for (cpu, MV) in (("znver4", 2), ("znver5", 2), ("icelake-server", 4), ("cascadelake", 4)), T in (Float64, Float32) - p = TargetProfile(:avx512, Sys.ARCH, cpu, 64, 32, CacheLevel(), CacheLevel(), CacheLevel()) + p = TargetProfile(:avx512, cpu, CacheLevel(), CacheLevel(), CacheLevel()) W = 64 ÷ sizeof(T) @test _derived_shape(p, T) === (MV * W, NR_DEFAULT, W) - @test _derived_shape(p, ComplexF64) === _derived_shape(synthetic(:avx512, 64), ComplexF64) + @test _derived_shape(p, ComplexF64) === _derived_shape(synthetic(:avx512), ComplexF64) end - @test _derived_shape(synthetic(:avx2, 32), Float64) === _fallback_shape(Float64) + @test _derived_shape(synthetic(:avx2), Float64) === _fallback_shape(Float64) # ISAs without a rule get the fallback shape, whatever their width. - for T in (Float64, Float32), vb in (0, 16, 32, 64) - @test _derived_shape(synthetic(:neon, vb), T) === _fallback_shape(T) - @test _derived_shape(synthetic(:unknown, vb), T) === _fallback_shape(T) + for T in (Float64, Float32) + @test _derived_shape(synthetic(:neon), T) === _fallback_shape(T) + @test _derived_shape(synthetic(:unknown), T) === _fallback_shape(T) end # No real override row on any ISA: the rule is the optimum. for T in (Float64, Float32), key in VALID_ISAS @@ -93,7 +93,7 @@ using StridedViews: StridedView @testset "_extent_shape: MV = 4 steps down to MV = 2 where it pads less" begin ext(shape, T, m_length) = QuasiStrided._extent_shape(shape, T, RealMethod(), m_length) for T in (Float64, Float32) - tall = _derived_shape(synthetic(:avx512, 64), T) + tall = _derived_shape(synthetic(:avx512), T) MR, NR, W = tall half = (2 * W, NR, W) @test ext(tall, T, 0) === tall @@ -104,9 +104,9 @@ using StridedViews: StridedView @test ext(tall, T, m_length) === tall end # Never below MV = 2, and complex methods are untouched. - avx2 = _derived_shape(synthetic(:avx2, 32), T) + avx2 = _derived_shape(synthetic(:avx2), T) @test ext(avx2, T, 3) === avx2 - cs = _derived_shape(synthetic(:avx512, 64), ComplexF64) + cs = _derived_shape(synthetic(:avx512), ComplexF64) @test QuasiStrided._extent_shape(cs, ComplexF64, PlanarMethod(), 3) === cs # Through `_default_kernel`, where the host's shape is the tall one. if _derived_shape(target_profile(), T) === tall @@ -121,7 +121,7 @@ using StridedViews: StridedView @testset "_store_shape: MV = 4 steps down where C's run breaks its slivers" begin st = QuasiStrided._store_shape for T in (Float64, Float32) - tall = _derived_shape(synthetic(:avx512, 64), T) + tall = _derived_shape(synthetic(:avx512), T) MR, NR, W = tall half = (2 * W, NR, W) m_length = 64 * MR @@ -137,7 +137,7 @@ using StridedViews: StridedView end @test st(half, T, RealMethod(), m_length, 3 * W) === half @test st(_fallback_shape(T), T, RealMethod(), m_length, 3 * W) === _fallback_shape(T) - cs = _derived_shape(synthetic(:avx512, 64), ComplexF64) + cs = _derived_shape(synthetic(:avx512), ComplexF64) @test st(cs, ComplexF64, PlanarMethod(), m_length, 3 * W) === cs if _derived_shape(target_profile(), T) === tall @test QuasiStrided._default_shape(T, m_length, 256, 2 * W)[1] === half @@ -179,7 +179,7 @@ end # W counts REAL lanes; the AVX-512 override replaces the spilling rule shape. @test QuasiStrided._rule_shape(64, T, 2) === (2 * W, NR_DEFAULT, W) @test _shape_override(Val(:avx512), T) === swept - @test _derived_shape(synthetic(:avx512, 64), T) === swept === first(kernel_shapes(T, PlanarMethod())) + @test _derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarMethod())) @test Set(kernel_shapes(T, PlanarMethod())) == Set( ( (2 * W, NR_DEFAULT, W), swept, (W, 8, W), _shape_override(Val(:avx2), T), @@ -195,24 +195,20 @@ end @test _fallback_shape(T) === (8, NR_DEFAULT, _fallback_shape(real(T))[3]) # Off :avx512 an override row wins whatever the width; without a row # the shape is fitted to the register budget. - for key in (:avx2, :neon), vb in (0, 16, 32, 64), nreg in (0, 16, 32) - @test !_rule_applies(Val(key), PlanarMethod()) - @test _derived_shape(synthetic(key, vb; nregisters = nreg), T) === _shape_override(Val(key), T) - end - for key in (:unknown, :somethingelse), vb in (0, 16, 32, 64), nreg in (0, 16, 32) + for key in (:avx2, :neon) @test !_rule_applies(Val(key), PlanarMethod()) - @test _shape_override(Val(key), T) === nothing - shape = _derived_shape(synthetic(key, vb; nregisters = nreg), T) - @test _planar_pressure(shape...) <= (nreg > 0 ? nreg : 16) + @test _derived_shape(synthetic(key), T) === _shape_override(Val(key), T) end - @test _derived_shape(synthetic(:avx512, 0), T) in kernel_shapes(T, PlanarMethod()) + @test !_rule_applies(Val(:unknown), PlanarMethod()) + @test _shape_override(Val(:unknown), T) === nothing + @test _planar_pressure(_derived_shape(synthetic(:unknown), T)...) <= 16 end @test _rule_applies(Val(:avx512), PlanarMethod()) @test _planar_pressure(_fallback_shape(ComplexF64)...) == 30 @test _planar_pressure(_fallback_shape(ComplexF32)...) == 16 # Complex menus fit AVX-512 and stay bounded. - nreg = _isa_nregisters(Val(:avx512)) + nreg = isa_nregisters(:avx512) target_profile().isa === :avx512 && @test target_profile().nregisters == nreg for T in (ComplexF64, ComplexF32), m in (PlanarMethod(), OneMMethod()) planes = accumulator_planes(m) @@ -249,10 +245,10 @@ end @test dmethod(R, Float32, Float64) === RealMethod() end # The real extent, store and small-M demotions carry over. - saved = QuasiStrided._TARGET[] + 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 key in (:avx512, :avx2, :neon) + QuasiStrided.TARGET[] = synthetic(key) for T in (ComplexF64, ComplexF32), m in (CR, RC), m_length in (1, 3, 17, 40, 4096), run in (0, 1, 8, m_length) real_problem = m === CR ? (2 * m_length, 50, 2 * run) : (m_length, 100, run) @test QuasiStrided._default_shape(T, m, m_length, 50, run) === @@ -260,7 +256,7 @@ end end end finally - QuasiStrided._TARGET[] = saved + QuasiStrided.TARGET[] = saved end end @@ -269,9 +265,8 @@ end (Float64, RealMethod()), (Float32, RealMethod()), (ComplexF64, PlanarMethod()), (ComplexF32, PlanarMethod()), (ComplexF64, OneMMethod()), (ComplexF32, OneMMethod()), ) - for (T, m) in methods, isakey in (VALID_ISAS..., :somethingelse), - vb in (0, 16, 32, 64, 128), nreg in (0, 16, 32) - shape = _derived_shape(synthetic(isakey, vb; nregisters = nreg), T, m) + for (T, m) in methods, isakey in VALID_ISAS + shape = _derived_shape(synthetic(isakey), T, m) @test shape in kernel_shapes(T, m) k = QuasiStrided._kernel_from_shape(shape, T, m) @test (tile_size(k)..., lanewidth(k)) === shape @@ -282,12 +277,12 @@ end KiB, MiB = 1024, 1024^2 # No SMT with an L3 shared by 4 cores; SMT 2 with an L3 shared by 8 cores. nosmt = TargetProfile( - :avx2, :x86_64, "znver2", 32, 16, CacheLevel(32KiB, 8, 64, 1), - CacheLevel(512KiB, 8, 64, 1), CacheLevel(16MiB, 16, 64, 4) + :avx2, "znver2", CacheLevel(32KiB, 64, 1), + CacheLevel(512KiB, 64, 1), CacheLevel(16MiB, 64, 4) ) smt2 = TargetProfile( - :avx512, :x86_64, "cascadelake", 64, 32, CacheLevel(32KiB, 8, 64, 2), - CacheLevel(1MiB, 16, 64, 2), CacheLevel(25952256, 11, 64, 16) + :avx512, "cascadelake", CacheLevel(32KiB, 64, 2), + CacheLevel(1MiB, 64, 2), CacheLevel(25952256, 64, 16) ) @test _modelled_blocking(nosmt, Float64, 8, 6) === Blocking(96, 341, 1728) @test _modelled_blocking(nosmt, Float32, 16, 6) === Blocking(96, 682, 1728) @@ -302,17 +297,17 @@ end end # No L3: the B panel is budgeted from the L2 alone. nol3 = TargetProfile( - :neon, :aarch64, "", 16, 32, CacheLevel(64KiB, 4, 64, 1), - CacheLevel(4MiB, 16, 64, 4), CacheLevel() + :neon, "", CacheLevel(64KiB, 64, 1), + CacheLevel(4MiB, 64, 4), CacheLevel() ) @test _modelled_blocking(nol3, Float64, 4, 6).n_block == (1MiB ÷ (682 * 8)) ÷ 6 * 6 tiny = TargetProfile( - :avx2, :x86_64, "", 32, 16, CacheLevel(64, 1, 64, 1), CacheLevel(64, 1, 64, 1), CacheLevel() + :avx2, "", CacheLevel(64, 64, 1), CacheLevel(64, 64, 1), CacheLevel() ) @test _modelled_blocking(tiny, Float64, 8, 6) === Blocking(8, 1, 6) @test _modelled_blocking(unknown_target(), Float64, 8, 6) === nothing @test _modelled_blocking( - TargetProfile(:avx2, :x86_64, "", 32, 16, CacheLevel(32KiB, 8, 64, 1), CacheLevel(), CacheLevel()), + TargetProfile(:avx2, "", CacheLevel(32KiB, 64, 1), CacheLevel(), CacheLevel()), Float64, 8, 6 ) === nothing end @@ -375,7 +370,7 @@ end end @testset "_small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin - sms(T, m_length, p = synthetic(:avx512, 64)) = + sms(T, m_length, p = synthetic(:avx512)) = QuasiStrided._small_m_shape(QuasiStrided._small_m_candidates(Val(p.isa), p, T), m_length) # Least padded rows, then the larger tile. @test sms(ComplexF64, 12) === (12, 8, 8) @@ -390,7 +385,7 @@ end end @test sms(Float64, 4) === nothing for isakey in (:avx2, :neon, :unknown), T in (ComplexF64, ComplexF32) - @test sms(T, 2, synthetic(isakey, 32)) === nothing + @test sms(T, 2, synthetic(isakey)) === nothing end end diff --git a/test/runtests.jl b/test/runtests.jl index c007900..5640280 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,10 +9,10 @@ using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescrip ncols, axis_offset_range, checked_tile_storage_bounds, pack_a!, pack_b!, zero_accumulator, accumulate, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, - target_profile, cache_topology, unknown_target, _detect_isa, _detect_target, + target_profile, cache_topology, unknown_target, detect_isa, detect_target, _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, - _fallback_blocking, kernel_shapes, _parse_size, _count_cpu_list, NR_DEFAULT, - _rule_applies, _isa_nregisters, sliver_width, realtype, + _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, + _rule_applies, isa_nregisters, sliver_width, realtype, complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, a_reals, b_reals, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, _planar_pressure, _pack_split, _NestPath From 23c4a238161d1918b8e2fa4da21ff9a77fce99bd Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 09:23:09 -0400 Subject: [PATCH 05/60] Chunk 2: simplify axis groups - `pair_group.jl` becomes the `AxisGroup(labels, (ind1, v1), (ind2, v2))` constructor. - `normalize_group` is gone and `offsets` moves to the test helpers (both were test-only). - `map_ramp_step` moves from dot.jl; `affine_ramp` is built on it. - `Base.setindex`, `prod` and `Base.Checked` replace hand-written versions. - `BlockDescriptor`/`describe_block` move to tiles.jl; `describe_block` checks `count <= length(buffer) - first`, so `first + count` cannot overflow past the check. - `fill_offsets!` drops its alias check. - No leading underscores in axis_group.jl. Full suite passes (30989 tests, 15 min; the drop is the removed normalize_group/offsets/alias tests). Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 19 +- src/QuasiStrided.jl | 2 +- src/execution/barrier.jl | 2 +- src/execution/dot.jl | 27 +-- src/execution/execute.jl | 2 +- src/execution/outer.jl | 4 +- src/layout/axis_group.jl | 236 ++++++--------------- src/layout/pair_group.jl | 24 --- src/layout/tiles.jl | 43 ++++ src/planning/labels.jl | 10 +- src/planning/pack_split.jl | 2 +- src/planning/plan.jl | 8 +- test/helpers.jl | 6 + test/layout/test_axis_group.jl | 24 +-- test/layout/test_stridedviews_axisgroup.jl | 9 +- test/planning/test_per_call_overhead.jl | 11 +- test/runtests.jl | 4 +- 17 files changed, 162 insertions(+), 271 deletions(-) delete mode 100644 src/layout/pair_group.jl diff --git a/CLEANUP.md b/CLEANUP.md index 233e1ab..87dcd80 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -14,7 +14,7 @@ TTFX when codegen is touched). |---|---|---|---| | 0 | Tour | `QuasiStrided.jl`, `plan.jl`, `execute.jl` call path | done | | 1 | Hardware | `hardware/target.jl` | done | -| 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | | +| 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | done | | 3 | Tiles | `layout/tiles.jl` | | | 4 | Packing formats | `packing/format.jl`, `panel.jl`, `transposed.jl` | | | 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | | @@ -90,6 +90,23 @@ time, CpuId is x86-only without cache sharing, Hwloc needs a C library and `vector_bytes`/`nregisters` from `isa`; `core_bytes(profile, level)` replaces three copies of the per-core share. +### D6. Axis groups (chunk 2, applied) + +`normalize_group` and `offsets` leave `src` (`offsets` is a test helper); +`map_ramp_step` moves from dot.jl and `affine_ramp` is built on it; +`Base.setindex`/`prod`/`Base.Checked` replace hand-written versions; +`BlockDescriptor` moves to tiles.jl; `pair_group.jl` becomes the +`AxisGroup(labels, (ind1, v1), (ind2, v2))` constructor; `fill_offsets!` +drops its alias check. Kept: generic `P` (a batch label in A, B and C would +be a `P = 3` group) and unconditional range checks in `fill_offsets!` +(`@boundscheck`/`@propagate_inbounds` only help when inlined). + +## Possible improvements + +- Piecewise-affine block descriptions instead of block-sized offset buffers: + https://github.com/lkdvos/QuasiStrided.jl/issues/13 (decide after chunks + 3, 5, 12, 14). + ## Open items (to revisit in their chunk) - Detected `l1d.line` is unused: `_K_LINE_BYTES = 64` in `labels.jl` (K-order diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index da4564e..ce92c5a 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -1,6 +1,7 @@ module QuasiStrided using StridedViews: StridedView, offset +using Base.Checked: checked_abs, checked_add, checked_mul # TensorOperations names are always qualified: a bare `using` collides on `scalartype`. import TensorOperations as TO @@ -15,7 +16,6 @@ include("hardware/target.jl") # --- Layout: zero-based strided/scattered addressing --- include("layout/axis_group.jl") -include("layout/pair_group.jl") include("layout/tiles.jl") # --- Packing: packed-panel formats and the packers that fill them --- diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index 4f4ef27..b29e767 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -44,7 +44,7 @@ end # 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 +@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. A split group is diff --git a/src/execution/dot.jl b/src/execution/dot.jl index eb440d8..d1ce4b1 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -20,29 +20,6 @@ end # (a complex output has two accumulator planes). _dot_group_width(::Type{T}) where {T} = T <: Complex ? 4 : 8 -# The affine step of map `p` of `g` alone, or `nothing` if it is not a single -# ramp; `affine_ramp` asks the same of all maps at once. -function _map_ramp_step(g::AxisGroup{D, P}, p::Int) where {D, P} - step = 0 - run = 1 - started = false - for d in 1:D - L = g.lengths[d] - L == 0 && return 0 - L == 1 && continue - S = g.strides[p][d] - if !started - step = S - run = L - started = true - else - Int128(run) * Int128(step) == Int128(S) || return nothing - run *= L - end - end - return step -end - # Whether the dot path applies (all but the workspace capacity): a degenerate # free extent, a matrix operand with unit-ramp K in dense storage (raw-pointer # loads), and at least one vector of K. @@ -52,10 +29,10 @@ function _dot_applicable(::Type{T}, Astorage, Bstorage, kgroup::AxisGroup, m_len k_length >= _dot_lanewidth(T) || return false if m_length == 1 Bstorage isa DenseVector{T} || return false - _map_ramp_step(kgroup, 2) == 1 || return false + map_ramp_step(kgroup, 2) == 1 || return false else Astorage isa DenseVector{T} || return false - _map_ramp_step(kgroup, 1) == 1 || return false + map_ramp_step(kgroup, 1) == 1 || return false end return true end diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 42235e6..7702f38 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -235,7 +235,7 @@ 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)) + dense = ntuple(d -> step * prod(ntuple(i -> i < d ? g.lengths[i] : 1, Val(D))), Val(D)) return AxisGroup(g.lengths, (g.strides[1], dense)) end diff --git a/src/execution/outer.jl b/src/execution/outer.jl index 80f7e91..3f444ed 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -15,8 +15,8 @@ function _outer_applicable(::Type{T}, Astorage, Cstorage, mgroup::AxisGroup, m_l T <: Real || return false (Astorage isa DenseVector{T} && Cstorage isa DenseVector{T}) || return false m_length >= _dot_lanewidth(T) || return false - _map_ramp_step(mgroup, 1) == 1 || return false - return _map_ramp_step(mgroup, 2) == 1 + map_ramp_step(mgroup, 1) == 1 || return false + return map_ramp_step(mgroup, 2) == 1 end function _execute_outer!( diff --git a/src/layout/axis_group.jl b/src/layout/axis_group.jl index 59ef8fd..4e107cc 100644 --- a/src/layout/axis_group.jl +++ b/src/layout/axis_group.jl @@ -14,7 +14,7 @@ struct AxisGroup{D, P} for (d, L) in enumerate(lengths) L >= 0 || throw(ArgumentError("AxisGroup lengths must be nonnegative, got lengths[$d] = $L")) end - _validate_axis_group_bounds(lengths, strides) + validate_axis_group_bounds(lengths, strides) return new{D, P}(lengths, strides) end end @@ -22,80 +22,68 @@ end AxisGroup(lengths::NTuple{D, Int}, strides::NTuple{P, NTuple{D, Int}}) where {D, P} = AxisGroup{D, P}(lengths, strides) -# Validation-time only; Int128 so nothing (incl. abs(typemin(Int))) can wrap. -function _checked_axis_length(lengths::NTuple{D, Int}) where {D} - any(==(0), lengths) && return 0 - q = one(Int128) - for L in lengths - q *= Int128(L) - q > Int128(typemax(Int)) && - throw(OverflowError("AxisGroup cardinality (product of lengths) exceeds typemax(Int)")) - end - return Int(q) -end - -function _validate_axis_group_bounds( +# The checks that make all later offset arithmetic safe in plain `Int`: the +# cardinality, and every map's offset range `sum((L[d]-1)*abs(S[d]))`. +function validate_axis_group_bounds( lengths::NTuple{D, Int}, strides::NTuple{P, NTuple{D, Int}} ) where {D, P} - Q = _checked_axis_length(lengths) - Q == 0 && return nothing - for (p, S) in enumerate(strides) - acc = zero(Int128) + any(iszero, lengths) && return nothing + foldl(checked_mul, lengths; init = 1) + for S in strides + acc = 0 for d in 1:D - acc += Int128(lengths[d] - 1) * abs(Int128(S[d])) - acc > Int128(typemax(Int)) && - throw( - OverflowError( - "AxisGroup map $p: sum((L[d]-1)*abs(S[d])) exceeds typemax(Int); " * - "this layout is not representable under the conservative offset-range bound" - ) - ) + lengths[d] > 1 || continue + acc = checked_add(acc, checked_mul(lengths[d] - 1, checked_abs(S[d]))) end end return nothing end -axis_length(g::AxisGroup) = _unchecked_axis_length(g.lengths) - -@inline function _unchecked_axis_length(lengths::NTuple{D, Int}) where {D} - q = 1 - for L in lengths - q *= L +""" + AxisGroup(labels, (ind1, v1), (ind2, v2)) + +The two-map group of `labels` in operands `v1` and `v2`, whose axes carry the +labels `ind1` and `ind2`. Throws a `DimensionMismatch` if a label's axis +lengths differ. +""" +@inline function AxisGroup( + labels::NTuple{D, Int}, + (ind1, v1)::Tuple{NTuple{N1, Int}, StridedView}, + (ind2, v2)::Tuple{NTuple{N2, Int}, StridedView} + ) where {D, N1, N2} + # `Base.strides` on a `StridedView` rebuilds a tuple: call it once. + st1 = Base.strides(v1) + st2 = Base.strides(v2) + pos1 = ntuple(d -> findfirst(==(@inbounds labels[d]), ind1)::Int, Val(D)) + pos2 = ntuple(d -> findfirst(==(@inbounds labels[d]), ind2)::Int, Val(D)) + lens = ntuple(Val(D)) do d + l1 = size(v1, pos1[d]) + l2 = size(v2, pos2[d]) + l1 == l2 || throw_label_length((@inbounds labels[d]), l1, l2) + l1 end - return q + s1 = ntuple(d -> st1[pos1[d]], Val(D)) + s2 = ntuple(d -> st2[pos2[d]], Val(D)) + return AxisGroup(lens, (s1, s2)) end -# Type-stable, stack-allocated "replace element i of an NTuple{N,Int}". -@inline _tupleset(t::NTuple{N, Int}, i::Int, v::Int) where {N} = - ntuple(j -> ifelse(j == i, v, t[j]), Val(N)) +@noinline throw_label_length(label::Int, l1::Int, l2::Int) = throw( + DimensionMismatch("label $label has mismatched axis length: $l1 vs $l2") +) + +axis_length(g::AxisGroup) = prod(g.lengths) # Not inline lambdas: a closure over loop-reassigned variables is boxed and # allocates every call. -@inline function _add_offsets(offs::NTuple{P, Int}, g::AxisGroup{D, P}, d::Int, x::Int) where {D, P} +@inline function add_offsets(offs::NTuple{P, Int}, g::AxisGroup{D, P}, d::Int, x::Int) where {D, P} return ntuple(p -> offs[p] + x * g.strides[p][d], Val(P)) end -@inline function _sub_reset_offsets(offs::NTuple{P, Int}, g::AxisGroup{D, P}, d::Int, Ld::Int) where {D, P} +@inline function sub_reset_offsets(offs::NTuple{P, Int}, g::AxisGroup{D, P}, d::Int, Ld::Int) where {D, P} return ntuple(p -> offs[p] - (Ld - 1) * g.strides[p][d], Val(P)) end -function offsets(g::AxisGroup{D, P}, q::Int) where {D, P} - Q = axis_length(g) - (0 <= q < Q) || throw(BoundsError(g, q)) - r = q - offs = ntuple(_ -> 0, Val(P)) - for d in 1:D - L = g.lengths[d] - x = r % L - r = r ÷ L - if x != 0 - offs = _add_offsets(offs, g, d, x) - end - end - return offs -end - function fill_offsets!( buffers::NTuple{P, Vector{Int}}, g::AxisGroup{D, P}, first::Int, count::Int @@ -112,10 +100,6 @@ function fill_offsets!( length(buffers[p]) >= count || throw(DimensionMismatch("buffer $p has length $(length(buffers[p])), need at least $count")) end - for i in 1:P, j in (i + 1):P - buffers[i] === buffers[j] && - throw(ArgumentError("buffers must be distinct Vector{Int} objects (buffers $i and $j alias)")) - end count == 0 && return buffers @@ -127,9 +111,9 @@ function fill_offsets!( L = g.lengths[d] xd = r % L r = r ÷ L - x = _tupleset(x, d, xd) + x = Base.setindex(x, xd, d) if xd != 0 - offs = _add_offsets(offs, g, d, xd) + offs = add_offsets(offs, g, d, xd) end end @@ -146,12 +130,12 @@ function fill_offsets!( L = g.lengths[d] xd = x[d] if xd < L - 1 - x = _tupleset(x, d, xd + 1) - offs = _add_offsets(offs, g, d, 1) + x = Base.setindex(x, xd + 1, d) + offs = add_offsets(offs, g, d, 1) break else - x = _tupleset(x, d, 0) - offs = _sub_reset_offsets(offs, g, d, L) + x = Base.setindex(x, 0, d) + offs = sub_reset_offsets(offs, g, d, L) d += 1 end end @@ -160,121 +144,37 @@ function fill_offsets!( return buffers end -# Whether every map is a single ramp `offsets(g, q)[p] == q * steps[p]` (so the -# offset buffer can be replaced by arithmetic). `steps` is meaningless when not. -function affine_ramp(g::AxisGroup{D, P}) where {D, P} - zerosteps = ntuple(_ -> 0, Val(P)) - steps = zerosteps +# The step of map `p` if it is a single ramp, `offsets(g, q)[p] == q * step` +# (0 for an empty group), else `nothing`. +function map_ramp_step(g::AxisGroup{D}, p::Int) where {D} + step = 0 run = 1 started = false for d in 1:D L = g.lengths[d] - L == 0 && return (true, zerosteps) # empty domain: vacuously a ramp. - L == 1 && continue # singleton: coordinate never advances. - Sd = ntuple(p -> g.strides[p][d], Val(P)) + L == 0 && return 0 + L == 1 && continue + S = g.strides[p][d] if !started - steps = Sd + step = S run = L started = true else - for p in 1:P - Int128(run) * Int128(steps[p]) == Int128(Sd[p]) || - return (false, zerosteps) - end - run *= L # a sub-product of the validated cardinality + Int128(run) * Int128(step) == Int128(S) || return nothing + run *= L end end - return (true, steps) -end - -# An offset interval of one map: `regular` iff `buffer[t+1] == base + t*stride` -# for all `t < count`; otherwise read the buffer (valid until it is refilled). -struct BlockDescriptor - base::Int - stride::Int - count::Int - regular::Bool -end - -# Classifies `buffer[first+1 : first+count]`; an overflowing difference is irregular. -function describe_block(buffer::Vector{Int}, first::Int, count::Int) - first >= 0 || throw(ArgumentError("first must be nonnegative, got $first")) - count >= 0 || throw(ArgumentError("count must be nonnegative, got $count")) - first + count <= length(buffer) || - throw( - DimensionMismatch( - "buffer length $(length(buffer)) is less than first+count = $(first + count)" - ) - ) - - count == 0 && return BlockDescriptor(0, 0, 0, true) - - @inbounds base = buffer[first + 1] - count == 1 && return BlockDescriptor(base, 0, 1, true) - - @inbounds stride, overflowed = Base.Checked.sub_with_overflow(buffer[first + 2], buffer[first + 1]) - overflowed && return BlockDescriptor(base, 0, count, false) - - @inbounds for t in 2:(count - 1) - diff, ovf = Base.Checked.sub_with_overflow(buffer[first + t + 1], buffer[first + t]) - (ovf || diff != stride) && return BlockDescriptor(base, 0, count, false) - end - - return BlockDescriptor(base, stride, count, true) -end - -describe_block(buffer::Vector{Int}, count::Int) = describe_block(buffer, 0, count) - -function block_descriptors!( - buffers::NTuple{P, Vector{Int}}, g::AxisGroup{D, P}, - first::Int, count::Int - ) where {D, P} - fill_offsets!(buffers, g, first, count) - return ntuple(p -> describe_block(buffers[p], count), Val(P)) + return step end -# Same offset sequence with singleton dims dropped and adjacent dims folded -# where `next_stride[p] == length * stride[p]` for every map. Never reorders. -function normalize_group(g::AxisGroup{D, P}) where {D, P} - axis_length(g) == 0 && return g - - dims = Tuple{Int, NTuple{P, Int}}[] - for d in 1:D - L = g.lengths[d] - if L != 1 - push!(dims, (L, ntuple(p -> g.strides[p][d], P))) - end - end - - if isempty(dims) - emptylengths = NTuple{0, Int}() - emptystrides = ntuple(_ -> NTuple{0, Int}(), P) - return AxisGroup(emptylengths, emptystrides) - end - - folded = Tuple{Int, NTuple{P, Int}}[] - curL, curS = dims[1] - for i in 2:length(dims) - nextL, nextS = dims[i] - foldable = true - for p in 1:P - if Int128(curL) * Int128(curS[p]) != Int128(nextS[p]) - foldable = false - break - end - end - if foldable - curL = curL * nextL - else - push!(folded, (curL, curS)) - curL, curS = nextL, nextS - end +# Whether every map is a single ramp, so that the offset buffer can be replaced +# by arithmetic, and the steps (meaningless when not). +function affine_ramp(g::AxisGroup{D, P}) where {D, P} + steps = ntuple(_ -> 0, Val(P)) + for p in 1:P + step = map_ramp_step(g, p) + step === nothing && return (false, ntuple(_ -> 0, Val(P))) + steps = Base.setindex(steps, step, p) end - push!(folded, (curL, curS)) - - newD = length(folded) - newlengths = ntuple(i -> folded[i][1], newD) - newstrides = ntuple(p -> ntuple(i -> folded[i][2][p], newD), P) - - return AxisGroup(newlengths, newstrides) + return (true, steps) end diff --git a/src/layout/pair_group.jl b/src/layout/pair_group.jl deleted file mode 100644 index ab5b878..0000000 --- a/src/layout/pair_group.jl +++ /dev/null @@ -1,24 +0,0 @@ -# The two-map AxisGroup for one of M/N/K from label positions in a pair of -# operands: (v1, v2) is (A, C), (B, C) or (A, B). `labels` is an `NTuple{D,Int}` -# so the group is a concrete `AxisGroup{D,2}`; a runtime-length list would be -# heap-boxed and allocate per group. -@inline function _build_pair_group( - labels::NTuple{D, Int}, - ind1::NTuple{N1, Int}, v1::StridedView, - ind2::NTuple{N2, Int}, v2::StridedView - ) where {D, N1, N2} - # `Base.strides` on a `StridedView` rebuilds a tuple: call it once. - st1 = Base.strides(v1) - st2 = Base.strides(v2) - pos1 = ntuple(d -> findfirst(==(@inbounds labels[d]), ind1)::Int, Val(D)) - pos2 = ntuple(d -> findfirst(==(@inbounds labels[d]), ind2)::Int, Val(D)) - lens = ntuple(Val(D)) do d - l1 = size(v1, pos1[d]) - l2 = size(v2, pos2[d]) - l1 == l2 || _throw_label_length((@inbounds labels[d]), l1, l2) - l1 - end - s1 = ntuple(d -> st1[pos1[d]], Val(D)) - s2 = ntuple(d -> st2[pos2[d]], Val(D)) - return AxisGroup(lens, (s1, s2)) -end diff --git a/src/layout/tiles.jl b/src/layout/tiles.jl index bca8045..658af2b 100644 --- a/src/layout/tiles.jl +++ b/src/layout/tiles.jl @@ -55,6 +55,49 @@ axis_length(ax::PtrScatterAxis) = ax.count @inline axis_offset(ax::PtrScatterAxis, t::Int) = unsafe_load(ax.offsets + sizeof(Int) * t) +# An offset interval of one map: `regular` iff `buffer[t+1] == base + t*stride` +# for all `t < count`; otherwise read the buffer (valid until it is refilled). +struct BlockDescriptor + base::Int + stride::Int + count::Int + regular::Bool +end + +# Classifies `buffer[first+1 : first+count]`. A difference that overflows is +# irregular: a regular block's offset range is computed as `base + t*stride`. +function describe_block(buffer::Vector{Int}, first::Int, count::Int) + first >= 0 || throw(ArgumentError("first must be nonnegative, got $first")) + count >= 0 || throw(ArgumentError("count must be nonnegative, got $count")) + count <= length(buffer) - first || + throw(DimensionMismatch("buffer length $(length(buffer)) is less than first + count = $first + $count")) + + count == 0 && return BlockDescriptor(0, 0, 0, true) + + @inbounds base = buffer[first + 1] + count == 1 && return BlockDescriptor(base, 0, 1, true) + + @inbounds stride, overflowed = Base.Checked.sub_with_overflow(buffer[first + 2], buffer[first + 1]) + overflowed && return BlockDescriptor(base, 0, count, false) + + @inbounds for t in 2:(count - 1) + diff, ovf = Base.Checked.sub_with_overflow(buffer[first + t + 1], buffer[first + t]) + (ovf || diff != stride) && return BlockDescriptor(base, 0, count, false) + end + + return BlockDescriptor(base, stride, count, true) +end + +describe_block(buffer::Vector{Int}, count::Int) = describe_block(buffer, 0, count) + +function block_descriptors!( + buffers::NTuple{P, Vector{Int}}, g::AxisGroup{D, P}, + first::Int, count::Int + ) where {D, P} + fill_offsets!(buffers, g, first, count) + return ntuple(p -> describe_block(buffers[p], count), Val(P)) +end + function axis_from_descriptor(descriptor::BlockDescriptor, buffer::Vector{Int}, first::Int) if descriptor.regular return AffineAxis(descriptor.base, descriptor.stride, descriptor.count) diff --git a/src/planning/labels.jl b/src/planning/labels.jl index 4f6713d..ffd73e5 100644 --- a/src/planning/labels.jl +++ b/src/planning/labels.jl @@ -20,7 +20,7 @@ # Set element `n` of `out`, ignoring an `n` past the end (input about to be rejected). @inline function _push_label(out::NTuple{D, Int}, n::Int, x::Int) where {D} - return n <= D ? _tupleset(out, n, x) : out + return n <= D ? Base.setindex(out, x, n) : out end # (mlabels, nlabels, klabels) in indA/indB order. Per (inA, inB, inC): @@ -101,10 +101,6 @@ function _classify_labels( return mlabels, nlabels, klabels end -@noinline _throw_label_length(lbl::Int, l1::Int, l2::Int) = throw( - DimensionMismatch("label $lbl has mismatched axis length: $l1 vs $l2") -) - # Stable insertion sort of the permutation `perm` by `key[perm[j]]`, ascending # (n <= ndims, and unlike `sortperm` it allocates nothing). @inline function _sort_perm(perm::NTuple{D, Int}, key::NTuple{D, Int}) where {D} @@ -114,10 +110,10 @@ end kx = key[x] j = i - 1 while j >= 1 && key[out[j]] > kx - out = _tupleset(out, j + 1, out[j]) + out = Base.setindex(out, out[j], j + 1) j -= 1 end - out = _tupleset(out, j + 1, x) + out = Base.setindex(out, x, j + 1) end return out end diff --git a/src/planning/pack_split.jl b/src/planning/pack_split.jl index 7e7c7a6..f48fef2 100644 --- a/src/planning/pack_split.jl +++ b/src/planning/pack_split.jl @@ -128,7 +128,7 @@ function _pack_split_blocks( (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)) + neweff = min(budget ÷ blk * blk, _roundup(prod(lengths), R)) return (neweff, PackSplit(q, dj, L, kinner, Eq, E)) end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 0ef9074..80fd0b4 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -185,14 +185,14 @@ function _planned( 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) - ngroup = _build_pair_group(norder, indB, B, indC, C) # maps: (B, C) + mgroup = AxisGroup(morder, (indA, A), (indC, C)) # maps: (A, C) + ngroup = AxisGroup(norder, (indB, B), (indC, C)) # maps: (B, C) m_length = axis_length(mgroup) n_length = axis_length(ngroup) korder = _order_contract_labels(klabels, indA, A, morder, indB, B, norder, m_length, n_length) - kgroup = _build_pair_group(korder, indA, A, indB, B) # maps: (A, B) + kgroup = AxisGroup(korder, (indA, A), (indB, B)) # maps: (A, B) k_length = axis_length(kgroup) @@ -201,7 +201,7 @@ function _planned( if T <: Real && _prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) # B takes the M role: groups, K maps, storages, run and transforms move # together; the sum is unchanged. - kgroup_swapped = _build_pair_group(korder, indB, B, indA, A) # maps: (B, A) + kgroup_swapped = AxisGroup(korder, (indB, B), (indA, A)) # maps: (B, A) req_swapped = _plan_request( T, f, ngroup, mgroup, kgroup_swapped, parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), diff --git a/test/helpers.jl b/test/helpers.jl index fa61167..9eadde8 100644 --- a/test/helpers.jl +++ b/test/helpers.jl @@ -11,6 +11,12 @@ const execute_tilewise! = QuasiStrided.execute_tilewise! import TensorOperations as TO +# Offsets of coordinate `q` of an `AxisGroup` in every map, by direct decoding. +function offsets(g::AxisGroup{D, P}, q::Int) where {D, P} + c = Tuple(CartesianIndices(g.lengths)[q + 1]) + return ntuple(p -> sum((c[d] - 1) * g.strides[p][d] for d in 1:D; init = 0), P) +end + # Plan for the matmul C[m,n] = sum_k A[m,k]*B[k,n]. function _mm_plan(Cmat, Amat, Bmat; kwargs...) return plan_contract( diff --git a/test/layout/test_axis_group.jl b/test/layout/test_axis_group.jl index 4cfcfbc..156a06a 100644 --- a/test/layout/test_axis_group.jl +++ b/test/layout/test_axis_group.jl @@ -1,5 +1,5 @@ # AxisGroup indexing against an oracle built on `CartesianIndices` (first axis -# fastest), never on offsets/fill_offsets!/block_descriptors! themselves. +# fastest), never on fill_offsets!/block_descriptors! themselves. using QuasiStrided: affine_ramp @@ -47,29 +47,20 @@ end full = oracle_offsets(lengths, strides) Q = axis_length(g) @test Q == length(full) - @test [offsets(g, q) for q in 0:(Q - 1)] == full - @test_throws BoundsError offsets(g, Q) - @test_throws BoundsError offsets(g, -1) isramp, steps = affine_ramp(g) if Q > 0 step1 = Q > 1 ? full[2] : ntuple(_ -> 0, P) @test isramp == (full == [step1 .* q for q in 0:(Q - 1)]) isramp && @test steps == step1 end - # normalize_group is planning-time only: check it at the pair arity. - gn = P == 2 ? normalize_group(g) : g - @test axis_length(gn) == Q for _ in 1:4 first = rand(rng, 0:Q) count = rand(rng, 0:(Q - first)) bufs = ntuple(_ -> fill(-1, count + 2), P) # the suffix must stay untouched descs = block_descriptors!(bufs, g, first, count) - bufsn = ntuple(_ -> zeros(Int, count), P) - fill_offsets!(bufsn, gn, first, count) for p in 1:P expected = [full[i + 1][p] for i in first:(first + count - 1)] @test bufs[p] == [expected; -1; -1] - @test bufsn[p] == expected d = descs[p] affine = count <= 1 || expected == [expected[1] + t * (expected[2] - expected[1]) for t in 0:(count - 1)] @test d.regular == affine && d.count == count @@ -92,17 +83,6 @@ end @test offsets(g1, 0) == (0,) end -@testset "normalize_group: folds only where every map folds" begin - folded = normalize_group(AxisGroup((3, 2), ((1, 3), (2, 6)))) - @test folded.lengths == (6,) && folded.strides == ((1,), (2,)) - @test normalize_group(AxisGroup((4, 3), ((-1, -4),))).lengths == (12,) - # map 1 alone would fold, map 2 does not: the group must not. - @test normalize_group(AxisGroup((3, 2), ((1, 3), (1, 10)))).lengths == (3, 2) - g = normalize_group(AxisGroup((1, 1, 1), ((5, 6, 7), (8, 9, 10)))) - @test g.lengths == () && length(g.strides) == 2 && offsets(g, 0) == (0, 0) - @test normalize_group(AxisGroup((3, 1, 2), ((1, 999, 3),))).lengths == (6,) -end - @testset "AxisGroup: constructor validation and overflow policy" begin @test_throws ArgumentError AxisGroup((-1,), ((1,),)) @test_throws ArgumentError AxisGroup((3, -2), ((1, 1),)) @@ -125,8 +105,6 @@ end short, ok = fill(-7, Q - 1), fill(-7, Q) @test_throws DimensionMismatch fill_offsets!((short, ok), g, 0, Q) @test ok == fill(-7, Q) - @test_throws ArgumentError fill_offsets!((ok, ok), g, 0, Q) # aliased buffers - @test ok == fill(-7, Q) fill_offsets!((ok, short), g, Q, 0) # empty interval at Q @test ok == fill(-7, Q) end diff --git a/test/layout/test_stridedviews_axisgroup.jl b/test/layout/test_stridedviews_axisgroup.jl index 2f648c4..b594374 100644 --- a/test/layout/test_stridedviews_axisgroup.jl +++ b/test/layout/test_stridedviews_axisgroup.jl @@ -2,7 +2,6 @@ # indexing. using StridedViews: StridedView, offset -using QuasiStrided: _build_pair_group @testset "AxisGroup over a StridedView reproduces its indexing ($name)" for (name, v) in ( ("permuted", permutedims(StridedView(reshape(collect(1.0:30.0), 3, 5, 2)), (3, 1, 2))), @@ -19,12 +18,12 @@ using QuasiStrided: _build_pair_group end end -@testset "_build_pair_group: worked example from label positions" begin +@testset "AxisGroup from labels: worked example" begin A, B, C = StridedView(randn(3, 5, 2)), StridedView(randn(5, 4)), StridedView(zeros(3, 4, 2)) indA, indB, indC = (1, 2, 3), (2, 4), (1, 4, 3) # A[a,k,b] B[k,n] C[a,n,b] - M = _build_pair_group((1, 3), indA, A, indC, C) + M = AxisGroup((1, 3), (indA, A), (indC, C)) @test (M.lengths, M.strides) == ((3, 2), ((1, 15), (1, 12))) - K = _build_pair_group((2,), indA, A, indB, B) + K = AxisGroup((2,), (indA, A), (indB, B)) @test (K.lengths, K.strides) == ((5,), ((3,), (1,))) - @test_throws DimensionMismatch _build_pair_group((2,), indA, A, indB, StridedView(randn(6, 4))) + @test_throws DimensionMismatch AxisGroup((2,), (indA, A), (indB, StridedView(randn(6, 4)))) end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index c928f65..094c93b 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -29,7 +29,7 @@ const _pcf_Plan = QuasiStrided.ContractPlan @test_throws ArgumentError QS._classify_labels((1, 2), (3, 4), (1, 2, 3, 9)) end -@testset "per-call floor: _build_pair_group matches a direct construction" begin +@testset "per-call floor: AxisGroup from labels matches a direct construction" begin Random.seed!(1234) for trial in 1:60 nd1 = rand(1:4) @@ -47,7 +47,7 @@ end ) v2 = StridedView(randn(dims2)) - g = QS._build_pair_group(Tuple(labels), ind1, v1, ind2, v2) + g = AxisGroup(Tuple(labels), (ind1, v1), (ind2, v2)) s1 = Base.strides(v1) s2 = Base.strides(v2) p1 = ntuple(d -> findfirst(==(labels[d]), ind1)::Int, shared) @@ -58,16 +58,16 @@ end @test g isa AxisGroup{shared, 2} @test isconcretetype(typeof(g)) @test Base.return_types( - QS._build_pair_group, (typeof(Tuple(labels)), typeof(ind1), typeof(v1), typeof(ind2), typeof(v2)) + AxisGroup, (typeof(Tuple(labels)), Tuple{typeof(ind1), typeof(v1)}, Tuple{typeof(ind2), typeof(v2)}) ) == [AxisGroup{shared, 2}] end v1 = StridedView(randn(3, 4)) v2 = StridedView(randn(5, 4)) - @test_throws DimensionMismatch QS._build_pair_group((1,), (1, 2), v1, (1, 2), v2) + @test_throws DimensionMismatch AxisGroup((1,), ((1, 2), v1), ((1, 2), v2)) # Rank zero: an outer product's K group. - g0 = QS._build_pair_group((), (1, 2), v1, (1, 2), v2) + g0 = AxisGroup((), ((1, 2), v1), ((1, 2), v2)) @test g0 isa AxisGroup{0, 2} @test axis_length(g0) == 1 end @@ -419,7 +419,6 @@ end s = offsets(g, 1) @test any(q -> offsets(g, q) != (q * s[1], q * s[2]), 0:(Q - 1)) end - Q > 0 && @test isramp == (length(normalize_group(g).lengths) <= 1) end end diff --git a/test/runtests.jl b/test/runtests.jl index 5640280..f2790d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,8 +2,8 @@ using Test using Random using QuasiStrided # Internal names the test files use unqualified. -using QuasiStrided: AxisGroup, axis_length, offsets, fill_offsets!, BlockDescriptor, - describe_block, block_descriptors!, normalize_group, KernelDescriptor, tile_size, +using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, + describe_block, block_descriptors!, KernelDescriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, AffineAxis, ScatterAxis, SourceTile, DestinationTile, axis_from_descriptor, nrows, ncols, axis_offset_range, checked_tile_storage_bounds, pack_a!, pack_b!, From 1b6229abc746c615cb750080b7ee490fe9186265 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 10:33:38 -0400 Subject: [PATCH 06/60] Chunk 3: tile axes as offset vectors, one Tile type - AffineAxis <: AbstractVector{Int} (one-based getindex, closed-form checked extrema); the scatter axis is a plain view of the offset buffer. ScatterAxis, PtrScatterAxis, axis_from_descriptor and the Axis union go. - Tile{S, R, C} replaces QSTile/SourceTile/DestinationTile; size and @propagate_inbounds getindex/setindex! replace nrows/ncols/tile_offset/ tile_load/tile_store!, with @inbounds at the former call sites. - extrema replaces axis_offset_range; is_unit_stride (with an AbstractVector fallback) replaces _unit_stride_rows. - Tile coordinates and axis_offset stay zero-based for now (temporary shims). - The nest passes axes to their consumers through _with_axis, one branch per axis type: a Union holding a view is heap-boxed (432 B to 1.6 KB per call on a contraction with scattered M, N and K). A zero-allocation test for that contraction is added. Tests: 30930 pass (30989 before; removed ScatterAxis/PtrScatterAxis/ axis_from_descriptor tests). TTFX probe, Float64 / ComplexF64 first call, interleaved runs: before 3.1-3.3 s / 2.1-2.3 s, after 3.4-3.7 s / 2.3-2.4 s. Reused-plan execute! medians (us, 3 interleaved runs, before -> after): matmul 256^3 Float64 404 -> 395, ComplexF64 1537 -> 1442; scattered-M 12x20x256x256 Float64 432 -> 431, ComplexF64 1460 -> 1456; scattered-M/N/K small Float64 17.2 -> 17.6, ComplexF64 40.4 -> 39.1. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_kernels.jl | 4 +- src/execution/barrier.jl | 2 +- src/execution/execute.jl | 102 +++++++----- src/execution/macrokernel.jl | 50 +++--- src/execution/unpackedb.jl | 30 ++-- src/layout/tiles.jl | 148 +++++------------- src/microkernels/fmaddsub.jl | 6 +- src/microkernels/interface.jl | 39 ++--- src/microkernels/mixed.jl | 4 +- src/microkernels/onem.jl | 4 +- src/microkernels/planar.jl | 12 +- src/microkernels/scalar.jl | 2 +- src/microkernels/simd.jl | 12 +- src/packing/pack.jl | 36 ++--- src/packing/pack_contiguous.jl | 8 +- src/packing/transposed.jl | 2 +- test/execution/test_execute.jl | 4 +- test/execution/test_manual_pipeline.jl | 31 ++-- test/layout/test_tiles.jl | 81 ++++------ test/microkernels/test_fmaddsub_kernel.jl | 6 +- test/microkernels/test_mixed_kernel.jl | 6 +- .../test_planar_store_fastpath.jl | 14 +- test/microkernels/test_scalar_kernel.jl | 28 ++-- test/microkernels/test_simd_kernel.jl | 20 +-- test/packing/test_pack_complex.jl | 28 ++-- test/packing/test_pack_complex_contiguous.jl | 11 +- test/packing/test_pack_real.jl | 63 ++++---- test/planning/test_per_call_overhead.jl | 29 ++-- test/runtests.jl | 3 +- 29 files changed, 352 insertions(+), 433 deletions(-) diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index e396431..3aa645f 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -23,7 +23,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: RealMethod, PlanarMethod, OneMMethod, FMAddSubMethod, kernel_shapes, _kernel_from_shape, default_blocking, packed_panel, unsafe_execute_tile!, - DestinationTile, AffineAxis, target_profile + Tile, AffineAxis, target_profile const RUN_DTYPES = parse_dtypes(argopt("dtypes", "Float64,Float32,ComplexF64,ComplexF32")) const SHAPES = let s = argval("shapes") @@ -60,7 +60,7 @@ function kernel_hot!(kernel, C, apack, bpack, k_block, reps) GC.@preserve apack bpack begin ap = packed_panel(apack, 1, length(apack)) bp = packed_panel(bpack, 1, length(bpack)) - dest = DestinationTile(C, 0, AffineAxis(0, 1, tile_size(kernel, 1)), AffineAxis(0, tile_size(kernel)...)) + dest = Tile(C, 0, AffineAxis(0, 1, tile_size(kernel, 1)), AffineAxis(0, tile_size(kernel)...)) for _ in 1:reps unsafe_execute_tile!(kernel, dest, ap, bp, k_block, one(eltype(C)), one(eltype(C))) end diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index b29e767..60ec9b3 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -40,7 +40,7 @@ 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. `SPLIT`: whether +# `AffineAxis` and the offset-buffer `view` 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 diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 7702f38..f402723 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -10,13 +10,15 @@ function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::In while m_tile_start < m_length m_tile_length = min(m_tile, m_length - m_tile_start) (_, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) - rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) n_tile_start = 0 while n_tile_start < n_length n_tile_length = min(n_tile, n_length - n_tile_start) (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) - _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) + _with_axis(dM_C, ws.tile_m_buf_C, 0) do rowsC + _with_axis(dN_C, ws.tile_n_buf_C, 0) do colsC + _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) + end + end n_tile_start += n_tile_length end m_tile_start += m_tile_length @@ -198,7 +200,7 @@ function _execute_path!( ) where {T} kernel = plan.kernel ws = plan.workspace - # Panels and `PtrScatterAxis`es borrow pointers into `ws`. + # Packed panels borrow pointers into `ws`. GC.@preserve ws begin _execute_nest!( plan, ws, kernel, tile_size(kernel)..., @@ -355,8 +357,6 @@ function _execute_nest!( descriptor_offset_range(dB, ws.k_buf_B, 0), ) end - colsA_k = _axis_of(dK_A, ws.k_buf_A, 0, aff_kA) - rowsB_k = _axis_of(dK_B, ws.k_buf_B, 0, aff_kB) # Hoisted bounds checks (B here, A and C per M block): each # rectangle is exactly the union of the per-sliver/per-tile @@ -366,21 +366,28 @@ function _execute_nest!( beta_eff = firstpanel ? betaT : one(T) if split_b - _pack_block_transposed!( - packed_b_plane_offset, b_format(kernel), - packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, - plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, - plan.npack - ) + _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k + _pack_block_transposed!( + packed_b_plane_offset, b_format(kernel), + packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, + plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, + plan.npack + ) + end elseif !UNPACKED_B for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile - colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - _pack_sliver!( - unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, - kernel, btransform - ) + _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k + @inline + _with_axis(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) do colsB + @inline + _pack_sliver!( + unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, + kernel, btransform + ) + end + end end end @@ -406,30 +413,39 @@ function _execute_nest!( checked_span_bounds(cplan.Cbase, rng_mC, rng_nC, lenC) if split_a - _pack_block_transposed!( - packed_a_plane_offset, a_format(kernel), - packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, - plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, - plan.mpack - ) + _with_axis(dK_A, ws.k_buf_A, 0, aff_kA) do colsA_k + _pack_block_transposed!( + packed_a_plane_offset, a_format(kernel), + packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, + plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, + plan.mpack + ) + end else for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile - rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - _pack_sliver!( - unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, - kernel, atransform - ) + _with_axis(dK_A, ws.k_buf_A, 0, aff_kA) do colsA_k + @inline + _with_axis(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) do rowsA + @inline + _pack_sliver!( + unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, + kernel, atransform + ) + end + end end end # --- loops over N tiles and M tiles --- if UNPACKED_B - _micro_tiles_unpacked_b!( - kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC - ) + _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k + _micro_tiles_unpacked_b!( + kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC + ) + end else _micro_tiles_packed_b!( kernel, cplan, ws, m_tiles, n_tiles, @@ -461,16 +477,20 @@ end ) where {K, MC, NC} for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile - colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - unsafe_execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bpanel, k_block_length, alphaT, beta_eff - ) + _with_axis(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) do colsC + @inline + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + _with_axis(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) do rowsC + @inline + unsafe_execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + apanel, bpanel, k_block_length, alphaT, beta_eff + ) + end + end end end return nothing diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index ab6f81a..e069b5b 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -1,22 +1,26 @@ # Macro-kernel helpers for the five-loop nest: function barriers over the # tile axis types, packed-sliver addressing, and sliver classification. -# GUARDRAIL: `_axis_of` returns a `Union{AffineAxis,PtrScatterAxis}` (both -# isbits, so unboxed), and each consumer below is a barrier specialised per -# concrete axis type, so no partially-typed, heap-boxed `QSTile` is ever -# built. Do not inline these helpers into their callers. -@inline function _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int) - return d.regular ? AffineAxis(d.base, d.stride, d.count) : - PtrScatterAxis(pointer(buffer, first + 1), d.count) -end +# The axis `d` describes: an `AffineAxis` when regular, else a `view` of its +# offsets in `buffer`. +@inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int) = + d.regular ? AffineAxis(d.base, d.stride, d.count) : view(buffer, (first + 1):(first + d.count)) -# The same with the axis type fixed by the path: `Val(true)` for a ramp map, -# whose descriptors are always regular (checked). -@inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{false}) = - _axis_of(d, buffer, first) -@inline function _axis_of(d::BlockDescriptor, ::Vector{Int}, ::Int, ::Val{true}) - d.regular || _throw_irregular_ramp_descriptor() - return AffineAxis(d.base, d.stride, d.count) +# GUARDRAIL: `f(_axis_of(d, buffer, first))` with each axis type in its own +# branch, which the nest uses instead of `_axis_of`: a `Union` holding a `view` +# (not `isbits`) is heap-boxed. Each consumer below binds each axis type as its +# own parameter, so it builds a concretely typed `Tile`. `Val(true)`: a ramp +# map, whose descriptors are always regular (checked), so only the affine +# branch is compiled. +@inline function _with_axis( + f::F, d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{AFF} = Val(false) + ) where {F, AFF} + if AFF + d.regular || _throw_irregular_ramp_descriptor() + elseif !d.regular + return f(view(buffer, (first + 1):(first + d.count))) + end + return f(AffineAxis(d.base, d.stride, d.count)) end @noinline _throw_irregular_ramp_descriptor() = throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) @@ -27,16 +31,16 @@ end @inline function _pack_sliver!( pack!::PF, packed::PK, storage::S, base::Int, rows::R, cols::C, kernel, transform::TF - ) where {PF, PK, S, R <: Axis, C <: Axis, TF} - pack!(packed, SourceTile(storage, base, rows, cols), kernel, transform) + ) where {PF, PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} + pack!(packed, Tile(storage, base, rows, cols), kernel, transform) return nothing end @inline function _execute_micro_tile!( kernel, storage::S, base::Int, rows::R, cols::C, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta - ) where {PA, PB, S, R <: Axis, C <: Axis} - destination = DestinationTile(storage, base, rows, cols) + ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} + destination = Tile(storage, base, rows, cols) execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) return nothing end @@ -45,16 +49,16 @@ end @inline function unsafe_execute_micro_tile!( kernel, storage::S, base::Int, rows::R, cols::C, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta - ) where {PA, PB, S, R <: Axis, C <: Axis} - destination = DestinationTile(storage, base, rows, cols) + ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} + destination = Tile(storage, base, rows, cols) unsafe_execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) return nothing end @inline function _scale_micro_tile!( storage::S, base::Int, rows::R, cols::C, beta - ) where {S, R <: Axis, C <: Axis} - destination = DestinationTile(storage, base, rows, cols) + ) where {S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} + destination = Tile(storage, base, rows, cols) scale_tile!(destination, beta) return nothing end diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 2baa77c..84386ac 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -10,7 +10,7 @@ # every address read is inside the hoisted B check; their results are never # stored. `length` is the packed-equivalent count that # `_execute_tile_prologue!` checks, not the storage length. -struct UnpackedBView{S, K <: Axis, NR, F} +struct UnpackedBView{S, K <: AbstractVector{Int}, NR, F} storage::S colbase::NTuple{NR, Int} ksteps::K @@ -18,7 +18,7 @@ struct UnpackedBView{S, K <: Axis, NR, F} transform::F end -Base.length(u::UnpackedBView) = u.per_k * axis_length(u.ksteps) +Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) # `@inbounds`: the caller has checked the whole B panel region. @inline function _unpacked_b_element(u::UnpackedBView{S, K, NR, F}, j::Int, p::Int) where {S, K, NR, F} @@ -72,7 +72,7 @@ const _UNPACKED_B_MMAX = 256 @inline function _unpacked_b_view( kernel::DescriptorKernel{MR, NR, T}, storage::S, base::Int, d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F - ) where {MR, NR, T, S, K <: Axis, F} + ) where {MR, NR, T, S, K <: AbstractVector{Int}, F} last = d.count - 1 colbase = ntuple(Val(NR)) do j1 jj = min(j1 - 1, last) @@ -90,24 +90,28 @@ end kernel::K, plan::ContractPlan, ws, rowsB_k::KA, m_tiles::Int, n_tiles::Int, m_tile::Int, n_tile::Int, a_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, aff_mC::Val{MC}, aff_nC::Val{NC} - ) where {K, KA <: Axis, MC, NC} + ) where {K, KA <: AbstractVector{Int}, MC, NC} Bstorage = plan.Bstorage Bbase = plan.Bbase btransform = plan.btransform for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile - colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) bview = _unpacked_b_view( kernel, Bstorage, Bbase, ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, rowsB_k, btransform ) - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - unsafe_execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bview, k_block_length, alphaT, beta_eff - ) + _with_axis(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) do colsC + @inline + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + _with_axis(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) do rowsC + @inline + unsafe_execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + apanel, bview, k_block_length, alphaT, beta_eff + ) + end + end end end return nothing diff --git a/src/layout/tiles.jl b/src/layout/tiles.jl index 658af2b..d3282a6 100644 --- a/src/layout/tiles.jl +++ b/src/layout/tiles.jl @@ -1,8 +1,10 @@ -# Tile axes and tiles. Coordinates and addresses are zero-based (storage index -# = address + 1); a tile's `base` is the address of its logical origin. +# Tile axes and tiles. Addresses are zero-based (storage index = address + 1); +# a tile's `base` is the address of its logical origin. A tile axis is the +# vector of offsets of its rows or columns: an `AffineAxis`, or a `view` of an +# offset buffer. -# `t -> base + t*stride`. -struct AffineAxis +# Offsets `base + (t - 1) * stride`. Not a `StepRange`: the stride may be zero. +struct AffineAxis <: AbstractVector{Int} base::Int stride::Int count::Int @@ -13,47 +15,22 @@ struct AffineAxis end end -# `t -> offsets[t+1]`, borrowing `offsets`. -struct ScatterAxis{V <: AbstractVector{Int}} - offsets::V - count::Int +Base.size(ax::AffineAxis) = (ax.count,) +Base.IndexStyle(::Type{AffineAxis}) = IndexLinear() - function ScatterAxis(offsets::V, count::Int) where {V <: AbstractVector{Int}} - count >= 0 || throw(ArgumentError("ScatterAxis count must be nonnegative, got $count")) - count <= length(offsets) || - throw( - DimensionMismatch( - "ScatterAxis: offsets has length $(length(offsets)), " * - "need at least count = $count" - ) - ) - return new{V}(offsets, count) - end +@inline function Base.getindex(ax::AffineAxis, t::Int) + @boundscheck checkbounds(ax, t) + return ax.base + (t - 1) * ax.stride end -# `ScatterAxis` over a borrowed raw pointer, so that it is `isbits` and -# `Union{AffineAxis, PtrScatterAxis}` needs no heap box in the driver. -struct PtrScatterAxis - offsets::Ptr{Int} - count::Int - - function PtrScatterAxis(offsets::Ptr{Int}, count::Int) - count >= 0 || - throw(ArgumentError("PtrScatterAxis count must be nonnegative, got $count")) - return new(offsets, count) - end +function Base.extrema(ax::AffineAxis) + isempty(ax) && throw(ArgumentError("extrema of an empty AffineAxis")) + last = Base.Checked.checked_add(ax.base, Base.Checked.checked_mul(ax.count - 1, ax.stride)) + return minmax(ax.base, last) end -const Axis = Union{AffineAxis, ScatterAxis, PtrScatterAxis} - -axis_length(ax::AffineAxis) = ax.count -axis_length(ax::ScatterAxis) = ax.count -axis_length(ax::PtrScatterAxis) = ax.count - -@inline axis_offset(ax::AffineAxis, t::Int) = ax.base + t * ax.stride -@inline axis_offset(ax::ScatterAxis, t::Int) = @inbounds ax.offsets[t + 1] -@inline axis_offset(ax::PtrScatterAxis, t::Int) = - unsafe_load(ax.offsets + sizeof(Int) * t) +# Zero-based until coordinates become one-based. +@inline axis_offset(ax::AbstractVector{Int}, t::Int) = @inbounds ax[t + 1] # An offset interval of one map: `regular` iff `buffer[t+1] == base + t*stride` # for all `t < count`; otherwise read the buffer (valid until it is refilled). @@ -98,84 +75,33 @@ function block_descriptors!( return ntuple(p -> describe_block(buffers[p], count), Val(P)) end -function axis_from_descriptor(descriptor::BlockDescriptor, buffer::Vector{Int}, first::Int) - if descriptor.regular - return AffineAxis(descriptor.base, descriptor.stride, descriptor.count) - else - return ScatterAxis(view(buffer, (first + 1):(first + descriptor.count)), descriptor.count) - end -end - -axis_from_descriptor(descriptor::BlockDescriptor, buffer::Vector{Int}) = - axis_from_descriptor(descriptor, buffer, 0) - -# Logical `(i, j)` addresses `base + row_offset(i) + col_offset(j)`. -struct QSTile{S, R <: Axis, C <: Axis} +# Logical `(i, j)` addresses `base + rows[i] + cols[j]`. +struct Tile{S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} storage::S base::Int rows::R cols::C end -const SourceTile = QSTile -const DestinationTile = QSTile +Base.size(tile::Tile) = (length(tile.rows), length(tile.cols)) -nrows(tile::QSTile) = axis_length(tile.rows) -ncols(tile::QSTile) = axis_length(tile.cols) +# Zero-based until coordinates become one-based. +Base.@propagate_inbounds Base.getindex(tile::Tile, i::Int, j::Int) = + tile.storage[tile.base + tile.rows[i + 1] + tile.cols[j + 1] + 1] -@inline function tile_offset(tile::QSTile, i::Int, j::Int) - return tile.base + axis_offset(tile.rows, i) + axis_offset(tile.cols, j) -end - -@inline function tile_load(tile::QSTile, i::Int, j::Int) - return @inbounds tile.storage[tile_offset(tile, i, j) + 1] -end - -@inline function tile_store!(tile::QSTile, i::Int, j::Int, v) - @inbounds tile.storage[tile_offset(tile, i, j) + 1] = v +Base.@propagate_inbounds function Base.setindex!(tile::Tile, v, i::Int, j::Int) + tile.storage[tile.base + tile.rows[i + 1] + tile.cols[j + 1] + 1] = v return tile end # Bounds checks run once per tile (or per macro block) before the unchecked # hot paths. Offset ranges are `(lo, hi)`, `(0, -1)` when empty. -function axis_offset_range(ax::AffineAxis) - ax.count == 0 && return (0, -1) - lo128 = Int128(ax.base) - hi128 = Int128(ax.base) + Int128(ax.count - 1) * Int128(ax.stride) - lo128, hi128 = minmax(lo128, hi128) - (typemin(Int) <= lo128 && hi128 <= typemax(Int)) || - throw(OverflowError("axis_offset_range: affine axis range not representable as Int")) - return (Int(lo128), Int(hi128)) -end - -function axis_offset_range(ax::ScatterAxis) - ax.count == 0 && return (0, -1) - prefix = view(ax.offsets, 1:ax.count) - return (Int(minimum(prefix)), Int(maximum(prefix))) -end - -function axis_offset_range(ax::PtrScatterAxis) - ax.count == 0 && return (0, -1) - lo = hi = unsafe_load(ax.offsets) - for t in 1:(ax.count - 1) - v = axis_offset(ax, t) - lo, hi = min(lo, v), max(hi, v) - end - return (lo, hi) -end - -# `axis_offset_range` of the axis `d` describes, without materializing it. +# `extrema` of the axis `d` describes, without materializing it. function descriptor_offset_range(d::BlockDescriptor, buffer::Vector{Int}, first::Int) d.count == 0 && return (0, -1) - d.regular && return axis_offset_range(AffineAxis(d.base, d.stride, d.count)) - lo = hi = buffer[first + 1] - for t in 1:(d.count - 1) - v = buffer[first + t + 1] - lo = min(lo, v) - hi = max(hi, v) - end - return (lo, hi) + d.regular && return extrema(AffineAxis(d.base, d.stride, d.count)) + return extrema(@view buffer[(first + 1):(first + d.count)]) end # Exact, not conservative, for any rectangular (row-set x column-set) region: @@ -194,18 +120,16 @@ function checked_span_bounds( return nothing end -function checked_tile_storage_bounds(base::Int, rows::Axis, cols::Axis, storage_length::Int) - (axis_length(rows) == 0 || axis_length(cols) == 0) && return nothing - return checked_span_bounds( - base, axis_offset_range(rows), axis_offset_range(cols), storage_length +function checked_tile_storage_bounds( + base::Int, rows::AbstractVector{Int}, cols::AbstractVector{Int}, storage_length::Int ) + (isempty(rows) || isempty(cols)) && return nothing + return checked_span_bounds(base, extrema(rows), extrema(cols), storage_length) end -checked_tile_storage_bounds(tile::QSTile) = +checked_tile_storage_bounds(tile::Tile) = checked_tile_storage_bounds(tile.base, tile.rows, tile.cols, length(tile.storage)) -# Whether an axis steps through storage one element at a time. Deliberately no -# fallback method: an unknown axis type must be a MethodError, not `false`. -_unit_stride_rows(ax::AffineAxis) = ax.stride == 1 -_unit_stride_rows(::ScatterAxis) = false -_unit_stride_rows(::PtrScatterAxis) = false +# Whether an axis steps through storage one element at a time. +is_unit_stride(ax::AffineAxis) = ax.stride == 1 +is_unit_stride(::AbstractVector{Int}) = false diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 19f7b26..09bf147 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -198,7 +198,7 @@ end # Same unroll and full-block / row-tail split as `_store_tile_planar_vector!`. @generated function _store_tile_fmaddsub_vector!( - destination::QSTile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, R}}, + destination::Tile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NV} @@ -249,7 +249,7 @@ end return quote storage = destination.storage cols = destination.cols - rowbase0 = destination.base + destination.rows.base + rowbase0 = destination.base + axis_offset(destination.rows, 0) ar = Vec{$W, $R}(real(alpha)) ai = Vec{$W, $R}(imag(alpha)) br = Vec{$W, $R}(real(beta)) @@ -267,7 +267,7 @@ end # Not `@inline`, unlike the real and planar stores: inlining it cost time # (code growth). function store_tile!( - destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::FMAddSubKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NV} m, n = _store_prologue!(destination, alpha, beta) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index af6ec15..55bc802 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -70,19 +70,19 @@ packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.de # unbound one makes the call dynamically dispatched and allocating on every # pack. `V` is unconstrained so a `PackedPanel` forwards too. pack_a!( - packed::V, source::QSTile, kernel::K, transform::F + packed::V, source::Tile, kernel::K, transform::F ) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = pack_a!(packed, source, kernel.descriptor, transform) pack_b!( - packed::V, source::QSTile, kernel::K, transform::F + packed::V, source::Tile, kernel::K, transform::F ) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = pack_b!(packed, source, kernel.descriptor, transform) @inline unsafe_pack_a!( - packed::V, source::QSTile, kernel::K, transform::F + packed::V, source::Tile, kernel::K, transform::F ) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = unsafe_pack_a!(packed, source, kernel.descriptor, transform) @inline unsafe_pack_b!( - packed::V, source::QSTile, kernel::K, transform::F + packed::V, source::Tile, kernel::K, transform::F ) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = unsafe_pack_b!(packed, source, kernel.descriptor, transform) @@ -113,31 +113,28 @@ _default_lanewidth(::Type{Float64}) = 4 _default_lanewidth(::Type{Float32}) = 8 # `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. -function scale_tile!(destination::QSTile, beta::T) where {T} - m = nrows(destination) - n = ncols(destination) +function scale_tile!(destination::Tile, beta::T) where {T} + m, n = size(destination) (m == 0 || n == 0) && return destination if isone(beta) return destination elseif iszero(beta) @inbounds for j in 0:(n - 1), i in 0:(m - 1) - tile_store!(destination, i, j, zero(T)) + destination[i, j] = zero(T) end else @inbounds for j in 0:(n - 1), i in 0:(m - 1) - tile_store!(destination, i, j, convert(T, tile_load(destination, i, j)) * beta) + destination[i, j] = convert(T, 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::T) where {T} = tile_store!( - dest, i, j, +@inline _axpby_tile!(dest, i::Int, j::Int, alpha, r, beta::T) where {T} = @inbounds dest[i, j] = iszero(beta) ? alpha * r : - isone(beta) ? muladd(alpha, r, convert(T, tile_load(dest, i, j))) : - muladd(alpha, r, beta * convert(T, tile_load(dest, i, j))) -) + isone(beta) ? muladd(alpha, r, convert(T, dest[i, j])) : + muladd(alpha, r, beta * convert(T, dest[i, j])) @inline _axpby_at!(storage, idx::Int, alpha, r, beta::T) where {T} = @inbounds storage[idx] = iszero(beta) ? alpha * r : @@ -146,9 +143,8 @@ end # The `(m, n)` to store over, or `(0, 0)` when already done (empty destination, # or `alpha == 0` handled by `scale_tile!`). -@inline function _store_prologue!(destination::QSTile, alpha, beta) - m = nrows(destination) - n = ncols(destination) +@inline function _store_prologue!(destination::Tile, alpha, beta) + m, n = size(destination) if m == 0 || n == 0 return (0, 0) elseif iszero(alpha) @@ -180,11 +176,10 @@ end # bounds check at compile time (`unsafe_execute_tile!` only). # GUARDRAIL: `@inline` and one bound type parameter per argument (hot path). @inline function _execute_tile_prologue!( - kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, + kernel::K, destination::Tile, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta, ::Val{BOUNDS} ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB, BOUNDS} - m = nrows(destination) - n = ncols(destination) + m, n = size(destination) m <= MR || _throw_tile_extent(:row, m, MR) n <= NR || _throw_tile_extent(:column, n, NR) k_block_length >= 0 || _throw_negative_k_block_length(:execute_tile!, k_block_length) @@ -211,7 +206,7 @@ end # One checked K panel: `zero_accumulator`, `accumulate`, `store_tile!`. function execute_tile!( - kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, + kernel::K, destination::Tile, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB} run, alphaT, betaT = _execute_tile_prologue!( @@ -228,7 +223,7 @@ end # `checked_span_bounds`, which is equivalent because the check only compares # range extremes. @inline function unsafe_execute_tile!( - kernel::K, destination::QSTile, packed_a::PA, packed_b::PB, + kernel::K, destination::Tile, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta ) where {K, PA, PB} run, alphaT, betaT = _execute_tile_prologue!( diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index 3a85395..4173a18 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -84,7 +84,7 @@ zero_accumulator(kernel::_MixedKernel) = zero_accumulator(kernel.inner) # Inlined or not as the fmaddsub and planar stores each reuses. function store_tile!( - destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + destination::Tile, 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) @@ -98,7 +98,7 @@ function store_tile!( end @inline function store_tile!( - destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + destination::Tile, 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) diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index 76a35db..cce67f8 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -65,7 +65,7 @@ end # and `2MR` a multiple of `W`, complex row `i = v*(W÷2) + u` is lanes `2u+1` # (re) and `2u+2` (im) of vector `v` -- never split across two vectors. @generated function _store_tile_lanepair!( - destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {MR, NR, T, W, R, NV} @@ -103,7 +103,7 @@ end # Scalar only: 1m has no vector store. Not `@inline`, as for fmaddsub. function store_tile!( - destination::QSTile, acc::NTuple{NV, Vec{W, R}}, + destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::OneMKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NV} m, n = _store_prologue!(destination, alpha, beta) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index b9655bf..10a7a76 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -145,7 +145,7 @@ 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}}, + destination::Tile, acc::NTuple{NA, Vec{W, R}}, alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {MR, NR, T, W, R, NA} @@ -184,8 +184,8 @@ end # Vector store: unit-stride rows into rank-1 dense `Complex` storage (so the # storage can be reinterpreted as `2W` consecutive reals per `W` rows), on an # ISA the complex fast paths ship for (shared with the complex pack fast path). -@inline _complex_vector_eligible(tile::QSTile, ::Type{T}) where {T} = - _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) && +@inline _complex_vector_eligible(tile::Tile, ::Type{T}) where {T} = + is_unit_stride(tile.rows) && _dense_lanes(tile.storage, T) && complex_fastpath_isa_eligible() # Shuffle patterns built from `W` at specialization time, never hardcoded to @@ -242,7 +242,7 @@ end # Same unroll and full-block / row-tail split as `_store_tile_vector!`, with a # 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}}, + destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NA} @@ -295,7 +295,7 @@ end Base.@_inline_meta storage = destination.storage cols = destination.cols - rowbase0 = destination.base + destination.rows.base + rowbase0 = destination.base + axis_offset(destination.rows, 0) ar = Vec{$W, $R}(real(alpha)) ai = Vec{$W, $R}(imag(alpha)) br = Vec{$W, $R}(real(beta)) @@ -312,7 +312,7 @@ end # `@inline` with the vector store, for the reason at the real `store_tile!`. @inline function store_tile!( - destination::QSTile, acc::NTuple{NA, Vec{W, R}}, + destination::Tile, acc::NTuple{NA, Vec{W, R}}, alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NA} m, n = _store_prologue!(destination, alpha, beta) diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl index b09dbae..362c986 100644 --- a/src/microkernels/scalar.jl +++ b/src/microkernels/scalar.jl @@ -33,7 +33,7 @@ function Base.accumulate( end function store_tile!( - destination::QSTile, acc::AbstractMatrix{T}, + destination::Tile, acc::AbstractMatrix{T}, alpha::T, beta::T, kernel::ScalarKernel ) where {T} m, n = _store_prologue!(destination, alpha, beta) diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index f780226..b8ca1e4 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -92,11 +92,11 @@ end # Vector store eligibility: unit-stride rows into rank-1 dense real storage, # exactly what SIMD.jl's array `vload`/`vstore` accept. Must admit `Memory{T}`: # that is the `parent` of an Array-backed `StridedView` on Julia >= 1.11. -@inline _vector_store_eligible(tile::QSTile, ::Type{T}) where {T} = - _unit_stride_rows(tile.rows) && _dense_lanes(tile.storage, T) +@inline _vector_store_eligible(tile::Tile, ::Type{T}) where {T} = + is_unit_stride(tile.rows) && _dense_lanes(tile.storage, T) @generated function _store_tile_scattered!( - destination::QSTile, acc::NTuple{NV, Vec{W, T}}, + destination::Tile, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, m::Int, n::Int ) where {MR, NR, T, W, NV} @@ -129,7 +129,7 @@ end # stored lane by lane, so nothing outside the valid rectangle is touched. # `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. @generated function _store_tile_vector!( - destination::QSTile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, T}}, + destination::Tile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, m::Int, n::Int ) where {S, MR, NR, T, W, NV} @@ -178,7 +178,7 @@ end Base.@_inline_meta storage = destination.storage cols = destination.cols - rowbase0 = destination.base + destination.rows.base + rowbase0 = destination.base + axis_offset(destination.rows, 0) @inbounds begin $(blocks...) end @@ -190,7 +190,7 @@ end # to the stack and reloaded on every micro-tile. The scattered store stays out # of line: it is scalar anyway, and inlining it bloats the tile function. @inline function store_tile!( - destination::QSTile, acc::NTuple{NV, Vec{W, T}}, + destination::Tile, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W} ) where {MR, NR, T, W, NV} m, n = _store_prologue!(destination, alpha, beta) diff --git a/src/packing/pack.jl b/src/packing/pack.jl index 8c35057..ba76d22 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -29,14 +29,14 @@ end # into the packed format; padding lanes are literal zeros and never read the # source or call `transform`. All validation happens before any write. function pack_a!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F ) where {V, MR, NR, T2, FA, FB, F} return _pack_a!(packed, source, kernel, transform, Val(true)) end function pack_b!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F ) where {V, MR, NR, T2, FA, FB, F} return _pack_b!(packed, source, kernel, transform, Val(true)) @@ -44,28 +44,27 @@ end # Skips only `checked_tile_storage_bounds(source)`: the caller (`_execute_nest!`) # has already validated the whole macro block the sliver belongs to. `@inline` -# because out of line each call marshals the `QSTile` through the stack. +# because out of line each call marshals the `Tile` through the stack. @inline function unsafe_pack_a!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F ) where {V, MR, NR, T2, FA, FB, F} return _pack_a!(packed, source, kernel, transform, Val(false)) end @inline function unsafe_pack_b!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F ) where {V, MR, NR, T2, FA, FB, F} return _pack_b!(packed, source, kernel, transform, Val(false)) end @inline function _pack_a!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F, ::Val{BOUNDS} ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} _check_packed_eltype(packed, kernel) - m = nrows(source) - k_block_length = ncols(source) + m, k_block_length = size(source) (0 <= m <= MR) || _throw_pack_extent(:pack_a, :row, m, MR) needed = packed_a_length(kernel, k_block_length) length(packed) >= needed || _throw_pack_short(:pack_a, length(packed), needed, k_block_length) @@ -77,12 +76,11 @@ end end @inline function _pack_b!( - packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2, FA, FB}, + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, transform::F, ::Val{BOUNDS} ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} _check_packed_eltype(packed, kernel) - k_block_length = nrows(source) - n = ncols(source) + k_block_length, n = size(source) (0 <= n <= NR) || _throw_pack_extent(:pack_b, :column, n, NR) needed = packed_b_length(kernel, k_block_length) length(packed) >= needed || _throw_pack_short(:pack_b, length(packed), needed, k_block_length) @@ -94,40 +92,40 @@ end # A's packed index runs along `source.rows`, B's along `source.cols`; each # contiguous fast path needs that lane axis to be unit-stride. @inline function _pack_a_sliver!( - format::FMT, packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2}, + format::FMT, packed::V, source::Tile, kernel::Descriptor{MR, NR, T2}, transform::F, m::Int, k_block_length::Int ) where {FMT, V, MR, NR, T2, F} if format isa RealFormat if _pack_a_contiguous_eligible(packed, source, transform, m, Val(MR), real(T2)) - rowbase = source.base + source.rows.base + rowbase = source.base + axis_offset(source.rows, 0) return _pack_a_contiguous!(packed, source.storage, rowbase, source.cols, Val(MR), k_block_length) end elseif _pack_complex_contiguous_eligible( packed, source.storage, source.rows, transform, format, m, Val(MR), T2 ) - elembase = source.base + source.rows.base + elembase = source.base + axis_offset(source.rows, 0) return _pack_complex_contiguous!( format, packed, source.storage, elembase, source.cols, Val(MR), k_block_length, transform ) end - load = (i, p) -> tile_load(source, i, p) + load = (i, p) -> @inbounds source[i, p] plane_offset = (plane, i, p) -> packed_a_plane_offset(kernel, plane, i, p) return _pack_panel!(packed, _element_type(format, T2), format, Val(MR), k_block_length, m, transform, load, plane_offset) end @inline function _pack_b_sliver!( - format::FMT, packed::V, source::QSTile, kernel::Descriptor{MR, NR, T2}, + format::FMT, packed::V, source::Tile, kernel::Descriptor{MR, NR, T2}, transform::F, n::Int, k_block_length::Int ) where {FMT, V, MR, NR, T2, F} if _pack_complex_contiguous_eligible( packed, source.storage, source.cols, transform, format, n, Val(NR), T2 ) - elembase = source.base + source.cols.base + elembase = source.base + axis_offset(source.cols, 0) return _pack_complex_contiguous!( format, packed, source.storage, elembase, source.rows, Val(NR), k_block_length, transform ) end - load = (j, p) -> tile_load(source, p, j) + load = (j, p) -> @inbounds source[p, j] plane_offset = (plane, j, p) -> packed_b_plane_offset(kernel, plane, j, p) return _pack_panel!(packed, _element_type(format, T2), format, Val(NR), k_block_length, n, transform, load, plane_offset) end @@ -168,7 +166,7 @@ end end # Per-lane stores of one element. `real`/`imag` on the loaded element are the -# only accessors: a `QSTile` may address scattered storage, so the source is +# only accessors: a `Tile` may address scattered storage, so the source is # never `reinterpret`ed. @inline function _pack_emit!( packed::V, ::RealFormat, plane_offset::P, t::Int, p::Int, z::T diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index d4e6a47..a716025 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -20,13 +20,11 @@ _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} + packed::V, source::Tile, transform::F, m::Int, ::Val{MR}, ::Type{T} ) where {V, F, MR, T} return packed isa PackedPanel{T} && _dense_lanes(source.storage, T) && - _copies_unchanged(transform, T) && m == MR && _unit_stride_rows(source.rows) + _copies_unchanged(transform, T) && m == MR && is_unit_stride(source.rows) end # Real A: one `Vec{MR}` load/convert/store per K step. @@ -62,7 +60,7 @@ end return !(format isa RealFormat) && packed isa PackedPanel{real(T)} && _dense_lanes(storage, T) && _complex_pack_transform_eligible(transform) && - valid == PD && _unit_stride_rows(lane_axis) && + valid == PD && is_unit_stride(lane_axis) && complex_fastpath_isa_eligible() end diff --git a/src/packing/transposed.jl b/src/packing/transposed.jl index c5575f6..3bf9854 100644 --- a/src/packing/transposed.jl +++ b/src/packing/transposed.jl @@ -8,7 +8,7 @@ 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, k_block_length::Int, split - ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: Axis, F} + ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: AbstractVector{Int}, 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`; diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 262e166..b9c1142 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -206,7 +206,7 @@ end _ws_union_members(t) = t isa Union ? (_ws_union_members(t.a)..., _ws_union_members(t.b)...) : (t,) -# Every non-concrete QSTile/ContractWorkspace type in `f`'s unoptimized typed +# Every non-concrete Tile/ContractWorkspace type in `f`'s unoptimized typed # IR, directly or as a union member (`Union{}` is a throw, not an instability). function _ws_nonconcrete_types(f, argtypes) bad = Any[] @@ -218,7 +218,7 @@ function _ws_nonconcrete_types(f, argtypes) t isa Type || continue for m in _ws_union_members(t) (m isa Type && m !== Union{}) || continue - if (m <: QuasiStrided.QSTile || m <: QuasiStrided.ContractWorkspace) && !isconcretetype(m) + if (m <: QuasiStrided.Tile || m <: QuasiStrided.ContractWorkspace) && !isconcretetype(m) push!(bad, t) break end diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index 631674b..99138e6 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -14,12 +14,13 @@ mb, nb, kb = bufs(M), bufs(N), bufs(K) (dM_A, dM_C) = block_descriptors!(mb, M, 0, 6) (dN_B, dN_C) = block_descriptors!(nb, N, 0, 4) - row_A, row_C = axis_from_descriptor(dM_A, mb[1]), axis_from_descriptor(dM_C, mb[2]) - col_B, col_C = axis_from_descriptor(dN_B, nb[1]), axis_from_descriptor(dN_C, nb[2]) + axis(d, buf) = QuasiStrided._axis_of(d, buf, 0) + row_A, row_C = axis(dM_A, mb[1]), axis(dM_C, mb[2]) + col_B, col_C = axis(dN_B, nb[1]), axis(dN_C, nb[2]) Cstorage = zeros(length(Cref)) - destination = DestinationTile(Cstorage, 0, row_C, col_C) - @test (nrows(destination), ncols(destination)) == (6, 4) + destination = Tile(Cstorage, 0, row_C, col_C) + @test size(destination) == (6, 4) packed_a, packed_b = zeros(tile_size(kernel, 1) * 5), zeros(tile_size(kernel, 2) * 5) Cstart = rand(MersenneTwister(1234), size(Cref)...) @@ -28,8 +29,8 @@ copyto!(Cstorage, vec(Cstart)) for (idx, (first, len)) in enumerate(panels) (dK_A, dK_B) = block_descriptors!(kb, K, first, len) - pack_a!(packed_a, SourceTile(vec(A), 0, row_A, axis_from_descriptor(dK_A, kb[1])), kernel, identity) - pack_b!(packed_b, SourceTile(vec(B), 0, axis_from_descriptor(dK_B, kb[2]), col_B), kernel, identity) + pack_a!(packed_a, Tile(vec(A), 0, row_A, axis(dK_A, kb[1])), kernel, identity) + pack_b!(packed_b, Tile(vec(B), 0, axis(dK_B, kb[2]), col_B), kernel, identity) execute_tile!(kernel, destination, packed_a, packed_b, len, alpha, idx == 1 ? beta : 1.0) end @test reshape(Cstorage, size(Cref)) ≈ alpha .* Cref .+ beta .* Cstart @@ -41,27 +42,23 @@ end # Out-of-bounds source or destination storage is rejected before any access. packed = zeros(4) for base in (10_000, -10_000) - src = SourceTile([1.0, 2.0, 3.0, 4.0], base, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) + src = Tile([1.0, 2.0, 3.0, 4.0], base, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) @test_throws BoundsError pack_a!(packed, src, k, identity) @test_throws BoundsError pack_b!(packed, src, k, identity) end @test all(iszero, packed) canary = fill(999.0, 4) - dst = DestinationTile(canary, 3, AffineAxis(0, 1, 2), AffineAxis(0, 0, 1)) + dst = Tile(canary, 3, AffineAxis(0, 1, 2), AffineAxis(0, 0, 1)) @test_throws BoundsError execute_tile!(k, dst, zeros(4), zeros(4), 0, 1.0, 2.0) @test all(==(999.0), canary) # Undersized packed buffers are rejected, except by the k_block_length = 0 / alpha = 0 # short-circuits, which never read them. - dst = DestinationTile(zeros(4), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) + dst = Tile(zeros(4), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) @test_throws DimensionMismatch execute_tile!(k, dst, zeros(4), zeros(4), 50, 1.0, 0.0) @test execute_tile!(k, dst, zeros(1), zeros(1), 0, 1.0, 0.0) === dst @test execute_tile!(k, dst, zeros(1), zeros(1), 5, 0.0, 1.0) === dst - @test axis_offset_range(AffineAxis(5, -2, 3)) == (1, 5) - @test axis_offset_range(AffineAxis(7, 3, 0)) == (0, -1) - @test axis_offset_range(ScatterAxis([4, 1, 9, 2], 3)) == (1, 9) - @test axis_offset_range(ScatterAxis(Int[], 0)) == (0, -1) @test checked_tile_storage_bounds(0, AffineAxis(0, 1, 3), AffineAxis(0, 3, 2), 6) === nothing @test checked_tile_storage_bounds(1_000_000, AffineAxis(0, 1, 0), AffineAxis(0, 1, 5), 1) === nothing @test_throws BoundsError checked_tile_storage_bounds(0, AffineAxis(0, 1, 3), AffineAxis(0, 3, 2), 5) @@ -69,7 +66,7 @@ end kernel = ScalarKernel(Val(3), Val(2), Float64) pa, pb = Float64[1, 2, 3, 4, 5, 6], Float64[10, 20, 30, 40] - tile(s) = DestinationTile(s, 0, AffineAxis(0, 1, 3), AffineAxis(0, 3, 2)) + tile(s) = Tile(s, 0, AffineAxis(0, 1, 3), AffineAxis(0, 3, 2)) s = fill(NaN, 6) execute_tile!(kernel, tile(s), pa, pb, 2, 0.0, 0.0) # alpha = beta = 0: zeros @test all(iszero, s) @@ -80,17 +77,17 @@ end execute_tile!(kernel, tile(s), pa, pb, 2, 1.0, 0.0) # beta = 0 never reads C @test all(isfinite, s) s = [42.0] - execute_tile!(kernel, DestinationTile(s, 0, AffineAxis(0, 1, 0), AffineAxis(0, 1, 0)), pa, pb, 2, 1.0, 0.0) + execute_tile!(kernel, Tile(s, 0, AffineAxis(0, 1, 0), AffineAxis(0, 1, 0)), pa, pb, 2, 1.0, 0.0) @test s == [42.0] @test_throws ArgumentError execute_tile!( - kernel, DestinationTile(zeros(20), 0, AffineAxis(0, 1, 4), AffineAxis(0, 4, 2)), pa, pb, 2, 1.0, 0.0 + kernel, Tile(zeros(20), 0, AffineAxis(0, 1, 4), AffineAxis(0, 4, 2)), pa, pb, 2, 1.0, 0.0 ) # Canaries outside the tile survive; scattered columns agree with affine ones. s = fill(-1.0, 8) execute_tile!(kernel, tile(s), pa, pb, 2, 1.0, 0.0) @test s[7] == -1.0 && s[8] == -1.0 s6 = zeros(6) - execute_tile!(kernel, DestinationTile(s6, 0, AffineAxis(0, 1, 3), ScatterAxis([0, 3], 2)), pa, pb, 2, 1.0, 0.0) + execute_tile!(kernel, Tile(s6, 0, AffineAxis(0, 1, 3), view([0, 3], 1:2)), pa, pb, 2, 1.0, 0.0) @test s6 == s[1:6] s32 = zeros(Float32, 6) execute_tile!(ScalarKernel(Val(3), Val(2), Float32), tile(s32), Float32.(pa), Float32.(pb), 2, 1.0f0, 0.0f0) diff --git a/test/layout/test_tiles.jl b/test/layout/test_tiles.jl index beb50e7..bdc2e74 100644 --- a/test/layout/test_tiles.jl +++ b/test/layout/test_tiles.jl @@ -1,69 +1,48 @@ -# Tile axes and QSTile addressing, against hand-computed addresses. +# Tile axes and Tile addressing, against hand-computed addresses. -using QuasiStrided: PtrScatterAxis, QSTile, axis_offset, tile_offset, tile_load, tile_store!, - checked_span_bounds, descriptor_offset_range +using QuasiStrided: checked_span_bounds, descriptor_offset_range, is_unit_stride -@testset "tile axes: addressing, borrowing, validation" begin - @test [axis_offset(AffineAxis(10, 3, 5), t) for t in 0:4] == [10, 13, 16, 19, 22] - @test [axis_offset(AffineAxis(20, -4, 4), t) for t in 0:3] == [20, 16, 12, 8] - @test [axis_offset(AffineAxis(7, 0, 3), t) for t in 0:2] == [7, 7, 7] - offs = [5, -3, 100, 0, 42] - ax = ScatterAxis(offs, 3) # a strict prefix of `offs` - @test axis_length(ax) == 3 - @test [axis_offset(ax, t) for t in 0:2] == offs[1:3] - offs[2] = 999 # borrowed, not copied - @test axis_offset(ax, 1) == 999 - GC.@preserve offs begin - pax = PtrScatterAxis(pointer(offs), 5) - @test [axis_offset(pax, t) for t in 0:4] == offs - @test axis_offset_range(pax) == (0, 999) - end +@testset "AffineAxis: offsets, validation" begin + @test AffineAxis(10, 3, 5) == [10, 13, 16, 19, 22] + @test AffineAxis(20, -4, 4) == [20, 16, 12, 8] + @test AffineAxis(7, 0, 3) == [7, 7, 7] + @test_throws BoundsError AffineAxis(7, 0, 3)[4] @test_throws ArgumentError AffineAxis(0, 1, -1) - @test_throws ArgumentError ScatterAxis(offs, -1) - @test_throws DimensionMismatch ScatterAxis(offs, 6) - @test_throws ArgumentError PtrScatterAxis(pointer(offs), -1) + @test is_unit_stride(AffineAxis(3, 1, 2)) && !is_unit_stride(AffineAxis(3, 2, 2)) + @test !is_unit_stride(view([0, 1, 2], 1:3)) end -@testset "axis_from_descriptor: regular -> AffineAxis, irregular -> borrowed ScatterAxis" begin - buf = [99, 99, 7, 10, 13, 16, -1] - ax = axis_from_descriptor(describe_block(buf, 2, 4), buf, 2) - @test ax isa AffineAxis && (ax.base, ax.stride, axis_length(ax)) == (7, 3, 4) - bufi = [-1, -1, 0, 5, 1, 6, -1] - d = describe_block(bufi, 2, 4) - axi = axis_from_descriptor(d, bufi, 2) - @test axi isa ScatterAxis && [axis_offset(axi, t) for t in 0:3] == bufi[3:6] - @test descriptor_offset_range(d, bufi, 2) == axis_offset_range(axi) == (0, 6) - bufi[3] = 555 - @test axis_offset(axi, 0) == 555 - @test axis_from_descriptor(describe_block(Int[], 0), Int[]) isa AffineAxis -end - -@testset "QSTile: base + row_offset(i) + col_offset(j), every axis combination" begin +@testset "Tile: base + rows[i] + cols[j], every axis combination" begin rowoffs, coloffs = [0, 100, 5, -20], [3, -7, 42] - for rows in (AffineAxis(2, 3, 4), ScatterAxis(rowoffs, 4)), cols in (AffineAxis(1, 10, 3), ScatterAxis(coloffs, 3)) + for rows in (AffineAxis(2, 3, 4), view(rowoffs, 1:4)), cols in (AffineAxis(1, 10, 3), view(coloffs, 1:3)) storage = zeros(200) - t = SourceTile(storage, 50, rows, cols) - @test (nrows(t), ncols(t)) == (4, 3) + t = Tile(storage, 50, rows, cols) + @test size(t) == (4, 3) for i in 0:3, j in 0:2 - addr = 50 + (rows isa AffineAxis ? 2 + 3i : rowoffs[i + 1]) + (cols isa AffineAxis ? 1 + 10j : coloffs[j + 1]) - @test tile_offset(t, i, j) == addr - tile_store!(t, i, j, 1000.0i + j) - @test storage[addr + 1] == tile_load(t, i, j) == 1000.0i + j + addr = 50 + rows[i + 1] + cols[j + 1] + t[i, j] = 1000.0i + j + @test storage[addr + 1] == t[i, j] == 1000.0i + j end end + @test_throws BoundsError Tile(zeros(3), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2))[1, 1] end @testset "storage bounds: offset ranges and span checks" begin - @test axis_offset_range(AffineAxis(10, -3, 4)) == (1, 10) - @test axis_offset_range(ScatterAxis([4, -2, 9], 2)) == (-2, 4) - @test axis_offset_range(AffineAxis(5, 1, 0)) == (0, -1) - @test_throws OverflowError axis_offset_range(AffineAxis(typemax(Int), 1, 2)) + @test extrema(AffineAxis(10, -3, 4)) == (1, 10) + @test_throws OverflowError extrema(AffineAxis(typemax(Int), 1, 2)) + @test_throws ArgumentError extrema(AffineAxis(5, 1, 0)) + buf = [99, 99, 7, 10, 13, 16, -1] + @test descriptor_offset_range(describe_block(buf, 2, 4), buf, 2) == (7, 16) + bufi = [-1, -1, 0, 5, 1, 6, -1] + @test descriptor_offset_range(describe_block(bufi, 2, 4), bufi, 2) == (0, 6) + @test descriptor_offset_range(describe_block(Int[], 0), Int[], 0) == (0, -1) storage = zeros(20) - @test checked_tile_storage_bounds(SourceTile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 5, 4))) === nothing - @test_throws BoundsError checked_tile_storage_bounds(SourceTile(storage, 2, AffineAxis(0, 1, 4), AffineAxis(0, 5, 4))) - @test_throws BoundsError checked_tile_storage_bounds(SourceTile(storage, 0, AffineAxis(-1, 1, 4), AffineAxis(0, 5, 4))) + @test checked_tile_storage_bounds(Tile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 5, 4))) === nothing + @test_throws BoundsError checked_tile_storage_bounds(Tile(storage, 2, AffineAxis(0, 1, 4), AffineAxis(0, 5, 4))) + @test_throws BoundsError checked_tile_storage_bounds(Tile(storage, 0, AffineAxis(-1, 1, 4), AffineAxis(0, 5, 4))) + @test_throws BoundsError checked_tile_storage_bounds(Tile(storage, 0, view([4, -2, 9], 1:2), AffineAxis(0, 5, 4))) # An empty axis reads nothing, whatever the other axis addresses. - @test checked_tile_storage_bounds(SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(500, 1, 3))) === nothing + @test checked_tile_storage_bounds(Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(500, 1, 3))) === nothing @test checked_span_bounds(0, (0, 9), (0, 10), 20) === nothing @test_throws BoundsError checked_span_bounds(0, (0, 9), (0, 11), 20) @test checked_span_bounds(-100, (0, -1), (0, 5), 20) === nothing diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 3c9db1a..2bdee37 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -1,5 +1,5 @@ using QuasiStrided: FMAddSubKernel, FMAddSubMethod, PlanarKernel, PlanarMethod, OneMKernel, - InterleavedFormat, SourceTile, complex_method, lanewidth, kernel_shapes, packed_panel, + InterleavedFormat, complex_method, lanewidth, kernel_shapes, packed_panel, _kernel_from_shape, _default_method, accumulator_planes, target_profile, PackedPanel using SIMD: Vec using InteractiveUtils: code_native @@ -118,7 +118,7 @@ const _QSF = QuasiStrided vals[5] = T(0, 0) vals[9] = T(R(-0.0), R(0)) for f in (identity, conj), m in (MR, max(1, MR - 1)) - src = SourceTile(vals, base, AffineAxis(0, 1, m), AffineAxis(0, lda, k_block_length)) + src = Tile(vals, base, AffineAxis(0, 1, m), AffineAxis(0, lda, k_block_length)) A = zeros(T, MR, k_block_length) A[1:m, :] = f.(reshape(vals[(base + 1):(base + lda * k_block_length)], lda, k_block_length)[1:m, :]) want = mk_pack_a(k, A) @@ -129,7 +129,7 @@ const _QSF = QuasiStrided GC.@preserve bufp pack_a!(packed_panel(bufp, 1, length(bufp)), src, k, f) @test isequal(bufp, want) end - src = SourceTile(vals, base, AffineAxis(0, 1, MR), AffineAxis(0, lda, k_block_length)) + src = Tile(vals, base, AffineAxis(0, 1, MR), AffineAxis(0, lda, k_block_length)) bufp = zeros(R, packed_a_length(k, k_block_length)) GC.@preserve bufp begin pp = packed_panel(bufp, 1, length(bufp)) diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index 158d05d..b7dc8aa 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -55,8 +55,8 @@ end k_block_length, lda = 6, MR + 2 Av, Bv = rand(rng, SA, lda * k_block_length + 1), rand(rng, SB, k_block_length * 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, k_block_length)) - srcB = SourceTile(Bv, 0, AffineAxis(0, NR, k_block_length), AffineAxis(0, 1, n)) + srcA = Tile(Av, 1, AffineAxis(0, 1, m), AffineAxis(0, lda, k_block_length)) + srcB = Tile(Bv, 0, AffineAxis(0, NR, k_block_length), AffineAxis(0, 1, n)) A = zeros(T, MR, k_block_length) B = zeros(T, k_block_length, NR) A[1:m, :] = f.(reshape(Av[2:(1 + lda * k_block_length)], lda, k_block_length)[1:m, :]) @@ -73,7 +73,7 @@ end 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, k_block_length, alpha, beta) + execute_tile!(k, Tile(got, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)), pa, pb, k_block_length, alpha, beta) @test mk_close(got, alpha .* vec((A * B)[1:m, 1:n]) .+ beta .* C0, T) end end diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 97a1b89..02ca01c 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -11,7 +11,7 @@ # The fast path ships for AVX-512 only; every expectation derives from the live # profile, so `test/forced_isa_runner.jl` checks the other ISAs too. -using QuasiStrided: PlanarKernel, FMAddSubKernel, KERNEL_SHAPES_C64_FMADDSUB, KERNEL_SHAPES_C32_FMADDSUB, PtrScatterAxis, TargetProfile, CacheLevel, +using QuasiStrided: PlanarKernel, FMAddSubKernel, KERNEL_SHAPES_C64_FMADDSUB, KERNEL_SHAPES_C32_FMADDSUB, TargetProfile, CacheLevel, target_profile, unknown_target, KERNEL_SHAPES_C64_PLANAR, KERNEL_SHAPES_C32_PLANAR const STORE_FASTPATH_ON = QuasiStrided.complex_fastpath_isa_eligible() @@ -64,11 +64,11 @@ function mk_store_fastpath(K, T, shapes; beta0_exact = false, S = T) 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)) + dfast = Tile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) @test QuasiStrided._complex_vector_eligible(dfast, T) == STORE_FASTPATH_ON 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) + store_tile!(Tile(scal, 0, view(collect(0:(m - 1)), 1:m), AffineAxis(0, m, n)), acc, alpha, beta, k) 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) @@ -92,13 +92,11 @@ end m, n = 16, 4 storage = zeros(T, 4 * m * n) eligible(s, rows, cols = AffineAxis(0, m, n)) = - QuasiStrided._complex_vector_eligible(DestinationTile(s, 0, rows, cols), T) + QuasiStrided._complex_vector_eligible(Tile(s, 0, rows, cols), T) @test eligible(storage, AffineAxis(0, 1, m)) == STORE_FASTPATH_ON @test !eligible(storage, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)) - @test !QuasiStrided._complex_vector_eligible(DestinationTile(storage, m - 1, AffineAxis(0, -1, m), AffineAxis(0, m, n)), T) - @test !eligible(storage, ScatterAxis(collect(0:(m - 1)), m)) - ptr_rows = collect(0:(m - 1)) - GC.@preserve ptr_rows @test !eligible(storage, PtrScatterAxis(pointer(ptr_rows), m)) + @test !QuasiStrided._complex_vector_eligible(Tile(storage, m - 1, AffineAxis(0, -1, m), AffineAxis(0, m, n)), T) + @test !eligible(storage, view(collect(0:(m - 1)), 1: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)) == STORE_FASTPATH_ON diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index fa92e8f..9a93282 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -6,7 +6,7 @@ # never through the engine's packers or offset accessors. using Random -using QuasiStrided: SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel, QSTile +using QuasiStrided: SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel using SIMD: Vec using StridedViews: StridedView @@ -89,8 +89,8 @@ function mk_contract_light(k) if !(k isa ScalarKernel) skip = VERSION < v"1.11" ab = (T(2), T(0.5)) - scat = DestinationTile(zeros(T, MR * NR), 0, ScatterAxis(collect(0:(MR - 1)), MR), AffineAxis(0, MR, NR)) - part = DestinationTile(mk_dense(zeros(T, m * NR)), 0, AffineAxis(0, 1, m), AffineAxis(0, m, NR)) + scat = Tile(zeros(T, MR * NR), 0, view(collect(0:(MR - 1)), 1:MR), AffineAxis(0, MR, NR)) + part = Tile(mk_dense(zeros(T, m * NR)), 0, AffineAxis(0, 1, m), AffineAxis(0, m, NR)) mk_run_acc(k, pa, pb, k_block_length) @test (@allocated mk_run_acc(k, pa, pb, k_block_length)) == 0 skip = skip for dst in (scat, part) @@ -137,8 +137,8 @@ function mk_contract_full(k) Cold = iszero(beta) ? fill(mk_nan(T), m * n) : rand(rng, T, m * n) storage = [fill(sentinel, pad); Cold; fill(sentinel, pad)] scattered || (storage = mk_dense(storage)) - rows = scattered ? ScatterAxis(collect(0:(m - 1)), m) : AffineAxis(0, 1, m) - execute_tile!(k, DestinationTile(storage, pad, rows, AffineAxis(0, m, n)), pa, pb, k_block_length, alpha, beta) + rows = scattered ? view(collect(0:(m - 1)), 1:m) : AffineAxis(0, 1, m) + execute_tile!(k, Tile(storage, pad, rows, AffineAxis(0, m, n)), pa, pb, k_block_length, alpha, beta) AB = (A * B)[1:m, 1:n] want = iszero(beta) ? alpha .* vec(AB) : alpha .* vec(AB) .+ beta .* Cold @test mk_close(storage[(pad + 1):(pad + m * n)], want, T) @@ -154,26 +154,26 @@ function mk_contract_full(k) 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)) - for (storage, rows) in ((mk_dense(fill(mk_nan(T), m * n)), AffineAxis(0, 1, m)), (fill(mk_nan(T), m * n), ScatterAxis(collect(0:(m - 1)), m))) - execute_tile!(k, DestinationTile(storage, 0, rows, AffineAxis(0, m, n)), pa, pb, 2, one(T), zero(T)) + for (storage, rows) in ((mk_dense(fill(mk_nan(T), m * n)), AffineAxis(0, 1, m)), (fill(mk_nan(T), m * n), view(collect(0:(m - 1)), 1:m))) + execute_tile!(k, Tile(storage, 0, rows, AffineAxis(0, m, n)), pa, pb, 2, one(T), zero(T)) @test mk_close(storage, vec(A[1:m, :] * B[:, 1:n]), T) end # alpha == 0 never reads acc or the panels; k_block_length == 0 never reads the panels. - fulltile(s) = DestinationTile(s, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, NR)) + fulltile(s) = Tile(s, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, NR)) storage = mk_dense(fill(T(2), MR * NR)) store_tile!(fulltile(storage), mk_fill_acc(k, R(NaN)), zero(T), T(3), k) @test all(==(T(6)), storage) execute_tile!(k, fulltile(storage), R[], R[], 5, zero(T), zero(T)) @test all(iszero, storage) - empty = DestinationTile(T[], 0, AffineAxis(0, 1, 0), AffineAxis(0, 0, 0)) + empty = Tile(T[], 0, AffineAxis(0, 1, 0), AffineAxis(0, 0, 0)) @test store_tile!(empty, acc, one(T), zero(T), k) === empty @test execute_tile!(k, empty, pa, pb, 2, one(T), one(T)) === empty # Rejected inputs. st = zeros(T, (MR + 1) * (NR + 1)) - @test_throws ArgumentError execute_tile!(k, DestinationTile(st, 0, AffineAxis(0, 1, MR + 1), AffineAxis(0, MR + 1, NR)), pa, pb, 2, one(T), zero(T)) - @test_throws ArgumentError execute_tile!(k, DestinationTile(st, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, NR + 1)), pa, pb, 2, one(T), zero(T)) + @test_throws ArgumentError execute_tile!(k, Tile(st, 0, AffineAxis(0, 1, MR + 1), AffineAxis(0, MR + 1, NR)), pa, pb, 2, one(T), zero(T)) + @test_throws ArgumentError execute_tile!(k, Tile(st, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, NR + 1)), pa, pb, 2, one(T), zero(T)) @test_throws ArgumentError execute_tile!(k, fulltile(st), pa, pb, -1, one(T), zero(T)) @test_throws DimensionMismatch execute_tile!(k, fulltile(st), pa[1:(end - 1)], pb, 2, one(T), zero(T)) @test_throws DimensionMismatch execute_tile!(k, fulltile(st), pa, pb[1:(end - 1)], 2, one(T), zero(T)) @@ -189,15 +189,15 @@ function mk_contract_full(k) nanbuf(len) = fill(mk_nan(T), len) layouts = ( # (storage, base, rows, cols, zero-based address of (i, j)) - (nanbuf(ld * n), 0, ScatterAxis(perm, m), AffineAxis(0, ld, n), (i, j) -> perm[i + 1] + ld * j), + (nanbuf(ld * n), 0, view(perm, 1:m), AffineAxis(0, ld, n), (i, j) -> perm[i + 1] + ld * j), (mk_dense(nanbuf(ld * n)), (m - 1) + ld * (n - 1), AffineAxis(0, -1, m), AffineAxis(0, -ld, n), (i, j) -> (m - 1 - i) + ld * (n - 1 - j)), (mk_dense(nanbuf(2ld * n)), 0, AffineAxis(0, 2, m), AffineAxis(0, 2ld, n), (i, j) -> 2i + 2ld * j), (mk_dense(nanbuf(ld * n)), 0, AffineAxis(0, 1, m), AffineAxis(0, ld, n), (i, j) -> i + ld * j), - (mk_dense(nanbuf(ld * n)), 0, AffineAxis(0, 1, m), ScatterAxis(cperm, n), (i, j) -> i + cperm[j + 1]), + (mk_dense(nanbuf(ld * n)), 0, AffineAxis(0, 1, m), view(cperm, 1:n), (i, j) -> i + cperm[j + 1]), (view(nanbuf(ld * n + 4), 3:(ld * n + 2)), 0, AffineAxis(0, 1, m), AffineAxis(0, ld, n), (i, j) -> i + ld * j), ) for (storage, base, rows, cols, addr) in layouts - execute_tile!(k, DestinationTile(storage, base, rows, cols), pa, pb, k_block_length, one(T), zero(T)) + execute_tile!(k, Tile(storage, base, rows, cols), pa, pb, k_block_length, one(T), zero(T)) got = [storage[addr(i, j) + 1] for i in 0:(m - 1), j in 0:(n - 1)] @test mk_close(got, AB[1:m, 1:n], T) @test count(!isnan, storage) == m * n diff --git a/test/microkernels/test_simd_kernel.jl b/test/microkernels/test_simd_kernel.jl index 7feecb6..9fbf142 100644 --- a/test/microkernels/test_simd_kernel.jl +++ b/test/microkernels/test_simd_kernel.jl @@ -23,14 +23,14 @@ using StridedViews: StridedView m, n = 8, 8 rows, cols = AffineAxis(0, 1, m), AffineAxis(0, m, n) mem = parent(StridedView(zeros(T, m * n))) # `Memory{T}` on Julia >= 1.11 - @test _vector_store_eligible(DestinationTile(mem, 0, rows, cols), T) - @test _vector_store_eligible(DestinationTile(zeros(T, m * n), 0, rows, cols), T) - @test !_vector_store_eligible(DestinationTile(view(zeros(T, m * n + 4), 3:(m * n + 2)), 0, rows, cols), T) - @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(zeros(complex(T), m * n), 0, rows, cols), T) + @test _vector_store_eligible(Tile(mem, 0, rows, cols), T) + @test _vector_store_eligible(Tile(zeros(T, m * n), 0, rows, cols), T) + @test !_vector_store_eligible(Tile(view(zeros(T, m * n + 4), 3:(m * n + 2)), 0, rows, cols), T) + @test !_vector_store_eligible(Tile(zeros(T, m, n), 0, rows, cols), T) + @test !_vector_store_eligible(Tile(zeros(T, 2m * n), 0, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)), T) + @test !_vector_store_eligible(Tile(mem, 0, view(collect(0:(m - 1)), 1:m), cols), T) + @test _vector_store_eligible(Tile(mem, 0, rows, cols), T === Float64 ? Float32 : Float64) + @test !_vector_store_eligible(Tile(zeros(complex(T), m * n), 0, rows, cols), T) end end @@ -41,11 +41,11 @@ using StridedViews: StridedView 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)) + dfast = Tile(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) + store_tile!(Tile(scal, 0, view(collect(0:(m - 1)), 1:m), AffineAxis(0, m, n)), acc, alpha, beta, k) @test mk_close(fast, scal, Float32) end end diff --git a/test/packing/test_pack_complex.jl b/test/packing/test_pack_complex.jl index 643e181..550f189 100644 --- a/test/packing/test_pack_complex.jl +++ b/test/packing/test_pack_complex.jl @@ -32,7 +32,7 @@ const COMPLEX_FORMATS = ( fixtures = ( (AffineAxis(0, 1, PD), AffineAxis(0, 97, k_block_length)), (AffineAxis(40, -1, PD), AffineAxis(3000, -97, k_block_length)), - (ScatterAxis(SCATTER_LANES, PD), ScatterAxis(SCATTER_STEPS, k_block_length)), + (view(SCATTER_LANES, 1:PD), view(SCATTER_STEPS, 1:k_block_length)), ) len = _ref_rpe(fmt) * PD * k_block_length for (lane, step) in fixtures, valid in (PD, 1, 0) @@ -53,24 +53,24 @@ end @testset "complex packing: validation before any write, k_block_length == 0 is a no-op" begin kernel = ComplexKernelDescriptor(Val(4), Val(3), ComplexF64, PlanarFormat(), PlanarFormat()) storage = fill(ComplexF64(3, 4), 100) - src = SourceTile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 8, 2)) + src = Tile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 8, 2)) # The buffer holds realtype(kernel), not scalartype(kernel). @test_throws ArgumentError pack_a!(zeros(ComplexF64, 1000), src, kernel, identity) @test_throws ArgumentError pack_b!(zeros(ComplexF64, 1000), src, kernel, identity) @test_throws ArgumentError pack_a!(zeros(Float32, 1000), src, kernel, identity) canary = fill(-42.0, 1000) - @test_throws ArgumentError pack_a!(canary, SourceTile(storage, 0, AffineAxis(0, 1, 5), src.cols), kernel, identity) - @test_throws ArgumentError pack_b!(canary, SourceTile(storage, 0, src.cols, AffineAxis(0, 1, 4)), kernel, identity) + @test_throws ArgumentError pack_a!(canary, Tile(storage, 0, AffineAxis(0, 1, 5), src.cols), kernel, identity) + @test_throws ArgumentError pack_b!(canary, Tile(storage, 0, src.cols, AffineAxis(0, 1, 4)), kernel, identity) # A buffer sized as if it held complex elements is half as long as needed. @test_throws DimensionMismatch pack_a!(zeros(Float64, 4 * 2), src, kernel, identity) - @test_throws BoundsError pack_a!(canary, SourceTile(storage, 90, src.rows, src.cols), kernel, identity) + @test_throws BoundsError pack_a!(canary, Tile(storage, 90, src.rows, src.cols), kernel, identity) @test all(==(-42.0), canary) read = Ref(false) spy = z -> (read[] = true; z) packed = fill(-42.0, 8) - @test pack_a!(packed, SourceTile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 8, 0)), kernel, spy) === packed - @test pack_b!(packed, SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 8, 3)), kernel, spy) === packed + @test pack_a!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 8, 0)), kernel, spy) === packed + @test pack_b!(packed, Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 8, 3)), kernel, spy) === packed @test all(==(-42.0), packed) @test !read[] end @@ -91,20 +91,20 @@ end GC.@preserve bufa bufb begin pa = packed_panel(bufa, 1, length(bufa)) pb = packed_panel(bufb, 1, length(bufb)) - full_a = SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)) - full_b = SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)) + full_a = Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)) + full_b = Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)) for f in (identity, conj) # the contiguous fast path, where the ISA has it push!(fast, steady_pack_allocs(pack_a!, pa, full_a, kernel, f)) push!(fast, steady_pack_allocs(pack_b!, pb, full_b, kernel, f)) end srcs_a = ( - full_a, SourceTile(storage, 0, ScatterAxis(ro, MR), ScatterAxis(co, k_block_length)), - SourceTile(storage, 0, AffineAxis(0, 1, MR - 1), AffineAxis(0, MR, k_block_length)), - SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), + full_a, Tile(storage, 0, view(ro, 1:MR), view(co, 1:k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, MR - 1), AffineAxis(0, MR, k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), ) srcs_b = ( - full_b, SourceTile(storage, 0, ScatterAxis(co, k_block_length), ScatterAxis(ro_b, NR)), - SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR - 1)), + full_b, Tile(storage, 0, view(co, 1:k_block_length), view(ro_b, 1:NR)), + Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR - 1)), ) for s in srcs_a, f in (identity, conj) push!(slow, steady_pack_allocs(pack_a!, bufa, s, kernel, f)) diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index 2424868..95650cb 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -28,10 +28,10 @@ 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) + len = _ref_rpe(fmt) * PD * length(step) got, canaries = pack_into(pack!, :panel, real(T), len, src, kernel, f) @test canaries - @test all(isequal.(got, ref_pack(fmt, T, PD, axis_length(step), axis_length(lane), g, f))) + @test all(isequal.(got, ref_pack(fmt, T, PD, length(step), length(lane), g, f))) return storage, src end @@ -61,7 +61,7 @@ end for fa in (PlanarFormat(), OneEFormat(), InterleavedFormat()), f in (identity, conj) kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, OneEFormat()) # The lane axis must be unit-stride; the step axis may scatter. - check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), ScatterAxis(koffs, k_block_length), f) + check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), view(koffs, 1:k_block_length), f) check_panel(:b, kernel, T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, k_block_length), f) # Ineligible shapes fall back to the scalar loop, into the same panel. for lane in (AffineAxis(20, -1, MR), AffineAxis(0, 3, MR), AffineAxis(0, 1, MR - 3)) @@ -76,7 +76,7 @@ end storage = complex_storage(T, 4000) buf = zeros(Float64, 64) offs = collect(0:(MR - 1)) - GC.@preserve buf offs begin + GC.@preserve buf begin pk = packed_panel(buf, 1, length(buf)) elig(; packed = pk, store = storage, ax = AffineAxis(0, 1, MR), tr = identity, @@ -95,8 +95,7 @@ end @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)) - @test !elig(ax = PtrScatterAxis(pointer(offs), MR)) + @test !elig(ax = view(offs, 1:MR)) @test !elig(valid = MR - 1) # padding lanes @test !elig(tr = z -> 2z) @test !elig(fmt = RealFormat()) diff --git a/test/packing/test_pack_real.jl b/test/packing/test_pack_real.jl index 9253790..57f9c87 100644 --- a/test/packing/test_pack_real.jl +++ b/test/packing/test_pack_real.jl @@ -3,7 +3,7 @@ # files included after this one. using QuasiStrided: RealFormat, PlanarFormat, OneEFormat, InterleavedFormat, PackedPanel, - packed_panel, PtrScatterAxis, ComplexKernelDescriptor, _copies_unchanged, + packed_panel, ComplexKernelDescriptor, _copies_unchanged, _pack_a_contiguous_eligible # Packed layouts written from the format definitions, not from the offset @@ -34,19 +34,14 @@ function ref_pack(fmt, ::Type{T}, vr, k_block_length, valid, g, f) where {T} return out end -ref_offs(ax::AffineAxis) = [ax.base + t * ax.stride for t in 0:(ax.count - 1)] -ref_offs(ax::ScatterAxis) = collect(ax.offsets[1:ax.count]) -ref_offs(ax::PtrScatterAxis) = [unsafe_load(ax.offsets, t + 1) for t in 0:(ax.count - 1)] resized(ax::AffineAxis, n) = AffineAxis(ax.base, ax.stride, n) -resized(ax::ScatterAxis, n) = ScatterAxis(ax.offsets, n) -resized(ax::PtrScatterAxis, n) = PtrScatterAxis(ax.offsets, n) +resized(ax::SubArray, n) = view(parent(ax), 1:n) # An A (`lane` = rows) or B (`lane` = cols) source over `storage`, and its # direct-indexing reader. function pack_fixture(side, storage, base, lane, step) - lo, so = ref_offs(lane), ref_offs(step) - src = side === :a ? SourceTile(storage, base, lane, step) : SourceTile(storage, base, step, lane) - return src, (t, p) -> storage[base + lo[t + 1] + so[p + 1] + 1] + src = side === :a ? Tile(storage, base, lane, step) : Tile(storage, base, step, lane) + return src, (t, p) -> storage[base + lane[t + 1] + step[p + 1] + 1] end # Packs `src` into a destination of kind `dst` (canaried unless a bare Vector) @@ -78,9 +73,9 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] counting = x -> (calls[] += 1; 3x + 1000) # nonzero at zero: padding must bypass it lanes = Any[ AffineAxis(0, 1, 8), AffineAxis(3, 7, 8), AffineAxis(40, -1, 8), AffineAxis(20, 0, 8), - ScatterAxis(SCATTER_LANES, 8), + view(SCATTER_LANES, 1:8), ] - steps = Any[AffineAxis(0, 64, k_block_length), AffineAxis(700, -97, k_block_length), ScatterAxis(SCATTER_STEPS, k_block_length)] + steps = Any[AffineAxis(0, 64, k_block_length), AffineAxis(700, -97, k_block_length), view(SCATTER_STEPS, 1:k_block_length)] for (side, pack!, PD) in ((:a, pack_a!, MR), (:b, pack_b!, NR)), lane in lanes, step in steps, valid in (PD, 1, 0), dst in (:vector, :view) src, g = pack_fixture(side, storage, 11, resized(lane, valid), step) @@ -108,22 +103,22 @@ end 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 - steps = Any[AffineAxis(0, MR, 5), AffineAxis(1800, -MR, 6), PtrScatterAxis(pointer(koffs), 4)] + for storage in storages + steps = Any[AffineAxis(0, MR, 5), AffineAxis(1800, -MR, 6), view(koffs, 1:4)] lanes = Any[ (AffineAxis(0, 1, MR), true), (AffineAxis(5, 1, MR), true), (AffineAxis(0, 2, MR), false), (AffineAxis(MR + 3, -1, MR), false), - (AffineAxis(0, 1, MR - 1), false), (PtrScatterAxis(pointer(contig), MR), false), + (AffineAxis(0, 1, MR - 1), false), (view(contig, 1:MR), false), ] for (lane, eligible) in lanes, step in steps, f in (identity, conj, x -> -x) src, g = pack_fixture(:a, storage, 11, lane, step) - k_block_length = axis_length(step) + k_block_length = length(step) pp = packed_panel(zeros(T, 1), 1, 1) - @test _pack_a_contiguous_eligible(pp, src, f, nrows(src), Val(MR), T) == + @test _pack_a_contiguous_eligible(pp, src, f, length(lane), Val(MR), T) == (eligible && (f === identity || f === conj)) - @test !_pack_a_contiguous_eligible(zeros(T, 1), src, f, nrows(src), Val(MR), T) + @test !_pack_a_contiguous_eligible(zeros(T, 1), src, f, length(lane), Val(MR), T) got, canaries = pack_into(pack_a!, :panel, T, MR * k_block_length, src, kernel, f) - @test got == ref_pack(RealFormat(), T, MR, k_block_length, nrows(src), g, f) + @test got == ref_pack(RealFormat(), T, MR, k_block_length, length(lane), g, f) @test canaries end end @@ -137,8 +132,8 @@ end read = Ref(false) spy = x -> (read[] = true; x) packed = fill(-42.0, 8) - @test pack_a!(packed, SourceTile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 1, 0)), kernel, spy) === packed - @test pack_b!(packed, SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 1, 3)), kernel, spy) === packed + @test pack_a!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 1, 0)), kernel, spy) === packed + @test pack_b!(packed, Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 1, 3)), kernel, spy) === packed @test packed == fill(-42.0, 8) @test !read[] end @@ -147,20 +142,20 @@ end kernel = KernelDescriptor(Val(4), Val(3), Float64) kernel32 = KernelDescriptor(Val(4), Val(3), Float32) storage = fill(9.0, 20) - tile(m, n, base = 0) = SourceTile(storage, base, AffineAxis(0, 1, m), AffineAxis(0, 1, n)) + tile(m, n, base = 0) = Tile(storage, base, AffineAxis(0, 1, m), AffineAxis(0, 1, n)) for (pack!, bad, ok, short) in ((pack_a!, tile(5, 3), tile(4, 3), 11), (pack_b!, tile(3, 4), tile(3, 3), 8)) canary = fill(-1.0, 100) @test_throws ArgumentError pack!(canary, bad, kernel, identity) # m > MR / n > NR @test_throws ArgumentError pack!(canary, ok, kernel32, identity) # eltype mismatch @test_throws DimensionMismatch pack!(fill(-1.0, short), ok, kernel, identity) - @test_throws BoundsError pack!(canary, tile(nrows(ok), ncols(ok), 18), kernel, identity) + @test_throws BoundsError pack!(canary, tile(size(ok)..., 18), kernel, identity) @test canary == fill(-1.0, 100) end end @testset "pack_a!/pack_b!: zero steady-state allocation" begin nontrivial(x) = 2x + 1 - # Every destination kind, affine/scattered/pointer-scattered axes, tails and + # Every destination kind, affine/scattered axes, tails and # k_block_length == 0; the ScalarKernel/SIMDKernel forwarding methods on a subset. function run(ctor, ::Type{T}, MR, NR, full) where {T} kernel = ctor(Val(MR), Val(NR), T) @@ -171,23 +166,23 @@ end bufa = zeros(T, MR * k_block_length + 8) bufb = zeros(T, NR * k_block_length + 8) bytes = Int[] - GC.@preserve bufa bufb koffs begin + GC.@preserve bufa bufb begin pa = packed_panel(bufa, 1, MR * k_block_length) dsts_a = (bufa, view(bufa, 3:(2 + MR * k_block_length)), pa) dsts_b = (bufb, view(bufb, 3:(2 + NR * k_block_length)), packed_panel(bufb, 1, NR * k_block_length)) - pk = PtrScatterAxis(pointer(koffs), k_block_length) + pk = view(koffs, 1:k_block_length) srcs_a = ( - SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)), - SourceTile(storage, 0, AffineAxis(0, 1, MR), pk), - SourceTile(storage, 0, AffineAxis(0, 2, MR - 1), AffineAxis(0, 2MR, k_block_length)), - SourceTile(storage, 0, ScatterAxis(lanes, MR), ScatterAxis(koffs, k_block_length)), - SourceTile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), + Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, MR), pk), + Tile(storage, 0, AffineAxis(0, 2, MR - 1), AffineAxis(0, 2MR, k_block_length)), + Tile(storage, 0, view(lanes, 1:MR), view(koffs, 1:k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), ) srcs_b = ( - SourceTile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)), - SourceTile(storage, 0, pk, AffineAxis(0, k_block_length, NR - 1)), - SourceTile(storage, 0, ScatterAxis(koffs, k_block_length), ScatterAxis(lanes, NR)), - SourceTile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, k_block_length, NR)), + Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)), + Tile(storage, 0, pk, AffineAxis(0, k_block_length, NR - 1)), + Tile(storage, 0, view(koffs, 1:k_block_length), view(lanes, 1:NR)), + Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, k_block_length, NR)), ) if !full srcs_a, srcs_b, dsts_a, dsts_b = srcs_a[1:1], srcs_b[1:1], dsts_a[1:2], dsts_b[1:2] diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 094c93b..3c69b95 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -97,12 +97,12 @@ end reg = rand(1:4) # sliver height nsliv = cld(nrow, reg) - cols = ScatterAxis(coloffs, ncol) + cols = view(coloffs, 1:ncol) persliver = true for s in 0:(nsliv - 1) first = s * reg cnt = min(reg, nrow - first) - rows = ScatterAxis(view(rowoffs, (first + 1):(first + cnt)), cnt) + rows = view(rowoffs, (first + 1):(first + cnt)) ok = try checked_tile_storage_bounds(base, rows, cols, len) true @@ -164,7 +164,7 @@ end @test QS.descriptor_offset_range(d1[3], buf1, 12) == (1, 1) end -@testset "per-call floor: descriptor_offset_range agrees with axis_offset_range" begin +@testset "per-call floor: descriptor_offset_range agrees with extrema" begin Random.seed!(7) for trial in 1:100 n = rand(1:8) @@ -172,7 +172,7 @@ end rand() < 0.4 && (buf = [3 + 5 * (t - 1) for t in 1:n]) # force a regular run d = describe_block(buf, 0, n) ax = QS._axis_of(d, buf, 0) - @test QS.descriptor_offset_range(d, buf, 0) == axis_offset_range(ax) + @test QS.descriptor_offset_range(d, buf, 0) == extrema(ax) end @test QS.descriptor_offset_range(BlockDescriptor(0, 0, 0, true), Int[], 0) == (0, -1) end @@ -348,7 +348,7 @@ end @testset "per-call floor: the unsafe_* packers keep every non-bounds check" begin kernel = SIMDKernel(Val(8), Val(6), Float64) - src = SourceTile(collect(1.0:64.0), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 4)) + src = Tile(collect(1.0:64.0), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 4)) packed = zeros(8 * 4) ref = zeros(8 * 4) @@ -359,10 +359,10 @@ end # Only the bounds checks are skipped. @test_throws DimensionMismatch QS.unsafe_pack_a!(zeros(3), src, kernel, identity) @test_throws ArgumentError QS.unsafe_pack_a!(zeros(Float32, 64), src, kernel, identity) - toowide = SourceTile(collect(1.0:200.0), 0, AffineAxis(0, 1, 9), AffineAxis(0, 16, 4)) + toowide = Tile(collect(1.0:200.0), 0, AffineAxis(0, 1, 9), AffineAxis(0, 16, 4)) @test_throws ArgumentError QS.unsafe_pack_a!(zeros(200), toowide, kernel, identity) - srcB = SourceTile(collect(1.0:64.0), 0, AffineAxis(0, 1, 4), AffineAxis(0, 4, 6)) + srcB = Tile(collect(1.0:64.0), 0, AffineAxis(0, 1, 4), AffineAxis(0, 4, 6)) refB = zeros(6 * 4) packedB = zeros(6 * 4) pack_b!(refB, srcB, kernel, identity) @@ -370,15 +370,15 @@ end @test packedB == refB @test_throws DimensionMismatch QS.unsafe_pack_b!(zeros(3), srcB, kernel, identity) - dest = DestinationTile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) + dest = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) @test_throws DimensionMismatch QS.unsafe_execute_tile!( kernel, dest, zeros(3), zeros(6 * 4), 4, 1.0, 0.0 ) @test_throws ArgumentError QS.unsafe_execute_tile!( kernel, dest, zeros(8 * 4), zeros(6 * 4), -1, 1.0, 0.0 ) - d1 = DestinationTile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) - d2 = DestinationTile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) + d1 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) + d2 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) execute_tile!(kernel, d1, ref, refB, 4, 1.0, 0.0) QS.unsafe_execute_tile!(kernel, d2, ref, refB, 4, 1.0, 0.0) @test d1.storage == d2.storage @@ -542,4 +542,13 @@ end allocs_s = @allocated _pcf_exec(ps, 1.0, 0.0) @test Cs ≈ Ap * Bn @test allocs_s == 0 skip = (VERSION < v"1.11") + + # M, N and K composites ordered differently per operand: scattered tile axes. + A4 = randn(5, 6, 7, 9); B4 = randn(7, 6, 11, 3); C4 = zeros(3, 9, 11, 5) + p4 = _pcf_plan( + StridedView(C4), StridedView(A4), (1, 2, 3, 4), StridedView(B4), (3, 2, 5, 6), (6, 4, 5, 1) + ) + _pcf_exec(p4, 1.0, 0.5) + _pcf_exec(p4, 1.0, 0.5) + @test (@allocated _pcf_exec(p4, 1.0, 0.5)) == 0 skip = (VERSION < v"1.11") end diff --git a/test/runtests.jl b/test/runtests.jl index f2790d8..1bdba6e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,8 +5,7 @@ using QuasiStrided using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, describe_block, block_descriptors!, KernelDescriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, - AffineAxis, ScatterAxis, SourceTile, DestinationTile, axis_from_descriptor, nrows, - ncols, axis_offset_range, checked_tile_storage_bounds, pack_a!, pack_b!, + AffineAxis, Tile, checked_tile_storage_bounds, pack_a!, pack_b!, zero_accumulator, accumulate, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, target_profile, cache_topology, unknown_target, detect_isa, detect_target, From d2ee74b658e5b4166c0cbe64f6e28fcd881d020c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 11:02:40 -0400 Subject: [PATCH 07/60] One-based tile, sliver and panel coordinates - Tile indexing, AffineAxis/view axes, packer lane and K-step loops, microkernel K steps, columns and row lanes (generated code included) and UnpackedBView columns/K steps are one-based; the two zero-based shims go. - packed_{a,b}_offset and packed_{a,b}_plane_offset take one-based coordinates and still return zero-based offsets: the panel layout is unchanged. Addresses, block starts and AxisGroup coordinates stay zero-based, as does the macro-block coordinate of the line-by-line packer. Tests: 30930 pass. TTFX probe, Float64 / ComplexF64 first call, interleaved with the chunk-2 commit: before 3.1-3.2 s / 2.1-2.2 s, after 3.3-3.4 s / 2.2-2.4 s. Reused-plan execute! medians (us, 3 interleaved runs, before -> after): matmul 256^3 Float64 411 -> 400, ComplexF64 1564 -> 1532; scattered-M 12x20x256x256 Float64 448 -> 431, ComplexF64 1503 -> 1499; scattered-M/N/K small Float64 17.9 -> 17.6, ComplexF64 42.0 -> 41.4. bench_kernels.jl kernel arm, Float64, median of 5 interleaved runs (GF/s): 16x6/W8 85.8 -> 84.7, 32x6/W8 107.0 -> 106.9. Co-Authored-By: Claude Opus 5.5 --- src/execution/unpackedb.jl | 12 ++--- src/layout/tiles.jl | 16 +++---- src/microkernels/fmaddsub.jl | 42 +++++++++--------- src/microkernels/interface.jl | 4 +- src/microkernels/mixed.jl | 2 +- src/microkernels/onem.jl | 19 ++++---- src/microkernels/planar.jl | 59 +++++++++++++------------ src/microkernels/scalar.jl | 13 +++--- src/microkernels/simd.jl | 45 ++++++++++--------- src/packing/format.jl | 21 +++++---- src/packing/pack.jl | 34 +++++++------- src/packing/pack_contiguous.jl | 20 ++++----- src/packing/transposed.jl | 15 ++++--- test/layout/test_tiles.jl | 6 +-- test/microkernels/test_planar_kernel.jl | 2 +- test/packing/test_kernel_descriptor.jl | 4 +- 16 files changed, 158 insertions(+), 156 deletions(-) diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 84386ac..92cf58b 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -5,11 +5,11 @@ # (N, K) block, which at small M costs as much as the kernel itself. # Stands in for a packed B micro-panel of `NR` columns: column `j`, K step -# `p` is storage address `colbase[j+1] + axis_offset(ksteps, p)` (base -# included). Padding columns of a tail sliver alias the last valid column, so -# every address read is inside the hoisted B check; their results are never -# stored. `length` is the packed-equivalent count that -# `_execute_tile_prologue!` checks, not the storage length. +# `p` is storage address `colbase[j] + ksteps[p]` (base included). Padding +# columns of a tail sliver alias the last valid column, so every address read +# is inside the hoisted B check; their results are never stored. `length` is +# the packed-equivalent count that `_execute_tile_prologue!` checks, not the +# storage length. struct UnpackedBView{S, K <: AbstractVector{Int}, NR, F} storage::S colbase::NTuple{NR, Int} @@ -22,7 +22,7 @@ Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) # `@inbounds`: the caller has checked the whole B panel region. @inline function _unpacked_b_element(u::UnpackedBView{S, K, NR, F}, j::Int, p::Int) where {S, K, NR, F} - @inbounds z = u.storage[u.colbase[j + 1] + axis_offset(u.ksteps, p) + 1] + @inbounds z = u.storage[u.colbase[j] + u.ksteps[p] + 1] return u.transform(z) end diff --git a/src/layout/tiles.jl b/src/layout/tiles.jl index d3282a6..5d98746 100644 --- a/src/layout/tiles.jl +++ b/src/layout/tiles.jl @@ -1,7 +1,7 @@ -# Tile axes and tiles. Addresses are zero-based (storage index = address + 1); -# a tile's `base` is the address of its logical origin. A tile axis is the -# vector of offsets of its rows or columns: an `AffineAxis`, or a `view` of an -# offset buffer. +# Tile axes and tiles. Addresses are zero-based (storage index = address + 1), +# coordinates one-based; a tile's `base` is the address its offsets are +# relative to. A tile axis is the vector of offsets of its rows or columns: an +# `AffineAxis`, or a `view` of an offset buffer. # Offsets `base + (t - 1) * stride`. Not a `StepRange`: the stride may be zero. struct AffineAxis <: AbstractVector{Int} @@ -29,9 +29,6 @@ function Base.extrema(ax::AffineAxis) return minmax(ax.base, last) end -# Zero-based until coordinates become one-based. -@inline axis_offset(ax::AbstractVector{Int}, t::Int) = @inbounds ax[t + 1] - # An offset interval of one map: `regular` iff `buffer[t+1] == base + t*stride` # for all `t < count`; otherwise read the buffer (valid until it is refilled). struct BlockDescriptor @@ -85,12 +82,11 @@ end Base.size(tile::Tile) = (length(tile.rows), length(tile.cols)) -# Zero-based until coordinates become one-based. Base.@propagate_inbounds Base.getindex(tile::Tile, i::Int, j::Int) = - tile.storage[tile.base + tile.rows[i + 1] + tile.cols[j + 1] + 1] + tile.storage[tile.base + tile.rows[i] + tile.cols[j] + 1] Base.@propagate_inbounds function Base.setindex!(tile::Tile, v, i::Int, j::Int) - tile.storage[tile.base + tile.rows[i + 1] + tile.cols[j + 1] + 1] = v + tile.storage[tile.base + tile.rows[i] + tile.cols[j] + 1] = v return tile end diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 09bf147..ccc33e7 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -1,5 +1,5 @@ # FMAddSub (interleaved-accumulator) complex microkernel. One accumulator plane -# in 1m's layout (lanes `(2u, 2u+1)` are `(re, im)` of one complex row), updated +# in 1m's layout (lanes `(2u-1, 2u)` are `(re, im)` of one complex row), updated # per (A-vector, B-column) pair by two chained x86 `vfmaddsub` ops: # # acc = fmaddsub(a, br, fmaddsub(swap(a), bi, acc)) @@ -109,8 +109,8 @@ end av = [Symbol(:a, v) for v in 0:(MV - 1)] sv = [Symbol(:s, v) for v in 0:(MV - 1)] - brv = [Symbol(:br, j) for j in 0:(NR - 1)] - biv = [Symbol(:bi, j) for j in 0:(NR - 1)] + brv = [Symbol(:br, j) for j in 1:NR] + biv = [Symbol(:bi, j) for j in 1:NR] load_a = Any[] for v in 0:(MV - 1) @@ -118,7 +118,7 @@ end load_a, :( $(av[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W), p) + Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W + 1), p) ) ) ) @@ -126,19 +126,19 @@ end end load_b = Any[] - for j in 0:(NR - 1) + for j in 1:NR push!(load_b, :((br_s, bi_s) = _b_step_load2(packed_b, kernel, $j, p))) - push!(load_b, :($(brv[j + 1]) = Vec{$W, $R}(br_s))) - push!(load_b, :($(biv[j + 1]) = Vec{$W, $R}(bi_s))) + push!(load_b, :($(brv[j]) = Vec{$W, $R}(br_s))) + push!(load_b, :($(biv[j]) = Vec{$W, $R}(bi_s))) end acc_exprs = Vector{Any}(undef, NA) - for j in 0:(NR - 1), v in 0:(MV - 1) - idx = v + MV * j + 1 + for j in 1:NR, v in 0:(MV - 1) + idx = v + MV * (j - 1) + 1 acc_exprs[idx] = :( _fmaddsub( - $(av[v + 1]), $(brv[j + 1]), - _fmaddsub($(sv[v + 1]), $(biv[j + 1]), acc[$idx]) + $(av[v + 1]), $(brv[j]), + _fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) ) ) end @@ -159,7 +159,7 @@ function Base.accumulate( ) where {MR, NR, T, W, R, NA, PA, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) - @inbounds for p in 0:(k_block_length - 1) + @inbounds for p in 1:k_block_length acc = _accumulate_step_fmaddsub(kernel, acc, packed_a, packed_b, p) end return acc @@ -212,10 +212,10 @@ end HW = W ÷ 2 blocks = Any[] - for j in 0:(NR - 1) + for j in 1:NR vblocks = Any[] for v in 0:(MV - 1) - idx = v + MV * j + 1 + idx = v + MV * (j - 1) + 1 push!( vblocks, quote vec = acc[$idx] @@ -224,12 +224,12 @@ end sp, 2 * (colbase + $(v * HW)), vec, ar, ai, br, bi, beta ) elseif $(v * HW) < m - for u in 0:$(HW - 1) + for u in 1:$HW i = $(v * HW) + u - i < m || break + i <= m || break _axpby_at!( - storage, colbase + i + 1, alpha, - Complex(vec[2 * u + 1], vec[2 * u + 2]), beta + storage, colbase + i, alpha, + Complex(vec[2 * u - 1], vec[2 * u]), beta ) end end @@ -238,8 +238,8 @@ end end push!( blocks, quote - if $j < n - colbase = rowbase0 + axis_offset(cols, $j) # zero-based (0, j) + if $j <= n + colbase = rowbase0 + cols[$j] # the address of (1, j) $(vblocks...) end end @@ -249,7 +249,7 @@ end return quote storage = destination.storage cols = destination.cols - rowbase0 = destination.base + axis_offset(destination.rows, 0) + rowbase0 = @inbounds destination.base + destination.rows[1] ar = Vec{$W, $R}(real(alpha)) ai = Vec{$W, $R}(imag(alpha)) br = Vec{$W, $R}(real(beta)) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 55bc802..eeaee0b 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -119,11 +119,11 @@ function scale_tile!(destination::Tile, beta::T) where {T} if isone(beta) return destination elseif iszero(beta) - @inbounds for j in 0:(n - 1), i in 0:(m - 1) + @inbounds for j in 1:n, i in 1:m destination[i, j] = zero(T) end else - @inbounds for j in 0:(n - 1), i in 0:(m - 1) + @inbounds for j in 1:n, i in 1:m destination[i, j] = convert(T, destination[i, j]) * beta end end diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index 4173a18..dd916dd 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -4,7 +4,7 @@ # * complex A x real B: interleaved A is a real `2MR x k_block_length` # panel; the accumulator is 1m/fmaddsub's interleaved layout. # * real A x complex B: interleaved B is a real `k_block_length x 2NR` panel; -# accumulator column `2j` holds the real and `2j+1` the imaginary part of +# accumulator column `2j-1` holds the real and `2j` the imaginary part of # column `j`. using SIMD: Vec diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index cce67f8..82ab64a 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -2,7 +2,7 @@ # `2*k_block_length` real K steps against 1e-packed A and planar B. At logical K # step `p`, `OneEFormat` A holds `(re_0, im_0, re_1, im_1, ...)` then `(-im_0, # re_0, -im_1, re_1, ...)`, and planar B `re_0..` then `im_0..`, so accumulator -# real row `2i` is the real and `2i+1` the imaginary part of complex row `i`. +# real row `2i-1` is the real and `2i` the imaginary part of complex row `i`. # There is deliberately no FMA loop here: 1m reuses the real kernel body # verbatim. @@ -62,8 +62,9 @@ function Base.accumulate( end # Scalar store for an interleaved accumulator (1m and fmaddsub): with `W` even -# and `2MR` a multiple of `W`, complex row `i = v*(W÷2) + u` is lanes `2u+1` -# (re) and `2u+2` (im) of vector `v` -- never split across two vectors. +# and `2MR` a multiple of `W`, complex row `i = v*(W÷2) + u` is lanes `2u-1` +# (re) and `2u` (im) of vector `v` (zero-based) -- never split across two +# vectors. @generated function _store_tile_lanepair!( destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, @@ -75,18 +76,18 @@ end HW = W ÷ 2 blocks = Any[] - for j in 0:(NR - 1), v in 0:(MV - 1) - idx = v + MV * j + 1 + for j in 1:NR, v in 0:(MV - 1) + idx = v + MV * (j - 1) + 1 push!( blocks, quote - if $j < n + if $j <= n vec = acc[$idx] - for u in 0:$(HW - 1) + for u in 1:$HW i = $(v * HW) + u - i < m || break + i <= m || break _axpby_tile!( destination, i, $j, alpha, - Complex(vec[2 * u + 1], vec[2 * u + 2]), beta + Complex(vec[2 * u - 1], vec[2 * u]), beta ) end end diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 10a7a76..747d74a 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -66,8 +66,8 @@ end arv = [Symbol(:ar, v) for v in 0:(MV - 1)] aiv = [Symbol(:ai, v) for v in 0:(MV - 1)] naiv = [Symbol(:nai, v) for v in 0:(MV - 1)] - brv = [Symbol(:br, j) for j in 0:(NR - 1)] - biv = [Symbol(:bi, j) for j in 0:(NR - 1)] + brv = [Symbol(:br, j) for j in 1:NR] + biv = [Symbol(:bi, j) for j in 1:NR] load_a = Any[] for v in 0:(MV - 1) @@ -75,7 +75,7 @@ end load_a, :( $(arv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W), p) + Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W + 1), p) ) ) ) @@ -83,7 +83,7 @@ end load_a, :( $(aiv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 1, $(v * W), p) + Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 1, $(v * W + 1), p) ) ) ) @@ -91,23 +91,23 @@ end end load_b = Any[] - for j in 0:(NR - 1) - push!(load_b, :(($(brv[j + 1]), $(biv[j + 1])) = _b_step_load2(packed_b, kernel, $j, p))) + for j in 1:NR + push!(load_b, :(($(brv[j]), $(biv[j])) = _b_step_load2(packed_b, kernel, $j, p))) end acc_exprs = Vector{Any}(undef, NA) - for j in 0:(NR - 1), v in 0:(MV - 1) - idx = v + MV * j + 1 + for j in 1:NR, v in 0:(MV - 1) + idx = v + MV * (j - 1) + 1 acc_exprs[idx] = :( # re: ar*br - ai*bi muladd( - $(naiv[v + 1]), $(biv[j + 1]), - muladd($(arv[v + 1]), $(brv[j + 1]), acc[$idx]) + $(naiv[v + 1]), $(biv[j]), + muladd($(arv[v + 1]), $(brv[j]), acc[$idx]) ) ) acc_exprs[NV + idx] = :( # im: ar*bi + ai*br muladd( - $(aiv[v + 1]), $(brv[j + 1]), - muladd($(arv[v + 1]), $(biv[j + 1]), acc[$(NV + idx)]) + $(aiv[v + 1]), $(brv[j]), + muladd($(arv[v + 1]), $(biv[j]), acc[$(NV + idx)]) ) ) end @@ -128,18 +128,19 @@ function Base.accumulate( ) where {MR, NR, T, W, R, NA, PA, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) - @inbounds for p in 0:(k_block_length - 1) + @inbounds for p in 1:k_block_length acc = _accumulate_step_planar(kernel, acc, packed_a, packed_b, p) end 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. +# Generator-time accumulator indices of the re/im vectors of block `v` +# (zero-based), 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) + kernel <: PlanarKernel && return (v + MV * (j - 1) + 1, MV * NR + v + MV * (j - 1) + 1) + kernel <: RealComplexKernel && return (v + MV * (2j - 2) + 1, v + MV * (2j - 1) + 1) throw(ArgumentError("_planar_acc_index: no planar accumulator layout for $kernel")) end @@ -154,16 +155,16 @@ end _check_acc(:_store_tile_planar!, R, T, NA, 2NV) blocks = Any[] - for j in 0:(NR - 1), v in 0:(MV - 1) + for j in 1:NR, v in 0:(MV - 1) ire, iim = _planar_acc_index(kernel, MV, NR, v, j) push!( blocks, quote - if $j < n + if $j <= n revec = acc[$ire] imvec = acc[$iim] for lane in 1:$W - i = $(v * W) + lane - 1 - i < m || break + i = $(v * W) + lane + i <= m || break _axpby_tile!( destination, i, $j, alpha, Complex(revec[lane], imvec[lane]), beta @@ -255,7 +256,7 @@ end _check_acc(:_store_tile_planar_vector!, R, T, NA, 2NV) blocks = Any[] - for j in 0:(NR - 1) + for j in 1:NR vblocks = Any[] for v in 0:(MV - 1) ire, iim = _planar_acc_index(kernel, MV, NR, v, j) @@ -270,10 +271,10 @@ end ) elseif $(v * W) < m for lane in 1:$W - i = $(v * W) + lane - 1 - i < m || break + i = $(v * W) + lane + i <= m || break _axpby_at!( - storage, colbase + i + 1, alpha, + storage, colbase + i, alpha, Complex(revec[lane], imvec[lane]), beta ) end @@ -283,8 +284,8 @@ end end push!( blocks, quote - if $j < n - colbase = rowbase0 + axis_offset(cols, $j) # zero-based (0, j) + if $j <= n + colbase = rowbase0 + cols[$j] # the address of (1, j) $(vblocks...) end end @@ -295,7 +296,7 @@ end Base.@_inline_meta storage = destination.storage cols = destination.cols - rowbase0 = destination.base + axis_offset(destination.rows, 0) + rowbase0 = @inbounds destination.base + destination.rows[1] ar = Vec{$W, $R}(real(alpha)) ai = Vec{$W, $R}(imag(alpha)) br = Vec{$W, $R}(real(beta)) diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl index 362c986..d0b023f 100644 --- a/src/microkernels/scalar.jl +++ b/src/microkernels/scalar.jl @@ -10,7 +10,6 @@ end ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = ScalarKernel(KernelDescriptor(Val(MR), Val(NR), T)) -# `acc[i+1, j+1]` for zero-based `(i, j)`. zero_accumulator(kernel::ScalarKernel{MR, NR, T}) where {MR, NR, T} = zeros(T, MR, NR) # Extends `Base.accumulate` so `using QuasiStrided` does not clash with it. @@ -20,12 +19,12 @@ function Base.accumulate( ) where {MR, NR, T, PA, PB} k_block_length == 0 && return acc k_block_length > 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) - @inbounds for p in 0:(k_block_length - 1) - for j in 0:(NR - 1) + @inbounds for p in 1:k_block_length + for j in 1:NR bj = panel_load(packed_b, packed_b_offset(kernel, j, p)) - for i in 0:(MR - 1) + for i in 1:MR ai = panel_load(packed_a, packed_a_offset(kernel, i, p)) - acc[i + 1, j + 1] = muladd(ai, bj, acc[i + 1, j + 1]) + acc[i, j] = muladd(ai, bj, acc[i, j]) end end end @@ -38,8 +37,8 @@ function store_tile!( ) where {T} m, n = _store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination - @inbounds for j in 0:(n - 1), i in 0:(m - 1) - _axpby_tile!(destination, i, j, alpha, acc[i + 1, j + 1], beta) + @inbounds for j in 1:n, i in 1:m + _axpby_tile!(destination, i, j, alpha, acc[i, j], beta) end return destination end diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index b8ca1e4..699a435 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -31,7 +31,8 @@ SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = lanewidth(::SIMDKernel{MR, NR, T, W}) where {MR, NR, T, W} = W -# `(MR÷W)*NR` vectors; vector `v` of column `j` is at `v + (MR÷W)*j + 1`. +# `(MR÷W)*NR` vectors; vector `v` (zero-based) of column `j` is at +# `v + (MR÷W)*(j - 1) + 1`. function zero_accumulator(kernel::SIMDKernel{MR, NR, T, W}) where {MR, NR, T, W} z = zero(Vec{W, T}) return ntuple(_ -> z, Val((MR ÷ W) * NR)) @@ -50,21 +51,21 @@ end _check_acc(:_accumulate_step, T, T, NV, NVECA * NR) avars = [Symbol(:a, v) for v in 0:(NVECA - 1)] - bvars = [Symbol(:b, j) for j in 0:(NR - 1)] + bvars = [Symbol(:b, j) for j in 1:NR] load_a = [ - :($(avars[v + 1]) = panel_vload(Vec{$W, $T}, packed_a, packed_a_offset(kernel, $(v * W), p))) + :($(avars[v + 1]) = panel_vload(Vec{$W, $T}, packed_a, packed_a_offset(kernel, $(v * W + 1), p))) for v in 0:(NVECA - 1) ] load_b = [ - :($(bvars[j + 1]) = _b_step_load(packed_b, kernel, $j, p)) - for j in 0:(NR - 1) + :($(bvars[j]) = _b_step_load(packed_b, kernel, $j, p)) + for j in 1:NR ] acc_exprs = Vector{Any}(undef, NV) - for j in 0:(NR - 1), v in 0:(NVECA - 1) - idx = v + NVECA * j + 1 - acc_exprs[idx] = :(muladd($(avars[v + 1]), $(bvars[j + 1]), acc[$idx])) + for j in 1:NR, v in 0:(NVECA - 1) + idx = v + NVECA * (j - 1) + 1 + acc_exprs[idx] = :(muladd($(avars[v + 1]), $(bvars[j]), acc[$idx])) end return quote @@ -83,7 +84,7 @@ end ) where {MR, NR, T, W, NV, PA, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) - @inbounds for p in 0:(k_block_length - 1) + @inbounds for p in 1:k_block_length acc = _accumulate_step(kernel, acc, packed_a, packed_b, p) end return acc @@ -102,15 +103,15 @@ end ) where {MR, NR, T, W, NV} NVECA = MR ÷ W blocks = Any[] - for j in 0:(NR - 1), v in 0:(NVECA - 1) - idx = v + NVECA * j + 1 + for j in 1:NR, v in 0:(NVECA - 1) + idx = v + NVECA * (j - 1) + 1 push!( blocks, quote - if $j < n + if $j <= n vec = acc[$idx] for lane in 1:$W - i = $(v * W) + lane - 1 - i < m || break + i = $(v * W) + lane + i <= m || break _axpby_tile!(destination, i, $j, alpha, vec[lane], beta) end end @@ -137,10 +138,10 @@ end RC = eltype(S) old = :(convert(Vec{$W, $T}, vload(Vec{$W, $RC}, storage, at))) blocks = Any[] - for j in 0:(NR - 1) + for j in 1:NR vblocks = Any[] for v in 0:(NVECA - 1) - idx = v + NVECA * j + 1 + idx = v + NVECA * (j - 1) + 1 push!( vblocks, quote vec = acc[$idx] @@ -157,9 +158,9 @@ end ) elseif $(v * W) < m for lane in 1:$W - i = $(v * W) + lane - 1 - i < m || break - _axpby_at!(storage, colbase + i + 1, alpha, vec[lane], beta) + i = $(v * W) + lane + i <= m || break + _axpby_at!(storage, colbase + i, alpha, vec[lane], beta) end end end @@ -167,8 +168,8 @@ end end push!( blocks, quote - if $j < n - colbase = rowbase0 + axis_offset(cols, $j) # zero-based (0, j) + if $j <= n + colbase = rowbase0 + cols[$j] # the address of (1, j) $(vblocks...) end end @@ -178,7 +179,7 @@ end Base.@_inline_meta storage = destination.storage cols = destination.cols - rowbase0 = destination.base + axis_offset(destination.rows, 0) + rowbase0 = @inbounds destination.base + destination.rows[1] @inbounds begin $(blocks...) end diff --git a/src/packing/format.jl b/src/packing/format.jl index eff728f..b190114 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -1,10 +1,12 @@ # The contract between packing and the microkernels: the packed-panel formats # and the descriptor fixing a kernel's register tile, element type and panel -# formats. Complex panels are "N planes of real(T)"; every format is addressed by +# formats. Complex panels are "N planes of real(T)"; every format puts lane `i` +# of plane `plane` at K step `p` (one-based `i`, `p`) at zero-based offset # -# p * MR * reals_per_element + plane * MR + i +# (p - 1) * MR * reals_per_element + plane * MR + (i - 1) # -# which at RealFormat (one real per element, plane 0) reduces to `i + MR*p`. +# which at RealFormat (one real per element, plane 0) reduces to +# `(i - 1) + MR * (p - 1)`. abstract type PackFormat end @@ -87,13 +89,14 @@ sliver_width(k, i::Int) = sliver_width(k)[i] packed_a_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 1) * k_block_length packed_b_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 2) * k_block_length -# Zero-based offsets, in reals. For 1e and interleaved, `i`/`j` runs over reals -# (0:2MR-1) and 1e's second region is `plane == 2`. +# Zero-based offsets, in reals, of one-based coordinates. For 1e and +# interleaved, `i`/`j` runs over reals (1:2MR) and 1e's second region is +# `plane == 2`. @inline packed_a_plane_offset(d::Descriptor{MR}, plane::Int, i::Int, p::Int) where {MR} = - p * sliver_width(d, 1) + plane * MR + i + (p - 1) * sliver_width(d, 1) + plane * MR + (i - 1) @inline packed_b_plane_offset(d::Descriptor{MR, NR}, plane::Int, j::Int, p::Int) where {MR, NR} = - p * sliver_width(d, 2) + plane * NR + j + (p - 1) * sliver_width(d, 2) + plane * NR + (j - 1) # Real descriptors only: a complex element has no single offset. -packed_a_offset(kernel::KernelDescriptor{MR}, i::Int, p::Int) where {MR} = i + MR * p -packed_b_offset(kernel::KernelDescriptor{MR, NR}, j::Int, p::Int) where {MR, NR} = j + NR * p +packed_a_offset(kernel::KernelDescriptor{MR}, i::Int, p::Int) where {MR} = (i - 1) + MR * (p - 1) +packed_b_offset(kernel::KernelDescriptor{MR, NR}, j::Int, p::Int) where {MR, NR} = (j - 1) + NR * (p - 1) diff --git a/src/packing/pack.jl b/src/packing/pack.jl index ba76d22..964c97b 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -97,13 +97,13 @@ end ) where {FMT, V, MR, NR, T2, F} if format isa RealFormat if _pack_a_contiguous_eligible(packed, source, transform, m, Val(MR), real(T2)) - rowbase = source.base + axis_offset(source.rows, 0) + rowbase = @inbounds source.base + source.rows[1] return _pack_a_contiguous!(packed, source.storage, rowbase, source.cols, Val(MR), k_block_length) end elseif _pack_complex_contiguous_eligible( packed, source.storage, source.rows, transform, format, m, Val(MR), T2 ) - elembase = source.base + axis_offset(source.rows, 0) + elembase = @inbounds source.base + source.rows[1] return _pack_complex_contiguous!( format, packed, source.storage, elembase, source.cols, Val(MR), k_block_length, transform ) @@ -120,7 +120,7 @@ end if _pack_complex_contiguous_eligible( packed, source.storage, source.cols, transform, format, n, Val(NR), T2 ) - elembase = source.base + axis_offset(source.cols, 0) + elembase = @inbounds source.base + source.cols[1] return _pack_complex_contiguous!( format, packed, source.storage, elembase, source.rows, Val(NR), k_block_length, transform ) @@ -145,19 +145,19 @@ end transform::F, load::L, plane_offset::P ) where {V, T, FMT <: PackFormat, PD, F, L, P} if valid == PD - @inbounds for p in 0:(k_block_length - 1) - for t in 0:(PD - 1) + @inbounds for p in 1:k_block_length + for t in 1:PD z = convert(T, transform(load(t, p)))::T _pack_emit!(packed, format, plane_offset, t, p, z) end end else - @inbounds for p in 0:(k_block_length - 1) - for t in 0:(valid - 1) + @inbounds for p in 1:k_block_length + for t in 1:valid z = convert(T, transform(load(t, p)))::T _pack_emit!(packed, format, plane_offset, t, p, z) end - for t in valid:(PD - 1) + for t in (valid + 1):PD _pack_emit_zero!(packed, format, plane_offset, t, p, real(T)) end end @@ -186,8 +186,8 @@ end @inline function _pack_emit!( packed::V, ::InterleavedFormat, plane_offset::P, t::Int, p::Int, z::T ) where {V, P, T} - panel_store!(packed, plane_offset(0, 2 * t, p), real(z)) - panel_store!(packed, plane_offset(0, 2 * t + 1, p), imag(z)) + panel_store!(packed, plane_offset(0, 2 * t - 1, p), real(z)) + panel_store!(packed, plane_offset(0, 2 * t, p), imag(z)) return nothing end @@ -196,10 +196,10 @@ end ) where {V, P, T} re = real(z) im = imag(z) - panel_store!(packed, plane_offset(0, 2 * t, p), re) - panel_store!(packed, plane_offset(0, 2 * t + 1, p), im) - panel_store!(packed, plane_offset(2, 2 * t, p), -im) - panel_store!(packed, plane_offset(2, 2 * t + 1, p), re) + panel_store!(packed, plane_offset(0, 2 * t - 1, p), re) + panel_store!(packed, plane_offset(0, 2 * t, p), im) + panel_store!(packed, plane_offset(2, 2 * t - 1, p), -im) + panel_store!(packed, plane_offset(2, 2 * t, p), re) return nothing end @@ -223,17 +223,17 @@ end @inline function _pack_emit_zero!( packed::V, ::InterleavedFormat, plane_offset::P, t::Int, p::Int, ::Type{R} ) where {V, P, R} + panel_store!(packed, plane_offset(0, 2 * t - 1, p), zero(R)) panel_store!(packed, plane_offset(0, 2 * t, p), zero(R)) - panel_store!(packed, plane_offset(0, 2 * t + 1, p), zero(R)) return nothing end @inline function _pack_emit_zero!( packed::V, ::OneEFormat, plane_offset::P, t::Int, p::Int, ::Type{R} ) where {V, P, R} + panel_store!(packed, plane_offset(0, 2 * t - 1, p), zero(R)) panel_store!(packed, plane_offset(0, 2 * t, p), zero(R)) - panel_store!(packed, plane_offset(0, 2 * t + 1, p), zero(R)) + panel_store!(packed, plane_offset(2, 2 * t - 1, p), zero(R)) panel_store!(packed, plane_offset(2, 2 * t, p), zero(R)) - panel_store!(packed, plane_offset(2, 2 * t + 1, p), zero(R)) return nothing end diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index a716025..0e98e2d 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -35,9 +35,9 @@ end GC.@preserve storage begin sp = pointer(storage) dp = packed.ptr - for p in 0:(k_block_length - 1) - v = vload(Vec{MR, S}, sp + sizeof(S) * (rowbase + axis_offset(cols, p))) - vstore(convert(Vec{MR, T}, v), dp + sizeof(T) * (MR * p)) + for p in 1:k_block_length + v = vload(Vec{MR, S}, sp + sizeof(S) * (rowbase + (@inbounds cols[p]))) + vstore(convert(Vec{MR, T}, v), dp + sizeof(T) * (MR * (p - 1))) end end return packed @@ -107,7 +107,7 @@ end _single_region_shuffle(::PlanarFormat) = _planar_pack_shuffle _single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a -# Lane `t` of K step `p` is element `elembase + axis_offset(steps, p) + t`; +# Lane `t` of K step `p` is element `elembase + steps[p] + (t - 1)`; # only the lane axis must be unit-stride, `steps` may be scattered. Pinning # `Complex{RS}` in the signature keeps the bitcast to `RS` lanes sound; the # lanes convert to `R` before the shuffle. @@ -120,14 +120,14 @@ _single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a GC.@preserve storage begin sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr - for p in 0:(k_block_length - 1) + for p in 1:k_block_length src = convert( Vec{2 * PD, R}, - vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) + vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + (@inbounds steps[p])))) ) vstore( shuffle(src, _pack_alt(src, transform), Val(PD)), - dp + sizeof(R) * (2 * PD * p) + dp + sizeof(R) * (2 * PD * (p - 1)) ) end end @@ -141,12 +141,12 @@ end GC.@preserve storage begin sp = reinterpret(Ptr{RS}, pointer(storage)) dp = packed.ptr - for p in 0:(k_block_length - 1) + for p in 1:k_block_length src = convert( Vec{2 * PD, R}, - vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + axis_offset(steps, p)))) + vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + (@inbounds steps[p])))) ) - at = 4 * PD * p + at = 4 * PD * (p - 1) vstore( _onee_pack_shuffle_a(src, _pack_alt(src, transform), Val(PD)), dp + sizeof(R) * at diff --git a/src/packing/transposed.jl b/src/packing/transposed.jl index 3bf9854..643f103 100644 --- a/src/packing/transposed.jl +++ b/src/packing/transposed.jl @@ -18,15 +18,15 @@ plane_offset, format, buffer, descriptor, Val(R), Rp * k_block_length, 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:(k_block_length - 1) - kbase = base + axis_offset(kaxis, p) + 1 + for g0 in 0:G:(fcount - 1), y1 in 0:(E - 1), p in 1:k_block_length + kbase = @inbounds base + 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:(k_block_length - 1) - kbase = base + axis_offset(kaxis, p) + 1 + for g0 in 0:G:(fcount - 1), p in 1:k_block_length + kbase = @inbounds base + kaxis[p] + 1 for y1 in 0:(E - 1), y2 in 0:(L - 1) emit(kbase, g0 + y1 + E * y2, p) end @@ -36,19 +36,20 @@ if valid != 0 rb = (fcount ÷ R) * Rp * k_block_length pad = (plane, idx, pp) -> rb + plane_offset(descriptor, plane, idx, pp) - for p in 0:(k_block_length - 1), t in valid:(R - 1) + for p in 1:k_block_length, t in (valid + 1):R _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`. +# Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane +# `i % R + 1`) 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(_element_type(format, scalartype(descriptor)), x)) + _pack_emit!(buffer, format, po, t + 1, p, convert(_element_type(format, scalartype(descriptor)), x)) return nothing end diff --git a/test/layout/test_tiles.jl b/test/layout/test_tiles.jl index bdc2e74..277a0bf 100644 --- a/test/layout/test_tiles.jl +++ b/test/layout/test_tiles.jl @@ -18,13 +18,13 @@ end storage = zeros(200) t = Tile(storage, 50, rows, cols) @test size(t) == (4, 3) - for i in 0:3, j in 0:2 - addr = 50 + rows[i + 1] + cols[j + 1] + for i in 1:4, j in 1:3 + addr = 50 + rows[i] + cols[j] t[i, j] = 1000.0i + j @test storage[addr + 1] == t[i, j] == 1000.0i + j end end - @test_throws BoundsError Tile(zeros(3), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2))[1, 1] + @test_throws BoundsError Tile(zeros(3), 0, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2))[2, 2] end @testset "storage bounds: offset ranges and span checks" begin diff --git a/test/microkernels/test_planar_kernel.jl b/test/microkernels/test_planar_kernel.jl index 9163919..bc08f18 100644 --- a/test/microkernels/test_planar_kernel.jl +++ b/test/microkernels/test_planar_kernel.jl @@ -18,7 +18,7 @@ using QuasiStrided: PlanarKernel, PlanarMethod, complex_method, lanewidth, @test_throws ArgumentError PlanarKernel(Val(8), Val(4), ComplexF64, Val(0)) @test_throws ArgumentError PlanarKernel(Val(8), Val(4), Float64) # A complex kernel is never asked for a single-plane offset. - @test_throws MethodError QuasiStrided.packed_a_offset(k, 0, 0) + @test_throws MethodError QuasiStrided.packed_a_offset(k, 1, 1) mk_e2e(ComplexF64, nothing) end diff --git a/test/packing/test_kernel_descriptor.jl b/test/packing/test_kernel_descriptor.jl index 0e8f06f..34bbda5 100644 --- a/test/packing/test_kernel_descriptor.jl +++ b/test/packing/test_kernel_descriptor.jl @@ -4,8 +4,8 @@ using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEForma @testset "KernelDescriptor" begin k = KernelDescriptor(Val(8), Val(6), Float64) @test (tile_size(k)..., scalartype(k), realtype(k)) == (8, 6, Float64, Float64) - @test packed_a_offset(k, 3, 2) == 3 + 8 * 2 - @test packed_b_offset(k, 5, 2) == 5 + 6 * 2 + @test packed_a_offset(k, 4, 3) == 3 + 8 * 2 + @test packed_b_offset(k, 6, 3) == 5 + 6 * 2 @test packed_a_length(k, 0) == 0 @test packed_a_length(k, 4) == 32 @test packed_b_length(k, 4) == 24 From 81cfaa31954c6dcac0c79d72e0af9464942d091b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 12:19:00 -0400 Subject: [PATCH 08/60] Pointer-backed ScatterAxis instead of nested axis closures A `view` scatter axis is not `isbits`, so the nest's `Union{AffineAxis, view}` was heap-boxed, which `_with_axis` worked around with about ten nested closures. `ScatterAxis` is again a borrowed pointer, now an `AbstractVector{Int}` indexed one-based, so `_axis_of` returns an unboxed union and the nest reads flat. The oracle and `_scale_all_of_C!` now `GC.@preserve` the workspace they borrow from. Full suite passes (30930 tests). Reused-plan execute! medians, closures -> pointer, us: matmul 256^3 F64 410 -> 410, C64 1578 -> 1559; scattered M F64 432 -> 441, C64 1545 -> 1531; scattered M/N/K F64 18.4 -> 18.1, C64 41.9 -> 42.0. Zero allocations per call in all six, under --check-bounds=yes. First call F64/C64: 3.5-3.8 s / 2.4-2.8 s -> 3.2 s / 2.1-2.2 s (back to the pre-chunk-3 3.1-3.3 s / 2.1-2.3 s). Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 29 ++++++- src/execution/execute.jl | 103 ++++++++++-------------- src/execution/macrokernel.jl | 29 +++---- src/execution/oracle.jl | 79 +++++++++--------- src/execution/unpackedb.jl | 22 +++-- src/layout/tiles.jl | 18 ++++- test/execution/test_manual_pipeline.jl | 37 +++++---- test/planning/test_per_call_overhead.jl | 3 +- 8 files changed, 167 insertions(+), 153 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 87dcd80..667217c 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -15,15 +15,15 @@ TTFX when codegen is touched). | 0 | Tour | `QuasiStrided.jl`, `plan.jl`, `execute.jl` call path | done | | 1 | Hardware | `hardware/target.jl` | done | | 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | done | -| 3 | Tiles | `layout/tiles.jl` | | -| 4 | Packing formats | `packing/format.jl`, `panel.jl`, `transposed.jl` | | +| 3 | Tiles | `layout/tiles.jl` | done | +| 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | in progress | | 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | | | 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | | | 7 | SIMD kernels | `simd.jl`, `planar.jl` | | | 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | | | 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | | 10 | Kernel selection | `planning/kernel_selection.jl` | | -| 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl` | | +| 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | | 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | @@ -101,11 +101,34 @@ drops its alias check. Kept: generic `P` (a batch label in A, B and C would be a `P = 3` group) and unconditional range checks in `fill_offsets!` (`@boundscheck`/`@propagate_inbounds` only help when inlined). +### D7. Tiles (chunk 3, applied) + +Addresses stay zero-based (storage index = address + 1), as do `AxisGroup` +coordinates and block starts; coordinates within a tile, sliver or packed +panel are one-based. Tile axes are `AbstractVector{Int}`s of offsets: +`AffineAxis` (zero stride allowed, so not a `StepRange`) and `ScatterAxis`, +a borrowed pointer. The pointer is kept because the nest holds a +`Union{AffineAxis, ScatterAxis}`, which is only unboxed when both members are +`isbits`; a `view` there allocated per call, and splitting by hand needed +nested closures. One `Tile` type with `getindex`/`setindex!`/`size` replaces +`QSTile`/`SourceTile`/`DestinationTile`/`tile_load`/`tile_store!`. + +### D8. Packing formats (chunk 4) + +Panels are only `PackedPanel` (the oracle and tests too; the `AbstractVector` +panel methods go); one `Descriptor` constructor plus a `RealDescriptor` +alias; one `packed_a_offset(d, i, p, plane = 0)` (and B); +`transposed.jl` merges into `panel.jl` and is reviewed with `pack_split.jl`. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: https://github.com/lkdvos/QuasiStrided.jl/issues/13 (decide after chunks 3, 5, 12, 14). +- One tile-axis type with a `regular` flag instead of a union of axis types: + https://github.com/lkdvos/QuasiStrided.jl/issues/14. +- `PackedPanel` without a raw pointer: + https://github.com/lkdvos/QuasiStrided.jl/issues/15. ## Open items (to revisit in their chunk) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index f402723..b65a958 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -14,11 +14,9 @@ function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::In while n_tile_start < n_length n_tile_length = min(n_tile, n_length - n_tile_start) (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - _with_axis(dM_C, ws.tile_m_buf_C, 0) do rowsC - _with_axis(dN_C, ws.tile_n_buf_C, 0) do colsC - _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) - end - end + rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) + colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) + GC.@preserve ws _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) n_tile_start += n_tile_length end m_tile_start += m_tile_length @@ -200,7 +198,7 @@ function _execute_path!( ) where {T} kernel = plan.kernel ws = plan.workspace - # Packed panels borrow pointers into `ws`. + # Packed panels and scatter axes borrow pointers into `ws`. GC.@preserve ws begin _execute_nest!( plan, ws, kernel, tile_size(kernel)..., @@ -358,6 +356,9 @@ function _execute_nest!( ) end + colsA_k = _axis_of(dK_A, ws.k_buf_A, 0, aff_kA) + rowsB_k = _axis_of(dK_B, ws.k_buf_B, 0, aff_kB) + # Hoisted bounds checks (B here, A and C per M block): each # rectangle is exactly the union of the per-sliver/per-tile # regions, so the `unsafe_*` calls below read nothing unchecked. @@ -366,28 +367,21 @@ function _execute_nest!( beta_eff = firstpanel ? betaT : one(T) if split_b - _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k - _pack_block_transposed!( - packed_b_plane_offset, b_format(kernel), - packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, - plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, - plan.npack - ) - end + _pack_block_transposed!( + packed_b_plane_offset, b_format(kernel), + packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, + plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, + plan.npack + ) elseif !UNPACKED_B for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k - @inline - _with_axis(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) do colsB - @inline - _pack_sliver!( - unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, - kernel, btransform - ) - end - end + colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) + _pack_sliver!( + unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, + kernel, btransform + ) end end @@ -413,39 +407,30 @@ function _execute_nest!( checked_span_bounds(cplan.Cbase, rng_mC, rng_nC, lenC) if split_a - _with_axis(dK_A, ws.k_buf_A, 0, aff_kA) do colsA_k - _pack_block_transposed!( - packed_a_plane_offset, a_format(kernel), - packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, - plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, - plan.mpack - ) - end + _pack_block_transposed!( + packed_a_plane_offset, a_format(kernel), + packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, + plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, + plan.mpack + ) else for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - _with_axis(dK_A, ws.k_buf_A, 0, aff_kA) do colsA_k - @inline - _with_axis(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) do rowsA - @inline - _pack_sliver!( - unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, - kernel, atransform - ) - end - end + rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) + _pack_sliver!( + unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, + kernel, atransform + ) end end # --- loops over N tiles and M tiles --- if UNPACKED_B - _with_axis(dK_B, ws.k_buf_B, 0, aff_kB) do rowsB_k - _micro_tiles_unpacked_b!( - kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC - ) - end + _micro_tiles_unpacked_b!( + kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC + ) else _micro_tiles_packed_b!( kernel, cplan, ws, m_tiles, n_tiles, @@ -478,19 +463,15 @@ end for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - _with_axis(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) do colsC - @inline - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - _with_axis(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) do rowsC - @inline - unsafe_execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bpanel, k_block_length, alphaT, beta_eff - ) - end - end + colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + unsafe_execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + apanel, bpanel, k_block_length, alphaT, beta_eff + ) end end return nothing diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index e069b5b..79d6e0f 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -1,26 +1,19 @@ # Macro-kernel helpers for the five-loop nest: function barriers over the # tile axis types, packed-sliver addressing, and sliver classification. -# The axis `d` describes: an `AffineAxis` when regular, else a `view` of its -# offsets in `buffer`. +# The axis `d` describes. GUARDRAIL: a `Union{AffineAxis, ScatterAxis}` of +# `isbits` types, never boxed; each consumer below binds each axis type as its +# own parameter, so it builds a concretely typed `Tile`. @inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int) = - d.regular ? AffineAxis(d.base, d.stride, d.count) : view(buffer, (first + 1):(first + d.count)) + d.regular ? AffineAxis(d.base, d.stride, d.count) : ScatterAxis(pointer(buffer, first + 1), d.count) -# GUARDRAIL: `f(_axis_of(d, buffer, first))` with each axis type in its own -# branch, which the nest uses instead of `_axis_of`: a `Union` holding a `view` -# (not `isbits`) is heap-boxed. Each consumer below binds each axis type as its -# own parameter, so it builds a concretely typed `Tile`. `Val(true)`: a ramp -# map, whose descriptors are always regular (checked), so only the affine -# branch is compiled. -@inline function _with_axis( - f::F, d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{AFF} = Val(false) - ) where {F, AFF} - if AFF - d.regular || _throw_irregular_ramp_descriptor() - elseif !d.regular - return f(view(buffer, (first + 1):(first + d.count))) - end - return f(AffineAxis(d.base, d.stride, d.count)) +# The same with the axis type fixed by the path: `Val(true)` for a ramp map, +# whose descriptors are always regular (checked). +@inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{false}) = + _axis_of(d, buffer, first) +@inline function _axis_of(d::BlockDescriptor, ::Vector{Int}, ::Int, ::Val{true}) + d.regular || _throw_irregular_ramp_descriptor() + return AffineAxis(d.base, d.stride, d.count) end @noinline _throw_irregular_ramp_descriptor() = throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl index 533588f..fe0c9d7 100644 --- a/src/execution/oracle.jl +++ b/src/execution/oracle.jl @@ -35,51 +35,54 @@ function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) w k_bufs = (ws.tw_k_buf_A, ws.tw_k_buf_B) k_block = plan.blocking.k_block - m_tile_start = 0 - while m_tile_start < m_length - m_tile_length = min(m_tile, m_length - m_tile_start) - (dM_A, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) - rowsA = _axis_of(dM_A, ws.tile_m_buf_A, 0) - rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) + # Scatter axes borrow pointers into `ws`. + GC.@preserve ws begin + m_tile_start = 0 + while m_tile_start < m_length + m_tile_length = min(m_tile, m_length - m_tile_start) + (dM_A, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) + rowsA = _axis_of(dM_A, ws.tile_m_buf_A, 0) + rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) - n_tile_start = 0 - while n_tile_start < n_length - n_tile_length = min(n_tile, n_length - n_tile_start) - (dN_B, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - colsB = _axis_of(dN_B, ws.tile_n_buf_B, 0) - colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) + n_tile_start = 0 + while n_tile_start < n_length + n_tile_length = min(n_tile, n_length - n_tile_start) + (dN_B, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) + colsB = _axis_of(dN_B, ws.tile_n_buf_B, 0) + colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) - k_block_start = 0 - firstpanel = true - while k_block_start < k_length - k_block_length = min(k_block, k_length - k_block_start) - (dK_A, dK_B) = block_descriptors!(k_bufs, plan.kgroup, k_block_start, k_block_length) - colsK_A = _axis_of(dK_A, ws.tw_k_buf_A, 0) - rowsK_B = _axis_of(dK_B, ws.tw_k_buf_B, 0) + k_block_start = 0 + firstpanel = true + while k_block_start < k_length + k_block_length = min(k_block, k_length - k_block_start) + (dK_A, dK_B) = block_descriptors!(k_bufs, plan.kgroup, k_block_start, k_block_length) + colsK_A = _axis_of(dK_A, ws.tw_k_buf_A, 0) + rowsK_B = _axis_of(dK_B, ws.tw_k_buf_B, 0) - # Whole single-sliver buffers, so no `_sliver_panel`. - _pack_sliver!( - pack_a!, ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, - kernel, plan.atransform - ) - _pack_sliver!( - pack_b!, ws.tw_packed_b, plan.Bstorage, plan.Bbase, rowsK_B, colsB, - kernel, plan.btransform - ) + # Whole single-sliver buffers, so no `_sliver_panel`. + _pack_sliver!( + pack_a!, ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, + kernel, plan.atransform + ) + _pack_sliver!( + pack_b!, ws.tw_packed_b, plan.Bstorage, plan.Bbase, rowsK_B, colsB, + kernel, plan.btransform + ) - beta_eff = firstpanel ? betaT : one(T) - _execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - ws.tw_packed_a, ws.tw_packed_b, k_block_length, alphaT, beta_eff - ) + beta_eff = firstpanel ? betaT : one(T) + _execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + ws.tw_packed_a, ws.tw_packed_b, k_block_length, alphaT, beta_eff + ) - firstpanel = false - k_block_start += k_block_length - end + firstpanel = false + k_block_start += k_block_length + end - n_tile_start += n_tile_length + n_tile_start += n_tile_length + end + m_tile_start += m_tile_length end - m_tile_start += m_tile_length end return plan.Cstorage diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 92cf58b..7f9e4c5 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -99,19 +99,15 @@ end bview = _unpacked_b_view( kernel, Bstorage, Bbase, ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, rowsB_k, btransform ) - _with_axis(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) do colsC - @inline - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - _with_axis(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) do rowsC - @inline - unsafe_execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bview, k_block_length, alphaT, beta_eff - ) - end - end + colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + unsafe_execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + apanel, bview, k_block_length, alphaT, beta_eff + ) end end return nothing diff --git a/src/layout/tiles.jl b/src/layout/tiles.jl index 5d98746..915cc7b 100644 --- a/src/layout/tiles.jl +++ b/src/layout/tiles.jl @@ -1,7 +1,7 @@ # Tile axes and tiles. Addresses are zero-based (storage index = address + 1), # coordinates one-based; a tile's `base` is the address its offsets are # relative to. A tile axis is the vector of offsets of its rows or columns: an -# `AffineAxis`, or a `view` of an offset buffer. +# `AffineAxis`, a `ScatterAxis`, or any `AbstractVector{Int}`. # Offsets `base + (t - 1) * stride`. Not a `StepRange`: the stride may be zero. struct AffineAxis <: AbstractVector{Int} @@ -29,6 +29,22 @@ function Base.extrema(ax::AffineAxis) return minmax(ax.base, last) end +# Offsets read from a borrowed buffer, valid while its owner is +# `GC.@preserve`d. A pointer rather than a `view`, so that it is `isbits` and a +# `Union{AffineAxis, ScatterAxis}` is never heap-boxed. +struct ScatterAxis <: AbstractVector{Int} + offsets::Ptr{Int} + count::Int +end + +Base.size(ax::ScatterAxis) = (ax.count,) +Base.IndexStyle(::Type{ScatterAxis}) = IndexLinear() + +@inline function Base.getindex(ax::ScatterAxis, t::Int) + @boundscheck checkbounds(ax, t) + return unsafe_load(ax.offsets, t) +end + # An offset interval of one map: `regular` iff `buffer[t+1] == base + t*stride` # for all `t < count`; otherwise read the buffer (valid until it is refilled). struct BlockDescriptor diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index 99138e6..c50c959 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -14,26 +14,29 @@ mb, nb, kb = bufs(M), bufs(N), bufs(K) (dM_A, dM_C) = block_descriptors!(mb, M, 0, 6) (dN_B, dN_C) = block_descriptors!(nb, N, 0, 4) - axis(d, buf) = QuasiStrided._axis_of(d, buf, 0) - row_A, row_C = axis(dM_A, mb[1]), axis(dM_C, mb[2]) - col_B, col_C = axis(dN_B, nb[1]), axis(dN_C, nb[2]) + # Scatter axes borrow pointers into the offset buffers. + GC.@preserve mb nb kb begin + axis(d, buf) = QuasiStrided._axis_of(d, buf, 0) + row_A, row_C = axis(dM_A, mb[1]), axis(dM_C, mb[2]) + col_B, col_C = axis(dN_B, nb[1]), axis(dN_C, nb[2]) - Cstorage = zeros(length(Cref)) - destination = Tile(Cstorage, 0, row_C, col_C) - @test size(destination) == (6, 4) - packed_a, packed_b = zeros(tile_size(kernel, 1) * 5), zeros(tile_size(kernel, 2) * 5) + Cstorage = zeros(length(Cref)) + destination = Tile(Cstorage, 0, row_C, col_C) + @test size(destination) == (6, 4) + packed_a, packed_b = zeros(tile_size(kernel, 1) * 5), zeros(tile_size(kernel, 2) * 5) - Cstart = rand(MersenneTwister(1234), size(Cref)...) - # One panel covering K, then panels of 2, 2, 1 with nontrivial alpha/beta. - for (panels, alpha, beta) in ((((0, 5),), 1.0, 0.0), (((0, 2), (2, 2), (4, 1)), 2.5, 0.75)) - copyto!(Cstorage, vec(Cstart)) - for (idx, (first, len)) in enumerate(panels) - (dK_A, dK_B) = block_descriptors!(kb, K, first, len) - pack_a!(packed_a, Tile(vec(A), 0, row_A, axis(dK_A, kb[1])), kernel, identity) - pack_b!(packed_b, Tile(vec(B), 0, axis(dK_B, kb[2]), col_B), kernel, identity) - execute_tile!(kernel, destination, packed_a, packed_b, len, alpha, idx == 1 ? beta : 1.0) + Cstart = rand(MersenneTwister(1234), size(Cref)...) + # One panel covering K, then panels of 2, 2, 1 with nontrivial alpha/beta. + for (panels, alpha, beta) in ((((0, 5),), 1.0, 0.0), (((0, 2), (2, 2), (4, 1)), 2.5, 0.75)) + copyto!(Cstorage, vec(Cstart)) + for (idx, (first, len)) in enumerate(panels) + (dK_A, dK_B) = block_descriptors!(kb, K, first, len) + pack_a!(packed_a, Tile(vec(A), 0, row_A, axis(dK_A, kb[1])), kernel, identity) + pack_b!(packed_b, Tile(vec(B), 0, axis(dK_B, kb[2]), col_B), kernel, identity) + execute_tile!(kernel, destination, packed_a, packed_b, len, alpha, idx == 1 ? beta : 1.0) + end + @test reshape(Cstorage, size(Cref)) ≈ alpha .* Cref .+ beta .* Cstart end - @test reshape(Cstorage, size(Cref)) ≈ alpha .* Cref .+ beta .* Cstart end end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 3c69b95..a6d1c2f 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -171,8 +171,7 @@ end buf = [rand(-30:30) for _ in 1:n] rand() < 0.4 && (buf = [3 + 5 * (t - 1) for t in 1:n]) # force a regular run d = describe_block(buf, 0, n) - ax = QS._axis_of(d, buf, 0) - @test QS.descriptor_offset_range(d, buf, 0) == extrema(ax) + GC.@preserve buf @test QS.descriptor_offset_range(d, buf, 0) == extrema(QS._axis_of(d, buf, 0)) end @test QS.descriptor_offset_range(BlockDescriptor(0, 0, 0, true), Int[], 0) == (0, -1) end From b0f1b687b399044518560ab16b2476feb66dcc29 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 13:05:59 -0400 Subject: [PATCH 09/60] Chunk 4: packing formats - Panels are only `PackedPanel`: the `AbstractVector` methods of `panel_vload`/`panel_load`/`panel_store!` are gone. The checked entry points `pack_a!`, `pack_b!` and `accumulate` take a `DenseVector` and borrow it as a `PackedPanel` under `GC.@preserve`; `execute_tile!` and the oracle reach vectors through `accumulate`. The kernels' `accumulate` methods constrain the A panel to `PackedPanel` (B may be an `UnpackedBView`), which keeps the vector method unambiguous. - One `Descriptor(Val(MR), Val(NR), T, a_format = RealFormat(), b_format = RealFormat())` constructor and a `RealDescriptor` alias replace `KernelDescriptor`/`ComplexKernelDescriptor`; `check_descriptor` names `Descriptor` in its messages. - `packed_a_offset(d, i, p, plane = 0)` / `packed_b_offset(d, j, p, plane = 0)` replace the plane and real-only offset functions; the layout is unchanged. The packers' `(plane, i, p)` closures adapt the order. - transposed.jl is merged into the end of panel.jl. - Tests: view destinations dropped (no longer supported), the unsafe packer test borrows its buffers, the "complex kernel has no single-plane offset" MethodError test is replaced by a plane-offset check. Measured against 81cfaa3 on a workstation: - Tests: 30931 pass (baseline 30930). - TTFX (load; first Float64, ComplexF64, K=1 outer call), 2 runs each, interleaved: before 0.5/2.9/2.0/0.6 s and 0.5/3.0/2.0/0.6 s, after 0.5/3.0/2.1/0.6 s and 0.5/2.9/2.0/0.6 s. - Reused-plan execute! medians over 5 interleaved processes (before -> after, us): matmul256 Float64 393.5 -> 391.8, ComplexF64 1472.7 -> 1487.2; scatterM Float64 414.9 -> 414.5, ComplexF64 1445.8 -> 1460.7; scatterMN Float64 17.5 -> 16.7, ComplexF64 39.7 -> 39.2 (minima equal). - 0 bytes allocated per call in every case under --check-bounds=yes. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 1 - src/execution/execute.jl | 4 +- src/microkernels/fmaddsub.jl | 10 +-- src/microkernels/interface.jl | 25 ++++--- src/microkernels/mixed.jl | 14 ++-- src/microkernels/onem.jl | 8 +-- src/microkernels/planar.jl | 16 ++--- src/microkernels/scalar.jl | 6 +- src/microkernels/simd.jl | 8 +-- src/packing/format.jl | 48 +++++-------- src/packing/pack.jl | 21 +++++- src/packing/panel.jl | 74 +++++++++++++++++--- src/packing/transposed.jl | 55 --------------- test/microkernels/test_onem_kernel.jl | 4 +- test/microkernels/test_planar_kernel.jl | 2 - test/microkernels/test_scalar_kernel.jl | 2 +- test/packing/test_kernel_descriptor.jl | 29 ++++---- test/packing/test_pack_complex.jl | 8 +-- test/packing/test_pack_complex_contiguous.jl | 6 +- test/packing/test_pack_real.jl | 26 +++---- test/planning/test_per_call_overhead.jl | 4 +- test/runtests.jl | 2 +- 22 files changed, 190 insertions(+), 183 deletions(-) delete mode 100644 src/packing/transposed.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index ce92c5a..159debe 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -23,7 +23,6 @@ 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") diff --git a/src/execution/execute.jl b/src/execution/execute.jl index b65a958..1b0433f 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -368,7 +368,7 @@ function _execute_nest!( if split_b _pack_block_transposed!( - packed_b_plane_offset, b_format(kernel), + packed_b_offset, b_format(kernel), packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, plan.npack @@ -408,7 +408,7 @@ function _execute_nest!( if split_a _pack_block_transposed!( - packed_a_plane_offset, a_format(kernel), + packed_a_offset, a_format(kernel), packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, plan.mpack diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index ccc33e7..5c868be 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -26,10 +26,10 @@ Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over complex contractions on AVX-512. """ struct FMAddSubKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} - descriptor::ComplexKernelDescriptor{MR, NR, T, InterleavedFormat, PlanarFormat} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} function FMAddSubKernel{MR, NR, T, W}( - descriptor::ComplexKernelDescriptor{MR, NR, T, InterleavedFormat, PlanarFormat} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} ) where {MR, NR, T, W} _check_vector_shape("FMAddSubKernel", 2 * MR, W, true) return new{MR, NR, T, W}(descriptor) @@ -40,7 +40,7 @@ function FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR T <: Complex || throw(ArgumentError("FMAddSubKernel requires a complex element type, got $T")) return FMAddSubKernel{MR, NR, T, W}( - ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) + Descriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) ) end function FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} @@ -118,7 +118,7 @@ end load_a, :( $(av[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W + 1), p) + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) ) ) ) @@ -156,7 +156,7 @@ end function Base.accumulate( kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NA, PA, PB} + ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) @inbounds for p in 1:k_block_length diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index eeaee0b..45c9374 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -55,20 +55,16 @@ a_format(k::DescriptorKernel) = a_format(k.descriptor) b_format(k::DescriptorKernel) = b_format(k.descriptor) tile_size(k::DescriptorKernel) = tile_size(k.descriptor) scalartype(k::DescriptorKernel) = scalartype(k.descriptor) -packed_a_offset(k::DescriptorKernel, i::Int, p::Int) = packed_a_offset(k.descriptor, i, p) -packed_b_offset(k::DescriptorKernel, j::Int, p::Int) = packed_b_offset(k.descriptor, j, p) +@inline packed_a_offset(k::DescriptorKernel, i::Int, p::Int, plane::Int = 0) = + packed_a_offset(k.descriptor, i, p, plane) +@inline packed_b_offset(k::DescriptorKernel, j::Int, p::Int, plane::Int = 0) = + packed_b_offset(k.descriptor, j, p, plane) packed_a_length(k::DescriptorKernel, k_block_length::Int) = packed_a_length(k.descriptor, k_block_length) packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.descriptor, k_block_length) -# Only these resolve for a complex descriptor; the single-plane offsets above -# are a MethodError there, by design. -@inline packed_a_plane_offset(k::DescriptorKernel, plane::Int, i::Int, p::Int) = - packed_a_plane_offset(k.descriptor, plane, i, p) -@inline packed_b_plane_offset(k::DescriptorKernel, plane::Int, j::Int, p::Int) = - packed_b_plane_offset(k.descriptor, plane, j, p) # GUARDRAIL: every forwarded argument needs its OWN bound type parameter; an # unbound one makes the call dynamically dispatched and allocating on every -# pack. `V` is unconstrained so a `PackedPanel` forwards too. +# pack. `V` is unconstrained so a `DenseVector` forwards too. pack_a!( packed::V, source::Tile, kernel::K, transform::F ) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = @@ -204,6 +200,17 @@ end return (true, alphaT, betaT) end +# The kernels read only `PackedPanel`s; the checked entry points borrow a +# `DenseVector` as one. +function Base.accumulate( + kernel::K, acc::A, packed_a::PA, packed_b::PB, k_block_length::Int + ) where {K <: DescriptorKernel, A, PA <: DenseVector, PB <: DenseVector} + return GC.@preserve packed_a packed_b accumulate( + kernel, acc, packed_panel(packed_a, 1, length(packed_a)), + packed_panel(packed_b, 1, length(packed_b)), k_block_length + ) +end + # One checked K panel: `zero_accumulator`, `accumulate`, `store_tile!`. function execute_tile!( kernel::K, destination::Tile, packed_a::PA, packed_b::PB, diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index dd916dd..dced3d8 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -17,11 +17,11 @@ Mixed-domain microkernel for complex `A` times real `B`, computing in complex 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} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} inner::KI function ComplexRealKernel{MR, NR, T, W, KI}( - descriptor::ComplexKernelDescriptor{MR, NR, T, InterleavedFormat, RealFormat}, + descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} _check_vector_shape("ComplexRealKernel", 2 * MR, W, true) @@ -40,11 +40,11 @@ Mixed-domain microkernel for real `A` times complex `B`, computing in complex 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} + descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} inner::KI function RealComplexKernel{MR, NR, T, W, KI}( - descriptor::ComplexKernelDescriptor{MR, NR, T, RealFormat, InterleavedFormat}, + descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} KI === SIMDKernel{MR, 2 * NR, real(T), W} || throw( @@ -57,13 +57,13 @@ 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()) + descriptor = Descriptor(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()) + descriptor = Descriptor(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 @@ -79,7 +79,7 @@ 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, k_block_length::Int) = +@inline Base.accumulate(kernel::_MixedKernel, acc::NTuple, packed_a::PackedPanel, packed_b, k_block_length::Int) = accumulate(kernel.inner, acc, packed_a, packed_b, k_block_length) # Inlined or not as the fmaddsub and planar stores each reuses. diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index 82ab64a..b0344ba 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -16,13 +16,13 @@ over `SIMD.Vec{W,real(T)}` lanes. `2MR` must be a multiple of `W`, and `W` must be even. Used only when named in `plan_contract(...; kernel = ...)`. """ struct OneMKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, NR, T} - descriptor::ComplexKernelDescriptor{MR, NR, T, OneEFormat, PlanarFormat} + descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} # `KI` stands in for `SIMDKernel{2MR,NR,real(T),W}`, which is not a legal # field type; the constructor pins it to exactly that. inner::KI function OneMKernel{MR, NR, T, W, KI}( - descriptor::ComplexKernelDescriptor{MR, NR, T, OneEFormat, PlanarFormat}, + descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} _check_vector_shape("OneMKernel", 2 * MR, W, true) @@ -36,7 +36,7 @@ end function OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} T <: Complex || throw(ArgumentError("OneMKernel requires a complex element type, got $T")) - descriptor = ComplexKernelDescriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) + descriptor = Descriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) inner = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) return OneMKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) end @@ -54,7 +54,7 @@ zero_accumulator(kernel::OneMKernel) = zero_accumulator(kernel.inner) function Base.accumulate( kernel::OneMKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NV, PA, PB} + ) where {MR, NR, T, W, R, NV, PA <: PackedPanel, PB} # Checked here so the message reports the logical k_block_length, not the # doubled one. k_block_length >= 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 747d74a..f0c1ba8 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -13,10 +13,10 @@ Split-complex microkernel for `T = ComplexF32/ComplexF64` over multiple of `W`, which counts reals. """ struct PlanarKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} - descriptor::ComplexKernelDescriptor{MR, NR, T, PlanarFormat, PlanarFormat} + descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} function PlanarKernel{MR, NR, T, W}( - descriptor::ComplexKernelDescriptor{MR, NR, T, PlanarFormat, PlanarFormat} + descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} ) where {MR, NR, T, W} _check_vector_shape("PlanarKernel", MR, W) return new{MR, NR, T, W}(descriptor) @@ -25,7 +25,7 @@ end function PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} return PlanarKernel{MR, NR, T, W}( - ComplexKernelDescriptor(Val(MR), Val(NR), T, PlanarFormat(), PlanarFormat()) + Descriptor(Val(MR), Val(NR), T, PlanarFormat(), PlanarFormat()) ) end function PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} @@ -47,8 +47,8 @@ end # B column `j` at K step `p` as `(re, im)`; `UnpackedBView` overrides it. # Shared with the fmaddsub kernel. @inline _b_step_load2(packed_b::PB, kernel, j::Int, p::Int) where {PB} = ( - panel_load(packed_b, packed_b_plane_offset(kernel, 0, j, p)), - panel_load(packed_b, packed_b_plane_offset(kernel, 1, j, p)), + panel_load(packed_b, packed_b_offset(kernel, j, p)), + panel_load(packed_b, packed_b_offset(kernel, j, p, 1)), ) # GUARDRAIL: the real part is `muladd(-ai, bi, muladd(ar, br, c))`. `c - ai*bi` @@ -75,7 +75,7 @@ end load_a, :( $(arv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 0, $(v * W + 1), p) + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) ) ) ) @@ -83,7 +83,7 @@ end load_a, :( $(aiv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_plane_offset(kernel, 1, $(v * W + 1), p) + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p, 1) ) ) ) @@ -125,7 +125,7 @@ end function Base.accumulate( kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NA, PA, PB} + ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) @inbounds for p in 1:k_block_length diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl index d0b023f..342d690 100644 --- a/src/microkernels/scalar.jl +++ b/src/microkernels/scalar.jl @@ -4,11 +4,11 @@ Scalar reference microkernel for register tile `(MR, NR)` and element type `T`. """ struct ScalarKernel{MR, NR, T} <: DescriptorKernel{MR, NR, T} - descriptor::KernelDescriptor{MR, NR, T} + descriptor::RealDescriptor{MR, NR, T} end ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = - ScalarKernel(KernelDescriptor(Val(MR), Val(NR), T)) + ScalarKernel(Descriptor(Val(MR), Val(NR), T)) zero_accumulator(kernel::ScalarKernel{MR, NR, T}) where {MR, NR, T} = zeros(T, MR, NR) @@ -16,7 +16,7 @@ zero_accumulator(kernel::ScalarKernel{MR, NR, T}) where {MR, NR, T} = zeros(T, M function Base.accumulate( kernel::ScalarKernel{MR, NR, T}, acc::AbstractMatrix{T}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, PA, PB} + ) where {MR, NR, T, PA <: PackedPanel, PB} k_block_length == 0 && return acc k_block_length > 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) @inbounds for p in 1:k_block_length diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index 699a435..50aede2 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -16,16 +16,16 @@ Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). """ struct SIMDKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} - descriptor::KernelDescriptor{MR, NR, T} + descriptor::RealDescriptor{MR, NR, T} - function SIMDKernel{MR, NR, T, W}(descriptor::KernelDescriptor{MR, NR, T}) where {MR, NR, T, W} + function SIMDKernel{MR, NR, T, W}(descriptor::RealDescriptor{MR, NR, T}) where {MR, NR, T, W} _check_vector_shape("SIMDKernel", MR, W) return new{MR, NR, T, W}(descriptor) end end SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} = - SIMDKernel{MR, NR, T, W}(KernelDescriptor(Val(MR), Val(NR), T)) + SIMDKernel{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T)) SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = SIMDKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(T))) @@ -81,7 +81,7 @@ end @inline function Base.accumulate( kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, NV, PA, PB} + ) where {MR, NR, T, W, NV, PA <: PackedPanel, PB} k_block_length == 0 && return acc k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) @inbounds for p in 1:k_block_length diff --git a/src/packing/format.jl b/src/packing/format.jl index b190114..f3f428e 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -1,12 +1,10 @@ # The contract between packing and the microkernels: the packed-panel formats # and the descriptor fixing a kernel's register tile, element type and panel -# formats. Complex panels are "N planes of real(T)"; every format puts lane `i` -# of plane `plane` at K step `p` (one-based `i`, `p`) at zero-based offset +# formats. Complex panels are "N planes of real(T)". A sliver of register +# extent `R` (`MR` for A, `NR` for B) puts lane `i` of plane `plane` at K step +# `p` (one-based `i`, `p`) at zero-based offset, in reals, # -# (p - 1) * MR * reals_per_element + plane * MR + (i - 1) -# -# which at RealFormat (one real per element, plane 0) reduces to -# `(i - 1) + MR * (p - 1)`. +# (p - 1) * R * reals_per_element + plane * R + (i - 1) abstract type PackFormat end @@ -30,44 +28,38 @@ reals_per_element(::OneEFormat) = 4 # length and offset below takes the logical `k_block_length` and counts reals. struct Descriptor{MR, NR, T, FA <: PackFormat, FB <: PackFormat} function Descriptor{MR, NR, T, FA, FB}() where {MR, NR, T, FA, FB} - _check_descriptor(MR, NR, T, FA, FB) + check_descriptor(MR, NR, T, FA, FB) return new{MR, NR, T, FA, FB}() end end -const KernelDescriptor{MR, NR, T} = Descriptor{MR, NR, T, RealFormat, RealFormat} -const ComplexKernelDescriptor{MR, NR, T <: Complex, FA, FB} = Descriptor{MR, NR, T, FA, FB} - -KernelDescriptor(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = KernelDescriptor{MR, NR, T}() +const RealDescriptor{MR, NR, T} = Descriptor{MR, NR, T, RealFormat, RealFormat} -function ComplexKernelDescriptor( - ::Val{MR}, ::Val{NR}, ::Type{T}, ::FA, ::FB - ) where {MR, NR, T, FA <: PackFormat, FB <: PackFormat} - return Descriptor{MR, NR, T, FA, FB}() -end +Descriptor( + ::Val{MR}, ::Val{NR}, ::Type{T}, a_format::PackFormat = RealFormat(), b_format::PackFormat = RealFormat() +) where {MR, NR, T} = Descriptor{MR, NR, T, typeof(a_format), typeof(b_format)}() -function _check_descriptor(MR, NR, T, FA, FB) +function check_descriptor(MR, NR, T, FA, FB) real_formats = FA === RealFormat && FB === RealFormat - name = real_formats ? "KernelDescriptor" : "ComplexKernelDescriptor" MR isa Int && NR isa Int || - throw(ArgumentError("$name requires Int type parameters MR, NR")) - MR > 0 || throw(ArgumentError("$name requires MR > 0, got MR = $MR")) - NR > 0 || throw(ArgumentError("$name requires NR > 0, got NR = $NR")) + throw(ArgumentError("Descriptor requires Int type parameters MR, NR")) + MR > 0 || throw(ArgumentError("Descriptor requires MR > 0, got MR = $MR")) + NR > 0 || throw(ArgumentError("Descriptor requires NR > 0, got NR = $NR")) if real_formats T === Float32 || T === Float64 || - throw(ArgumentError("KernelDescriptor requires T ∈ (Float32, Float64), got $T")) + throw(ArgumentError("Descriptor with real formats requires T ∈ (Float32, Float64), got $T")) else # 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, " * + "Descriptor requires complex formats on both operands, " * "or (InterleavedFormat, RealFormat) / (RealFormat, InterleavedFormat), got ($FA, $FB)" ) ) T === ComplexF32 || T === ComplexF64 || throw( - ArgumentError("ComplexKernelDescriptor requires T in (ComplexF32, ComplexF64), got $T") + ArgumentError("Descriptor with complex formats requires T in (ComplexF32, ComplexF64), got $T") ) end return nothing @@ -92,11 +84,7 @@ packed_b_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 2) * k_blo # Zero-based offsets, in reals, of one-based coordinates. For 1e and # interleaved, `i`/`j` runs over reals (1:2MR) and 1e's second region is # `plane == 2`. -@inline packed_a_plane_offset(d::Descriptor{MR}, plane::Int, i::Int, p::Int) where {MR} = +@inline packed_a_offset(d::Descriptor{MR}, i::Int, p::Int, plane::Int = 0) where {MR} = (p - 1) * sliver_width(d, 1) + plane * MR + (i - 1) -@inline packed_b_plane_offset(d::Descriptor{MR, NR}, plane::Int, j::Int, p::Int) where {MR, NR} = +@inline packed_b_offset(d::Descriptor{MR, NR}, j::Int, p::Int, plane::Int = 0) where {MR, NR} = (p - 1) * sliver_width(d, 2) + plane * NR + (j - 1) - -# Real descriptors only: a complex element has no single offset. -packed_a_offset(kernel::KernelDescriptor{MR}, i::Int, p::Int) where {MR} = (i - 1) + MR * (p - 1) -packed_b_offset(kernel::KernelDescriptor{MR, NR}, j::Int, p::Int) where {MR, NR} = (j - 1) + NR * (p - 1) diff --git a/src/packing/pack.jl b/src/packing/pack.jl index 964c97b..79dcc9b 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -42,6 +42,23 @@ function pack_b!( return _pack_b!(packed, source, kernel, transform, Val(true)) end +# The packers write only `PackedPanel`s; a `DenseVector` is borrowed as one. +function pack_a!( + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, + transform::F + ) where {V <: DenseVector, MR, NR, T2, FA, FB, F} + GC.@preserve packed pack_a!(packed_panel(packed, 1, length(packed)), source, kernel, transform) + return packed +end + +function pack_b!( + packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, + transform::F + ) where {V <: DenseVector, MR, NR, T2, FA, FB, F} + GC.@preserve packed pack_b!(packed_panel(packed, 1, length(packed)), source, kernel, transform) + return packed +end + # Skips only `checked_tile_storage_bounds(source)`: the caller (`_execute_nest!`) # has already validated the whole macro block the sliver belongs to. `@inline` # because out of line each call marshals the `Tile` through the stack. @@ -109,7 +126,7 @@ end ) end load = (i, p) -> @inbounds source[i, p] - plane_offset = (plane, i, p) -> packed_a_plane_offset(kernel, plane, i, p) + plane_offset = (plane, i, p) -> packed_a_offset(kernel, i, p, plane) return _pack_panel!(packed, _element_type(format, T2), format, Val(MR), k_block_length, m, transform, load, plane_offset) end @@ -126,7 +143,7 @@ end ) end load = (j, p) -> @inbounds source[p, j] - plane_offset = (plane, j, p) -> packed_b_plane_offset(kernel, plane, j, p) + plane_offset = (plane, j, p) -> packed_b_offset(kernel, j, p, plane) return _pack_panel!(packed, _element_type(format, T2), format, Val(NR), k_block_length, n, transform, load, plane_offset) end diff --git a/src/packing/panel.jl b/src/packing/panel.jl index 59146f5..1a065e7 100644 --- a/src/packing/panel.jl +++ b/src/packing/panel.jl @@ -1,6 +1,7 @@ -# GUARDRAIL: packed micro-panels reach the kernel as borrowed pointers, never -# `view`s: a `SubArray`'s address arithmetic evicts the accumulator at large -# register tiles and heap-allocates per `execute!`. +# GUARDRAIL: packers and kernels read and write packed micro-panels only as +# `PackedPanel`s, borrowed pointers, never `view`s: a `SubArray`'s address +# arithmetic evicts the accumulator at large register tiles and heap-allocates +# per `execute!`. Only the checked entry points also take a `DenseVector`. using SIMD: Vec, vload, vstore @@ -17,16 +18,67 @@ Base.eltype(::Type{PackedPanel{T}}) where {T} = T @inline packed_panel(buffer::AbstractVector{T}, first1::Int, len::Int) where {T} = PackedPanel{T}(pointer(buffer, first1), len) -# Zero-based panel element access; the `AbstractVector` methods serve `Vector` -# and `view` destinations. +# Zero-based panel element access. @inline panel_vload(::Type{Vec{W, T}}, p::PackedPanel{T}, o::Int) where {W, T} = vload(Vec{W, T}, p.ptr + sizeof(T) * o) -@inline panel_vload(::Type{Vec{W, T}}, v::AbstractVector{T}, o::Int) where {W, T} = - vload(Vec{W, T}, v, o + 1) - @inline panel_load(p::PackedPanel{T}, o::Int) where {T} = unsafe_load(p.ptr + sizeof(T) * o) -@inline panel_load(v::AbstractVector, o::Int) = @inbounds v[o + 1] - @inline panel_store!(p::PackedPanel{T}, o::Int, x::T) where {T} = unsafe_store!(p.ptr + sizeof(T) * o, x) -@inline panel_store!(v::AbstractVector{T}, o::Int, x::T) where {T} = @inbounds v[o + 1] = x + +# --- Line-by-line packing of a split macro block --- + +# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated by +# a split group, whole `E x L` groups) by `k_block_length` 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, k_block_length::Int, split + ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: AbstractVector{Int}, 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 * k_block_length, 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 1:k_block_length + kbase = @inbounds base + 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 1:k_block_length + kbase = @inbounds base + 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 * k_block_length + pad = (plane, idx, pp) -> rb + plane_offset(descriptor, idx, pp, plane) + for p in 1:k_block_length, t in (valid + 1):R + _pack_emit_zero!(buffer, format, pad, t, p, realtype(descriptor)) + end + end + return nothing +end + +# Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane +# `i % R + 1`) 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, idx, pp, plane) + _pack_emit!(buffer, format, po, t + 1, p, convert(_element_type(format, scalartype(descriptor)), x)) + return nothing +end diff --git a/src/packing/transposed.jl b/src/packing/transposed.jl deleted file mode 100644 index 643f103..0000000 --- a/src/packing/transposed.jl +++ /dev/null @@ -1,55 +0,0 @@ -# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated by -# a split group, whole `E x L` groups) by `k_block_length` 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, k_block_length::Int, split - ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: AbstractVector{Int}, 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 * k_block_length, 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 1:k_block_length - kbase = @inbounds base + 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 1:k_block_length - kbase = @inbounds base + 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 * k_block_length - pad = (plane, idx, pp) -> rb + plane_offset(descriptor, plane, idx, pp) - for p in 1:k_block_length, t in (valid + 1):R - _pack_emit_zero!(buffer, format, pad, t, p, realtype(descriptor)) - end - end - return nothing -end - -# Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane -# `i % R + 1`) 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 + 1, p, convert(_element_type(format, scalartype(descriptor)), x)) - return nothing -end diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 8f40fd8..0d33402 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -1,5 +1,5 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEFormat, - PlanarFormat, ComplexKernelDescriptor, complex_method, lanewidth, kernel_shapes, + PlanarFormat, complex_method, lanewidth, kernel_shapes, _kernel_from_shape, _default_method @testset "OneMKernel" begin @@ -24,7 +24,7 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEForm @test_throws ArgumentError OneMKernel(Val(8), Val(4), Float64) @test_throws ArgumentError OneMKernel(Val(8), Val(4), Float64, Val(4)) wrong = SIMDKernel(Val(12), Val(8), Float64, Val(4)) - d = ComplexKernelDescriptor(Val(12), Val(8), ComplexF64, OneEFormat(), PlanarFormat()) + d = Descriptor(Val(12), Val(8), ComplexF64, OneEFormat(), PlanarFormat()) @test_throws ArgumentError OneMKernel{12, 8, ComplexF64, 8, typeof(wrong)}(d, wrong) # The error reports the logical k_block_length, not the doubled real one. err = try diff --git a/test/microkernels/test_planar_kernel.jl b/test/microkernels/test_planar_kernel.jl index bc08f18..5035fe5 100644 --- a/test/microkernels/test_planar_kernel.jl +++ b/test/microkernels/test_planar_kernel.jl @@ -17,8 +17,6 @@ using QuasiStrided: PlanarKernel, PlanarMethod, complex_method, lanewidth, @test_throws ArgumentError PlanarKernel(Val(6), Val(4), ComplexF64, Val(4)) @test_throws ArgumentError PlanarKernel(Val(8), Val(4), ComplexF64, Val(0)) @test_throws ArgumentError PlanarKernel(Val(8), Val(4), Float64) - # A complex kernel is never asked for a single-plane offset. - @test_throws MethodError QuasiStrided.packed_a_offset(k, 1, 1) mk_e2e(ComplexF64, nothing) end diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index 9a93282..77c8a9f 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -122,7 +122,7 @@ function mk_contract_full(k) la, lb = length(pa) ÷ k_block_length, length(pb) ÷ k_block_length accs = zero_accumulator(k) for p in 0:(k_block_length - 1) - accs = accumulate(k, accs, view(pa, (p * la + 1):((p + 1) * la)), view(pb, (p * lb + 1):((p + 1) * lb)), 1) + accs = accumulate(k, accs, pa[(p * la + 1):((p + 1) * la)], pb[(p * lb + 1):((p + 1) * lb)], 1) end @test all(mk_read(k, accs, i, j) ≈ mk_read(k, acc, i, j) for i in 0:(MR - 1), j in 0:(NR - 1)) diff --git a/test/packing/test_kernel_descriptor.jl b/test/packing/test_kernel_descriptor.jl index 34bbda5..96c1659 100644 --- a/test/packing/test_kernel_descriptor.jl +++ b/test/packing/test_kernel_descriptor.jl @@ -1,8 +1,8 @@ -using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEFormat, +using QuasiStrided: RealFormat, PlanarFormat, OneEFormat, InterleavedFormat, realtype, sliver_width -@testset "KernelDescriptor" begin - k = KernelDescriptor(Val(8), Val(6), Float64) +@testset "Descriptor" begin + k = Descriptor(Val(8), Val(6), Float64) @test (tile_size(k)..., scalartype(k), realtype(k)) == (8, 6, Float64, Float64) @test packed_a_offset(k, 4, 3) == 3 + 8 * 2 @test packed_b_offset(k, 6, 3) == 5 + 6 * 2 @@ -10,27 +10,28 @@ using QuasiStrided: ComplexKernelDescriptor, RealFormat, PlanarFormat, OneEForma @test packed_a_length(k, 4) == 32 @test packed_b_length(k, 4) == 24 - @test_throws ArgumentError KernelDescriptor(Val(0), Val(6), Float64) - @test_throws ArgumentError KernelDescriptor(Val(8), Val(-1), Float64) - @test_throws ArgumentError KernelDescriptor(Val(8), Val(6), Int) + @test_throws ArgumentError Descriptor(Val(0), Val(6), Float64) + @test_throws ArgumentError Descriptor(Val(8), Val(-1), Float64) + @test_throws ArgumentError Descriptor(Val(8), Val(6), Int) end -@testset "ComplexKernelDescriptor: lengths count reals at a logical k_block_length ($T)" for +@testset "Descriptor: lengths count reals at a logical k_block_length ($T)" for T in (ComplexF64, ComplexF32) MR, NR = 8, 6 - planar = ComplexKernelDescriptor(Val(MR), Val(NR), T, PlanarFormat(), PlanarFormat()) - onem = ComplexKernelDescriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) - fmas = ComplexKernelDescriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) + planar = Descriptor(Val(MR), Val(NR), T, PlanarFormat(), PlanarFormat()) + onem = Descriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) + fmas = Descriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) @test realtype(planar) === real(T) @test (sliver_width(planar, 1), sliver_width(onem, 1), sliver_width(fmas, 1)) == (2MR, 4MR, 2MR) @test sliver_width(planar, 2) == sliver_width(onem, 2) == sliver_width(fmas, 2) == 2NR @test packed_a_length(onem, 7) == 4MR * 7 @test packed_b_length(planar, 7) == 2NR * 7 + @test packed_a_offset(planar, 3, 2, 1) == 2MR + MR + 2 - @test_throws ArgumentError ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), PlanarFormat()) - @test sliver_width(ComplexKernelDescriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()))[2] == 2NR - @test_throws ArgumentError ComplexKernelDescriptor(Val(MR), Val(NR), Float64, PlanarFormat(), PlanarFormat()) - @test_throws ArgumentError ComplexKernelDescriptor(Val(0), Val(NR), T, PlanarFormat(), PlanarFormat()) + @test_throws ArgumentError Descriptor(Val(MR), Val(NR), T, RealFormat(), PlanarFormat()) + @test sliver_width(Descriptor(Val(MR), Val(NR), T, RealFormat(), InterleavedFormat()))[2] == 2NR + @test_throws ArgumentError Descriptor(Val(MR), Val(NR), Float64, PlanarFormat(), PlanarFormat()) + @test_throws ArgumentError Descriptor(Val(0), Val(NR), T, PlanarFormat(), PlanarFormat()) end diff --git a/test/packing/test_pack_complex.jl b/test/packing/test_pack_complex.jl index 550f189..2eb59ca 100644 --- a/test/packing/test_pack_complex.jl +++ b/test/packing/test_pack_complex.jl @@ -23,7 +23,7 @@ const COMPLEX_FORMATS = ( T in (ComplexF64, ComplexF32), (fa, fb) in COMPLEX_FORMATS MR, NR = 4, 3 - kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, fb) + kernel = Descriptor(Val(MR), Val(NR), T, fa, fb) R = real(T) storage = complex_storage(T, 4000) calls = Ref(0) @@ -42,7 +42,7 @@ const COMPLEX_FORMATS = ( @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, f))) end calls[] = 0 - got, canaries = pack_into(pack!, :view, R, len, src, kernel, counting) + got, canaries = pack_into(pack!, :panel, R, len, src, kernel, counting) @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, z -> 2z + one(T)))) @test calls[] == valid * k_block_length # once per element, not per real half @test canaries @@ -51,7 +51,7 @@ const COMPLEX_FORMATS = ( end @testset "complex packing: validation before any write, k_block_length == 0 is a no-op" begin - kernel = ComplexKernelDescriptor(Val(4), Val(3), ComplexF64, PlanarFormat(), PlanarFormat()) + kernel = Descriptor(Val(4), Val(3), ComplexF64, PlanarFormat(), PlanarFormat()) storage = fill(ComplexF64(3, 4), 100) src = Tile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 8, 2)) # The buffer holds realtype(kernel), not scalartype(kernel). @@ -80,7 +80,7 @@ end function run(::Type{T}, fa, fb) where {T} MR, NR, k_block_length = 8, 6, 4 - kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, fb) + kernel = Descriptor(Val(MR), Val(NR), T, fa, fb) R = real(T) storage = rand(T, 4000) bufa = zeros(R, packed_a_length(kernel, k_block_length)) diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index 95650cb..7dab653 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -36,7 +36,7 @@ function check_panel(side, kernel, T, PD, fmt, lane, step, f; S = T) 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()) + kernel = Descriptor(Val(MR), Val(3), T, fa, PlanarFormat()) 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) @@ -46,7 +46,7 @@ end 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()) + kernel = Descriptor(Val(4), Val(NR), T, PlanarFormat(), PlanarFormat()) 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) @@ -59,7 +59,7 @@ end MR, NR, k_block_length = 8, 6, 6 koffs = [((p * 5) % 7) * 1013 + 3p for p in 0:(k_block_length - 1)] for fa in (PlanarFormat(), OneEFormat(), InterleavedFormat()), f in (identity, conj) - kernel = ComplexKernelDescriptor(Val(MR), Val(NR), T, fa, OneEFormat()) + kernel = Descriptor(Val(MR), Val(NR), T, fa, OneEFormat()) # The lane axis must be unit-stride; the step axis may scatter. check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), view(koffs, 1:k_block_length), f) check_panel(:b, kernel, T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, k_block_length), f) diff --git a/test/packing/test_pack_real.jl b/test/packing/test_pack_real.jl index 57f9c87..053f35b 100644 --- a/test/packing/test_pack_real.jl +++ b/test/packing/test_pack_real.jl @@ -3,7 +3,7 @@ # files included after this one. using QuasiStrided: RealFormat, PlanarFormat, OneEFormat, InterleavedFormat, PackedPanel, - packed_panel, ComplexKernelDescriptor, _copies_unchanged, + packed_panel, _copies_unchanged, _pack_a_contiguous_eligible # Packed layouts written from the format definitions, not from the offset @@ -54,7 +54,7 @@ function pack_into(pack!, dst, R, len, src, kernel, f) return v, true end GC.@preserve buf begin - d = dst === :view ? view(buf, 5:(4 + len)) : packed_panel(buf, 5, len) + d = packed_panel(buf, 5, len) @test pack!(d, src, kernel, f) === d end return buf[5:(4 + len)], all(==(R(-777)), buf[1:4]) && all(==(R(-777)), buf[(5 + len):end]) @@ -67,7 +67,7 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] T in (Float64, Float32) MR, NR, k_block_length = 8, 6, 5 - kernel = KernelDescriptor(Val(MR), Val(NR), T) + kernel = Descriptor(Val(MR), Val(NR), T) storage = T.(collect(1.0:2000.0)) calls = Ref(0) counting = x -> (calls[] += 1; 3x + 1000) # nonzero at zero: padding must bypass it @@ -77,7 +77,7 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] ] steps = Any[AffineAxis(0, 64, k_block_length), AffineAxis(700, -97, k_block_length), view(SCATTER_STEPS, 1:k_block_length)] for (side, pack!, PD) in ((:a, pack_a!, MR), (:b, pack_b!, NR)), - lane in lanes, step in steps, valid in (PD, 1, 0), dst in (:vector, :view) + lane in lanes, step in steps, valid in (PD, 1, 0), dst in (:vector, :panel) src, g = pack_fixture(side, storage, 11, resized(lane, valid), step) calls[] = 0 got, canaries = pack_into(pack!, dst, T, PD * k_block_length, src, kernel, counting) @@ -97,7 +97,7 @@ end @testset "pack_a! contiguous fast path: fires exactly when eligible ($T, MR=$MR)" for T in (Float64, Float32), MR in (4, 16) - kernel = KernelDescriptor(Val(MR), Val(3), T) + kernel = Descriptor(Val(MR), Val(3), T) vals = T.(collect(1.0:2000.0)) 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) @@ -127,7 +127,7 @@ end end @testset "pack_a!/pack_b!: k_block_length == 0 reads and writes nothing" begin - kernel = KernelDescriptor(Val(4), Val(3), Float64) + kernel = Descriptor(Val(4), Val(3), Float64) storage = fill(3.0, 10) read = Ref(false) spy = x -> (read[] = true; x) @@ -139,8 +139,8 @@ end end @testset "pack_a!/pack_b!: invalid metadata rejected before mutation" begin - kernel = KernelDescriptor(Val(4), Val(3), Float64) - kernel32 = KernelDescriptor(Val(4), Val(3), Float32) + kernel = Descriptor(Val(4), Val(3), Float64) + kernel32 = Descriptor(Val(4), Val(3), Float32) storage = fill(9.0, 20) tile(m, n, base = 0) = Tile(storage, base, AffineAxis(0, 1, m), AffineAxis(0, 1, n)) for (pack!, bad, ok, short) in ((pack_a!, tile(5, 3), tile(4, 3), 11), (pack_b!, tile(3, 4), tile(3, 3), 8)) @@ -168,8 +168,8 @@ end bytes = Int[] GC.@preserve bufa bufb begin pa = packed_panel(bufa, 1, MR * k_block_length) - dsts_a = (bufa, view(bufa, 3:(2 + MR * k_block_length)), pa) - dsts_b = (bufb, view(bufb, 3:(2 + NR * k_block_length)), packed_panel(bufb, 1, NR * k_block_length)) + dsts_a = (bufa, pa) + dsts_b = (bufb, packed_panel(bufb, 1, NR * k_block_length)) pk = view(koffs, 1:k_block_length) srcs_a = ( Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)), @@ -185,7 +185,7 @@ end Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, k_block_length, NR)), ) if !full - srcs_a, srcs_b, dsts_a, dsts_b = srcs_a[1:1], srcs_b[1:1], dsts_a[1:2], dsts_b[1:2] + srcs_a, srcs_b = srcs_a[1:1], srcs_b[1:1] end for d in dsts_a for s in srcs_a @@ -201,8 +201,8 @@ end end return bytes end - @test all(iszero, run(KernelDescriptor, Float64, 16, 6, true)) - @test all(iszero, run(KernelDescriptor, Float32, 32, 6, true)) + @test all(iszero, run(Descriptor, Float64, 16, 6, true)) + @test all(iszero, run(Descriptor, Float32, 32, 6, true)) @test all(iszero, run(ScalarKernel, Float64, 8, 6, false)) @test all(iszero, run(SIMDKernel, Float64, 8, 6, false)) end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index a6d1c2f..cd6b16d 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -352,7 +352,7 @@ end ref = zeros(8 * 4) pack_a!(ref, src, kernel, identity) - QS.unsafe_pack_a!(packed, src, kernel, identity) + GC.@preserve packed QS.unsafe_pack_a!(QS.packed_panel(packed, 1, length(packed)), src, kernel, identity) @test packed == ref # Only the bounds checks are skipped. @@ -365,7 +365,7 @@ end refB = zeros(6 * 4) packedB = zeros(6 * 4) pack_b!(refB, srcB, kernel, identity) - QS.unsafe_pack_b!(packedB, srcB, kernel, identity) + GC.@preserve packedB QS.unsafe_pack_b!(QS.packed_panel(packedB, 1, length(packedB)), srcB, kernel, identity) @test packedB == refB @test_throws DimensionMismatch QS.unsafe_pack_b!(zeros(3), srcB, kernel, identity) diff --git a/test/runtests.jl b/test/runtests.jl index 1bdba6e..9e71cc9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,7 +3,7 @@ using Random using QuasiStrided # Internal names the test files use unqualified. using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, - describe_block, block_descriptors!, KernelDescriptor, tile_size, + describe_block, block_descriptors!, Descriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, AffineAxis, Tile, checked_tile_storage_bounds, pack_a!, pack_b!, zero_accumulator, accumulate, store_tile!, execute_tile!, lanewidth, contract!, From 2d3cb24e64fcf198247a97144bc55e392eca2a0e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 15:17:25 -0400 Subject: [PATCH 10/60] Chunk 5: one symmetric packer A and B are packed by one code path; the kernel's descriptor decides each operand's packing through a `SliverSpec{I, L, F, T}` (side, lanes, format, scalar type), from `sliver_spec(kernel, i)`. - `pack!(panel, tile, spec, transform)` / `unsafe_pack!` replace pack_a!/pack_b!/unsafe_pack_a!/unsafe_pack_b! and the DescriptorKernel forwarding methods. Lanes run along `tile.rows`, K steps along `tile.cols`; B is packed as `transpose` of its tile (the nest and oracle pass B's axes swapped). `Base.transpose(::Tile)` added. - `panel_offset(spec, t, p, plane)` is the one offset formula; `packed_a_offset`/`packed_b_offset`/`packed_*_length` are defined on it (layout unchanged). - Fast paths are symmetric: real B with unit-stride columns now takes the `Vec{L}` load/convert/store path. - `emit!` takes `(re, im, neg_im)`; the value and padding callers supply them (padding stays literal +0.0). The block packer takes the spec. - One complex contiguous K loop with a per-format `store_step!` (LLVM IR of the shuffles/loads/stores identical to before for all formats x transforms). - Dead `PackedPanel` eligibility conditions dropped; eltype error built in an `@noinline` helper; `is_unit_stride`, `dense_lanes`, `lane_convertible`, `complex_fastpath_isa_eligible` collected in pack_contiguous.jl; no leading underscores in pack.jl / pack_contiguous.jl. Tests: 30931 -> 31537. Real A/B halves merged into sweeps over specs, plus "B packs as A packs its transposed tile" and a real-B fast-path sweep (L = 6); removed the duplicate ISA-gate testset (planar store), the PackedPanel-gate eligibility checks, the fmaddsub Vector-vs-panel duplicate and the unsafe_pack_b! duplicates. TTFX (base / new, 2 interleaved runs): first call Float64 4.5, 3.4 / 4.2, 4.3 s; ComplexF64 3.1, 3.0 / 2.7, 2.7 s; load 0.5-0.6 s both. A/B, reused-plan execute! medians (5 rounds, us, base -> new): matmul256 F64 513 -> 521, scatterM F64 600 -> 568, scatterMN F64 31.6 -> 32.0, matmul256 C64 2005 -> 1981, scatterM C64 2011 -> 1963, scatterMN C64 62.1 -> 61.7. Row-major B (7 rounds): F64 571 -> 482, F32 282 -> 265 (5-round run: 605 -> 566, 232 -> 269; workstation noise is +-30%). Scattered K (7 rounds): F64 582 -> 590, C64 2069 -> 1717. Isolated real B sliver pack, fast vs scalar, kc = 256: 2-3x faster for every L in 4, 6, 8, 12, 16 (Float32 and Float64), so no power-of-two restriction. Per-call allocations 0 for all cases under --check-bounds=yes. Co-Authored-By: Claude Opus 5.5 --- src/execution/execute.jl | 14 +- src/execution/macrokernel.jl | 10 +- src/execution/oracle.jl | 8 +- src/hardware/target.jl | 5 - src/layout/tiles.jl | 5 +- src/microkernels/fmaddsub.jl | 2 +- src/microkernels/interface.jl | 21 +- src/microkernels/planar.jl | 4 +- src/microkernels/simd.jl | 2 +- src/packing/format.jl | 27 +- src/packing/pack.jl | 293 ++++++------------ src/packing/pack_contiguous.jl | 220 +++++++------ src/packing/panel.jl | 21 +- test/execution/test_manual_pipeline.jl | 10 +- test/helpers.jl | 6 +- test/microkernels/test_fmaddsub_kernel.jl | 18 +- test/microkernels/test_mixed_kernel.jl | 4 +- .../test_planar_store_fastpath.jl | 8 - test/packing/test_pack_complex.jl | 72 ++--- test/packing/test_pack_complex_contiguous.jl | 78 ++--- test/packing/test_pack_real.jl | 151 +++++---- test/planning/test_per_call_overhead.jl | 26 +- test/runtests.jl | 2 +- 23 files changed, 402 insertions(+), 605 deletions(-) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 1b0433f..23a3771 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -368,8 +368,7 @@ function _execute_nest!( if split_b _pack_block_transposed!( - packed_b_offset, b_format(kernel), - packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), kernel, Val(tile_size(kernel, 2)), b_sliver_width, + packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), sliver_spec(kernel, 2), plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, plan.npack ) @@ -379,8 +378,8 @@ function _execute_nest!( bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) _pack_sliver!( - unsafe_pack_b!, bpanel, plan.Bstorage, plan.Bbase, rowsB_k, colsB, - kernel, btransform + unsafe_pack!, bpanel, plan.Bstorage, plan.Bbase, colsB, rowsB_k, + sliver_spec(kernel, 2), btransform ) end end @@ -408,8 +407,7 @@ function _execute_nest!( if split_a _pack_block_transposed!( - packed_a_offset, a_format(kernel), - packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), kernel, Val(tile_size(kernel, 1)), a_sliver_width, + packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), sliver_spec(kernel, 1), plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, plan.mpack ) @@ -419,8 +417,8 @@ function _execute_nest!( apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) _pack_sliver!( - unsafe_pack_a!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, - kernel, atransform + unsafe_pack!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, + sliver_spec(kernel, 1), atransform ) end end diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index 79d6e0f..82f7684 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -18,14 +18,14 @@ end @noinline _throw_irregular_ramp_descriptor() = throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) -# `pack!` is one of pack_a!/pack_b!/unsafe_pack_a!/unsafe_pack_b!; call sites -# spell the unsafe name out. GUARDRAIL: `transform` needs its own bound type -# parameter, or it costs a dynamic dispatch per call (see `pack_a!`). +# `pack!` is `pack!` or `unsafe_pack!`; B passes its axes swapped (its +# `transpose`). GUARDRAIL: `transform` needs its own bound type parameter, or +# it costs a dynamic dispatch per call. @inline function _pack_sliver!( pack!::PF, packed::PK, storage::S, base::Int, - rows::R, cols::C, kernel, transform::TF + lanes::R, steps::C, spec, transform::TF ) where {PF, PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} - pack!(packed, Tile(storage, base, rows, cols), kernel, transform) + pack!(packed, Tile(storage, base, lanes, steps), spec, transform) return nothing end diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl index fe0c9d7..5041196 100644 --- a/src/execution/oracle.jl +++ b/src/execution/oracle.jl @@ -61,12 +61,12 @@ function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) w # Whole single-sliver buffers, so no `_sliver_panel`. _pack_sliver!( - pack_a!, ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, - kernel, plan.atransform + pack!, ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, + sliver_spec(kernel, 1), plan.atransform ) _pack_sliver!( - pack_b!, ws.tw_packed_b, plan.Bstorage, plan.Bbase, rowsK_B, colsB, - kernel, plan.btransform + pack!, ws.tw_packed_b, plan.Bstorage, plan.Bbase, colsB, rowsK_B, + sliver_spec(kernel, 2), plan.btransform ) beta_eff = firstpanel ? betaT : one(T) diff --git a/src/hardware/target.jl b/src/hardware/target.jl index 928c5c8..ed1e7ae 100644 --- a/src/hardware/target.jl +++ b/src/hardware/target.jl @@ -178,8 +178,3 @@ init_target!() = (TARGET[] = detect_target(); nothing) # core, which are the CPUs sharing L1d. core_bytes(profile::TargetProfile, level::CacheLevel) = level.bytes ÷ max(1, level.sharing ÷ max(1, profile.l1d.sharing)) - -# Shared by the deinterleaving complex packer and the planar store: both pay -# off only with 512-bit vector registers. -@inline complex_fastpath_isa_eligible(profile::TargetProfile) = profile.isa === :avx512 -@inline complex_fastpath_isa_eligible() = complex_fastpath_isa_eligible(target_profile()) diff --git a/src/layout/tiles.jl b/src/layout/tiles.jl index 915cc7b..5a50401 100644 --- a/src/layout/tiles.jl +++ b/src/layout/tiles.jl @@ -97,6 +97,7 @@ struct Tile{S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} end Base.size(tile::Tile) = (length(tile.rows), length(tile.cols)) +Base.transpose(tile::Tile) = Tile(tile.storage, tile.base, tile.cols, tile.rows) Base.@propagate_inbounds Base.getindex(tile::Tile, i::Int, j::Int) = tile.storage[tile.base + tile.rows[i] + tile.cols[j] + 1] @@ -141,7 +142,3 @@ end checked_tile_storage_bounds(tile::Tile) = checked_tile_storage_bounds(tile.base, tile.rows, tile.cols, length(tile.storage)) - -# Whether an axis steps through storage one element at a time. -is_unit_stride(ax::AffineAxis) = ax.stride == 1 -is_unit_stride(::AbstractVector{Int}) = false diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 5c868be..a809b24 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -203,7 +203,7 @@ end m::Int, n::Int ) where {S, MR, NR, T, W, R, NV} # The pointer reinterpretation is only sound on dense rank-1 complex storage. - S <: DenseVector && _lane_convertible(eltype(S), T) || + 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")) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 45c9374..15442d1 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -55,6 +55,7 @@ a_format(k::DescriptorKernel) = a_format(k.descriptor) b_format(k::DescriptorKernel) = b_format(k.descriptor) tile_size(k::DescriptorKernel) = tile_size(k.descriptor) scalartype(k::DescriptorKernel) = scalartype(k.descriptor) +@inline sliver_spec(k::DescriptorKernel, i::Int) = sliver_spec(k.descriptor, i) @inline packed_a_offset(k::DescriptorKernel, i::Int, p::Int, plane::Int = 0) = packed_a_offset(k.descriptor, i, p, plane) @inline packed_b_offset(k::DescriptorKernel, j::Int, p::Int, plane::Int = 0) = @@ -62,26 +63,6 @@ scalartype(k::DescriptorKernel) = scalartype(k.descriptor) packed_a_length(k::DescriptorKernel, k_block_length::Int) = packed_a_length(k.descriptor, k_block_length) packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.descriptor, k_block_length) -# GUARDRAIL: every forwarded argument needs its OWN bound type parameter; an -# unbound one makes the call dynamically dispatched and allocating on every -# pack. `V` is unconstrained so a `DenseVector` forwards too. -pack_a!( - packed::V, source::Tile, kernel::K, transform::F -) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = - pack_a!(packed, source, kernel.descriptor, transform) -pack_b!( - packed::V, source::Tile, kernel::K, transform::F -) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = - pack_b!(packed, source, kernel.descriptor, transform) -@inline unsafe_pack_a!( - packed::V, source::Tile, kernel::K, transform::F -) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = - unsafe_pack_a!(packed, source, kernel.descriptor, transform) -@inline unsafe_pack_b!( - packed::V, source::Tile, kernel::K, transform::F -) where {V, MR, NR, T, K <: DescriptorKernel{MR, NR, T}, F} = - unsafe_pack_b!(packed, source, kernel.descriptor, transform) - # Constructor check: `rows` reals per sliver must split into whole `W`-vectors. # `even`: the lane-pair kernels keep one element's re/im in adjacent lanes, and # an odd `W` can divide `2MR` while splitting an element across two vectors -- diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index f0c1ba8..180ebcf 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -186,7 +186,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::Tile, ::Type{T}) where {T} = - is_unit_stride(tile.rows) && _dense_lanes(tile.storage, T) && + is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) && complex_fastpath_isa_eligible() # Shuffle patterns built from `W` at specialization time, never hardcoded to @@ -248,7 +248,7 @@ end m::Int, n::Int ) where {S, MR, NR, T, W, R, NA} # The pointer reinterpretation is only sound on dense rank-1 complex storage. - S <: DenseVector && _lane_convertible(eltype(S), T) || + 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 diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index 50aede2..f594376 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -94,7 +94,7 @@ end # 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::Tile, ::Type{T}) where {T} = - is_unit_stride(tile.rows) && _dense_lanes(tile.storage, T) + is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) @generated function _store_tile_scattered!( destination::Tile, acc::NTuple{NV, Vec{W, T}}, diff --git a/src/packing/format.jl b/src/packing/format.jl index f3f428e..cf1c977 100644 --- a/src/packing/format.jl +++ b/src/packing/format.jl @@ -78,13 +78,24 @@ sliver_width(d::Descriptor{MR, NR}) where {MR, NR} = (MR * reals_per_element(a_format(d)), NR * reals_per_element(b_format(d))) sliver_width(k, i::Int) = sliver_width(k)[i] -packed_a_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 1) * k_block_length -packed_b_length(d::Descriptor, k_block_length::Int) = sliver_width(d, 2) * k_block_length +# How one operand's slivers are packed: side `I` (1 for A, 2 for B), `L` lanes +# (`MR` or `NR`) in format `F`, for a kernel of scalar type `T`. +struct SliverSpec{I, L, F <: PackFormat, T} end -# Zero-based offsets, in reals, of one-based coordinates. For 1e and -# interleaved, `i`/`j` runs over reals (1:2MR) and 1e's second region is +@inline sliver_spec(::Descriptor{MR, NR, T, FA, FB}, i::Int) where {MR, NR, T, FA, FB} = + i == 1 ? SliverSpec{1, MR, FA, T}() : SliverSpec{2, NR, FB, T}() + +realtype(::SliverSpec{I, L, F, T}) where {I, L, F, T} = real(T) +sliver_width(::SliverSpec{I, L, F}) where {I, L, F} = L * reals_per_element(F()) +packed_length(spec::SliverSpec, k_block_length::Int) = sliver_width(spec) * k_block_length + +# Zero-based offset, in reals, of lane `t` of plane `plane` at K step `p`. For +# 1e and interleaved, `t` runs over reals (1:2L) and 1e's second region is # `plane == 2`. -@inline packed_a_offset(d::Descriptor{MR}, i::Int, p::Int, plane::Int = 0) where {MR} = - (p - 1) * sliver_width(d, 1) + plane * MR + (i - 1) -@inline packed_b_offset(d::Descriptor{MR, NR}, j::Int, p::Int, plane::Int = 0) where {MR, NR} = - (p - 1) * sliver_width(d, 2) + plane * NR + (j - 1) +@inline panel_offset(spec::SliverSpec{I, L}, t::Int, p::Int, plane::Int = 0) where {I, L} = + (p - 1) * sliver_width(spec) + plane * L + (t - 1) + +packed_a_length(d::Descriptor, k_block_length::Int) = packed_length(sliver_spec(d, 1), k_block_length) +packed_b_length(d::Descriptor, k_block_length::Int) = packed_length(sliver_spec(d, 2), k_block_length) +@inline packed_a_offset(d::Descriptor, i::Int, p::Int, plane::Int = 0) = panel_offset(sliver_spec(d, 1), i, p, plane) +@inline packed_b_offset(d::Descriptor, j::Int, p::Int, plane::Int = 0) = panel_offset(sliver_spec(d, 2), j, p, plane) diff --git a/src/packing/pack.jl b/src/packing/pack.jl index 79dcc9b..edca708 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -1,256 +1,137 @@ -# Packers: copy one sliver of A (up to MR logical rows) or B (up to NR logical -# columns) over a K range into a contiguous panel in the descriptor's format. +# The packer: copies one sliver, up to `L` lanes along `tile.rows` by the K +# steps along `tile.cols`, into a contiguous panel in the spec's format. A is +# packed as its tile, B as `transpose` of its tile. + +operand_name(i::Int) = i == 1 ? "A" : "B" # A runtime check rather than dispatch, so a mismatch is an ArgumentError. -@inline function _check_packed_eltype(packed, kernel::Descriptor{MR, NR, T2}) where {MR, NR, T2} - R = realtype(kernel) - eltype(packed) === R && return nothing - msg = R === T2 ? - "packed buffer eltype $(eltype(packed)) does not match kernel scalar type $T2" : - "packed buffer eltype $(eltype(packed)) does not match kernel real type $R " * - "(scalar type $T2)" - throw(ArgumentError(msg)) +@inline function check_packed_eltype(panel, spec::SliverSpec{I, L, F, T}) where {I, L, F, T} + eltype(panel) === real(T) || throw_packed_eltype(eltype(panel), T) + return nothing end -# Out of line so the packers carry no string formatting (and no GC frame). -@noinline _throw_pack_extent(which::Symbol, name::Symbol, got::Int, limit::Int) = throw( +# Out of line so the packer carries no string formatting (and no GC frame). +@noinline throw_packed_eltype(got::Type, ::Type{T}) where {T} = throw( + ArgumentError( + real(T) === T ? + "packed buffer eltype $got does not match kernel scalar type $T" : + "packed buffer eltype $got does not match kernel real type $(real(T)) (scalar type $T)" + ) +) +@noinline throw_pack_extent(i::Int, got::Int, limit::Int) = throw( ArgumentError( - "$(which)!: source $(name) count $got must satisfy 0 <= $got <= tile_size(kernel, $(which === :pack_a ? 1 : 2))=$limit" + "pack!: $(operand_name(i)) sliver has $got lanes, need 0 <= $got <= tile_size(kernel, $i) = $limit" ) ) -@noinline _throw_pack_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( +@noinline throw_pack_short(i::Int, got::Int, need::Int, k_block_length::Int) = throw( DimensionMismatch( - "$(which)!: packed buffer has length $got, " * - "need at least packed_$(which === :pack_a ? "a" : "b")_length(kernel, k_block_length=$k_block_length) = $need" + "pack!: packed $(operand_name(i)) buffer has length $got, need at least " * + "packed_$(i == 1 ? "a" : "b")_length(kernel, k_block_length=$k_block_length) = $need" ) ) # `transform` is applied to each loaded (complex) element before it is split # into the packed format; padding lanes are literal zeros and never read the # source or call `transform`. All validation happens before any write. -function pack_a!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V, MR, NR, T2, FA, FB, F} - return _pack_a!(packed, source, kernel, transform, Val(true)) -end - -function pack_b!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V, MR, NR, T2, FA, FB, F} - return _pack_b!(packed, source, kernel, transform, Val(true)) -end +pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V, F} = + pack_tile!(panel, tile, spec, transform, Val(true)) -# The packers write only `PackedPanel`s; a `DenseVector` is borrowed as one. -function pack_a!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V <: DenseVector, MR, NR, T2, FA, FB, F} - GC.@preserve packed pack_a!(packed_panel(packed, 1, length(packed)), source, kernel, transform) - return packed +# The packer writes only `PackedPanel`s; a `DenseVector` is borrowed as one. +function pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V <: DenseVector, F} + GC.@preserve panel pack!(packed_panel(panel, 1, length(panel)), tile, spec, transform) + return panel end -function pack_b!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V <: DenseVector, MR, NR, T2, FA, FB, F} - GC.@preserve packed pack_b!(packed_panel(packed, 1, length(packed)), source, kernel, transform) - return packed -end - -# Skips only `checked_tile_storage_bounds(source)`: the caller (`_execute_nest!`) +# Skips only `checked_tile_storage_bounds(tile)`: the caller (`_execute_nest!`) # has already validated the whole macro block the sliver belongs to. `@inline` # because out of line each call marshals the `Tile` through the stack. -@inline function unsafe_pack_a!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V, MR, NR, T2, FA, FB, F} - return _pack_a!(packed, source, kernel, transform, Val(false)) -end - -@inline function unsafe_pack_b!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F - ) where {V, MR, NR, T2, FA, FB, F} - return _pack_b!(packed, source, kernel, transform, Val(false)) -end - -@inline function _pack_a!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F, ::Val{BOUNDS} - ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} - _check_packed_eltype(packed, kernel) - m, k_block_length = size(source) - (0 <= m <= MR) || _throw_pack_extent(:pack_a, :row, m, MR) - needed = packed_a_length(kernel, k_block_length) - length(packed) >= needed || _throw_pack_short(:pack_a, length(packed), needed, k_block_length) +@inline unsafe_pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V, F} = + pack_tile!(panel, tile, spec, transform, Val(false)) + +@inline function pack_tile!( + panel::V, tile::Tile, spec::SliverSpec{I, L}, transform::F, ::Val{BOUNDS} + ) where {V, I, L, F, BOUNDS} + check_packed_eltype(panel, spec) + lanes, k_block_length = size(tile) + (0 <= lanes <= L) || throw_pack_extent(I, lanes, L) + needed = packed_length(spec, k_block_length) + length(panel) >= needed || throw_pack_short(I, length(panel), needed, k_block_length) # `<= 0`, not `== 0` (counts are never negative): LLVM may then assume # `k_block_length >= 1` in the loops below. - k_block_length <= 0 && return packed - BOUNDS && checked_tile_storage_bounds(source) - return _pack_a_sliver!(FA(), packed, source, kernel, transform, m, k_block_length) -end - -@inline function _pack_b!( - packed::V, source::Tile, kernel::Descriptor{MR, NR, T2, FA, FB}, - transform::F, ::Val{BOUNDS} - ) where {V, MR, NR, T2, FA, FB, F, BOUNDS} - _check_packed_eltype(packed, kernel) - k_block_length, n = size(source) - (0 <= n <= NR) || _throw_pack_extent(:pack_b, :column, n, NR) - needed = packed_b_length(kernel, k_block_length) - length(packed) >= needed || _throw_pack_short(:pack_b, length(packed), needed, k_block_length) - k_block_length <= 0 && return packed - BOUNDS && checked_tile_storage_bounds(source) - return _pack_b_sliver!(FB(), packed, source, kernel, transform, n, k_block_length) -end - -# A's packed index runs along `source.rows`, B's along `source.cols`; each -# contiguous fast path needs that lane axis to be unit-stride. -@inline function _pack_a_sliver!( - format::FMT, packed::V, source::Tile, kernel::Descriptor{MR, NR, T2}, - transform::F, m::Int, k_block_length::Int - ) where {FMT, V, MR, NR, T2, F} - if format isa RealFormat - if _pack_a_contiguous_eligible(packed, source, transform, m, Val(MR), real(T2)) - rowbase = @inbounds source.base + source.rows[1] - return _pack_a_contiguous!(packed, source.storage, rowbase, source.cols, Val(MR), k_block_length) - end - elseif _pack_complex_contiguous_eligible( - packed, source.storage, source.rows, transform, format, m, Val(MR), T2 - ) - elembase = @inbounds source.base + source.rows[1] - return _pack_complex_contiguous!( - format, packed, source.storage, elembase, source.cols, Val(MR), k_block_length, transform - ) - end - load = (i, p) -> @inbounds source[i, p] - plane_offset = (plane, i, p) -> packed_a_offset(kernel, i, p, plane) - return _pack_panel!(packed, _element_type(format, T2), format, Val(MR), k_block_length, m, transform, load, plane_offset) -end - -@inline function _pack_b_sliver!( - format::FMT, packed::V, source::Tile, kernel::Descriptor{MR, NR, T2}, - transform::F, n::Int, k_block_length::Int - ) where {FMT, V, MR, NR, T2, F} - if _pack_complex_contiguous_eligible( - packed, source.storage, source.cols, transform, format, n, Val(NR), T2 - ) - elembase = @inbounds source.base + source.cols[1] - return _pack_complex_contiguous!( - format, packed, source.storage, elembase, source.rows, Val(NR), k_block_length, transform - ) + k_block_length <= 0 && return panel + BOUNDS && checked_tile_storage_bounds(tile) + if real_contiguous_eligible(tile, spec, transform, lanes) + return pack_real_contiguous!(panel, tile, spec, k_block_length) + elseif complex_contiguous_eligible(tile, spec, transform, lanes) + return pack_complex_contiguous!(panel, tile, spec, k_block_length, transform) end - load = (j, p) -> @inbounds source[p, j] - plane_offset = (plane, j, p) -> packed_b_offset(kernel, j, p, plane) - return _pack_panel!(packed, _element_type(format, T2), format, Val(NR), k_block_length, n, transform, load, plane_offset) + return pack_scalar!(panel, tile, spec, transform, lanes, k_block_length) 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 `k_block_length` 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 `t < valid ? load : zero` would compile to a branch around the load -# and keep the loop scalar. -@inline function _pack_panel!( - packed::V, ::Type{T}, format::FMT, ::Val{PD}, k_block_length::Int, valid::Int, - transform::F, load::L, plane_offset::P - ) where {V, T, FMT <: PackFormat, PD, F, L, P} - if valid == PD +element_type(::SliverSpec{I, L, RealFormat, T}) where {I, L, T} = real(T) +element_type(::SliverSpec{I, L, F, T}) where {I, L, F, T} = T + +# 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 +# `t < lanes ? load : zero` would compile to a branch around the load and keep +# the loop scalar. +@inline function pack_scalar!( + panel::V, tile::Tile, spec::SliverSpec{I, L}, transform::F, lanes::Int, k_block_length::Int + ) where {V, I, L, F} + T = element_type(spec) + if lanes == L @inbounds for p in 1:k_block_length - for t in 1:PD - z = convert(T, transform(load(t, p)))::T - _pack_emit!(packed, format, plane_offset, t, p, z) + for t in 1:L + emit_value!(panel, spec, 0, t, p, convert(T, transform(tile[t, p]))::T) end end else @inbounds for p in 1:k_block_length - for t in 1:valid - z = convert(T, transform(load(t, p)))::T - _pack_emit!(packed, format, plane_offset, t, p, z) + for t in 1:lanes + emit_value!(panel, spec, 0, t, p, convert(T, transform(tile[t, p]))::T) end - for t in (valid + 1):PD - _pack_emit_zero!(packed, format, plane_offset, t, p, real(T)) + for t in (lanes + 1):L + emit_padding!(panel, spec, 0, t, p) end end end - return packed -end - -# Per-lane stores of one element. `real`/`imag` on the loaded element are the -# only accessors: a `Tile` may address scattered storage, so the source is -# never `reinterpret`ed. -@inline function _pack_emit!( - packed::V, ::RealFormat, plane_offset::P, t::Int, p::Int, z::T - ) where {V, P, T} - panel_store!(packed, plane_offset(0, t, p), z) - return nothing -end - -@inline function _pack_emit!( - packed::V, ::PlanarFormat, plane_offset::P, t::Int, p::Int, z::T - ) where {V, P, T} - panel_store!(packed, plane_offset(0, t, p), real(z)) - panel_store!(packed, plane_offset(1, t, p), imag(z)) - return nothing + return panel end -@inline function _pack_emit!( - packed::V, ::InterleavedFormat, plane_offset::P, t::Int, p::Int, z::T - ) where {V, P, T} - panel_store!(packed, plane_offset(0, 2 * t - 1, p), real(z)) - panel_store!(packed, plane_offset(0, 2 * t, p), imag(z)) - return nothing -end - -@inline function _pack_emit!( - packed::V, ::OneEFormat, plane_offset::P, t::Int, p::Int, z::T - ) where {V, P, T} - re = real(z) - im = imag(z) - panel_store!(packed, plane_offset(0, 2 * t - 1, p), re) - panel_store!(packed, plane_offset(0, 2 * t, p), im) - panel_store!(packed, plane_offset(2, 2 * t - 1, p), -im) - panel_store!(packed, plane_offset(2, 2 * t, p), re) - return nothing +# Stores of lane `t` at K step `p`, `base` reals into `panel`. `real`/`imag` +# on the loaded element are the only accessors: a `Tile` may address scattered +# storage, so the source is never `reinterpret`ed. +@inline emit_value!(panel, spec::SliverSpec, base::Int, t::Int, p::Int, z) = + emit!(panel, spec, base, t, p, real(z), imag(z), -imag(z)) +# Literal zeros: the value path would write 1e's `-im` as `-0.0`. +@inline function emit_padding!(panel, spec::SliverSpec, base::Int, t::Int, p::Int) + z = zero(realtype(spec)) + return emit!(panel, spec, base, t, p, z, z, z) end -# Padding stores literal zeros; `_pack_emit!(.., zero(T))` would write 1e's -# `-im` as `-0.0`. -@inline function _pack_emit_zero!( - packed::V, ::RealFormat, plane_offset::P, t::Int, p::Int, ::Type{R} - ) where {V, P, R} - panel_store!(packed, plane_offset(0, t, p), zero(R)) +@inline function emit!(panel, spec::SliverSpec{I, L, RealFormat}, base, t, p, re, im, neg_im) where {I, L} + panel_store!(panel, base + panel_offset(spec, t, p), re) return nothing end -@inline function _pack_emit_zero!( - packed::V, ::PlanarFormat, plane_offset::P, t::Int, p::Int, ::Type{R} - ) where {V, P, R} - panel_store!(packed, plane_offset(0, t, p), zero(R)) - panel_store!(packed, plane_offset(1, t, p), zero(R)) +@inline function emit!(panel, spec::SliverSpec{I, L, PlanarFormat}, base, t, p, re, im, neg_im) where {I, L} + panel_store!(panel, base + panel_offset(spec, t, p), re) + panel_store!(panel, base + panel_offset(spec, t, p, 1), im) return nothing end -@inline function _pack_emit_zero!( - packed::V, ::InterleavedFormat, plane_offset::P, t::Int, p::Int, ::Type{R} - ) where {V, P, R} - panel_store!(packed, plane_offset(0, 2 * t - 1, p), zero(R)) - panel_store!(packed, plane_offset(0, 2 * t, p), zero(R)) +@inline function emit!(panel, spec::SliverSpec{I, L, InterleavedFormat}, base, t, p, re, im, neg_im) where {I, L} + panel_store!(panel, base + panel_offset(spec, 2t - 1, p), re) + panel_store!(panel, base + panel_offset(spec, 2t, p), im) return nothing end -@inline function _pack_emit_zero!( - packed::V, ::OneEFormat, plane_offset::P, t::Int, p::Int, ::Type{R} - ) where {V, P, R} - panel_store!(packed, plane_offset(0, 2 * t - 1, p), zero(R)) - panel_store!(packed, plane_offset(0, 2 * t, p), zero(R)) - panel_store!(packed, plane_offset(2, 2 * t - 1, p), zero(R)) - panel_store!(packed, plane_offset(2, 2 * t, p), zero(R)) +@inline function emit!(panel, spec::SliverSpec{I, L, OneEFormat}, base, t, p, re, im, neg_im) where {I, L} + panel_store!(panel, base + panel_offset(spec, 2t - 1, p), re) + panel_store!(panel, base + panel_offset(spec, 2t, p), im) + panel_store!(panel, base + panel_offset(spec, 2t - 1, p, 2), neg_im) + panel_store!(panel, base + panel_offset(spec, 2t, p, 2), re) return nothing end diff --git a/src/packing/pack_contiguous.jl b/src/packing/pack_contiguous.jl index 0e98e2d..d2410ff 100644 --- a/src/packing/pack_contiguous.jl +++ b/src/packing/pack_contiguous.jl @@ -1,66 +1,80 @@ # Contiguous packing fast paths, each for the one sliver shape it can serve: a -# full sliver whose lane axis is unit-stride, dense storage, a `PackedPanel` -# destination and the driver's `identity`/`conj` transform. They read exactly -# the addresses the scalar loop would, so they skip no validation. +# full sliver whose lane axis is unit-stride, dense storage and the driver's +# `identity`/`conj` transform. They read exactly the addresses the scalar loop +# would, so they skip no validation. using SIMD: shufflevector -# `conj` is the identity on a real element type. -@inline _copies_unchanged(::typeof(identity), ::Type) = true -@inline _copies_unchanged(::typeof(conj), ::Type{T}) where {T <: Real} = true -@inline _copies_unchanged(::Any, ::Type) = false +# --- Vector-path predicates (the microkernels' vector stores use them too) --- + +# Whether an axis steps through storage one element at a time. +is_unit_stride(ax::AffineAxis) = ax.stride == 1 +is_unit_stride(::AbstractVector{Int}) = 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. -@inline function _pack_a_contiguous_eligible( - packed::V, source::Tile, transform::F, m::Int, ::Val{MR}, ::Type{T} - ) where {V, F, MR, T} - return packed isa PackedPanel{T} && _dense_lanes(source.storage, T) && - _copies_unchanged(transform, T) && m == MR && is_unit_stride(source.rows) +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 + +# The deinterleaving complex packer and the planar store pay off only with +# 512-bit vector registers. +@inline complex_fastpath_isa_eligible(profile::TargetProfile) = profile.isa === :avx512 +@inline complex_fastpath_isa_eligible() = complex_fastpath_isa_eligible(target_profile()) + +# --- Real --- + +# `conj` is the identity on a real element type. +@inline copies_unchanged(::typeof(identity), ::Type) = true +@inline copies_unchanged(::typeof(conj), ::Type{T}) where {T <: Real} = true +@inline copies_unchanged(::Any, ::Type) = false + +# Everything but `lanes == L` and the stride test folds at compile time. +@inline function real_contiguous_eligible( + tile::Tile, spec::SliverSpec{I, L, F}, transform, lanes::Int + ) where {I, L, F} + T = realtype(spec) + return F === RealFormat && dense_lanes(tile.storage, T) && + copies_unchanged(transform, T) && lanes == L && is_unit_stride(tile.rows) end -# Real A: one `Vec{MR}` load/convert/store per K step. -@inline function _pack_a_contiguous!( - packed::PackedPanel{T}, storage::DenseVector{S}, rowbase::Int, cols::C, - ::Val{MR}, k_block_length::Int - ) where {T, S, C, MR} +# One `Vec{L}` load/convert/store per K step. +@inline function pack_real_contiguous!( + panel::PackedPanel{T}, tile::Tile{<:DenseVector{S}}, ::SliverSpec{I, L}, k_block_length::Int + ) where {T, S, I, L} + storage = tile.storage + lanebase = @inbounds tile.base + tile.rows[1] GC.@preserve storage begin sp = pointer(storage) - dp = packed.ptr + dp = panel.ptr for p in 1:k_block_length - v = vload(Vec{MR, S}, sp + sizeof(S) * (rowbase + (@inbounds cols[p]))) - vstore(convert(Vec{MR, T}, v), dp + sizeof(T) * (MR * (p - 1))) + v = vload(Vec{L, S}, sp + sizeof(S) * (lanebase + (@inbounds tile.cols[p]))) + vstore(convert(Vec{L, T}, v), dp + sizeof(T) * (L * (p - 1))) end end - return packed + return panel end -# The transforms `_pack_alt` covers; anything else must take the scalar path. -@inline _complex_pack_transform_eligible(::typeof(identity)) = true -@inline _complex_pack_transform_eligible(::typeof(conj)) = true -@inline _complex_pack_transform_eligible(::Any) = false - -# Complex: a deinterleave of the source's native `[re, im, ...]` layout plus, -# for `conj`, a sign flip. `lane_axis` is the axis the packed index runs along -# (`source.rows` for A, `source.cols` for B); unit stride over dense rank-1 -# storage is what makes reading `PD` elements as `2PD` reals a sound bitcast. -# `valid == PD` keeps padding (and its no-`-0.0` rule) on the scalar path. -@inline function _pack_complex_contiguous_eligible( - packed::V, storage::S, lane_axis::AX, transform::F, format::FMT, - valid::Int, ::Val{PD}, ::Type{T} - ) where {V, S, AX, F, FMT, PD, T} - return !(format isa RealFormat) && - packed isa PackedPanel{real(T)} && _dense_lanes(storage, T) && - _complex_pack_transform_eligible(transform) && - valid == PD && is_unit_stride(lane_axis) && +# --- Complex --- + +# The transforms `pack_alt` covers; anything else must take the scalar path. +@inline complex_pack_transform_eligible(::typeof(identity)) = true +@inline complex_pack_transform_eligible(::typeof(conj)) = true +@inline complex_pack_transform_eligible(::Any) = false + +# A deinterleave of the source's native `[re, im, ...]` layout plus, for +# `conj`, a sign flip. Unit stride over dense rank-1 storage is what makes +# reading `L` elements as `2L` reals a sound bitcast. `lanes == L` keeps +# padding (and its no-`-0.0` rule) on the scalar path. +@inline function complex_contiguous_eligible( + tile::Tile, spec::SliverSpec{I, L, F, T}, transform, lanes::Int + ) where {I, L, F, T} + return F !== RealFormat && dense_lanes(tile.storage, T) && + complex_pack_transform_eligible(transform) && + lanes == L && is_unit_stride(tile.rows) && complex_fastpath_isa_eligible() end @@ -68,94 +82,72 @@ end # transform is entirely the choice of `alt`. `-src` is a sign-bit flip, # bit-identical to the scalar `imag(conj(z))` (a multiply by -1 would not be, # on NaN payloads). -@inline _pack_alt(src::Vec, ::typeof(identity)) = src -@inline _pack_alt(src::Vec, ::typeof(conj)) = -src +@inline pack_alt(src::Vec, ::typeof(identity)) = src +@inline pack_alt(src::Vec, ::typeof(conj)) = -src # 1e's second region stores `-im`, so it wants the opposite choice. -@inline _pack_alt_flipped(src::Vec, ::typeof(identity)) = -src -@inline _pack_alt_flipped(src::Vec, ::typeof(conj)) = src +@inline pack_alt_flipped(src::Vec, ::typeof(identity)) = -src +@inline pack_alt_flipped(src::Vec, ::typeof(conj)) = src # `@generated` because `shufflevector` needs a literal `Val` index tuple, which -# `Val(ntuple(...))` is not reliably. `src` holds `PD` elements of one K step. - -# Planar, one whole K step: `[re_0 .. re_{PD-1} | im_0 .. im_{PD-1}]`. -@generated function _planar_pack_shuffle( - src::Vec{N, R}, alt::Vec{N, R}, ::Val{PD} - ) where {N, R, PD} - N == 2 * PD || return :(throw(ArgumentError("_planar_pack_shuffle: expected N == 2PD"))) - idx = ntuple(k -> (k - 1) < PD ? 2 * (k - 1) : N + 2 * ((k - 1) - PD) + 1, 2 * PD) +# `Val(ntuple(...))` is not reliably. `src` holds `L` elements of one K step. + +# Planar, one whole K step: `[re_0 .. re_{L-1} | im_0 .. im_{L-1}]`. +@generated function planar_pack_shuffle( + src::Vec{N, R}, alt::Vec{N, R}, ::Val{L} + ) where {N, R, L} + N == 2 * L || return :(throw(ArgumentError("planar_pack_shuffle: expected N == 2L"))) + idx = ntuple(k -> (k - 1) < L ? 2 * (k - 1) : N + 2 * ((k - 1) - L) + 1, 2 * L) return :(shufflevector(src, alt, Val($idx))) end # 1e's first region (and interleaved): `[re_0, ±im_0, re_1, ±im_1, ...]`. -@generated function _onee_pack_shuffle_a( - src::Vec{N, R}, alt::Vec{N, R}, ::Val{PD} - ) where {N, R, PD} - N == 2 * PD || return :(throw(ArgumentError("_onee_pack_shuffle_a: expected N == 2PD"))) - idx = ntuple(k -> iseven(k) ? N + (k - 1) : (k - 1), 2 * PD) +@generated function onee_pack_shuffle_a( + src::Vec{N, R}, alt::Vec{N, R}, ::Val{L} + ) where {N, R, L} + N == 2 * L || return :(throw(ArgumentError("onee_pack_shuffle_a: expected N == 2L"))) + idx = ntuple(k -> iseven(k) ? N + (k - 1) : (k - 1), 2 * L) return :(shufflevector(src, alt, Val($idx))) end # 1e's second region: `[∓im_0, re_0, ∓im_1, re_1, ...]`. -@generated function _onee_pack_shuffle_b( - src::Vec{N, R}, alt::Vec{N, R}, ::Val{PD} - ) where {N, R, PD} - N == 2 * PD || return :(throw(ArgumentError("_onee_pack_shuffle_b: expected N == 2PD"))) - idx = ntuple(k -> isodd(k) ? N + k : k - 2, 2 * PD) +@generated function onee_pack_shuffle_b( + src::Vec{N, R}, alt::Vec{N, R}, ::Val{L} + ) where {N, R, L} + N == 2 * L || return :(throw(ArgumentError("onee_pack_shuffle_b: expected N == 2L"))) + idx = ntuple(k -> isodd(k) ? N + k : k - 2, 2 * L) return :(shufflevector(src, alt, Val($idx))) end -_single_region_shuffle(::PlanarFormat) = _planar_pack_shuffle -_single_region_shuffle(::InterleavedFormat) = _onee_pack_shuffle_a - -# Lane `t` of K step `p` is element `elembase + steps[p] + (t - 1)`; -# only the lane axis must be unit-stride, `steps` may be scattered. Pinning -# `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{RS}}, elembase::Int, steps::C, ::Val{PD}, k_block_length::Int, - transform::F - ) where {R, RS, C, PD, F} - shuffle = _single_region_shuffle(format) - GC.@preserve storage begin - sp = reinterpret(Ptr{RS}, pointer(storage)) - dp = packed.ptr - for p in 1:k_block_length - src = convert( - Vec{2 * PD, R}, - vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + (@inbounds steps[p])))) - ) - vstore( - shuffle(src, _pack_alt(src, transform), Val(PD)), - dp + sizeof(R) * (2 * PD * (p - 1)) - ) - end - end - return packed +# One K step's `2L` reals `src`, stored at `dp`. +@inline store_step!(dp::Ptr, src::Vec, ::SliverSpec{I, L, PlanarFormat}, transform) where {I, L} = + vstore(planar_pack_shuffle(src, pack_alt(src, transform), Val(L)), dp) +@inline store_step!(dp::Ptr, src::Vec, ::SliverSpec{I, L, InterleavedFormat}, transform) where {I, L} = + vstore(onee_pack_shuffle_a(src, pack_alt(src, transform), Val(L)), dp) +@inline function store_step!(dp::Ptr{R}, src::Vec, ::SliverSpec{I, L, OneEFormat}, transform) where {R, I, L} + vstore(onee_pack_shuffle_a(src, pack_alt(src, transform), Val(L)), dp) + vstore(onee_pack_shuffle_b(src, pack_alt_flipped(src, transform), Val(L)), dp + sizeof(R) * 2L) + return nothing end -@inline function _pack_complex_contiguous!( - ::OneEFormat, packed::PackedPanel{R}, storage::DenseVector{Complex{RS}}, - elembase::Int, steps::C, ::Val{PD}, k_block_length::Int, transform::F - ) where {R, RS, C, PD, F} +# Lane `t` of K step `p` is element `tile.base + tile.rows[1] + tile.cols[p] + +# (t - 1)`; only the lane axis must be unit-stride, the K steps may be +# scattered. Pinning `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!( + panel::PackedPanel{R}, tile::Tile{<:DenseVector{Complex{RS}}}, spec::SliverSpec{I, L}, + k_block_length::Int, transform::F + ) where {R, RS, I, L, F} + storage = tile.storage + elembase = @inbounds tile.base + tile.rows[1] GC.@preserve storage begin sp = reinterpret(Ptr{RS}, pointer(storage)) - dp = packed.ptr for p in 1:k_block_length src = convert( - Vec{2 * PD, R}, - vload(Vec{2 * PD, RS}, sp + sizeof(RS) * (2 * (elembase + (@inbounds steps[p])))) - ) - at = 4 * PD * (p - 1) - vstore( - _onee_pack_shuffle_a(src, _pack_alt(src, transform), Val(PD)), - dp + sizeof(R) * at - ) - vstore( - _onee_pack_shuffle_b(src, _pack_alt_flipped(src, transform), Val(PD)), - dp + sizeof(R) * (at + 2 * PD) + Vec{2 * L, R}, + vload(Vec{2 * L, RS}, sp + sizeof(RS) * (2 * (elembase + (@inbounds tile.cols[p])))) ) + store_step!(panel.ptr + sizeof(R) * panel_offset(spec, 1, p), src, spec, transform) end end - return packed + return panel end diff --git a/src/packing/panel.jl b/src/packing/panel.jl index 1a065e7..661673f 100644 --- a/src/packing/panel.jl +++ b/src/packing/panel.jl @@ -34,17 +34,16 @@ Base.eltype(::Type{PackedPanel{T}}) where {T} = T # 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, k_block_length::Int, split - ) where {PO, FMT <: PackFormat, PK, DS, R, ST, KA <: AbstractVector{Int}, F} + buffer::PK, spec::SliverSpec{I, R}, storage::ST, base::Int, fbuf::Vector{Int}, + kaxis::KA, transform::F, fcount::Int, k_block_length::Int, split + ) where {PK, I, R, ST, KA <: AbstractVector{Int}, 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 * k_block_length, transform(storage[kbase + fbuf[i + 1]]), i, p + buffer, spec, sliver_width(spec) * k_block_length, 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 1:k_block_length @@ -63,10 +62,9 @@ Base.eltype(::Type{PackedPanel{T}}) where {T} = T end valid = fcount % R if valid != 0 - rb = (fcount ÷ R) * Rp * k_block_length - pad = (plane, idx, pp) -> rb + plane_offset(descriptor, idx, pp, plane) + rb = (fcount ÷ R) * sliver_width(spec) * k_block_length for p in 1:k_block_length, t in (valid + 1):R - _pack_emit_zero!(buffer, format, pad, t, p, realtype(descriptor)) + emit_padding!(buffer, spec, rb, t, p) end end return nothing @@ -75,10 +73,9 @@ end # Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane # `i % R + 1`) 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} + buffer::PK, spec::SliverSpec{I, R}, panel::Int, x, i::Int, p::Int + ) where {PK, I, R} r, t = divrem(i, R) - po = (plane, idx, pp) -> r * panel + plane_offset(descriptor, idx, pp, plane) - _pack_emit!(buffer, format, po, t + 1, p, convert(_element_type(format, scalartype(descriptor)), x)) + emit_value!(buffer, spec, r * panel, t + 1, p, convert(element_type(spec), x)) return nothing end diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index c50c959..c0cf5af 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -1,5 +1,5 @@ # The stages driven by hand, end to end: AxisGroup -> BlockDescriptor -> tile -# axes -> pack_a!/pack_b! -> ScalarKernel execute_tile!, with beta applied on +# axes -> pack! -> ScalarKernel execute_tile!, with beta applied on # the first K panel only. @testset "manual pipeline: AxisGroup -> tiles -> packing -> ScalarKernel" begin @@ -31,8 +31,8 @@ copyto!(Cstorage, vec(Cstart)) for (idx, (first, len)) in enumerate(panels) (dK_A, dK_B) = block_descriptors!(kb, K, first, len) - pack_a!(packed_a, Tile(vec(A), 0, row_A, axis(dK_A, kb[1])), kernel, identity) - pack_b!(packed_b, Tile(vec(B), 0, axis(dK_B, kb[2]), col_B), kernel, identity) + pack!(packed_a, Tile(vec(A), 0, row_A, axis(dK_A, kb[1])), sliver_spec(kernel, 1), identity) + pack!(packed_b, Tile(vec(B), 0, col_B, axis(dK_B, kb[2])), sliver_spec(kernel, 2), identity) execute_tile!(kernel, destination, packed_a, packed_b, len, alpha, idx == 1 ? beta : 1.0) end @test reshape(Cstorage, size(Cref)) ≈ alpha .* Cref .+ beta .* Cstart @@ -46,8 +46,8 @@ end packed = zeros(4) for base in (10_000, -10_000) src = Tile([1.0, 2.0, 3.0, 4.0], base, AffineAxis(0, 1, 2), AffineAxis(0, 2, 2)) - @test_throws BoundsError pack_a!(packed, src, k, identity) - @test_throws BoundsError pack_b!(packed, src, k, identity) + @test_throws BoundsError pack!(packed, src, sliver_spec(k, 1), identity) + @test_throws BoundsError pack!(packed, src, sliver_spec(k, 2), identity) end @test all(iszero, packed) canary = fill(999.0, 4) diff --git a/test/helpers.jl b/test/helpers.jl index 9eadde8..7cf9e68 100644 --- a/test/helpers.jl +++ b/test/helpers.jl @@ -66,7 +66,7 @@ end # Steady-state allocation of one `pack!` call. `::F where {F}` forces # specialization on the pass-through `transform`; without it the dynamic call # allocates on Julia 1.10. -function steady_pack_allocs(pack!::P, packed, source, kernel, transform::F) where {P, F} - pack!(packed, source, kernel, transform) - return @allocated pack!(packed, source, kernel, transform) +function steady_pack_allocs(packed, source, spec, transform::F) where {F} + pack!(packed, source, spec, transform) + return @allocated pack!(packed, source, spec, transform) end diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 2bdee37..6e964a4 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -108,7 +108,7 @@ const _QSF = QuasiStrided end end - @testset "pack_a! InterleavedFormat: scalar loop and fast path vs local layout" begin + @testset "pack! InterleavedFormat A: scalar loop and fast path vs local layout" begin # Bitwise (`isequal` separates -0.0): packing is a copy, or a sign flip under conj. for T in (ComplexF64, ComplexF32), (MR, NR, W) in (kernel_shapes(T, FMAddSubMethod())..., (3, 2, 2)) R = real(T) @@ -122,21 +122,13 @@ const _QSF = QuasiStrided A = zeros(T, MR, k_block_length) A[1:m, :] = f.(reshape(vals[(base + 1):(base + lda * k_block_length)], lda, k_block_length)[1:m, :]) want = mk_pack_a(k, A) - bufv = fill(R(-777), packed_a_length(k, k_block_length)) # a Vector: the scalar loop - pack_a!(bufv, src, k, f) - @test isequal(bufv, want) - bufp = fill(R(-777), packed_a_length(k, k_block_length)) # a PackedPanel: the fast path where gated in - GC.@preserve bufp pack_a!(packed_panel(bufp, 1, length(bufp)), src, k, f) + bufp = fill(R(-777), packed_a_length(k, k_block_length)) + pack!(bufp, src, sliver_spec(k, 1), f) @test isequal(bufp, want) end src = Tile(vals, base, AffineAxis(0, 1, MR), AffineAxis(0, lda, k_block_length)) - bufp = zeros(R, packed_a_length(k, k_block_length)) - GC.@preserve bufp begin - pp = packed_panel(bufp, 1, length(bufp)) - @test _QSF._pack_complex_contiguous_eligible( - pp, vals, src.rows, identity, InterleavedFormat(), MR, Val(MR), T - ) == _QSF.complex_fastpath_isa_eligible() - end + @test _QSF.complex_contiguous_eligible(src, sliver_spec(k, 1), identity, MR) == + _QSF.complex_fastpath_isa_eligible() end end diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index b7dc8aa..e92dcb8 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -66,8 +66,8 @@ end 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) + pack!(dpa, srcA, sliver_spec(k, 1), f) + pack!(dpb, transpose(srcB), sliver_spec(k, 2), f) end @test isequal((pa, pb), mk_pack(k, TA.(A), TB.(B))) alpha, beta = T(1.5, -0.5), T(0.25, 1) diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 02ca01c..2fb8900 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -102,14 +102,6 @@ end @test eligible(zeros(ComplexF32, m * n), AffineAxis(0, 1, m)) == STORE_FASTPATH_ON @test !eligible(zeros(Float64, m * n), AffineAxis(0, 1, m)) end - - @testset "ISA gate: AVX-512 only" begin - @test QuasiStrided.complex_fastpath_isa_eligible(synthetic(:avx512)) - for key in (:avx2, :neon, :unknown) - @test !QuasiStrided.complex_fastpath_isa_eligible(synthetic(key)) - end - @test STORE_FASTPATH_ON == (target_profile().isa === :avx512) - end end diff --git a/test/packing/test_pack_complex.jl b/test/packing/test_pack_complex.jl index 2eb59ca..98bc69a 100644 --- a/test/packing/test_pack_complex.jl +++ b/test/packing/test_pack_complex.jl @@ -19,7 +19,7 @@ const COMPLEX_FORMATS = ( (InterleavedFormat(), PlanarFormat()), (OneEFormat(), OneEFormat()), ) -@testset "complex pack_a!/pack_b! ($T, $(typeof(fa)), $(typeof(fb))) vs reference" for +@testset "complex pack! ($T, $(typeof(fa)), $(typeof(fb))) vs reference" for T in (ComplexF64, ComplexF32), (fa, fb) in COMPLEX_FORMATS MR, NR = 4, 3 @@ -28,7 +28,7 @@ const COMPLEX_FORMATS = ( storage = complex_storage(T, 4000) calls = Ref(0) counting = z -> (calls[] += 1; 2z + one(T)) # nonzero at zero - for (side, pack!, PD, fmt) in ((:a, pack_a!, MR, fa), (:b, pack_b!, NR, fb)), k_block_length in (1, 5) + for (i, PD, fmt) in ((1, MR, fa), (2, NR, fb)), k_block_length in (1, 5) fixtures = ( (AffineAxis(0, 1, PD), AffineAxis(0, 97, k_block_length)), (AffineAxis(40, -1, PD), AffineAxis(3000, -97, k_block_length)), @@ -36,13 +36,13 @@ const COMPLEX_FORMATS = ( ) len = _ref_rpe(fmt) * PD * k_block_length for (lane, step) in fixtures, valid in (PD, 1, 0) - src, g = pack_fixture(side, storage, 500, resized(lane, valid), step) + src, g = pack_fixture(storage, 500, resized(lane, valid), step) for f in (identity, conj) - got, _ = pack_into(pack!, :vector, R, len, src, kernel, f) + got, _ = pack_into(:vector, R, len, src, sliver_spec(kernel, i), f) @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, f))) end calls[] = 0 - got, canaries = pack_into(pack!, :panel, R, len, src, kernel, counting) + got, canaries = pack_into(:panel, R, len, src, sliver_spec(kernel, i), counting) @test all(isequal.(got, ref_pack(fmt, T, PD, k_block_length, valid, g, z -> 2z + one(T)))) @test calls[] == valid * k_block_length # once per element, not per real half @test canaries @@ -53,24 +53,26 @@ end @testset "complex packing: validation before any write, k_block_length == 0 is a no-op" begin kernel = Descriptor(Val(4), Val(3), ComplexF64, PlanarFormat(), PlanarFormat()) storage = fill(ComplexF64(3, 4), 100) - src = Tile(storage, 0, AffineAxis(0, 1, 4), AffineAxis(0, 8, 2)) + src = Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 8, 2)) + a, b = sliver_spec(kernel, 1), sliver_spec(kernel, 2) # The buffer holds realtype(kernel), not scalartype(kernel). - @test_throws ArgumentError pack_a!(zeros(ComplexF64, 1000), src, kernel, identity) - @test_throws ArgumentError pack_b!(zeros(ComplexF64, 1000), src, kernel, identity) - @test_throws ArgumentError pack_a!(zeros(Float32, 1000), src, kernel, identity) + @test_throws ArgumentError pack!(zeros(ComplexF64, 1000), src, a, identity) + @test_throws ArgumentError pack!(zeros(ComplexF64, 1000), src, b, identity) + @test_throws ArgumentError pack!(zeros(Float32, 1000), src, a, identity) canary = fill(-42.0, 1000) - @test_throws ArgumentError pack_a!(canary, Tile(storage, 0, AffineAxis(0, 1, 5), src.cols), kernel, identity) - @test_throws ArgumentError pack_b!(canary, Tile(storage, 0, src.cols, AffineAxis(0, 1, 4)), kernel, identity) + @test_throws ArgumentError pack!(canary, Tile(storage, 0, AffineAxis(0, 1, 5), src.cols), a, identity) + @test_throws ArgumentError pack!(canary, Tile(storage, 0, AffineAxis(0, 1, 4), src.cols), b, identity) # A buffer sized as if it held complex elements is half as long as needed. - @test_throws DimensionMismatch pack_a!(zeros(Float64, 4 * 2), src, kernel, identity) - @test_throws BoundsError pack_a!(canary, Tile(storage, 90, src.rows, src.cols), kernel, identity) + @test_throws DimensionMismatch pack!(zeros(Float64, 4 * 2), src, a, identity) + @test_throws BoundsError pack!(canary, Tile(storage, 90, src.rows, src.cols), a, identity) @test all(==(-42.0), canary) read = Ref(false) spy = z -> (read[] = true; z) packed = fill(-42.0, 8) - @test pack_a!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 8, 0)), kernel, spy) === packed - @test pack_b!(packed, Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 8, 3)), kernel, spy) === packed + for spec in (a, b) + @test pack!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 8, 0)), spec, spy) === packed + end @test all(==(-42.0), packed) @test !read[] end @@ -81,38 +83,26 @@ end function run(::Type{T}, fa, fb) where {T} MR, NR, k_block_length = 8, 6, 4 kernel = Descriptor(Val(MR), Val(NR), T, fa, fb) - R = real(T) storage = rand(T, 4000) - bufa = zeros(R, packed_a_length(kernel, k_block_length)) - bufb = zeros(R, packed_b_length(kernel, k_block_length)) - ro, co, ro_b = collect(0:(MR - 1)) .* 3, [0, 100, 250, 400], collect(0:(NR - 1)) .* 5 + steps, lanes = [0, 100, 250, 400], collect(0:(max(MR, NR) - 1)) .* 3 fast = Int[] slow = Int[] - GC.@preserve bufa bufb begin - pa = packed_panel(bufa, 1, length(bufa)) - pb = packed_panel(bufb, 1, length(bufb)) - full_a = Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)) - full_b = Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)) - for f in (identity, conj) # the contiguous fast path, where the ISA has it - push!(fast, steady_pack_allocs(pack_a!, pa, full_a, kernel, f)) - push!(fast, steady_pack_allocs(pack_b!, pb, full_b, kernel, f)) - end - srcs_a = ( - full_a, Tile(storage, 0, view(ro, 1:MR), view(co, 1:k_block_length)), - Tile(storage, 0, AffineAxis(0, 1, MR - 1), AffineAxis(0, MR, k_block_length)), - Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), - ) - srcs_b = ( - full_b, Tile(storage, 0, view(co, 1:k_block_length), view(ro_b, 1:NR)), - Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR - 1)), + for (i, L) in ((1, MR), (2, NR)) + spec = sliver_spec(kernel, i) + buf = zeros(real(T), packed_length(spec, k_block_length)) + full = Tile(storage, 0, AffineAxis(0, 1, L), AffineAxis(0, L, k_block_length)) + srcs = ( + full, Tile(storage, 0, view(lanes, 1:L), view(steps, 1:k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, L - 1), AffineAxis(0, L, k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, L), AffineAxis(0, L, 0)), ) - for s in srcs_a, f in (identity, conj) - push!(slow, steady_pack_allocs(pack_a!, bufa, s, kernel, f)) + GC.@preserve buf for f in (identity, conj) # the contiguous fast path, where the ISA has it + push!(fast, steady_pack_allocs(packed_panel(buf, 1, length(buf)), full, spec, f)) end - for s in srcs_b, f in (identity, conj) - push!(slow, steady_pack_allocs(pack_b!, bufb, s, kernel, f)) + for s in srcs, f in (identity, conj) + push!(slow, steady_pack_allocs(buf, s, spec, f)) end - push!(slow, steady_pack_allocs(pack_a!, bufa, full_a, kernel, z -> 2z + oneunit(z))) + push!(slow, steady_pack_allocs(buf, full, spec, z -> 2z + oneunit(z))) end return fast, slow end diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index 7dab653..e379bed 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -1,13 +1,13 @@ -# The vectorized complex packing fast path (`_pack_complex_contiguous!`), +# The vectorized complex packing fast path (`pack_complex_contiguous!`), # checked bitwise against test_pack_real.jl's reference layouts. Gate # expectations derive from the live profile, so test/forced_isa_runner.jl # checks the same outputs with the fast path off. using QuasiStrided: target_profile, unknown_target, TargetProfile, CacheLevel, - kernel_shapes, PlanarMethod, OneMMethod, FMAddSubMethod + kernel_shapes, PlanarMethod, OneMMethod, FMAddSubMethod, SliverSpec, complex_contiguous_eligible, + complex_fastpath_isa_eligible -const QS = QuasiStrided -const FASTPATH_ON = QS.complex_fastpath_isa_eligible() +const FASTPATH_ON = complex_fastpath_isa_eligible() # The extremes of each method's menu plus extents that are not multiples of any # lane width (interleaved shares 1e's shuffle); B is planar under every method. @@ -24,12 +24,11 @@ const FAST_B_NRS = (1, 3, 7, 8) # Packs into a `PackedPanel` (the only destination the gate admits). # `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) +function check_panel(spec, 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! + src, g = pack_fixture(storage, 17, lane, step) len = _ref_rpe(fmt) * PD * length(step) - got, canaries = pack_into(pack!, :panel, real(T), len, src, kernel, f) + got, canaries = pack_into(:panel, real(T), len, src, spec, f) @test canaries @test all(isequal.(got, ref_pack(fmt, T, PD, length(step), length(lane), g, f))) return storage, src @@ -38,20 +37,16 @@ end @testset "complex pack fast path: A panel, $T / $(typeof(fa)) / MR=$MR" for (T, fa, MR) in FAST_A_CASES kernel = Descriptor(Val(MR), Val(3), T, fa, PlanarFormat()) 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 + storage, src = check_panel(sliver_spec(kernel, 1), T, MR, fa, AffineAxis(0, 1, MR), AffineAxis(0, 997, 5), f; S) + @test complex_contiguous_eligible(src, sliver_spec(kernel, 1), f, MR) == FASTPATH_ON end end @testset "complex pack fast path: B panel, $T / NR=$NR" for T in (ComplexF64, ComplexF32), NR in FAST_B_NRS kernel = Descriptor(Val(4), Val(NR), T, PlanarFormat(), PlanarFormat()) 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 + storage, src = check_panel(sliver_spec(kernel, 2), T, NR, PlanarFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, 5), f; S) + @test complex_contiguous_eligible(src, sliver_spec(kernel, 2), f, NR) == FASTPATH_ON end end @@ -61,11 +56,11 @@ end for fa in (PlanarFormat(), OneEFormat(), InterleavedFormat()), f in (identity, conj) kernel = Descriptor(Val(MR), Val(NR), T, fa, OneEFormat()) # The lane axis must be unit-stride; the step axis may scatter. - check_panel(:a, kernel, T, MR, fa, AffineAxis(0, 1, MR), view(koffs, 1:k_block_length), f) - check_panel(:b, kernel, T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, k_block_length), f) + check_panel(sliver_spec(kernel, 1), T, MR, fa, AffineAxis(0, 1, MR), view(koffs, 1:k_block_length), f) + check_panel(sliver_spec(kernel, 2), T, NR, OneEFormat(), AffineAxis(0, 1, NR), AffineAxis(0, 997, k_block_length), f) # Ineligible shapes fall back to the scalar loop, into the same panel. for lane in (AffineAxis(20, -1, MR), AffineAxis(0, 3, MR), AffineAxis(0, 1, MR - 3)) - check_panel(:a, kernel, T, MR, fa, lane, AffineAxis(3000, -97, k_block_length), f) + check_panel(sliver_spec(kernel, 1), T, MR, fa, lane, AffineAxis(3000, -97, k_block_length), f) end end end @@ -74,38 +69,31 @@ end T = ComplexF64 MR = 8 storage = complex_storage(T, 4000) - buf = zeros(Float64, 64) offs = collect(0:(MR - 1)) - GC.@preserve buf begin - pk = packed_panel(buf, 1, length(buf)) - elig(; - packed = pk, store = storage, ax = AffineAxis(0, 1, MR), tr = identity, - fmt = PlanarFormat(), valid = MR, - ) = QS._pack_complex_contiguous_eligible(packed, store, ax, tr, fmt, valid, Val(MR), T) + elig(; + store = storage, ax = AffineAxis(0, 1, MR), tr = identity, fmt = PlanarFormat(), valid = MR, + ) = complex_contiguous_eligible(Tile(store, 0, ax, AffineAxis(0, MR, 2)), SliverSpec{1, MR, typeof(fmt), T}(), tr, valid) - @test elig() == FASTPATH_ON - @test elig(tr = conj) == FASTPATH_ON - @test elig(fmt = OneEFormat()) == FASTPATH_ON - @test elig(fmt = InterleavedFormat()) == FASTPATH_ON - @test !elig(packed = buf) # not a PackedPanel - @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 = view(offs, 1:MR)) - @test !elig(valid = MR - 1) # padding lanes - @test !elig(tr = z -> 2z) - @test !elig(fmt = RealFormat()) - end + @test elig() == FASTPATH_ON + @test elig(tr = conj) == FASTPATH_ON + @test elig(fmt = OneEFormat()) == FASTPATH_ON + @test elig(fmt = InterleavedFormat()) == FASTPATH_ON + @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 = view(offs, 1:MR)) + @test !elig(valid = MR - 1) # padding lanes + @test !elig(tr = z -> 2z) + @test !elig(fmt = RealFormat()) end @testset "complex pack fast path: ISA gate is a register-width question" begin - @test QS.complex_fastpath_isa_eligible(synthetic(:avx512)) + @test complex_fastpath_isa_eligible(synthetic(:avx512)) for key in (:avx2, :neon, :unknown) - @test !QS.complex_fastpath_isa_eligible(synthetic(key)) + @test !complex_fastpath_isa_eligible(synthetic(key)) end @test FASTPATH_ON == (target_profile().isa === :avx512) end diff --git a/test/packing/test_pack_real.jl b/test/packing/test_pack_real.jl index 053f35b..081e372 100644 --- a/test/packing/test_pack_real.jl +++ b/test/packing/test_pack_real.jl @@ -1,10 +1,9 @@ -# pack_a!/pack_b! for real kernels against direct storage indexing. The -# reference layouts and fixtures here are shared with the two complex packing -# files included after this one. +# pack! for real kernels against direct storage indexing. The reference +# layouts and fixtures here are shared with the two complex packing files +# included after this one. using QuasiStrided: RealFormat, PlanarFormat, OneEFormat, InterleavedFormat, PackedPanel, - packed_panel, _copies_unchanged, - _pack_a_contiguous_eligible + packed_panel, packed_length, copies_unchanged, real_contiguous_eligible # Packed layouts written from the format definitions, not from the offset # helpers. `g(t, p)` is the source element at lane `t`, K step `p`; lanes @@ -37,25 +36,24 @@ end resized(ax::AffineAxis, n) = AffineAxis(ax.base, ax.stride, n) resized(ax::SubArray, n) = view(parent(ax), 1:n) -# An A (`lane` = rows) or B (`lane` = cols) source over `storage`, and its -# direct-indexing reader. -function pack_fixture(side, storage, base, lane, step) - src = side === :a ? Tile(storage, base, lane, step) : Tile(storage, base, step, lane) - return src, (t, p) -> storage[base + lane[t + 1] + step[p + 1] + 1] +# A sliver source with lanes along `lane` and K steps along `step` (B's tile +# transposed), and its direct-indexing reader. +function pack_fixture(storage, base, lane, step) + return Tile(storage, base, lane, step), (t, p) -> storage[base + lane[t + 1] + step[p + 1] + 1] end # Packs `src` into a destination of kind `dst` (canaried unless a bare Vector) # and returns the packed prefix and whether the canaries survived. -function pack_into(pack!, dst, R, len, src, kernel, f) +function pack_into(dst, R, len, src, spec, f) buf = fill(R(-777), len + 12) if dst === :vector v = buf[1:len] - @test pack!(v, src, kernel, f) === v + @test pack!(v, src, spec, f) === v return v, true end GC.@preserve buf begin d = packed_panel(buf, 5, len) - @test pack!(d, src, kernel, f) === d + @test pack!(d, src, spec, f) === d end return buf[5:(4 + len)], all(==(R(-777)), buf[1:4]) && all(==(R(-777)), buf[(5 + len):end]) end @@ -63,9 +61,7 @@ end const SCATTER_LANES = [5, 30, 1, 17, 9, 44, 2, 23, 11, 38, 7, 60, 14, 51, 3, 29] const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] -@testset "pack_a!/pack_b! ($T): every stride kind and tail width vs direct indexing" for - T in (Float64, Float32) - +@testset "pack! ($T): every stride kind and tail width vs direct indexing" for T in (Float64, Float32) MR, NR, k_block_length = 8, 6, 5 kernel = Descriptor(Val(MR), Val(NR), T) storage = T.(collect(1.0:2000.0)) @@ -76,84 +72,94 @@ const SCATTER_STEPS = [7, 900, 300, 1500, 60, 1210, 420] view(SCATTER_LANES, 1:8), ] steps = Any[AffineAxis(0, 64, k_block_length), AffineAxis(700, -97, k_block_length), view(SCATTER_STEPS, 1:k_block_length)] - for (side, pack!, PD) in ((:a, pack_a!, MR), (:b, pack_b!, NR)), - lane in lanes, step in steps, valid in (PD, 1, 0), dst in (:vector, :panel) - src, g = pack_fixture(side, storage, 11, resized(lane, valid), step) + for (spec, L) in ((sliver_spec(kernel, 1), MR), (sliver_spec(kernel, 2), NR)), + lane in lanes, step in steps, valid in (L, 1, 0), dst in (:vector, :panel) + src, g = pack_fixture(storage, 11, resized(lane, valid), step) calls[] = 0 - got, canaries = pack_into(pack!, dst, T, PD * k_block_length, src, kernel, counting) - @test got == ref_pack(RealFormat(), T, PD, k_block_length, valid, g, x -> 3x + 1000) + got, canaries = pack_into(dst, T, L * k_block_length, src, spec, counting) + @test got == ref_pack(RealFormat(), T, L, k_block_length, valid, g, x -> 3x + 1000) @test calls[] == valid * k_block_length @test canaries end # Every tail width, at the minimal and a longer K depth. negate = x -> -x - for (side, pack!, PD) in ((:a, pack_a!, MR), (:b, pack_b!, NR)), valid in 0:PD, kc1 in (1, 5) - src, g = pack_fixture(side, storage, 0, AffineAxis(2, 3, valid), AffineAxis(0, 40, kc1)) - got, _ = pack_into(pack!, :vector, T, PD * kc1, src, kernel, negate) - @test got == ref_pack(RealFormat(), T, PD, kc1, valid, g, negate) + for (spec, L) in ((sliver_spec(kernel, 1), MR), (sliver_spec(kernel, 2), NR)), valid in 0:L, kc1 in (1, 5) + src, g = pack_fixture(storage, 0, AffineAxis(2, 3, valid), AffineAxis(0, 40, kc1)) + got, _ = pack_into(:vector, T, L * kc1, src, spec, negate) + @test got == ref_pack(RealFormat(), T, L, kc1, valid, g, negate) end end -@testset "pack_a! contiguous fast path: fires exactly when eligible ($T, MR=$MR)" for - T in (Float64, Float32), MR in (4, 16) +@testset "pack!: B packs as A packs its transposed tile" begin + kernel = Descriptor(Val(6), Val(6), Float64) + storage = collect(1.0:2000.0) + tile_b = Tile(storage, 3, view(SCATTER_STEPS, 1:5), AffineAxis(0, 1, 6)) # K x N + got_b, _ = pack_into(:panel, Float64, 30, transpose(tile_b), sliver_spec(kernel, 2), identity) + got_a, _ = pack_into(:panel, Float64, 30, transpose(tile_b), sliver_spec(kernel, 1), identity) + @test got_b == got_a +end + +@testset "pack! contiguous fast path: fires exactly when eligible ($T, L=$L, side $i)" for + T in (Float64, Float32), (L, i) in ((4, 1), (16, 1), (6, 2)) - kernel = Descriptor(Val(MR), Val(3), T) + kernel = Descriptor(Val(L), Val(L), T) vals = T.(collect(1.0:2000.0)) 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)) + contig = collect(0:(L - 1)) + spec = sliver_spec(kernel, i) for storage in storages - steps = Any[AffineAxis(0, MR, 5), AffineAxis(1800, -MR, 6), view(koffs, 1:4)] + steps = Any[AffineAxis(0, L, 5), AffineAxis(1800, -L, 6), view(koffs, 1:4)] lanes = Any[ - (AffineAxis(0, 1, MR), true), (AffineAxis(5, 1, MR), true), - (AffineAxis(0, 2, MR), false), (AffineAxis(MR + 3, -1, MR), false), - (AffineAxis(0, 1, MR - 1), false), (view(contig, 1:MR), false), + (AffineAxis(0, 1, L), true), (AffineAxis(5, 1, L), true), + (AffineAxis(0, 2, L), false), (AffineAxis(L + 3, -1, L), false), + (AffineAxis(0, 1, L - 1), false), (view(contig, 1:L), false), ] for (lane, eligible) in lanes, step in steps, f in (identity, conj, x -> -x) - src, g = pack_fixture(:a, storage, 11, lane, step) + src, g = pack_fixture(storage, 11, lane, step) k_block_length = length(step) - pp = packed_panel(zeros(T, 1), 1, 1) - @test _pack_a_contiguous_eligible(pp, src, f, length(lane), Val(MR), T) == + @test real_contiguous_eligible(src, spec, f, length(lane)) == (eligible && (f === identity || f === conj)) - @test !_pack_a_contiguous_eligible(zeros(T, 1), src, f, length(lane), Val(MR), T) - got, canaries = pack_into(pack_a!, :panel, T, MR * k_block_length, src, kernel, f) - @test got == ref_pack(RealFormat(), T, MR, k_block_length, length(lane), g, f) + got, canaries = pack_into(:panel, T, L * k_block_length, src, spec, f) + @test got == ref_pack(RealFormat(), T, L, k_block_length, length(lane), g, f) @test canaries end end - @test _copies_unchanged(conj, T) - @test !_copies_unchanged(conj, complex(T)) # conj is not the identity on complex + @test copies_unchanged(conj, T) + @test !copies_unchanged(conj, complex(T)) # conj is not the identity on complex end -@testset "pack_a!/pack_b!: k_block_length == 0 reads and writes nothing" begin +@testset "pack!: k_block_length == 0 reads and writes nothing" begin kernel = Descriptor(Val(4), Val(3), Float64) storage = fill(3.0, 10) read = Ref(false) spy = x -> (read[] = true; x) packed = fill(-42.0, 8) - @test pack_a!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 1, 0)), kernel, spy) === packed - @test pack_b!(packed, Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, 1, 3)), kernel, spy) === packed + for i in 1:2 + @test pack!(packed, Tile(storage, 0, AffineAxis(0, 1, 3), AffineAxis(0, 1, 0)), sliver_spec(kernel, i), spy) === packed + end @test packed == fill(-42.0, 8) @test !read[] end -@testset "pack_a!/pack_b!: invalid metadata rejected before mutation" begin +@testset "pack!: invalid metadata rejected before mutation" begin kernel = Descriptor(Val(4), Val(3), Float64) kernel32 = Descriptor(Val(4), Val(3), Float32) storage = fill(9.0, 20) tile(m, n, base = 0) = Tile(storage, base, AffineAxis(0, 1, m), AffineAxis(0, 1, n)) - for (pack!, bad, ok, short) in ((pack_a!, tile(5, 3), tile(4, 3), 11), (pack_b!, tile(3, 4), tile(3, 3), 8)) + for (i, bad, ok, short) in ((1, tile(5, 3), tile(4, 3), 11), (2, tile(4, 3), tile(3, 3), 8)) + spec = sliver_spec(kernel, i) canary = fill(-1.0, 100) - @test_throws ArgumentError pack!(canary, bad, kernel, identity) # m > MR / n > NR - @test_throws ArgumentError pack!(canary, ok, kernel32, identity) # eltype mismatch - @test_throws DimensionMismatch pack!(fill(-1.0, short), ok, kernel, identity) - @test_throws BoundsError pack!(canary, tile(size(ok)..., 18), kernel, identity) + @test_throws ArgumentError pack!(canary, bad, spec, identity) # lanes > L + @test_throws ArgumentError pack!(canary, ok, sliver_spec(kernel32, i), identity) # eltype mismatch + @test_throws DimensionMismatch pack!(fill(-1.0, short), ok, spec, identity) + @test_throws BoundsError pack!(canary, tile(size(ok)..., 18), spec, identity) @test canary == fill(-1.0, 100) end end -@testset "pack_a!/pack_b!: zero steady-state allocation" begin +@testset "pack!: zero steady-state allocation" begin nontrivial(x) = 2x + 1 # Every destination kind, affine/scattered axes, tails and # k_block_length == 0; the ScalarKernel/SIMDKernel forwarding methods on a subset. @@ -163,41 +169,24 @@ end storage = rand(T, 4000) koffs = [0, MR, 3MR, 2MR, 5MR, 4MR, 6MR] lanes = collect(0:(max(MR, NR) - 1)) .* 3 - bufa = zeros(T, MR * k_block_length + 8) - bufb = zeros(T, NR * k_block_length + 8) bytes = Int[] - GC.@preserve bufa bufb begin - pa = packed_panel(bufa, 1, MR * k_block_length) - dsts_a = (bufa, pa) - dsts_b = (bufb, packed_panel(bufb, 1, NR * k_block_length)) + for (i, L) in ((1, MR), (2, NR)) + spec = sliver_spec(kernel, i) + buf = zeros(T, L * k_block_length + 8) pk = view(koffs, 1:k_block_length) - srcs_a = ( - Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, k_block_length)), - Tile(storage, 0, AffineAxis(0, 1, MR), pk), - Tile(storage, 0, AffineAxis(0, 2, MR - 1), AffineAxis(0, 2MR, k_block_length)), - Tile(storage, 0, view(lanes, 1:MR), view(koffs, 1:k_block_length)), - Tile(storage, 0, AffineAxis(0, 1, MR), AffineAxis(0, MR, 0)), + srcs = ( + Tile(storage, 0, AffineAxis(0, 1, L), AffineAxis(0, L, k_block_length)), + Tile(storage, 0, AffineAxis(0, 1, L), pk), + Tile(storage, 0, AffineAxis(0, 2, L - 1), AffineAxis(0, 2L, k_block_length)), + Tile(storage, 0, view(lanes, 1:L), pk), + Tile(storage, 0, AffineAxis(0, 1, L), AffineAxis(0, L, 0)), ) - srcs_b = ( - Tile(storage, 0, AffineAxis(0, 1, k_block_length), AffineAxis(0, k_block_length, NR)), - Tile(storage, 0, pk, AffineAxis(0, k_block_length, NR - 1)), - Tile(storage, 0, view(koffs, 1:k_block_length), view(lanes, 1:NR)), - Tile(storage, 0, AffineAxis(0, 1, 0), AffineAxis(0, k_block_length, NR)), - ) - if !full - srcs_a, srcs_b = srcs_a[1:1], srcs_b[1:1] - end - for d in dsts_a - for s in srcs_a - push!(bytes, steady_pack_allocs(pack_a!, d, s, kernel, identity)) + GC.@preserve buf for d in (buf, packed_panel(buf, 1, L * k_block_length)) + for s in (full ? srcs : srcs[1:1]), f in (identity, conj) + push!(bytes, steady_pack_allocs(d, s, spec, f)) end - push!(bytes, steady_pack_allocs(pack_a!, d, srcs_a[1], kernel, nontrivial)) - end - for d in dsts_b, s in srcs_b - push!(bytes, steady_pack_allocs(pack_b!, d, s, kernel, identity)) + push!(bytes, steady_pack_allocs(d, srcs[1], spec, nontrivial)) end - push!(bytes, steady_pack_allocs(pack_b!, dsts_b[1], srcs_b[1], kernel, nontrivial)) - push!(bytes, steady_pack_allocs(pack_a!, pa, srcs_a[1], kernel, conj)) end return bytes end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index cd6b16d..baa239a 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -345,29 +345,22 @@ end @test_throws BoundsError _pcf_exec(plow, 1.0, 0.0) end -@testset "per-call floor: the unsafe_* packers keep every non-bounds check" begin +@testset "per-call floor: unsafe_pack! keeps every non-bounds check" begin kernel = SIMDKernel(Val(8), Val(6), Float64) src = Tile(collect(1.0:64.0), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 4)) packed = zeros(8 * 4) + a = sliver_spec(kernel, 1) ref = zeros(8 * 4) - pack_a!(ref, src, kernel, identity) - GC.@preserve packed QS.unsafe_pack_a!(QS.packed_panel(packed, 1, length(packed)), src, kernel, identity) + pack!(ref, src, a, identity) + GC.@preserve packed unsafe_pack!(QS.packed_panel(packed, 1, length(packed)), src, a, identity) @test packed == ref # Only the bounds checks are skipped. - @test_throws DimensionMismatch QS.unsafe_pack_a!(zeros(3), src, kernel, identity) - @test_throws ArgumentError QS.unsafe_pack_a!(zeros(Float32, 64), src, kernel, identity) + @test_throws DimensionMismatch unsafe_pack!(zeros(3), src, a, identity) + @test_throws ArgumentError unsafe_pack!(zeros(Float32, 64), src, a, identity) toowide = Tile(collect(1.0:200.0), 0, AffineAxis(0, 1, 9), AffineAxis(0, 16, 4)) - @test_throws ArgumentError QS.unsafe_pack_a!(zeros(200), toowide, kernel, identity) - - srcB = Tile(collect(1.0:64.0), 0, AffineAxis(0, 1, 4), AffineAxis(0, 4, 6)) - refB = zeros(6 * 4) - packedB = zeros(6 * 4) - pack_b!(refB, srcB, kernel, identity) - GC.@preserve packedB QS.unsafe_pack_b!(QS.packed_panel(packedB, 1, length(packedB)), srcB, kernel, identity) - @test packedB == refB - @test_throws DimensionMismatch QS.unsafe_pack_b!(zeros(3), srcB, kernel, identity) + @test_throws ArgumentError unsafe_pack!(zeros(200), toowide, a, identity) dest = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) @test_throws DimensionMismatch QS.unsafe_execute_tile!( @@ -378,8 +371,9 @@ end ) d1 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) d2 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) - execute_tile!(kernel, d1, ref, refB, 4, 1.0, 0.0) - QS.unsafe_execute_tile!(kernel, d2, ref, refB, 4, 1.0, 0.0) + packedB = collect(1.0:24.0) + execute_tile!(kernel, d1, ref, packedB, 4, 1.0, 0.0) + QS.unsafe_execute_tile!(kernel, d2, ref, packedB, 4, 1.0, 0.0) @test d1.storage == d2.storage end diff --git a/test/runtests.jl b/test/runtests.jl index 9e71cc9..34c8e37 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,7 +5,7 @@ using QuasiStrided using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, describe_block, block_descriptors!, Descriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, - AffineAxis, Tile, checked_tile_storage_bounds, pack_a!, pack_b!, + AffineAxis, Tile, checked_tile_storage_bounds, pack!, unsafe_pack!, sliver_spec, zero_accumulator, accumulate, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, target_profile, cache_topology, unknown_target, detect_isa, detect_target, From 505e43342a1dbf932278c6a27a032db10339803c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 15:17:59 -0400 Subject: [PATCH 11/60] Decision log: chunks 4 and 5 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 667217c..cbbf1fd 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -16,8 +16,8 @@ TTFX when codegen is touched). | 1 | Hardware | `hardware/target.jl` | done | | 2 | Axis groups | `layout/axis_group.jl`, `pair_group.jl` | done | | 3 | Tiles | `layout/tiles.jl` | done | -| 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | in progress | -| 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | | +| 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | done | +| 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | done | | 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | | | 7 | SIMD kernels | `simd.jl`, `planar.jl` | | | 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | | @@ -113,13 +113,24 @@ a borrowed pointer. The pointer is kept because the nest holds a nested closures. One `Tile` type with `getindex`/`setindex!`/`size` replaces `QSTile`/`SourceTile`/`DestinationTile`/`tile_load`/`tile_store!`. -### D8. Packing formats (chunk 4) +### D8. Packing formats (chunk 4, applied) Panels are only `PackedPanel` (the oracle and tests too; the `AbstractVector` panel methods go); one `Descriptor` constructor plus a `RealDescriptor` alias; one `packed_a_offset(d, i, p, plane = 0)` (and B); `transposed.jl` merges into `panel.jl` and is reviewed with `pack_split.jl`. +### D9. Packers (chunk 5, applied) + +A and B on the same footing: one `pack!(panel, tile, sliver_spec(kernel, i), transform)` +with lanes along the tile's rows; B is packed as `transpose(tile)`. A kernel +needing asymmetric packing gets it through its per-operand spec (format, +extent), or overloads `pack!` for its own spec. Fast paths are symmetric, so +real B with unit-stride columns takes the vector path (measured before +keeping). One `emit!` per format taking the values to store; padding passes +literal zeros. Vector-path predicates (`dense_lanes`, `is_unit_stride`, +`complex_fastpath_isa_eligible`) live in pack_contiguous.jl. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: @@ -135,8 +146,6 @@ alias; one `packed_a_offset(d, i, p, plane = 0)` (and B); - Detected `l1d.line` is unused: `_K_LINE_BYTES = 64` in `labels.jl` (K-order cost model) and `pack_split.jl`; thread the detected line size through planning (chunks 9/11). -- `complex_fastpath_isa_eligible` is policy, not detection: move it next to - its consumers (chunks 5/7). - 16-argument `ContractPlan(...)` spelled out five times, `_PlanRequest` twice more: one "copy with fields replaced" helper (chunk 13/14). - `execute.jl` mixes public API, barrier/hint machinery, the C-panel path and From 984e6eb7afda129b8cd316ced9837955d41d5fad Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 16:54:45 -0400 Subject: [PATCH 12/60] Chunk 6: kernel interface - `Base.accumulate` -> `add_tile`, a function of our own, in every kernel, the `DenseVector` method, callers and tests. - `KernelMethod` (was `ComplexMethod`) singletons declare their packing formats once with `pack_formats`; `a_reals`/`b_reals` go, blocking uses `reals_per_element` of the formats. A test checks every kernel's descriptor agrees with `pack_formats(complex_method(k))`. - `DescriptorKernel` -> `Microkernel`; leading underscores dropped from the names in interface.jl and scalar.jl. - One `@inline execute_tile!` and one `@inline pack!`, whose storage check is `@boundscheck`; every other validation stays unconditional. The nest calls `@inbounds _pack_sliver!`/`@inbounds _execute_micro_tile!` (both `@propagate_inbounds`, replacing `unsafe_execute_micro_tile!` and the function-valued `pack!` argument); the oracle calls them checked. LLVM of `_execute_nest!`, `_micro_tiles_packed_b!` and `_micro_tiles_unpacked_b!` (Float64/ComplexF64, affine and scattered) has the same line count as before and no per-tile storage check; under `--check-bounds=yes` the per-tile checks now run, so the hoisting test counts them. - `ScalarKernel` accumulates into an `NTuple{MR*NR,T}` (column-major) and its store branches on `beta` once per tile; it now allocates nothing. - Contract moved from the header comment into docstrings (`Microkernel`, `KernelMethod`, `zero_accumulator`, `add_tile`, `store_tile!`, `execute_tile!`, `pack!`). Tests: 31548 pass (was 31537). TTFX (2 runs each, interleaved), before and after alike: load 0.5 s; first call Float64 3.2-3.3 s, ComplexF64 2.3 s; K=1 outer 0.6 s. A/B medians of 5 interleaved rounds, after/before: matmul256 F64 0.997, C64 0.991; scatterM 1.017/0.994; scatterMN 1.003/1.001; rowmajB256 F64 1.009, F32 0.996; scatterK F64 0.996, C64 1.009. Allocations 0 per call under --check-bounds=yes. ScalarKernel 4x4 nest, 20x30x25 output, K = 16: Float64 155 -> 50 us, Float32 159 -> 46-55 us; allocations per call 218400 / 151200 bytes -> 0. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_kernels.jl | 6 +- src/QuasiStrided.jl | 2 +- src/execution/dot.jl | 4 +- src/execution/execute.jl | 14 +- src/execution/macrokernel.jl | 23 +- src/execution/oracle.jl | 4 +- src/execution/unpackedb.jl | 7 +- src/microkernels/fmaddsub.jl | 20 +- src/microkernels/interface.jl | 264 +++++++++--------- src/microkernels/mixed.jl | 18 +- src/microkernels/onem.jl | 20 +- src/microkernels/planar.jl | 28 +- src/microkernels/scalar.jl | 58 ++-- src/microkernels/simd.jl | 18 +- src/packing/pack.jl | 40 ++- src/planning/blocking.jl | 6 +- src/planning/defaults.jl | 2 +- src/planning/kernel_selection.jl | 12 +- test/execution/test_execute.jl | 3 +- test/microkernels/test_fmaddsub_kernel.jl | 4 +- test/microkernels/test_mixed_kernel.jl | 2 +- test/microkernels/test_onem_kernel.jl | 2 +- .../test_planar_store_fastpath.jl | 4 +- test/microkernels/test_scalar_kernel.jl | 34 ++- test/planning/test_kernel_selection.jl | 17 +- test/planning/test_per_call_overhead.jl | 37 +-- test/planning/test_plan_contract.jl | 2 +- test/runtests.jl | 8 +- 28 files changed, 340 insertions(+), 319 deletions(-) diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index 3aa645f..beb58d1 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -5,7 +5,7 @@ # [--shapes 64x64x64,2048x2048x2048] [--reps 11] [--outdir DIR] # # Arms: -# kernel (real dtypes) the ceiling of each menu tile: `unsafe_execute_tile!` +# kernel (real dtypes) the ceiling of each menu tile: `@inbounds execute_tile!` # on one L1-resident packed A and B sliver at the tile's default `k_block`. # engine per dtype and shape: OpenBLAS, the default plan, and every menu # kernel (of every complex method, for complex dtypes) at its own @@ -22,7 +22,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: RealMethod, PlanarMethod, OneMMethod, FMAddSubMethod, kernel_shapes, - _kernel_from_shape, default_blocking, packed_panel, unsafe_execute_tile!, + _kernel_from_shape, default_blocking, packed_panel, execute_tile!, Tile, AffineAxis, target_profile const RUN_DTYPES = parse_dtypes(argopt("dtypes", "Float64,Float32,ComplexF64,ComplexF32")) @@ -62,7 +62,7 @@ function kernel_hot!(kernel, C, apack, bpack, k_block, reps) bp = packed_panel(bpack, 1, length(bpack)) dest = Tile(C, 0, AffineAxis(0, 1, tile_size(kernel, 1)), AffineAxis(0, tile_size(kernel)...)) for _ in 1:reps - unsafe_execute_tile!(kernel, dest, ap, bp, k_block, one(eltype(C)), one(eltype(C))) + @inbounds execute_tile!(kernel, dest, ap, bp, k_block, one(eltype(C)), one(eltype(C))) end end return nothing diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 159debe..39e705f 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -24,7 +24,7 @@ include("packing/panel.jl") include("packing/pack.jl") include("packing/pack_contiguous.jl") -# --- Microkernels: accumulate over one packed K panel, store into C --- +# --- Microkernels: add one packed K block to a register tile, store into C --- include("microkernels/interface.jl") include("microkernels/scalar.jl") include("microkernels/simd.jl") diff --git a/src/execution/dot.jl b/src/execution/dot.jl index d1ce4b1..3ad3230 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -13,7 +13,7 @@ const _DOT_MODE = Ref{Symbol}(:auto) @inline function _dot_lanewidth(::Type{T}) where {T} R = real(T) vb = target_profile().vector_bytes - return (vb > 0 && vb % sizeof(R) == 0) ? vb ÷ sizeof(R) : _default_lanewidth(R) + return (vb > 0 && vb % sizeof(R) == 0) ? vb ÷ sizeof(R) : default_lanewidth(R) end # Outputs computed together: 8 independent FMA chains cover the FMA latency @@ -175,7 +175,7 @@ end sums = _dot_group(mptr, gptr, mbases, k_block_length, Val(W), Val(NB)) @inbounds for j in 1:nvalid s = MCONJ ? conj(sums[j]) : sums[j] - _axpby_at!(Cstorage, cbase + bufC[q + j] + 1, alpha, s, beta) + axpby_at!(Cstorage, cbase + bufC[q + j] + 1, alpha, s, beta) end q += nvalid end diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 23a3771..33f1ed7 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -361,7 +361,8 @@ function _execute_nest!( # Hoisted bounds checks (B here, A and C per M block): each # rectangle is exactly the union of the per-sliver/per-tile - # regions, so the `unsafe_*` calls below read nothing unchecked. + # regions, and the storage check compares only range extremes, so + # the `@inbounds` pack and tile calls below touch nothing unchecked. checked_span_bounds(plan.Bbase, rng_kB, rng_nB, lenB) beta_eff = firstpanel ? betaT : one(T) @@ -377,8 +378,8 @@ function _execute_nest!( n_tile_start = n_tile_index * n_tile bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) - _pack_sliver!( - unsafe_pack!, bpanel, plan.Bstorage, plan.Bbase, colsB, rowsB_k, + @inbounds _pack_sliver!( + bpanel, plan.Bstorage, plan.Bbase, colsB, rowsB_k, sliver_spec(kernel, 2), btransform ) end @@ -416,8 +417,8 @@ function _execute_nest!( m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) - _pack_sliver!( - unsafe_pack!, apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, + @inbounds _pack_sliver!( + apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, sliver_spec(kernel, 1), atransform ) end @@ -466,7 +467,8 @@ end m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) - unsafe_execute_micro_tile!( + # Inside the caller's C check. + @inbounds _execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, apanel, bpanel, k_block_length, alphaT, beta_eff ) diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index 82f7684..2a48f4d 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -18,18 +18,17 @@ end @noinline _throw_irregular_ramp_descriptor() = throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) -# `pack!` is `pack!` or `unsafe_pack!`; B passes its axes swapped (its -# `transpose`). GUARDRAIL: `transform` needs its own bound type parameter, or -# it costs a dynamic dispatch per call. -@inline function _pack_sliver!( - pack!::PF, packed::PK, storage::S, base::Int, +# B passes its axes swapped (its `transpose`). GUARDRAIL: `transform` needs its +# own bound type parameter, or it costs a dynamic dispatch per call. +Base.@propagate_inbounds function _pack_sliver!( + packed::PK, storage::S, base::Int, lanes::R, steps::C, spec, transform::TF - ) where {PF, PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} + ) where {PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} pack!(packed, Tile(storage, base, lanes, steps), spec, transform) return nothing end -@inline function _execute_micro_tile!( +Base.@propagate_inbounds function _execute_micro_tile!( kernel, storage::S, base::Int, rows::R, cols::C, packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} @@ -38,16 +37,6 @@ end return nothing end -# The caller has bounds-checked the whole macro block this tile belongs to. -@inline function unsafe_execute_micro_tile!( - kernel, storage::S, base::Int, rows::R, cols::C, - packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta - ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} - destination = Tile(storage, base, rows, cols) - unsafe_execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) - return nothing -end - @inline function _scale_micro_tile!( storage::S, base::Int, rows::R, cols::C, beta ) where {S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl index 5041196..cbb7303 100644 --- a/src/execution/oracle.jl +++ b/src/execution/oracle.jl @@ -61,11 +61,11 @@ function execute_tilewise!(plan::ContractPlan{T}, alpha::Number, beta::Number) w # Whole single-sliver buffers, so no `_sliver_panel`. _pack_sliver!( - pack!, ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, + ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, sliver_spec(kernel, 1), plan.atransform ) _pack_sliver!( - pack!, ws.tw_packed_b, plan.Bstorage, plan.Bbase, colsB, rowsK_B, + ws.tw_packed_b, plan.Bstorage, plan.Bbase, colsB, rowsK_B, sliver_spec(kernel, 2), plan.btransform ) diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 7f9e4c5..64b3a3b 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -8,7 +8,7 @@ # `p` is storage address `colbase[j] + ksteps[p]` (base included). Padding # columns of a tail sliver alias the last valid column, so every address read # is inside the hoisted B check; their results are never stored. `length` is -# the packed-equivalent count that `_execute_tile_prologue!` checks, not the +# the packed-equivalent count that `execute_tile!` checks, not the # storage length. struct UnpackedBView{S, K <: AbstractVector{Int}, NR, F} storage::S @@ -70,7 +70,7 @@ const _UNPACKED_B_MMAX = 256 # The view for one N sliver; `buf`/`n_tile_start` locate an irregular sliver's # offsets in `ws.n_buf_B`. @inline function _unpacked_b_view( - kernel::DescriptorKernel{MR, NR, T}, storage::S, base::Int, + kernel::Microkernel{MR, NR, T}, storage::S, base::Int, d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F ) where {MR, NR, T, S, K <: AbstractVector{Int}, F} last = d.count - 1 @@ -104,7 +104,8 @@ end m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) - unsafe_execute_micro_tile!( + # Inside the caller's C check. + @inbounds _execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, apanel, bview, k_block_length, alphaT, beta_eff ) diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index a809b24..2b07e8a 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -25,13 +25,13 @@ Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over `2MR` must be a multiple of `W`, and `W` must be even. The default for small-M complex contractions on AVX-512. """ -struct FMAddSubKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} +struct FMAddSubKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} function FMAddSubKernel{MR, NR, T, W}( descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} ) where {MR, NR, T, W} - _check_vector_shape("FMAddSubKernel", 2 * MR, W, true) + check_vector_shape("FMAddSubKernel", 2 * MR, W, true) return new{MR, NR, T, W}(descriptor) end end @@ -46,7 +46,7 @@ end function FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} T <: Complex || throw(ArgumentError("FMAddSubKernel requires a complex element type, got $T")) - return FMAddSubKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) + return FMAddSubKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) end complex_method(::FMAddSubKernel) = FMAddSubMethod() @@ -105,7 +105,7 @@ end packed_a::PA, packed_b::PB, p::Int ) where {MR, NR, T, W, R, NA, PA, PB} MV = (2 * MR) ÷ W - _check_acc(:_accumulate_step_fmaddsub, R, T, NA, MV * NR) + check_acc(:_accumulate_step_fmaddsub, R, T, NA, MV * NR) av = [Symbol(:a, v) for v in 0:(MV - 1)] sv = [Symbol(:s, v) for v in 0:(MV - 1)] @@ -153,12 +153,12 @@ end end end -function Base.accumulate( +function add_tile( kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} k_block_length == 0 && return acc - k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) @inbounds for p in 1:k_block_length acc = _accumulate_step_fmaddsub(kernel, acc, packed_a, packed_b, p) end @@ -199,7 +199,7 @@ end # Same unroll and full-block / row-tail split as `_store_tile_planar_vector!`. @generated function _store_tile_fmaddsub_vector!( destination::Tile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, + alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NV} # The pointer reinterpretation is only sound on dense rank-1 complex storage. @@ -208,7 +208,7 @@ end 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) + check_acc(:_store_tile_fmaddsub_vector!, R, T, NV, MV * NR) HW = W ÷ 2 blocks = Any[] @@ -227,7 +227,7 @@ end for u in 1:$HW i = $(v * HW) + u i <= m || break - _axpby_at!( + axpby_at!( storage, colbase + i, alpha, Complex(vec[2 * u - 1], vec[2 * u]), beta ) @@ -270,7 +270,7 @@ function store_tile!( destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::FMAddSubKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NV} - m, n = _store_prologue!(destination, alpha, beta) + m, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination if _complex_vector_eligible(destination, T) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 15442d1..5dc6488 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -1,73 +1,103 @@ -# What every microkernel shares. Each kernel file implements only -# `zero_accumulator`, `accumulate` and `store_tile!`, with one contract: -# * `accumulate(kernel, acc, packed_a, packed_b, k_block_length)` adds -# `k_block_length` K steps to `acc`; `k_block_length == 0` returns `acc` -# without reading the panels. -# * `store_tile!(dest, acc, alpha, beta, kernel)` writes `alpha*acc + beta*C` -# over the valid rectangle only: `alpha == 0` never reads `acc`, `beta == 0` -# never reads old `C`, padding lanes are never read. -# The vector kernels match `ScalarKernel` only to within rounding (FMA grouping -# differs), never bitwise. - -# Kernels wrapping a descriptor as `.descriptor`; the accessors forward once here. -abstract type DescriptorKernel{MR, NR, T} end - -# Which complex-arithmetic method a kernel implements, as singletons so blocking -# and the shape menus dispatch on it. No method ranking is hardcoded: planar is -# the default, 1m and fmaddsub are used only when named (plus fmaddsub for the -# AVX-512 small-M demotion, `_small_m_shape` in src/planning/kernel_selection.jl). -abstract type ComplexMethod end -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 `m_block`/`n_block` -# by these, so every method gets the same packed byte budget. -a_reals(::RealMethod) = 1 -b_reals(::RealMethod) = 1 -a_reals(::PlanarMethod) = 2 -b_reals(::PlanarMethod) = 2 -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 +# The microkernel contract, the complex-arithmetic methods, and the validation +# and store helpers the kernels share. + +""" + Microkernel{MR, NR, T} + +A register-tile microkernel: computes an `MR x NR` tile of `C`, element type +`T`, from one packed A sliver and one packed B sliver of a K block. A kernel +wraps a `Descriptor` (tile size, element type, packing formats) as +`.descriptor` and implements [`zero_accumulator`](@ref), [`add_tile`](@ref) and +[`store_tile!`](@ref); [`execute_tile!`](@ref) composes them. The vector +kernels match `ScalarKernel` only to within rounding (FMA grouping differs), +never bitwise. +""" +abstract type Microkernel{MR, NR, T} end + +""" + KernelMethod + +How a kernel maps its multiply-adds onto real arithmetic, as singletons so +blocking and the shape menus dispatch on it: + + - `RealMethod`: a real kernel. + - `PlanarMethod`: split re/im planes, 4 real FMAs per complex MAC. + - `OneMMethod`: Van Zee's 1m, a real `2MR x NR` kernel. + - `FMAddSubMethod`: interleaved A, x86 `vfmaddsub`. + - `ComplexRealMethod`: complex A, real B; the real kernel on `2MR` rows. + - `RealComplexMethod`: real A, complex B; the real kernel on `2NR` columns. + +No method ranking is hardcoded: planar is the default, 1m and fmaddsub are used +only when named (plus fmaddsub for the AVX-512 small-M demotion, +`_small_m_shape` in src/planning/kernel_selection.jl). +""" +abstract type KernelMethod end +struct RealMethod <: KernelMethod end +struct PlanarMethod <: KernelMethod end +struct OneMMethod <: KernelMethod end +struct FMAddSubMethod <: KernelMethod end +struct ComplexRealMethod <: KernelMethod end +struct RealComplexMethod <: KernelMethod end + +# The (A, B) packing formats of every kernel of a method. +pack_formats(::RealMethod) = (RealFormat(), RealFormat()) +pack_formats(::PlanarMethod) = (PlanarFormat(), PlanarFormat()) +pack_formats(::OneMMethod) = (OneEFormat(), PlanarFormat()) +pack_formats(::FMAddSubMethod) = (InterleavedFormat(), PlanarFormat()) +pack_formats(::ComplexRealMethod) = (InterleavedFormat(), RealFormat()) +pack_formats(::RealComplexMethod) = (RealFormat(), InterleavedFormat()) # Accumulator planes held live: planar keeps separate re/im planes. -accumulator_planes(::RealMethod) = 1 +accumulator_planes(::KernelMethod) = 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() -realtype(k::DescriptorKernel) = realtype(k.descriptor) -sliver_width(k::DescriptorKernel) = sliver_width(k.descriptor) -a_format(k::DescriptorKernel) = a_format(k.descriptor) -b_format(k::DescriptorKernel) = b_format(k.descriptor) -tile_size(k::DescriptorKernel) = tile_size(k.descriptor) -scalartype(k::DescriptorKernel) = scalartype(k.descriptor) -@inline sliver_spec(k::DescriptorKernel, i::Int) = sliver_spec(k.descriptor, i) -@inline packed_a_offset(k::DescriptorKernel, i::Int, p::Int, plane::Int = 0) = +realtype(k::Microkernel) = realtype(k.descriptor) +sliver_width(k::Microkernel) = sliver_width(k.descriptor) +a_format(k::Microkernel) = a_format(k.descriptor) +b_format(k::Microkernel) = b_format(k.descriptor) +tile_size(k::Microkernel) = tile_size(k.descriptor) +scalartype(k::Microkernel) = scalartype(k.descriptor) +@inline sliver_spec(k::Microkernel, i::Int) = sliver_spec(k.descriptor, i) +@inline packed_a_offset(k::Microkernel, i::Int, p::Int, plane::Int = 0) = packed_a_offset(k.descriptor, i, p, plane) -@inline packed_b_offset(k::DescriptorKernel, j::Int, p::Int, plane::Int = 0) = +@inline packed_b_offset(k::Microkernel, j::Int, p::Int, plane::Int = 0) = packed_b_offset(k.descriptor, j, p, plane) -packed_a_length(k::DescriptorKernel, k_block_length::Int) = packed_a_length(k.descriptor, k_block_length) -packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.descriptor, k_block_length) +packed_a_length(k::Microkernel, k_block_length::Int) = packed_a_length(k.descriptor, k_block_length) +packed_b_length(k::Microkernel, k_block_length::Int) = packed_b_length(k.descriptor, k_block_length) + +""" + zero_accumulator(kernel::Microkernel) -> acc + +The kernel's zero register tile, an immutable value. +""" +function zero_accumulator end + +""" + add_tile(kernel::Microkernel, acc, packed_a, packed_b, k_block_length) -> acc + +Add the products of `k_block_length` K steps of the packed A and B slivers to +`acc` and return the result. `k_block_length == 0` returns `acc` without +reading the panels; a negative one throws an `ArgumentError`. +""" +function add_tile end + +""" + store_tile!(destination::Tile, acc, alpha, beta, kernel::Microkernel) -> destination + +Write `alpha * acc + beta * C` over the valid rectangle `size(destination)` +only: `alpha == 0` never reads `acc`, `beta == 0` never reads old `C`, and +padding lanes of `acc` are never read. `alpha` and `beta` are of the kernel's +element type. +""" +function store_tile! end # Constructor check: `rows` reals per sliver must split into whole `W`-vectors. # `even`: the lane-pair kernels keep one element's re/im in adjacent lanes, and # an odd `W` can divide `2MR` while splitting an element across two vectors -- # a wrong answer, not an error, hence rejected here. -@inline function _check_vector_shape(name, rows::Int, W, even::Bool = false) +@inline function check_vector_shape(name, rows::Int, W, even::Bool = false) W isa Int && W > 0 || throw(ArgumentError("$name requires an Int vector width W > 0, got W = $W")) even && isodd(W) && @@ -79,15 +109,15 @@ packed_b_length(k::DescriptorKernel, k_block_length::Int) = packed_b_length(k.de end # Generator-time check of an accumulator tuple against the kernel's shape. -function _check_acc(f::Symbol, R, T, NA::Int, want::Int) +function check_acc(f::Symbol, R, T, NA::Int, want::Int) R === real(T) || throw(ArgumentError("$f: accumulator lane type $R is not real($T)")) NA == want || throw(ArgumentError("$f: accumulator length $NA, expected $want")) return nothing end # Default `SIMD.Vec` lane count: one 256-bit register. -_default_lanewidth(::Type{Float64}) = 4 -_default_lanewidth(::Type{Float32}) = 8 +default_lanewidth(::Type{Float64}) = 4 +default_lanewidth(::Type{Float32}) = 8 # `beta == 0` writes zeros without reading `C`; `beta == 1` is a no-op. function scale_tile!(destination::Tile, beta::T) where {T} @@ -108,19 +138,19 @@ function scale_tile!(destination::Tile, beta::T) where {T} 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::T) where {T} = @inbounds dest[i, j] = +@inline axpby_tile!(dest, i::Int, j::Int, alpha, r, beta::T) where {T} = @inbounds dest[i, j] = iszero(beta) ? alpha * r : isone(beta) ? muladd(alpha, r, convert(T, dest[i, j])) : muladd(alpha, r, beta * convert(T, dest[i, j])) -@inline _axpby_at!(storage, idx::Int, alpha, r, beta::T) where {T} = @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, 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!`). -@inline function _store_prologue!(destination::Tile, alpha, beta) +@inline function store_prologue!(destination::Tile, alpha, beta) m, n = size(destination) if m == 0 || n == 0 return (0, 0) @@ -131,93 +161,73 @@ end return (m, n) end -# GUARDRAIL: every throw reachable from the per-tile prologue sits behind a +# GUARDRAIL: every throw reachable from `execute_tile!` sits behind a # `@noinline` helper. An inline interpolated message pulls `print_to_string`, a # GC frame and ~1 KB of stack into the hot tile function. -@noinline _throw_packed_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( +@noinline throw_packed_short(which::Symbol, got::Int, need::Int, k_block_length::Int) = throw( DimensionMismatch( "execute_tile!: packed_$which has length $got, " * "need at least packed_$(which)_length(kernel, k_block_length=$k_block_length) = $need" ) ) -@noinline _throw_tile_extent(which::Symbol, got::Int, limit::Int) = throw( +@noinline throw_tile_extent(which::Symbol, got::Int, limit::Int) = throw( ArgumentError( "destination valid $which extent $got exceeds tile_size(kernel, $(which === :row ? 1 : 2)) = $limit" ) ) -@noinline _throw_negative_k_block_length(where::Symbol, k_block_length::Int) = +@noinline throw_negative_k_block_length(where::Symbol, k_block_length::Int) = throw(ArgumentError("$where requires k_block_length >= 0, got k_block_length = $k_block_length")) -# `execute_tile!`'s validation, in order. Returns `(run, alphaT, betaT)`; -# `run == false` means the call is finished. `Val(false)` drops the storage -# bounds check at compile time (`unsafe_execute_tile!` only). -# GUARDRAIL: `@inline` and one bound type parameter per argument (hot path). -@inline function _execute_tile_prologue!( +# The kernels read only `PackedPanel`s; a `DenseVector` is borrowed as one. +function add_tile( + kernel::K, acc::A, packed_a::PA, packed_b::PB, k_block_length::Int + ) where {K <: Microkernel, A, PA <: DenseVector, PB <: DenseVector} + return GC.@preserve packed_a packed_b add_tile( + kernel, acc, packed_panel(packed_a, 1, length(packed_a)), + packed_panel(packed_b, 1, length(packed_b)), k_block_length + ) +end + +""" + execute_tile!(kernel::Microkernel, destination::Tile, packed_a, packed_b, + k_block_length, alpha, beta) -> destination + +One K block of one register tile: `zero_accumulator`, `add_tile`, +`store_tile!`, i.e. `C = alpha * A * B + beta * C` over `destination`'s valid +rectangle. Validates, in order: the destination extents against +`tile_size(kernel)` (`ArgumentError`), `k_block_length >= 0` +(`ArgumentError`), `alpha` and `beta` converting to the kernel's element type; +then returns at once for an empty destination; checks the +destination's storage bounds (`BoundsError`); applies only `beta` when +`k_block_length == 0` or `alpha == 0`, without reading the panels; and checks +the panel lengths (`DimensionMismatch`). + +Under `@inbounds` (the call inlined) the storage bounds check is skipped, so an +out-of-range destination is a silent out-of-bounds write; every other check +stays. +""" +@inline function execute_tile!( kernel::K, destination::Tile, packed_a::PA, packed_b::PB, - k_block_length::Int, alpha, beta, ::Val{BOUNDS} - ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB, BOUNDS} + k_block_length::Int, alpha, beta + ) where {MR, NR, T, K <: Microkernel{MR, NR, T}, PA, PB} m, n = size(destination) - m <= MR || _throw_tile_extent(:row, m, MR) - n <= NR || _throw_tile_extent(:column, n, NR) - k_block_length >= 0 || _throw_negative_k_block_length(:execute_tile!, k_block_length) - + m <= MR || throw_tile_extent(:row, m, MR) + n <= NR || throw_tile_extent(:column, n, NR) + k_block_length >= 0 || throw_negative_k_block_length(:execute_tile!, k_block_length) alphaT = convert(T, alpha) betaT = convert(T, beta) - - (m == 0 || n == 0) && return (false, alphaT, betaT) - - BOUNDS && checked_tile_storage_bounds(destination) - + (m == 0 || n == 0) && return destination + @boundscheck checked_tile_storage_bounds(destination) if k_block_length == 0 || iszero(alphaT) scale_tile!(destination, betaT) - return (false, alphaT, betaT) + return destination end need_a = packed_a_length(kernel, k_block_length) - length(packed_a) >= need_a || _throw_packed_short(:a, length(packed_a), need_a, k_block_length) + length(packed_a) >= need_a || throw_packed_short(:a, length(packed_a), need_a, k_block_length) need_b = packed_b_length(kernel, k_block_length) - length(packed_b) >= need_b || _throw_packed_short(:b, length(packed_b), need_b, k_block_length) - - return (true, alphaT, betaT) -end - -# The kernels read only `PackedPanel`s; the checked entry points borrow a -# `DenseVector` as one. -function Base.accumulate( - kernel::K, acc::A, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {K <: DescriptorKernel, A, PA <: DenseVector, PB <: DenseVector} - return GC.@preserve packed_a packed_b accumulate( - kernel, acc, packed_panel(packed_a, 1, length(packed_a)), - packed_panel(packed_b, 1, length(packed_b)), k_block_length - ) -end + length(packed_b) >= need_b || throw_packed_short(:b, length(packed_b), need_b, k_block_length) -# One checked K panel: `zero_accumulator`, `accumulate`, `store_tile!`. -function execute_tile!( - kernel::K, destination::Tile, packed_a::PA, packed_b::PB, - k_block_length::Int, alpha, beta - ) where {MR, NR, T, K <: DescriptorKernel{MR, NR, T}, PA, PB} - run, alphaT, betaT = _execute_tile_prologue!( - kernel, destination, packed_a, packed_b, k_block_length, alpha, beta, Val(true) - ) - run || return destination - acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, k_block_length) - return store_tile!(destination, acc, alphaT, betaT, kernel) -end - -# `execute_tile!` without `checked_tile_storage_bounds(destination)`: an -# out-of-range destination is a silent out-of-bounds WRITE. `_execute_nest!` -# checks the union of a whole macro block's tiles once with -# `checked_span_bounds`, which is equivalent because the check only compares -# range extremes. -@inline function unsafe_execute_tile!( - kernel::K, destination::Tile, packed_a::PA, packed_b::PB, - k_block_length::Int, alpha, beta - ) where {K, PA, PB} - run, alphaT, betaT = _execute_tile_prologue!( - kernel, destination, packed_a, packed_b, k_block_length, alpha, beta, Val(false) - ) - run || return destination - acc = accumulate(kernel, zero_accumulator(kernel), packed_a, packed_b, k_block_length) + acc = add_tile(kernel, zero_accumulator(kernel), packed_a, packed_b, k_block_length) return store_tile!(destination, acc, alphaT, betaT, kernel) end diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index dced3d8..14137f9 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -16,7 +16,7 @@ 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} +struct ComplexRealKernel{MR, NR, T, W, KI <: SIMDKernel} <: Microkernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} inner::KI @@ -24,7 +24,7 @@ struct ComplexRealKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} - _check_vector_shape("ComplexRealKernel", 2 * MR, W, true) + 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") ) @@ -39,7 +39,7 @@ 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} +struct RealComplexKernel{MR, NR, T, W, KI <: SIMDKernel} <: Microkernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} inner::KI @@ -69,9 +69,9 @@ function RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, end ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = - ComplexRealKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(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)))) + RealComplexKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) complex_method(::ComplexRealKernel) = ComplexRealMethod() complex_method(::RealComplexKernel) = RealComplexMethod() @@ -79,15 +79,15 @@ lanewidth(kernel::_MixedKernel) = lanewidth(kernel.inner) zero_accumulator(kernel::_MixedKernel) = zero_accumulator(kernel.inner) -@inline Base.accumulate(kernel::_MixedKernel, acc::NTuple, packed_a::PackedPanel, packed_b, k_block_length::Int) = - accumulate(kernel.inner, acc, packed_a, packed_b, k_block_length) +@inline add_tile(kernel::_MixedKernel, acc::NTuple, packed_a::PackedPanel, packed_b, k_block_length::Int) = + add_tile(kernel.inner, acc, packed_a, packed_b, k_block_length) # Inlined or not as the fmaddsub and planar stores each reuses. function store_tile!( destination::Tile, 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, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination if _complex_vector_eligible(destination, T) @@ -101,7 +101,7 @@ end destination::Tile, 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, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination if _complex_vector_eligible(destination, T) diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index b0344ba..9a7c48d 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -15,7 +15,7 @@ Van Zee's induced 1m complex microkernel: 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. Used only when named in `plan_contract(...; kernel = ...)`. """ -struct OneMKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, NR, T} +struct OneMKernel{MR, NR, T, W, KI <: SIMDKernel} <: Microkernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} # `KI` stands in for `SIMDKernel{2MR,NR,real(T),W}`, which is not a legal # field type; the constructor pins it to exactly that. @@ -25,7 +25,7 @@ struct OneMKernel{MR, NR, T, W, KI <: SIMDKernel} <: DescriptorKernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat}, inner::KI ) where {MR, NR, T, W, KI <: SIMDKernel} - _check_vector_shape("OneMKernel", 2 * MR, W, true) + check_vector_shape("OneMKernel", 2 * MR, W, true) KI === SIMDKernel{2 * MR, NR, real(T), W} || throw( ArgumentError("OneMKernel's inner kernel must be SIMDKernel{$(2 * MR),$NR,$(real(T)),$W}, got $KI") ) @@ -43,7 +43,7 @@ end function OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} T <: Complex || throw(ArgumentError("OneMKernel requires a complex element type, got $T")) - return OneMKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) + return OneMKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) end complex_method(::OneMKernel) = OneMMethod() @@ -51,14 +51,14 @@ lanewidth(::OneMKernel{MR, NR, T, W}) where {MR, NR, T, W} = W zero_accumulator(kernel::OneMKernel) = zero_accumulator(kernel.inner) -function Base.accumulate( +function add_tile( kernel::OneMKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NV, PA <: PackedPanel, PB} # Checked here so the message reports the logical k_block_length, not the # doubled one. - k_block_length >= 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) - return accumulate(kernel.inner, acc, packed_a, packed_b, 2 * k_block_length) + k_block_length >= 0 || throw_negative_k_block_length(:add_tile, k_block_length) + return add_tile(kernel.inner, acc, packed_a, packed_b, 2 * k_block_length) end # Scalar store for an interleaved accumulator (1m and fmaddsub): with `W` even @@ -67,12 +67,12 @@ end # vectors. @generated function _store_tile_lanepair!( destination::Tile, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, + alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, m::Int, n::Int ) where {MR, NR, T, W, R, NV} iseven(W) || throw(ArgumentError("_store_tile_lanepair!: requires an even W, got $W")) MV = (2 * MR) ÷ W - _check_acc(:_store_tile_lanepair!, R, T, NV, MV * NR) + check_acc(:_store_tile_lanepair!, R, T, NV, MV * NR) HW = W ÷ 2 blocks = Any[] @@ -85,7 +85,7 @@ end for u in 1:$HW i = $(v * HW) + u i <= m || break - _axpby_tile!( + axpby_tile!( destination, i, $j, alpha, Complex(vec[2 * u - 1], vec[2 * u]), beta ) @@ -107,7 +107,7 @@ function store_tile!( destination::Tile, acc::NTuple{NV, Vec{W, R}}, alpha::T, beta::T, kernel::OneMKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NV} - m, n = _store_prologue!(destination, alpha, beta) + m, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination return _store_tile_lanepair!(destination, acc, alpha, beta, kernel, m, n) end diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 180ebcf..f18051c 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -12,13 +12,13 @@ Split-complex microkernel for `T = ComplexF32/ComplexF64` over `SIMD.Vec{W,real(T)}` lanes. `MR`/`NR` are complex extents; `MR` must be a multiple of `W`, which counts reals. """ -struct PlanarKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} +struct PlanarKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} function PlanarKernel{MR, NR, T, W}( descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} ) where {MR, NR, T, W} - _check_vector_shape("PlanarKernel", MR, W) + check_vector_shape("PlanarKernel", MR, W) return new{MR, NR, T, W}(descriptor) end end @@ -31,7 +31,7 @@ end function PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} T <: Complex || throw(ArgumentError("PlanarKernel requires a complex element type, got $T")) - return PlanarKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(real(T)))) + return PlanarKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) end complex_method(::PlanarKernel) = PlanarMethod() @@ -61,7 +61,7 @@ end ) where {MR, NR, T, W, R, NA, PA, PB} MV = MR ÷ W NV = MV * NR - _check_acc(:_accumulate_step_planar, R, T, NA, 2NV) + check_acc(:_accumulate_step_planar, R, T, NA, 2NV) arv = [Symbol(:ar, v) for v in 0:(MV - 1)] aiv = [Symbol(:ai, v) for v in 0:(MV - 1)] @@ -122,12 +122,12 @@ end end end -function Base.accumulate( +function add_tile( kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} k_block_length == 0 && return acc - k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) @inbounds for p in 1:k_block_length acc = _accumulate_step_planar(kernel, acc, packed_a, packed_b, p) end @@ -147,12 +147,12 @@ end # Scalar fallback store, for every destination the vector store cannot take. @generated function _store_tile_planar!( destination::Tile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, + alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, m::Int, n::Int ) where {MR, NR, T, W, R, NA} MV = MR ÷ W NV = MV * NR - _check_acc(:_store_tile_planar!, R, T, NA, 2NV) + check_acc(:_store_tile_planar!, R, T, NA, 2NV) blocks = Any[] for j in 1:NR, v in 0:(MV - 1) @@ -165,7 +165,7 @@ end for lane in 1:$W i = $(v * W) + lane i <= m || break - _axpby_tile!( + axpby_tile!( destination, i, $j, alpha, Complex(revec[lane], imvec[lane]), beta ) @@ -210,7 +210,7 @@ end # One full `W`-row block; `at` is the zero-based index of its first real. # The arithmetic transcribes Base's `Complex` expression trees (the ones -# `_axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, +# `axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, # `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, # muladd(zi, wr, xi)))`. Against the scalar fallback it is exact at # `beta == 0/1` and ~1 ULP otherwise, because LLVM contracts Base's scalar @@ -244,7 +244,7 @@ end # block of `2W` reals. The raw pointer is only dereferenced inside `GC.@preserve`. @generated function _store_tile_planar_vector!( destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::DescriptorKernel{MR, NR, T}, + alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, m::Int, n::Int ) where {S, MR, NR, T, W, R, NA} # The pointer reinterpretation is only sound on dense rank-1 complex storage. @@ -253,7 +253,7 @@ end RC = real(eltype(S)) MV = MR ÷ W NV = MV * NR - _check_acc(:_store_tile_planar_vector!, R, T, NA, 2NV) + check_acc(:_store_tile_planar_vector!, R, T, NA, 2NV) blocks = Any[] for j in 1:NR @@ -273,7 +273,7 @@ end for lane in 1:$W i = $(v * W) + lane i <= m || break - _axpby_at!( + axpby_at!( storage, colbase + i, alpha, Complex(revec[lane], imvec[lane]), beta ) @@ -316,7 +316,7 @@ end destination::Tile, acc::NTuple{NA, Vec{W, R}}, alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W} ) where {MR, NR, T, W, R, NA} - m, n = _store_prologue!(destination, alpha, beta) + m, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination if _complex_vector_eligible(destination, T) diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl index 342d690..634efdf 100644 --- a/src/microkernels/scalar.jl +++ b/src/microkernels/scalar.jl @@ -2,43 +2,59 @@ ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) Scalar reference microkernel for register tile `(MR, NR)` and element type `T`. +Its accumulator is an `NTuple{MR * NR, T}` holding the tile column-major. """ -struct ScalarKernel{MR, NR, T} <: DescriptorKernel{MR, NR, T} +struct ScalarKernel{MR, NR, T} <: Microkernel{MR, NR, T} descriptor::RealDescriptor{MR, NR, T} end ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = ScalarKernel(Descriptor(Val(MR), Val(NR), T)) -zero_accumulator(kernel::ScalarKernel{MR, NR, T}) where {MR, NR, T} = zeros(T, MR, NR) +zero_accumulator(::ScalarKernel{MR, NR, T}) where {MR, NR, T} = ntuple(_ -> zero(T), Val(MR * NR)) -# Extends `Base.accumulate` so `using QuasiStrided` does not clash with it. -function Base.accumulate( - kernel::ScalarKernel{MR, NR, T}, acc::AbstractMatrix{T}, +function add_tile( + kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, PA <: PackedPanel, PB} + ) where {MR, NR, T, N, PA <: PackedPanel, PB} k_block_length == 0 && return acc - k_block_length > 0 || throw(ArgumentError("accumulate requires k_block_length >= 0, got k_block_length = $k_block_length")) - @inbounds for p in 1:k_block_length - for j in 1:NR - bj = panel_load(packed_b, packed_b_offset(kernel, j, p)) - for i in 1:MR - ai = panel_load(packed_a, packed_a_offset(kernel, i, p)) - acc[i, j] = muladd(ai, bj, acc[i, j]) - end - end + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) + for p in 1:k_block_length + acc = add_step(kernel, acc, packed_a, packed_b, p) end return acc end +# One K step: element `(i, j)` of the tile is `acc[i + MR * (j - 1)]`. +@inline function add_step( + kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, packed_a, packed_b, p::Int + ) where {MR, NR, T, N} + return ntuple(Val(N)) do q + i, j = mod1(q, MR), cld(q, MR) + a = panel_load(packed_a, packed_a_offset(kernel, i, p)) + b = panel_load(packed_b, packed_b_offset(kernel, j, p)) + muladd(a, b, acc[q]) + end +end + function store_tile!( - destination::Tile, acc::AbstractMatrix{T}, - alpha::T, beta::T, kernel::ScalarKernel - ) where {T} - m, n = _store_prologue!(destination, alpha, beta) + destination::Tile, acc::NTuple{N, T}, + alpha::T, beta::T, kernel::ScalarKernel{MR} + ) where {MR, N, T} + m, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination - @inbounds for j in 1:n, i in 1:m - _axpby_tile!(destination, i, j, alpha, acc[i, j], beta) + if iszero(beta) + @inbounds for j in 1:n, i in 1:m + destination[i, j] = alpha * acc[i + MR * (j - 1)] + end + elseif isone(beta) + @inbounds for j in 1:n, i in 1:m + destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], convert(T, destination[i, j])) + end + else + @inbounds for j in 1:n, i in 1:m + destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], beta * convert(T, destination[i, j])) + end end return destination end diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index f594376..6040cc1 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -15,11 +15,11 @@ using SIMD: Vec, vload, vstore Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). """ -struct SIMDKernel{MR, NR, T, W} <: DescriptorKernel{MR, NR, T} +struct SIMDKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} descriptor::RealDescriptor{MR, NR, T} function SIMDKernel{MR, NR, T, W}(descriptor::RealDescriptor{MR, NR, T}) where {MR, NR, T, W} - _check_vector_shape("SIMDKernel", MR, W) + check_vector_shape("SIMDKernel", MR, W) return new{MR, NR, T, W}(descriptor) end end @@ -27,7 +27,7 @@ end SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} = SIMDKernel{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T)) SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = - SIMDKernel(Val(MR), Val(NR), T, Val(_default_lanewidth(T))) + SIMDKernel(Val(MR), Val(NR), T, Val(default_lanewidth(T))) lanewidth(::SIMDKernel{MR, NR, T, W}) where {MR, NR, T, W} = W @@ -48,7 +48,7 @@ end packed_a::PA, packed_b::PB, p::Int ) where {MR, NR, T, W, NV, PA, PB} NVECA = MR ÷ W - _check_acc(:_accumulate_step, T, T, NV, NVECA * NR) + check_acc(:_accumulate_step, T, T, NV, NVECA * NR) avars = [Symbol(:a, v) for v in 0:(NVECA - 1)] bvars = [Symbol(:b, j) for j in 1:NR] @@ -78,12 +78,12 @@ end end end -@inline function Base.accumulate( +@inline function add_tile( kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, packed_a::PA, packed_b::PB, k_block_length::Int ) where {MR, NR, T, W, NV, PA <: PackedPanel, PB} k_block_length == 0 && return acc - k_block_length > 0 || _throw_negative_k_block_length(:accumulate, k_block_length) + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) @inbounds for p in 1:k_block_length acc = _accumulate_step(kernel, acc, packed_a, packed_b, p) end @@ -112,7 +112,7 @@ end for lane in 1:$W i = $(v * W) + lane i <= m || break - _axpby_tile!(destination, i, $j, alpha, vec[lane], beta) + axpby_tile!(destination, i, $j, alpha, vec[lane], beta) end end end @@ -160,7 +160,7 @@ end for lane in 1:$W i = $(v * W) + lane i <= m || break - _axpby_at!(storage, colbase + i, alpha, vec[lane], beta) + axpby_at!(storage, colbase + i, alpha, vec[lane], beta) end end end @@ -194,7 +194,7 @@ end destination::Tile, acc::NTuple{NV, Vec{W, T}}, alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W} ) where {MR, NR, T, W, NV} - m, n = _store_prologue!(destination, alpha, beta) + m, n = store_prologue!(destination, alpha, beta) (m == 0 || n == 0) && return destination if _vector_store_eligible(destination, T) diff --git a/src/packing/pack.jl b/src/packing/pack.jl index edca708..c3c2ff2 100644 --- a/src/packing/pack.jl +++ b/src/packing/pack.jl @@ -30,27 +30,17 @@ end ) ) -# `transform` is applied to each loaded (complex) element before it is split -# into the packed format; padding lanes are literal zeros and never read the -# source or call `transform`. All validation happens before any write. -pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V, F} = - pack_tile!(panel, tile, spec, transform, Val(true)) +""" + pack!(panel, tile::Tile, spec::SliverSpec, transform) -> panel -# The packer writes only `PackedPanel`s; a `DenseVector` is borrowed as one. -function pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V <: DenseVector, F} - GC.@preserve panel pack!(packed_panel(panel, 1, length(panel)), tile, spec, transform) - return panel -end - -# Skips only `checked_tile_storage_bounds(tile)`: the caller (`_execute_nest!`) -# has already validated the whole macro block the sliver belongs to. `@inline` -# because out of line each call marshals the `Tile` through the stack. -@inline unsafe_pack!(panel::V, tile::Tile, spec::SliverSpec, transform::F) where {V, F} = - pack_tile!(panel, tile, spec, transform, Val(false)) - -@inline function pack_tile!( - panel::V, tile::Tile, spec::SliverSpec{I, L}, transform::F, ::Val{BOUNDS} - ) where {V, I, L, F, BOUNDS} +Pack one sliver of `tile` (lanes along its rows, K steps along its columns) +into `panel` in `spec`'s format. `transform` is applied to each loaded +element before it is split into the format; padding lanes are literal zeros +and never read the source or call `transform`. All validation happens before +any write. Under `@inbounds` (the call inlined) the storage bounds check of +`tile` is skipped; the eltype, lane count and panel length checks stay. +""" +@inline function pack!(panel::V, tile::Tile, spec::SliverSpec{I, L}, transform::F) where {V, I, L, F} check_packed_eltype(panel, spec) lanes, k_block_length = size(tile) (0 <= lanes <= L) || throw_pack_extent(I, lanes, L) @@ -59,7 +49,7 @@ end # `<= 0`, not `== 0` (counts are never negative): LLVM may then assume # `k_block_length >= 1` in the loops below. k_block_length <= 0 && return panel - BOUNDS && checked_tile_storage_bounds(tile) + @boundscheck checked_tile_storage_bounds(tile) if real_contiguous_eligible(tile, spec, transform, lanes) return pack_real_contiguous!(panel, tile, spec, k_block_length) elseif complex_contiguous_eligible(tile, spec, transform, lanes) @@ -68,6 +58,14 @@ end return pack_scalar!(panel, tile, spec, transform, lanes, k_block_length) end +# The packer writes only `PackedPanel`s; a `DenseVector` is borrowed as one. +Base.@propagate_inbounds function pack!( + panel::V, tile::Tile, spec::SliverSpec, transform::F + ) where {V <: DenseVector, F} + GC.@preserve panel pack!(packed_panel(panel, 1, length(panel)), tile, spec, transform) + return panel +end + # What a packed element converts to: the real operand of a mixed-domain kernel # packs `real(T)`. element_type(::SliverSpec{I, L, RealFormat, T}) where {I, L, T} = real(T) diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index 884b822..509595f 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -22,8 +22,10 @@ end # Complex blocking is the real row divided by the packed reals per element of # each operand, so every method gets the same packed BYTE budget (1m's # `m_block` is half planar's). -@inline _scale_blocking(base::Blocking, m::ComplexMethod) = - Blocking(max(1, base.m_block ÷ a_reals(m)), base.k_block, max(1, base.n_block ÷ b_reals(m))) +@inline function _scale_blocking(base::Blocking, m::KernelMethod) + a_reals, b_reals = map(reals_per_element, pack_formats(m)) + return Blocking(max(1, base.m_block ÷ a_reals), base.k_block, max(1, base.n_block ÷ b_reals)) +end # For a host whose L1d or L2 size is undetected. _fallback_blocking(::Type{Float64}) = Blocking(128, 256, 768) diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index 32450b0..c869f5e 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -77,7 +77,7 @@ _default_kernel(::Type{T}) where {T} = end # The automatic `(shape, method)` under `method`, `_default_method(T, TA, TB)`. -@inline _default_shape(::Type{T}, ::ComplexMethod, m_length::Int, n_length::Int, run::Int) where {T} = +@inline _default_shape(::Type{T}, ::KernelMethod, m_length::Int, n_length::Int, run::Int) where {T} = _default_shape(T, m_length, n_length, run) # The real default shape of the real problem, mapped: the real extent, store and diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index 69ba47a..0204cc4 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -22,7 +22,7 @@ _default_method(::Type{T}, ::Type{<:Real}, ::Type{<:Complex}) where {T <: Comple const NR_DEFAULT = 6 # An (8, 6) tile at one 256-bit register's lane width: the shape wherever no rule applies. -_fallback_shape(::Type{T}) where {T} = (8, NR_DEFAULT, _default_lanewidth(real(T))) +_fallback_shape(::Type{T}) where {T} = (8, NR_DEFAULT, default_lanewidth(real(T))) # Closed menus of `(MR, NR, W)` shapes, so the compiled specializations stay # bounded. Complex `MR` counts complex rows and `W` real lanes; 1m runs a real @@ -45,7 +45,7 @@ const KERNEL_SHAPES_C64_FMADDSUB = ((12, 8, 8), (8, 8, 8), (4, 6, 4), (4, 5, 4)) const KERNEL_SHAPES_C32_FMADDSUB = ((24, 8, 16), (16, 8, 16), (8, 6, 8), (8, 5, 8)) """ - kernel_shapes(T, method::ComplexMethod = _default_method(T)) -> NTuple{<:Any,NTuple{3,Int}} + kernel_shapes(T, method::KernelMethod = _default_method(T)) -> NTuple{<:Any,NTuple{3,Int}} The closed menu of `(MR, NR, W)` register shapes the engine may build for element type `T` under `method`. @@ -130,7 +130,7 @@ _kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = # Per-ISA shapes, consulted first. None for real types: the rule below is the optimum. _shape_override(key::Val, ::Type{T}) where {T} = _shape_override(key, T, _default_method(T)) -_shape_override(::Val, ::Type, ::ComplexMethod) = nothing +_shape_override(::Val, ::Type, ::KernelMethod) = nothing # On AVX-512 the rule's planar `MR = 2W, NR = 6` spills; `NR = 3` tiles do not. _shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::PlanarMethod) = (24, 3, 8) _shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::PlanarMethod) = (48, 3, 16) @@ -150,7 +150,7 @@ _shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::OneMMethod) = (4, 6, 4) _rule_applies(::Val{:avx512}, ::RealMethod) = true _rule_applies(::Val{:avx512}, ::Union{PlanarMethod, OneMMethod}) = true _rule_applies(::Val{:avx2}, ::RealMethod) = true -_rule_applies(::Val, ::ComplexMethod) = false +_rule_applies(::Val, ::KernelMethod) = false # `W` real lanes per register, `MV` A vectors per column. _rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = @@ -160,7 +160,7 @@ _rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = # cores with 2 FMA ports, and 4*6 accumulators + 4 A vectors + 2 still fit 32 # registers. Complex methods keep MV = 2 (planar already spills there). _rule_mv(::Val{:avx512}, ::RealMethod) = 4 -_rule_mv(::Val, ::ComplexMethod) = 2 +_rule_mv(::Val, ::KernelMethod) = 2 # AMD's AVX-512 cores double-pump 512-bit FMAs, so their MV = 2 tile is not # front-end bound and the taller tile only adds edge and store cost. @@ -204,7 +204,7 @@ _fitted_shape(::TargetProfile, ::Type{T}, method::Union{OneMMethod, FMAddSubMeth function _fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) where {T} R = real(T) vb = profile.vector_bytes - lanes = (vb > 0 && vb % sizeof(R) == 0) ? vb ÷ sizeof(R) : _default_lanewidth(R) + lanes = (vb > 0 && vb % sizeof(R) == 0) ? vb ÷ sizeof(R) : default_lanewidth(R) budget = profile.nregisters > 0 ? profile.nregisters : 16 best = nothing for shape in kernel_shapes(T, method) diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index b9c1142..92b9808 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -199,9 +199,8 @@ end allocs_tw = _steady_allocs!(execute_tilewise!, plans[4], Cmat) @test Cmat ≈ Amat * Bmat @test allocs_tw == 0 skip = (VERSION < v"1.11") - # ScalarKernel's accumulator is a Matrix: bounded, one per tile call. ps = _mm_plan(Cmat, Amat, Bmat; kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 6, n_block = 7) - @test _steady_allocs!(execute!, ps, Cmat) <= 176 * cld(Ma, 4) * cld(Na, 3) * cld(Ka, 6) + 1 + @test _steady_allocs!(execute!, ps, Cmat) == 0 skip = (VERSION < v"1.11") end _ws_union_members(t) = t isa Union ? (_ws_union_members(t.a)..., _ws_union_members(t.b)...) : (t,) diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 6e964a4..78d6b26 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -68,7 +68,7 @@ const _QSF = QuasiStrided end @testset "instruction selection: vfmaddsub, no separate mul/add/sub, no spill" begin - # Exact instruction counts in the hot loop of `accumulate`, so only on + # Exact instruction counts in the hot loop of `add_tile`, so only on # an FMA3 x86 host and not on hosted CI, whose virtualized CPU feature # sets do not reliably match. isa = target_profile().isa @@ -93,7 +93,7 @@ const _QSF = QuasiStrided R = real(T) asm = sprint() do io code_native( - io, Base.accumulate, + io, QuasiStrided.add_tile, (typeof(k), typeof(zero_accumulator(k)), PackedPanel{R}, PackedPanel{R}, Int); debuginfo = :none, syntax = :intel ) diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index e92dcb8..3dbb5bb 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -43,7 +43,7 @@ end 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 ( + @testset "pack -> add_tile -> 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), diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 0d33402..fcd7e3c 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -28,7 +28,7 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEForm @test_throws ArgumentError OneMKernel{12, 8, ComplexF64, 8, typeof(wrong)}(d, wrong) # The error reports the logical k_block_length, not the doubled real one. err = try - accumulate(k, zero_accumulator(k), Float64[], Float64[], -3) + add_tile(k, zero_accumulator(k), Float64[], Float64[], -3) catch e e end diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 2fb8900..0a041f9 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -3,7 +3,7 @@ # # Values are pinned two ways. On the elements the fast path vectorizes (full # row blocks) it must match, bitwise, an independent transcription of Base's -# `Complex` `*`/`muladd` expression trees (the arithmetic `_axpby_tile!` does). +# `Complex` `*`/`muladd` expression trees (the arithmetic `axpby_tile!` does). # Against the scalar store it is compared with a tolerance only: LLVM contracts # Base's scalar complex `muladd` depending on inlining context, so the scalar # path is not bit-reproducible even against itself. @@ -36,7 +36,7 @@ function ref_axpby(alpha::T, rr::R, ri::R, beta::T, cold::T) where {T, R} return Complex(fma(ar, rr, barrier(-fma(ai, ri, -xr))), fma(ar, ri, barrier(fma(ai, rr, xi)))) end -# Every `_axpby_tile!` branch, plus purely imaginary and real non-unit beta. +# Every `axpby_tile!` branch, plus purely imaginary and real non-unit beta. mk_store_ab(T) = ( (one(T), zero(T)), (T(-0.5, 0.25), zero(T)), (one(T), one(T)), (T(2, -1), one(T)), (T(2.5, -1), T(-1.75, 0.5)), (one(T), T(2, 0)), (T(1, 1), T(0, -3)), diff --git a/test/microkernels/test_scalar_kernel.jl b/test/microkernels/test_scalar_kernel.jl index 77c8a9f..3af1d1a 100644 --- a/test/microkernels/test_scalar_kernel.jl +++ b/test/microkernels/test_scalar_kernel.jl @@ -25,7 +25,7 @@ mk_pack_b(k, B) = mk_cols(c -> [real(c); imag(c)], permutedims(B)) mk_pack(k, A, B) = (mk_pack_a(k, A), mk_pack_b(k, B)) # Element (i, j), zero-based, of an accumulator. -mk_read(::ScalarKernel, acc, i, j) = acc[i + 1, j + 1] +mk_read(k::ScalarKernel, acc, i, j) = acc[i + tile_size(k, 1) * j + 1] function mk_read(k::SIMDKernel, acc, i, j) W = lanewidth(k) return acc[i ÷ W + (tile_size(k, 1) ÷ W) * j + 1][i % W + 1] @@ -42,7 +42,7 @@ function mk_read(k::Union{OneMKernel, FMAddSubKernel}, acc, i, j) return Complex(vec[2u + 1], vec[2u + 2]) end -mk_fill_acc(k, x) = (acc = zero_accumulator(k); acc isa AbstractMatrix ? fill(x, size(acc)) : map(v -> typeof(v)(x), acc)) +mk_fill_acc(k, x) = map(v -> typeof(v)(x), zero_accumulator(k)) mk_nan(T) = T <: Complex ? T(NaN, NaN) : T(NaN) mk_inf(T) = T <: Complex ? T(Inf, Inf) : T(Inf) mk_tol(T) = real(T) === Float64 ? 1.0e-11 : 2.0f-4 @@ -57,10 +57,10 @@ mk_dense(v::AbstractVector{T}) where {T} = copyto!(parent(StridedView(zeros(T, l # 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, k_block_length) = accumulate(k, zero_accumulator(k), pa, pb, k_block_length) +mk_run_acc(k, pa, pb, k_block_length) = add_tile(k, zero_accumulator(k), pa, pb, k_block_length) mk_run_exec(k, dst, pa, pb, k_block_length, alpha, beta) = (execute_tile!(k, dst, pa, pb, k_block_length, alpha, beta); nothing) -# `full = false` checks only accumulate against the reference and allocations. +# `full = false` checks only add_tile against the reference and allocations. # Separate functions, so the light check does not compile the full one. function mk_contract(k; full::Bool = true) return @testset "$(nameof(typeof(k))){$(tile_size(k, 1)),$(tile_size(k, 2)),$(scalartype(k))}" begin @@ -86,17 +86,15 @@ function mk_contract_light(k) # with a partial row block. Julia 1.10 does not keep the tuple # accumulator in registers. m = max(1, MR - 1) - if !(k isa ScalarKernel) - skip = VERSION < v"1.11" - ab = (T(2), T(0.5)) - scat = Tile(zeros(T, MR * NR), 0, view(collect(0:(MR - 1)), 1:MR), AffineAxis(0, MR, NR)) - part = Tile(mk_dense(zeros(T, m * NR)), 0, AffineAxis(0, 1, m), AffineAxis(0, m, NR)) - mk_run_acc(k, pa, pb, k_block_length) - @test (@allocated mk_run_acc(k, pa, pb, k_block_length)) == 0 skip = skip - for dst in (scat, part) - mk_run_exec(k, dst, pa, pb, k_block_length, ab...) - @test (@allocated mk_run_exec(k, dst, pa, pb, k_block_length, ab...)) == 0 skip = skip - end + skip = VERSION < v"1.11" + ab = (T(2), T(0.5)) + scat = Tile(zeros(T, MR * NR), 0, view(collect(0:(MR - 1)), 1:MR), AffineAxis(0, MR, NR)) + part = Tile(mk_dense(zeros(T, m * NR)), 0, AffineAxis(0, 1, m), AffineAxis(0, m, NR)) + mk_run_acc(k, pa, pb, k_block_length) + @test (@allocated mk_run_acc(k, pa, pb, k_block_length)) == 0 skip = skip + for dst in (scat, part) + mk_run_exec(k, dst, pa, pb, k_block_length, ab...) + @test (@allocated mk_run_exec(k, dst, pa, pb, k_block_length, ab...)) == 0 skip = skip end return nothing end @@ -116,13 +114,13 @@ function mk_contract_full(k) acc0 = zero_accumulator(k) @test all(iszero(mk_read(k, acc0, i, j)) for i in 0:(MR - 1), j in 0:(NR - 1)) - @test accumulate(k, acc0, R[], R[], 0) === acc0 - @test_throws ArgumentError accumulate(k, acc0, R[], R[], -1) + @test add_tile(k, acc0, R[], R[], 0) === acc0 + @test_throws ArgumentError add_tile(k, acc0, R[], R[], -1) # Splitting k_block_length across calls composes. la, lb = length(pa) ÷ k_block_length, length(pb) ÷ k_block_length accs = zero_accumulator(k) for p in 0:(k_block_length - 1) - accs = accumulate(k, accs, pa[(p * la + 1):((p + 1) * la)], pb[(p * lb + 1):((p + 1) * lb)], 1) + accs = add_tile(k, accs, pa[(p * la + 1):((p + 1) * la)], pb[(p * lb + 1):((p + 1) * lb)], 1) end @test all(mk_read(k, accs, i, j) ≈ mk_read(k, acc, i, j) for i in 0:(MR - 1), j in 0:(NR - 1)) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index 7775764..6c46416 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -312,13 +312,26 @@ end ) === nothing end +@testset "pack_formats agrees with every kernel's descriptor" begin + QS = QuasiStrided + kernels = ( + ScalarKernel(Val(4), Val(2), Float64), SIMDKernel(Val(8), Val(4), Float32), + QS.PlanarKernel(Val(4), Val(2), ComplexF64), QS.OneMKernel(Val(4), Val(2), ComplexF64), + QS.FMAddSubKernel(Val(4), Val(2), ComplexF64), QS.ComplexRealKernel(Val(4), Val(2), ComplexF64), + QS.RealComplexKernel(Val(4), Val(2), ComplexF64), + ) + for k in kernels + @test pack_formats(complex_method(k)) === (QS.a_format(k), QS.b_format(k)) + end +end + @testset "complex blocking is the real row scaled by packed reals" begin for base in (_fallback_blocking(Float64), _fallback_blocking(Float32), Blocking(1, 5, 1)) for m in (PlanarMethod(), OneMMethod(), FMAddSubMethod()) b = _scale_blocking(base, m) @test b.k_block === base.k_block - @test b.m_block === max(1, base.m_block ÷ a_reals(m)) - @test b.n_block === max(1, base.n_block ÷ b_reals(m)) + @test b.m_block === max(1, base.m_block ÷ reals_per_element(pack_formats(m)[1])) + @test b.n_block === max(1, base.n_block ÷ reals_per_element(pack_formats(m)[2])) end end @test _scale_blocking(_fallback_blocking(Float64), PlanarMethod()) === Blocking(64, 256, 384) diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index baa239a..713f8fa 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -219,8 +219,11 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() _pcf_exec(p, 1.0, 0.0) @test reshape(cstore.data, Ma, Na) ≈ Amat * Bmat - @test cstore.n == 1 - @test astore.n == 1 + # `--check-bounds=yes` (as under `Pkg.test`) also runs the per-tile and + # per-sliver checks that `@inbounds` skips. + forced = Base.JLOptions().check_bounds == 1 + @test cstore.n == 1 + forced * ntiles + @test astore.n == 1 + forced * m_tiles @test bstore.n == 1 # The oracle checks per tile. @@ -345,36 +348,26 @@ end @test_throws BoundsError _pcf_exec(plow, 1.0, 0.0) end -@testset "per-call floor: unsafe_pack! keeps every non-bounds check" begin +# `@inbounds` reaches the storage check only through an inlined call. +pcf_inbounds_pack!(args...) = @inbounds pack!(args...) +pcf_inbounds_execute_tile!(args...) = @inbounds execute_tile!(args...) + +@testset "per-call floor: @inbounds skips only the storage bounds check" begin kernel = SIMDKernel(Val(8), Val(6), Float64) src = Tile(collect(1.0:64.0), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 4)) - packed = zeros(8 * 4) - a = sliver_spec(kernel, 1) - ref = zeros(8 * 4) - pack!(ref, src, a, identity) - GC.@preserve packed unsafe_pack!(QS.packed_panel(packed, 1, length(packed)), src, a, identity) - @test packed == ref - - # Only the bounds checks are skipped. - @test_throws DimensionMismatch unsafe_pack!(zeros(3), src, a, identity) - @test_throws ArgumentError unsafe_pack!(zeros(Float32, 64), src, a, identity) + @test_throws DimensionMismatch pcf_inbounds_pack!(zeros(3), src, a, identity) + @test_throws ArgumentError pcf_inbounds_pack!(zeros(Float32, 64), src, a, identity) toowide = Tile(collect(1.0:200.0), 0, AffineAxis(0, 1, 9), AffineAxis(0, 16, 4)) - @test_throws ArgumentError unsafe_pack!(zeros(200), toowide, a, identity) + @test_throws ArgumentError pcf_inbounds_pack!(zeros(200), toowide, a, identity) dest = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) - @test_throws DimensionMismatch QS.unsafe_execute_tile!( + @test_throws DimensionMismatch pcf_inbounds_execute_tile!( kernel, dest, zeros(3), zeros(6 * 4), 4, 1.0, 0.0 ) - @test_throws ArgumentError QS.unsafe_execute_tile!( + @test_throws ArgumentError pcf_inbounds_execute_tile!( kernel, dest, zeros(8 * 4), zeros(6 * 4), -1, 1.0, 0.0 ) - d1 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) - d2 = Tile(zeros(8 * 6), 0, AffineAxis(0, 1, 8), AffineAxis(0, 8, 6)) - packedB = collect(1.0:24.0) - execute_tile!(kernel, d1, ref, packedB, 4, 1.0, 0.0) - QS.unsafe_execute_tile!(kernel, d2, ref, packedB, 4, 1.0, 0.0) - @test d1.storage == d2.storage end @testset "per-call floor: affine_ramp classifies exactly the rank-<=1 folds" begin diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 78caa65..0b7c2fa 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -596,7 +596,7 @@ end # the flag XOR A's `conj` op. d = 4 for T in (ComplexF64, ComplexF32) - W = QuasiStrided._default_lanewidth(real(T)) + W = QuasiStrided.default_lanewidth(real(T)) kernel = QuasiStrided.PlanarKernel(Val(W), Val(8), T, Val(W)) @test tile_size(kernel, 1) <= 16 (indA, indB, indC), _ = _lo_labels(_LO_CASES[3][2], _LO_CASES[3][3]) diff --git a/test/runtests.jl b/test/runtests.jl index 34c8e37..e8c9411 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,15 +5,15 @@ using QuasiStrided using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, describe_block, block_descriptors!, Descriptor, tile_size, scalartype, packed_a_offset, packed_b_offset, packed_a_length, packed_b_length, - AffineAxis, Tile, checked_tile_storage_bounds, pack!, unsafe_pack!, sliver_spec, - zero_accumulator, accumulate, store_tile!, execute_tile!, lanewidth, contract!, + AffineAxis, Tile, checked_tile_storage_bounds, pack!, sliver_spec, + zero_accumulator, add_tile, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, target_profile, cache_topology, unknown_target, detect_isa, detect_target, _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, _rule_applies, isa_nregisters, sliver_width, realtype, - complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, a_reals, - b_reals, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, + complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, pack_formats, + reals_per_element, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, _planar_pressure, _pack_split, _NestPath # plan_contract, execute! and ContractPlan are bound in helpers.jl instead. From c8ffbd3a3f7247f18c92b6663e962e143ab65dca Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 16:55:23 -0400 Subject: [PATCH 13/60] Decision log: chunk 6 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/CLEANUP.md b/CLEANUP.md index cbbf1fd..b4f0d8e 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -18,7 +18,7 @@ TTFX when codegen is touched). | 3 | Tiles | `layout/tiles.jl` | done | | 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | done | | 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | done | -| 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | | +| 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | done | | 7 | SIMD kernels | `simd.jl`, `planar.jl` | | | 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | | | 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | @@ -131,6 +131,17 @@ keeping). One `emit!` per format taking the values to store; padding passes literal zeros. Vector-path predicates (`dense_lanes`, `is_unit_stride`, `complex_fastpath_isa_eligible`) live in pack_contiguous.jl. +### D10. Kernel interface (chunk 6, applied) + +`add_tile` (not an extension of `Base.accumulate`); `Microkernel` and +`KernelMethod` replace `DescriptorKernel`/`ComplexMethod`; each method +declares `pack_formats`, from which the reals per packed element follow. One +`execute_tile!` and one `pack!`, whose per-tile storage check is a +`@boundscheck` the nest skips with `@inbounds` (the macro block is span-checked; +under `--check-bounds=yes`, as in the tests, the checks run). `ScalarKernel` +uses a tuple accumulator (allocation-free) and branches on `beta` once per +tile. The contract lives in docstrings. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: From 95728d00ddebf77f6722764d9fed4345a3b5a2a7 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 17:21:11 -0400 Subject: [PATCH 14/60] Decision log: chunks 7 and 8 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index b4f0d8e..cbd33ed 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -19,8 +19,8 @@ TTFX when codegen is touched). | 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | done | | 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | done | | 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | done | -| 7 | SIMD kernels | `simd.jl`, `planar.jl` | | -| 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | | +| 7 | SIMD kernels | `simd.jl`, `planar.jl` | in progress (with 8) | +| 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | in progress (with 7) | | 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | | 10 | Kernel selection | `planning/kernel_selection.jl` | | | 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | @@ -142,6 +142,32 @@ under `--check-bounds=yes`, as in the tests, the checks run). `ScalarKernel` uses a tuple accumulator (allocation-free) and branches on `beta` once per tile. The contract lives in docstrings. +### D11. Vector kernels (chunk 7, decided; applied together with chunk 8) + +One abstract `VectorKernel{MR, NR, T, W}` with generic `add_tile`, +`store_tile!` and scattered store; each kernel supplies its K step, its +full-block vector store and a generator-time accumulator-index function. All +kernel structs in one types file before the implementations, so layout +methods dispatch instead of `<:` tests (a generator only sees methods defined +before it). Kept as is: `@generated` (not `ntuple`); `muladd(-a, b, c)` (it +already compiles to one `vfnmadd`, no separate negation); the per-block `beta` tests in +the vector stores are left for now (LLVM does not unswitch them: about 4 +compares per block). Hoisting is decided after chunk 8, leaning towards doing +it: each executed variant is smaller and branch-free (less code fetched per +store), at the cost of compile time and code size; measure both. `b_step_load*` hooks get descriptive +names. + +### D12. Kernels = K step + accumulator layout (chunks 7–8) + +Three K steps (real, planar, fmaddsub) and three accumulator layouts (real, +split re/im, lane pairs). `accumulator_layout(kernel)` selects one generic +`store_tile!` (vector or scalar store, shared generated block/tail loop; each +layout supplies its full-block store and lane value); one generic `add_tile` +over `accumulate_step` with `k_steps` (1m: 2k). `inner(k)` is computed, not a +field (no `KI` parameter). One generic constructor without `W`; generic +`complex_method`/`lanewidth`. 1m gets the lane-pair vector store. The `beta` +hoist is a separate commit on top, measured on its own. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: From c9a9f3aa67c4d9caa3b0f1ed7d01469146083c25 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 18:35:36 -0400 Subject: [PATCH 15/60] Kernels: one K step plus an accumulator layout All vector kernel structs move to src/microkernels/types.jl under an abstract VectorKernel{MR, NR, T, W}, with three accumulator layouts (RealLayout, SplitLayout, LanePairLayout) and the generator-time pieces of the stores (index, lane value, full-block store, store prelude), so they dispatch instead of `<:` tests. One generic add_tile loops accumulate_step over k_steps (1m: 2k); 1m and the mixed-domain kernels forward their step to a computed inner(kernel) (no inner field, no KI parameter; the call folds to a singleton). One generic store_tile!, vector store and scalar store parameterised by the layout; @inline per layout as before (real/split yes, lane pairs no). One generic constructor with the complex-T check and vector-shape check, generic lanewidth and zero_accumulator. 1m gets the lane-pair vector store. Renames: b_scalar/b_complex (UnpackedBView hooks), accumulate_step, vector_store!/scalar_store!/vector_store_eligible, split_index, split_store_block!/lanepair_store_block!, fmaddsub/swap_pairs/addsub. Codegen (AVX-512, default shapes): K-loop instruction counts identical for all six kernels; store_tile!/vector/scalar store instruction counts identical (SIMD 2015/1953/1317, Planar 1511/1494/907, FMAddSub 101/3619/2423, ComplexReal 101/3612/2417, RealComplex 1976/1900/1202); 1m now has the 3612-instruction lane-pair vector store. add_tile is now @inline for every kernel (before: only SIMD and mixed). A/B vs 95728d0, medians of 5 interleaved runs: small-K (M = N = 200, K = 1..8) ComplexF64 -23..-35%, Float64 -3..+2%; 256^3 named kernels: SIMD -1%, Planar -2.5%, 1m -11%, FMAddSub -7%, ComplexReal +3%, RealComplex 0%; earlier matmul/scattered/row-major/scattered-K cases -4..+2%. Allocations 0 under --check-bounds=yes. TTFX: Float64 first call 3.1 s (before 3.1/3.4 s), ComplexF64 2.1 s (before 2.1/2.2 s). Tests: 31967 pass (+419 for the 1m vector store), suite ~18 min. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 1 + src/execution/outer.jl | 2 +- src/execution/unpackedb.jl | 6 +- src/microkernels/fmaddsub.jl | 174 ++------------ src/microkernels/mixed.jl | 105 +------- src/microkernels/onem.jl | 107 +-------- src/microkernels/planar.jl | 225 +----------------- src/microkernels/simd.jl | 205 ++++++++-------- src/microkernels/types.jl | 214 +++++++++++++++++ test/microkernels/test_fmaddsub_kernel.jl | 12 +- test/microkernels/test_mixed_kernel.jl | 4 +- test/microkernels/test_onem_kernel.jl | 11 +- .../test_planar_store_fastpath.jl | 15 +- test/microkernels/test_simd_kernel.jl | 22 +- 14 files changed, 383 insertions(+), 720 deletions(-) create mode 100644 src/microkernels/types.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 39e705f..da194d8 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -26,6 +26,7 @@ include("packing/pack_contiguous.jl") # --- Microkernels: add one packed K block to a register tile, store into C --- include("microkernels/interface.jl") +include("microkernels/types.jl") include("microkernels/scalar.jl") include("microkernels/simd.jl") include("microkernels/planar.jl") diff --git a/src/execution/outer.jl b/src/execution/outer.jl index 3f444ed..ce3bf14 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -1,7 +1,7 @@ # The outer-product path for `K == 1`, real `T`, M unit-stride in A and C: # `C[:, n] = alpha * A * B[n] + beta * C[:, n]` as a streaming write. Through # the nest a K = 1 contraction is all per-tile fixed cost. The arithmetic is -# the nest's term for term (`r = a * b`, then `_store_tile_vector!`'s +# the nest's term for term (`r = a * b`, then `vector_store!`'s # expressions), so the two agree bitwise except on a signed zero. Real only: # complex K = 1 tiles are not the bottleneck, and the path would need its own # interleaved arithmetic. diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 64b3a3b..5f5536d 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -26,16 +26,16 @@ Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) return u.transform(z) end -@inline _b_step_load(u::UnpackedBView, kernel, j::Int, p::Int) = +@inline b_scalar(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) +@inline function b_complex(u::UnpackedBView, kernel, j::Int, p::Int) z = convert(scalartype(kernel), _unpacked_b_element(u, j, p)) return (real(z), imag(z)) end -# The kernels whose K step reads B through `_b_step_load`/`_b_step_load2` +# The kernels whose K step reads B through `b_scalar`/`b_complex` # (not 1m, whose inner kernel walks the planar panel, nor the scalar one). @inline _unpacked_b_kernel_eligible(::SIMDKernel) = true @inline _unpacked_b_kernel_eligible(::PlanarKernel) = true diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 2b07e8a..8712b88 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -17,50 +17,15 @@ using SIMD: Vec, shufflevector -""" - FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over -`SIMD.Vec{W,real(T)}` lanes with `InterleavedFormat` A and `PlanarFormat` B. -`2MR` must be a multiple of `W`, and `W` must be even. The default for small-M -complex contractions on AVX-512. -""" -struct FMAddSubKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} - descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} - - function FMAddSubKernel{MR, NR, T, W}( - descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} - ) where {MR, NR, T, W} - check_vector_shape("FMAddSubKernel", 2 * MR, W, true) - return new{MR, NR, T, W}(descriptor) - end -end - -function FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} - T <: Complex || - throw(ArgumentError("FMAddSubKernel requires a complex element type, got $T")) - return FMAddSubKernel{MR, NR, T, W}( - Descriptor(Val(MR), Val(NR), T, InterleavedFormat(), PlanarFormat()) - ) -end -function FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} - T <: Complex || - throw(ArgumentError("FMAddSubKernel requires a complex element type, got $T")) - return FMAddSubKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) -end - -complex_method(::FMAddSubKernel) = FMAddSubMethod() -lanewidth(::FMAddSubKernel{MR, NR, T, W}) where {MR, NR, T, W} = W - # `x*y - c` in even lanes, `x*y + c` in odd lanes, each one fused rounding. # Generic IR (`fneg` + two `llvm.fma` + a blend), which the X86 backend folds # into one `vfmaddsub` wherever FMA3 exists (and elsewhere stays correct). # GUARDRAIL: the SIMD.jl spelling `shufflevector(muladd(x, y, -c), muladd(x, y, # c), ...)` also selects `vfmaddsub` but leaves a dead stack store of the # accumulator in the K loop, every step. -@generated function _fmaddsub(x::Vec{N, R}, y::Vec{N, R}, c::Vec{N, R}) where {N, R} +@generated function fmaddsub(x::Vec{N, R}, y::Vec{N, R}, c::Vec{N, R}) where {N, R} R === Float64 || R === Float32 || - return :(throw(ArgumentError("_fmaddsub: unsupported lane type $R"))) + return :(throw(ArgumentError("fmaddsub: unsupported lane type $R"))) ty = "<$N x $(R === Float64 ? "double" : "float")>" fn = "llvm.fma.v$(N)$(R === Float64 ? "f64" : "f32")" # lane k: even -> `s` (x*y - c); odd -> `d` (x*y + c), index N + k. @@ -87,25 +52,20 @@ lanewidth(::FMAddSubKernel{MR, NR, T, W}) where {MR, NR, T, W} = W end # `[x1, x0, x3, x2, ...]`: one in-lane `vshufpd`/`vpermilps`. -@generated function _swap_pairs(x::Vec{N, R}) where {N, R} - iseven(N) || return :(throw(ArgumentError("_swap_pairs: expected even N, got $N"))) +@generated function swap_pairs(x::Vec{N, R}) where {N, R} + iseven(N) || return :(throw(ArgumentError("swap_pairs: expected even N, got $N"))) idx = ntuple(k -> isodd(k) ? k : k - 2, N) return :(Base.@_inline_meta; shufflevector(x, Val($idx))) end -function zero_accumulator(kernel::FMAddSubKernel{MR, NR, T, W}) where {MR, NR, T, W} - z = zero(Vec{W, real(T)}) - return ntuple(_ -> z, Val(((2 * MR) ÷ W) * NR)) -end - # The INNER op must be the `swap(a) * bi` one: only the inner product is # subtracted in the real lanes. Reversed, the real part is `ai*bi - ar*br`. -@generated function _accumulate_step_fmaddsub( +@generated function accumulate_step( kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, p::Int ) where {MR, NR, T, W, R, NA, PA, PB} MV = (2 * MR) ÷ W - check_acc(:_accumulate_step_fmaddsub, R, T, NA, MV * NR) + check_acc(:accumulate_step, R, T, NA, MV * NR) av = [Symbol(:a, v) for v in 0:(MV - 1)] sv = [Symbol(:s, v) for v in 0:(MV - 1)] @@ -122,23 +82,23 @@ end ) ) ) - push!(load_a, :($(sv[v + 1]) = _swap_pairs($(av[v + 1])))) + push!(load_a, :($(sv[v + 1]) = swap_pairs($(av[v + 1])))) end load_b = Any[] for j in 1:NR - push!(load_b, :((br_s, bi_s) = _b_step_load2(packed_b, kernel, $j, p))) + push!(load_b, :((br_s, bi_s) = b_complex(packed_b, kernel, $j, p))) push!(load_b, :($(brv[j]) = Vec{$W, $R}(br_s))) push!(load_b, :($(biv[j]) = Vec{$W, $R}(bi_s))) end acc_exprs = Vector{Any}(undef, NA) for j in 1:NR, v in 0:(MV - 1) - idx = v + MV * (j - 1) + 1 + idx = acc_index(MV, v, j) acc_exprs[idx] = :( - _fmaddsub( + fmaddsub( $(av[v + 1]), $(brv[j]), - _fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) + fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) ) ) end @@ -153,129 +113,33 @@ end end end -function add_tile( - kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - @inbounds for p in 1:k_block_length - acc = _accumulate_step_fmaddsub(kernel, acc, packed_a, packed_b, p) - end - return acc -end - # `p - q` in even lanes, `p + q` in odd lanes: Base's unfused complex `*` on # interleaved data. -@generated function _addsub(p::Vec{N, R}, q::Vec{N, R}) where {N, R} - iseven(N) || return :(throw(ArgumentError("_addsub: expected even N, got $N"))) +@generated function addsub(p::Vec{N, R}, q::Vec{N, R}) where {N, R} + iseven(N) || return :(throw(ArgumentError("addsub: expected even N, got $N"))) idx = ntuple(k -> iseven(k - 1) ? k - 1 : N + k - 1, N) return :(Base.@_inline_meta; shufflevector(p - q, p + q, Val($idx))) end # One full `W÷2`-row block, already in `Complex`'s memory order, so no # interleave shuffle. Base's `Complex` expression trees in lanes, as planar's -# `_planar_store_block!`: +# `split_store_block!`: # beta == 0: addsub(ar*r, ai*swap(r)) # beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) # otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) -@inline function _fmaddsub_store_block!( +@inline function lanepair_store_block!( 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 {RC, R, W} - s = _swap_pairs(r) + s = swap_pairs(r) if iszero(beta) - new = _addsub(ar * r, ai * s) + new = addsub(ar * r, ai * s) else 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)) + x = isone(beta) ? old : addsub(br * old, bi * swap_pairs(old)) + new = fmaddsub(ar, r, fmaddsub(ai, s, x)) end 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::Tile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, - m::Int, n::Int - ) where {S, MR, NR, T, W, R, NV} - # The pointer reinterpretation is only sound on dense rank-1 complex storage. - 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) - HW = W ÷ 2 - - blocks = Any[] - for j in 1:NR - vblocks = Any[] - for v in 0:(MV - 1) - idx = v + MV * (j - 1) + 1 - push!( - vblocks, quote - vec = acc[$idx] - if $((v + 1) * HW) <= m - _fmaddsub_store_block!( - sp, 2 * (colbase + $(v * HW)), vec, ar, ai, br, bi, beta - ) - elseif $(v * HW) < m - for u in 1:$HW - i = $(v * HW) + u - i <= m || break - axpby_at!( - storage, colbase + i, alpha, - Complex(vec[2 * u - 1], vec[2 * u]), beta - ) - end - end - end - ) - end - push!( - blocks, quote - if $j <= n - colbase = rowbase0 + cols[$j] # the address of (1, j) - $(vblocks...) - end - end - ) - end - - return quote - storage = destination.storage - cols = destination.cols - rowbase0 = @inbounds destination.base + destination.rows[1] - ar = Vec{$W, $R}(real(alpha)) - ai = Vec{$W, $R}(imag(alpha)) - br = Vec{$W, $R}(real(beta)) - bi = Vec{$W, $R}(imag(beta)) - GC.@preserve storage begin - sp = reinterpret(Ptr{$RC}, pointer(storage)) - @inbounds begin - $(blocks...) - end - end - return destination - end -end - -# Not `@inline`, unlike the real and planar stores: inlining it cost time -# (code growth). -function store_tile!( - destination::Tile, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::FMAddSubKernel{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 diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl index 14137f9..138f52a 100644 --- a/src/microkernels/mixed.jl +++ b/src/microkernels/mixed.jl @@ -7,106 +7,5 @@ # accumulator column `2j-1` holds the real and `2j` 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} <: Microkernel{MR, NR, T} - descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} - inner::KI - - function ComplexRealKernel{MR, NR, T, W, KI}( - descriptor::Descriptor{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} <: Microkernel{MR, NR, T} - descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} - inner::KI - - function RealComplexKernel{MR, NR, T, W, KI}( - descriptor::Descriptor{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 = Descriptor(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 = Descriptor(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 add_tile(kernel::_MixedKernel, acc::NTuple, packed_a::PackedPanel, packed_b, k_block_length::Int) = - add_tile(kernel.inner, acc, packed_a, packed_b, k_block_length) - -# Inlined or not as the fmaddsub and planar stores each reuses. -function store_tile!( - destination::Tile, 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::Tile, 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 +inner(::ComplexRealKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) +inner(::RealComplexKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(MR), Val(2 * NR), real(T), Val(W)) diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl index 9a7c48d..ee9d9b8 100644 --- a/src/microkernels/onem.jl +++ b/src/microkernels/onem.jl @@ -6,108 +6,5 @@ # There is deliberately no FMA loop here: 1m reuses the real kernel body # verbatim. -using SIMD: Vec - -""" - OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Van Zee's induced 1m complex microkernel: 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. Used only when named in `plan_contract(...; kernel = ...)`. -""" -struct OneMKernel{MR, NR, T, W, KI <: SIMDKernel} <: Microkernel{MR, NR, T} - descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} - # `KI` stands in for `SIMDKernel{2MR,NR,real(T),W}`, which is not a legal - # field type; the constructor pins it to exactly that. - inner::KI - - function OneMKernel{MR, NR, T, W, KI}( - descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat}, - inner::KI - ) where {MR, NR, T, W, KI <: SIMDKernel} - check_vector_shape("OneMKernel", 2 * MR, W, true) - KI === SIMDKernel{2 * MR, NR, real(T), W} || throw( - ArgumentError("OneMKernel'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 - -function OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} - T <: Complex || - throw(ArgumentError("OneMKernel requires a complex element type, got $T")) - descriptor = Descriptor(Val(MR), Val(NR), T, OneEFormat(), PlanarFormat()) - inner = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) - return OneMKernel{MR, NR, T, W, typeof(inner)}(descriptor, inner) -end -function OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} - T <: Complex || - throw(ArgumentError("OneMKernel requires a complex element type, got $T")) - return OneMKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) -end - -complex_method(::OneMKernel) = OneMMethod() -lanewidth(::OneMKernel{MR, NR, T, W}) where {MR, NR, T, W} = W - -zero_accumulator(kernel::OneMKernel) = zero_accumulator(kernel.inner) - -function add_tile( - kernel::OneMKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, R}}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NV, PA <: PackedPanel, PB} - # Checked here so the message reports the logical k_block_length, not the - # doubled one. - k_block_length >= 0 || throw_negative_k_block_length(:add_tile, k_block_length) - return add_tile(kernel.inner, acc, packed_a, packed_b, 2 * k_block_length) -end - -# Scalar store for an interleaved accumulator (1m and fmaddsub): with `W` even -# and `2MR` a multiple of `W`, complex row `i = v*(W÷2) + u` is lanes `2u-1` -# (re) and `2u` (im) of vector `v` (zero-based) -- never split across two -# vectors. -@generated function _store_tile_lanepair!( - destination::Tile, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, - m::Int, n::Int - ) where {MR, NR, T, W, R, NV} - iseven(W) || throw(ArgumentError("_store_tile_lanepair!: requires an even W, got $W")) - MV = (2 * MR) ÷ W - check_acc(:_store_tile_lanepair!, R, T, NV, MV * NR) - HW = W ÷ 2 - - blocks = Any[] - for j in 1:NR, v in 0:(MV - 1) - idx = v + MV * (j - 1) + 1 - push!( - blocks, quote - if $j <= n - vec = acc[$idx] - for u in 1:$HW - i = $(v * HW) + u - i <= m || break - axpby_tile!( - destination, i, $j, alpha, - Complex(vec[2 * u - 1], vec[2 * u]), beta - ) - end - end - end - ) - end - return quote - @inbounds begin - $(blocks...) - end - return destination - end -end - -# Scalar only: 1m has no vector store. Not `@inline`, as for fmaddsub. -function store_tile!( - destination::Tile, acc::NTuple{NV, Vec{W, R}}, - alpha::T, beta::T, kernel::OneMKernel{MR, NR, T, W} - ) where {MR, NR, T, W, R, NV} - m, n = store_prologue!(destination, alpha, beta) - (m == 0 || n == 0) && return destination - return _store_tile_lanepair!(destination, acc, alpha, beta, kernel, m, n) -end +inner(::OneMKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) +k_steps(::OneMKernel, k_block_length::Int) = 2 * k_block_length diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index f18051c..17c969d 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -5,63 +5,17 @@ using SIMD: Vec, vload, vstore, shufflevector -""" - PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Split-complex microkernel for `T = ComplexF32/ComplexF64` over -`SIMD.Vec{W,real(T)}` lanes. `MR`/`NR` are complex extents; `MR` must be a -multiple of `W`, which counts reals. -""" -struct PlanarKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} - descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} - - function PlanarKernel{MR, NR, T, W}( - descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} - ) where {MR, NR, T, W} - check_vector_shape("PlanarKernel", MR, W) - return new{MR, NR, T, W}(descriptor) - end -end - -function PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} - return PlanarKernel{MR, NR, T, W}( - Descriptor(Val(MR), Val(NR), T, PlanarFormat(), PlanarFormat()) - ) -end -function PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} - T <: Complex || - throw(ArgumentError("PlanarKernel requires a complex element type, got $T")) - return PlanarKernel(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) -end - -complex_method(::PlanarKernel) = PlanarMethod() -lanewidth(::PlanarKernel{MR, NR, T, W}) where {MR, NR, T, W} = W - -# One flat tuple: the real plane at `1:NV`, the imaginary plane at `NV+1:2NV`, -# each laid out as the real kernel's. Flat keeps the accumulator in registers. -function zero_accumulator(kernel::PlanarKernel{MR, NR, T, W}) where {MR, NR, T, W} - z = zero(Vec{W, real(T)}) - return ntuple(_ -> z, Val(2 * (MR ÷ W) * NR)) -end - -# B column `j` at K step `p` as `(re, im)`; `UnpackedBView` overrides it. -# Shared with the fmaddsub kernel. -@inline _b_step_load2(packed_b::PB, kernel, j::Int, p::Int) where {PB} = ( - panel_load(packed_b, packed_b_offset(kernel, j, p)), - panel_load(packed_b, packed_b_offset(kernel, j, p, 1)), -) - # GUARDRAIL: the real part is `muladd(-ai, bi, muladd(ar, br, c))`. `c - ai*bi` # does NOT fuse (Julia sets no LLVM `contract` flag): two instructions, and # different rounding. `-ai` is hoisted out of the `j` loop and folds into # `vfnmadd`, so the negation is free. -@generated function _accumulate_step_planar( +@generated function accumulate_step( kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, p::Int ) where {MR, NR, T, W, R, NA, PA, PB} MV = MR ÷ W NV = MV * NR - check_acc(:_accumulate_step_planar, R, T, NA, 2NV) + check_acc(:accumulate_step, R, T, NA, 2NV) arv = [Symbol(:ar, v) for v in 0:(MV - 1)] aiv = [Symbol(:ai, v) for v in 0:(MV - 1)] @@ -92,12 +46,12 @@ end load_b = Any[] for j in 1:NR - push!(load_b, :(($(brv[j]), $(biv[j])) = _b_step_load2(packed_b, kernel, $j, p))) + push!(load_b, :(($(brv[j]), $(biv[j])) = b_complex(packed_b, kernel, $j, p))) end acc_exprs = Vector{Any}(undef, NA) for j in 1:NR, v in 0:(MV - 1) - idx = v + MV * (j - 1) + 1 + idx = acc_index(MV, v, j) acc_exprs[idx] = :( # re: ar*br - ai*bi muladd( $(naiv[v + 1]), $(biv[j]), @@ -122,88 +76,21 @@ end end end -function add_tile( - kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, R, NA, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - @inbounds for p in 1:k_block_length - acc = _accumulate_step_planar(kernel, acc, packed_a, packed_b, p) - end - return acc -end - -# Generator-time accumulator indices of the re/im vectors of block `v` -# (zero-based), 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) + 1, MV * NR + v + MV * (j - 1) + 1) - kernel <: RealComplexKernel && return (v + MV * (2j - 2) + 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::Tile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, - m::Int, n::Int - ) where {MR, NR, T, W, R, NA} - MV = MR ÷ W - NV = MV * NR - check_acc(:_store_tile_planar!, R, T, NA, 2NV) - - blocks = Any[] - for j in 1:NR, v in 0:(MV - 1) - ire, iim = _planar_acc_index(kernel, MV, NR, v, j) - push!( - blocks, quote - if $j <= n - revec = acc[$ire] - imvec = acc[$iim] - for lane in 1:$W - i = $(v * W) + lane - i <= m || break - axpby_tile!( - destination, i, $j, alpha, - Complex(revec[lane], imvec[lane]), beta - ) - end - end - end - ) - end - return quote - @inbounds begin - $(blocks...) - end - return destination - end -end - -# Vector store: unit-stride rows into rank-1 dense `Complex` storage (so the -# storage can be reinterpreted as `2W` consecutive reals per `W` rows), on an -# ISA the complex fast paths ship for (shared with the complex pack fast path). -@inline _complex_vector_eligible(tile::Tile, ::Type{T}) where {T} = - is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) && - complex_fastpath_isa_eligible() - # Shuffle patterns built from `W` at specialization time, never hardcoded to # one ISA. `v` holds `W` complex values `[re_0, im_0, ..., re_{W-1}, im_{W-1}]`. -@generated function _deinterleave_re(v::Vec{N, R}, ::Val{W}) where {N, R, W} - N == 2 * W || return :(throw(ArgumentError("_deinterleave_re: expected N == 2W"))) +@generated function deinterleave_re(v::Vec{N, R}, ::Val{W}) where {N, R, W} + N == 2 * W || return :(throw(ArgumentError("deinterleave_re: expected N == 2W"))) idx = ntuple(k -> 2 * (k - 1), W) return :(shufflevector(v, Val($idx))) end -@generated function _deinterleave_im(v::Vec{N, R}, ::Val{W}) where {N, R, W} - N == 2 * W || return :(throw(ArgumentError("_deinterleave_im: expected N == 2W"))) +@generated function deinterleave_im(v::Vec{N, R}, ::Val{W}) where {N, R, W} + N == 2 * W || return :(throw(ArgumentError("deinterleave_im: expected N == 2W"))) idx = ntuple(k -> 2 * (k - 1) + 1, W) return :(shufflevector(v, Val($idx))) end -@generated function _interleave_planes(re::Vec{W, R}, im::Vec{W, R}, ::Val{W}) where {W, R} +@generated function interleave_planes(re::Vec{W, R}, im::Vec{W, R}, ::Val{W}) where {W, R} idx = ntuple(k -> isodd(k) ? (k - 1) ÷ 2 : W + (k - 1) ÷ 2, 2 * W) return :(shufflevector(re, im, Val($idx))) end @@ -215,7 +102,7 @@ end # muladd(zi, wr, xi)))`. Against the scalar fallback it is exact at # `beta == 0/1` and ~1 ULP otherwise, because LLVM contracts Base's scalar # complex `muladd` depending on inlining context. -@inline function _planar_store_block!( +@inline function split_store_block!( 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} @@ -225,8 +112,8 @@ end newim = ar * imv + ai * rev else 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)) + orv = deinterleave_re(old, Val(W)) + oiv = deinterleave_im(old, Val(W)) if isone(beta) xr, xi = orv, oiv else @@ -236,92 +123,6 @@ end newre = muladd(ar, rev, -muladd(ai, imv, -xr)) newim = muladd(ar, imv, muladd(ai, rev, xi)) end - vstore(convert(Vec{2 * W, RC}, _interleave_planes(newre, newim, Val(W))), sp + sizeof(RC) * at) + vstore(convert(Vec{2 * W, RC}, interleave_planes(newre, newim, Val(W))), sp + sizeof(RC) * at) return nothing end - -# Same unroll and full-block / row-tail split as `_store_tile_vector!`, with a -# block of `2W` reals. The raw pointer is only dereferenced inside `GC.@preserve`. -@generated function _store_tile_planar_vector!( - destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::Microkernel{MR, NR, T}, - m::Int, n::Int - ) where {S, MR, NR, T, W, R, NA} - # The pointer reinterpretation is only sound on dense rank-1 complex storage. - 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) - - blocks = Any[] - for j in 1:NR - vblocks = Any[] - for v in 0:(MV - 1) - ire, iim = _planar_acc_index(kernel, MV, NR, v, j) - push!( - vblocks, quote - revec = acc[$ire] - imvec = acc[$iim] - if $((v + 1) * W) <= m - _planar_store_block!( - sp, 2 * (colbase + $(v * W)), revec, imvec, - ar, ai, br, bi, beta, Val($W) - ) - elseif $(v * W) < m - for lane in 1:$W - i = $(v * W) + lane - i <= m || break - axpby_at!( - storage, colbase + i, alpha, - Complex(revec[lane], imvec[lane]), beta - ) - end - end - end - ) - end - push!( - blocks, quote - if $j <= n - colbase = rowbase0 + cols[$j] # the address of (1, j) - $(vblocks...) - end - end - ) - end - - return quote - Base.@_inline_meta - storage = destination.storage - cols = destination.cols - rowbase0 = @inbounds destination.base + destination.rows[1] - ar = Vec{$W, $R}(real(alpha)) - ai = Vec{$W, $R}(imag(alpha)) - br = Vec{$W, $R}(real(beta)) - bi = Vec{$W, $R}(imag(beta)) - GC.@preserve storage begin - sp = reinterpret(Ptr{$RC}, pointer(storage)) - @inbounds begin - $(blocks...) - end - end - return destination - end -end - -# `@inline` with the vector store, for the reason at the real `store_tile!`. -@inline function store_tile!( - destination::Tile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::PlanarKernel{MR, NR, T, W} - ) where {MR, NR, T, W, R, NA} - 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/simd.jl b/src/microkernels/simd.jl index 6040cc1..093ec4c 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -1,6 +1,6 @@ -# Explicit-SIMD real microkernel. The accumulator is an immutable tuple of -# `Vec{W,T}` so it stays in registers, and every K step and store is -# `@generated` straight-line code. +# The vector kernels' shared `add_tile` and stores, and the explicit-SIMD real +# K step. The accumulator is an immutable tuple of `Vec`s so it stays in +# registers, and every K step and store is `@generated` straight-line code. # # GUARDRAIL (all kernels): every `acc[...]` must be a *literal* tuple index. # Indexing an `NTuple` dynamically forces it to memory, and above 16 vectors @@ -9,46 +9,42 @@ using SIMD: Vec, vload, vstore -""" - SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a -multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). -""" -struct SIMDKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} - descriptor::RealDescriptor{MR, NR, T} - - function SIMDKernel{MR, NR, T, W}(descriptor::RealDescriptor{MR, NR, T}) where {MR, NR, T, W} - check_vector_shape("SIMDKernel", MR, W) - return new{MR, NR, T, W}(descriptor) - end +function zero_accumulator(kernel::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} + layout = accumulator_layout(kernel) + z = zero(Vec{W, real(T)}) + return ntuple(_ -> z, Val(accumulator_length(layout, MR ÷ rows_per_vector(layout, W), NR))) end -SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {MR, NR, T, W} = - SIMDKernel{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T)) -SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = - SIMDKernel(Val(MR), Val(NR), T, Val(default_lanewidth(T))) +k_steps(::VectorKernel, k_block_length::Int) = k_block_length -lanewidth(::SIMDKernel{MR, NR, T, W}) where {MR, NR, T, W} = W - -# `(MR÷W)*NR` vectors; vector `v` (zero-based) of column `j` is at -# `v + (MR÷W)*(j - 1) + 1`. -function zero_accumulator(kernel::SIMDKernel{MR, NR, T, W}) where {MR, NR, T, W} - z = zero(Vec{W, T}) - return ntuple(_ -> z, Val((MR ÷ W) * NR)) +@inline function add_tile( + kernel::VectorKernel, acc::NTuple{NA, Vec{W, R}}, + packed_a::PA, packed_b::PB, k_block_length::Int + ) where {NA, W, R, PA <: PackedPanel, PB} + k_block_length == 0 && return acc + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) + @inbounds for p in 1:k_steps(kernel, k_block_length) + acc = accumulate_step(kernel, acc, packed_a, packed_b, p) + end + return acc end -# B column `j` at K step `p`; `UnpackedBView` (src/execution/unpackedb.jl) -# overrides it to read B in place, resolved at compile time. -@inline _b_step_load(packed_b::PB, kernel, j::Int, p::Int) where {PB} = +# B column `j` at K step `p`, as a real or as `(re, im)`; `UnpackedBView` +# (src/execution/unpackedb.jl) overrides both to read B in place, resolved at +# compile time. +@inline b_scalar(packed_b::PB, kernel, j::Int, p::Int) where {PB} = panel_load(packed_b, packed_b_offset(kernel, j, p)) +@inline b_complex(packed_b::PB, kernel, j::Int, p::Int) where {PB} = ( + panel_load(packed_b, packed_b_offset(kernel, j, p)), + panel_load(packed_b, packed_b_offset(kernel, j, p, 1)), +) -@generated function _accumulate_step( +@generated function accumulate_step( kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, packed_a::PA, packed_b::PB, p::Int ) where {MR, NR, T, W, NV, PA, PB} NVECA = MR ÷ W - check_acc(:_accumulate_step, T, T, NV, NVECA * NR) + check_acc(:accumulate_step, T, T, NV, NVECA * NR) avars = [Symbol(:a, v) for v in 0:(NVECA - 1)] bvars = [Symbol(:b, j) for j in 1:NR] @@ -58,13 +54,13 @@ end for v in 0:(NVECA - 1) ] load_b = [ - :($(bvars[j]) = _b_step_load(packed_b, kernel, $j, p)) + :($(bvars[j]) = b_scalar(packed_b, kernel, $j, p)) for j in 1:NR ] acc_exprs = Vector{Any}(undef, NV) for j in 1:NR, v in 0:(NVECA - 1) - idx = v + NVECA * (j - 1) + 1 + idx = acc_index(NVECA, v, j) acc_exprs[idx] = :(muladd($(avars[v + 1]), $(bvars[j]), acc[$idx])) end @@ -78,41 +74,60 @@ end end end -@inline function add_tile( - kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, W, NV, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - @inbounds for p in 1:k_block_length - acc = _accumulate_step(kernel, acc, packed_a, packed_b, p) - end - return acc -end +# A kernel without a K step of its own runs its `inner(kernel)`'s. +@inline accumulate_step(kernel::VectorKernel, acc::NTuple, packed_a::PA, packed_b::PB, p::Int) where {PA, PB} = + accumulate_step(inner(kernel), acc, packed_a, packed_b, p) -# 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::Tile, ::Type{T}) where {T} = +# Vector store eligibility: unit-stride rows into rank-1 dense 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. The complex +# layouts reinterpret `W` rows as `2W` consecutive reals, on an ISA the complex +# fast paths ship for (shared with the complex pack fast path). +@inline vector_store_eligible(::RealLayout, tile::Tile, ::Type{T}) where {T} = is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) +@inline vector_store_eligible(::AccumulatorLayout, tile::Tile, ::Type{T}) where {T} = + is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) && + complex_fastpath_isa_eligible() + +@generated function store_tile!( + destination::Tile, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W} + ) where {MR, NR, T, W, R, NA} + layout = accumulator_layout(kernel) + return quote + $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) + m, n = store_prologue!(destination, alpha, beta) + (m == 0 || n == 0) && return destination -@generated function _store_tile_scattered!( - destination::Tile, acc::NTuple{NV, Vec{W, T}}, - alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, + if vector_store_eligible($layout, destination, T) + return vector_store!(destination, acc, alpha, beta, kernel, m, n) + end + + return scalar_store!(destination, acc, alpha, beta, kernel, m, n) + end +end + +# Element by element, for every destination the vector store cannot take. +@generated function scalar_store!( + destination::Tile, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, m::Int, n::Int - ) where {MR, NR, T, W, NV} - NVECA = MR ÷ W + ) where {MR, NR, T, W, R, NA} + layout = accumulator_layout(kernel) + rows = rows_per_vector(layout, W) + MV = MR ÷ rows + check_acc(:scalar_store!, R, T, NA, accumulator_length(layout, MV, NR)) + blocks = Any[] - for j in 1:NR, v in 0:(NVECA - 1) - idx = v + NVECA * (j - 1) + 1 + for j in 1:NR, v in 0:(MV - 1) push!( blocks, quote if $j <= n - vec = acc[$idx] - for lane in 1:$W - i = $(v * W) + lane + $(acc_bindings(layout, kernel, MV, NR, v, j)) + for lane in 1:$rows + i = $(v * rows) + lane i <= m || break - axpby_tile!(destination, i, $j, alpha, vec[lane], beta) + axpby_tile!(destination, i, $j, alpha, $(lane_value(layout)), beta) end end end @@ -126,41 +141,34 @@ end end end -# Whole `W`-row blocks are one vector load/store; a block straddling `m` is -# stored lane by lane, so nothing outside the valid rectangle is touched. +# Whole row blocks are one vector load/store; a block straddling `m` is stored +# lane by lane, so nothing outside the valid rectangle is touched. # `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. -@generated function _store_tile_vector!( - destination::Tile{S, <:AffineAxis}, acc::NTuple{NV, Vec{W, T}}, - alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W}, +@generated function vector_store!( + destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, m::Int, n::Int - ) where {S, MR, NR, T, W, NV} - NVECA = MR ÷ W - RC = eltype(S) - old = :(convert(Vec{$W, $T}, vload(Vec{$W, $RC}, storage, at))) + ) where {S, MR, NR, T, W, R, NA} + layout = accumulator_layout(kernel) + rows = rows_per_vector(layout, W) + MV = MR ÷ rows + check_acc(:vector_store!, R, T, NA, accumulator_length(layout, MV, NR)) + RC = store_lanetype(layout, S, T) + blocks = Any[] for j in 1:NR vblocks = Any[] - for v in 0:(NVECA - 1) - idx = v + NVECA * (j - 1) + 1 + for v in 0:(MV - 1) push!( vblocks, quote - vec = acc[$idx] - if $((v + 1) * W) <= m - at = colbase + $(v * W) + 1 - vstore( - 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 - for lane in 1:$W - i = $(v * W) + lane + $(acc_bindings(layout, kernel, MV, NR, v, j)) + if $((v + 1) * rows) <= m + $(block_store(layout, :(colbase + $(v * rows)), W, R, RC)) + elseif $(v * rows) < m + for lane in 1:$rows + i = $(v * rows) + lane i <= m || break - axpby_at!(storage, colbase + i, alpha, vec[lane], beta) + axpby_at!(storage, colbase + i, alpha, $(lane_value(layout)), beta) end end end @@ -176,30 +184,11 @@ end ) end return quote - Base.@_inline_meta + $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) storage = destination.storage cols = destination.cols rowbase0 = @inbounds destination.base + destination.rows[1] - @inbounds begin - $(blocks...) - end + $(store_body(layout, W, R, RC, blocks)) return destination end end - -# `@inline` with the vector store: out of line, the whole accumulator is spilled -# to the stack and reloaded on every micro-tile. The scattered store stays out -# of line: it is scalar anyway, and inlining it bloats the tile function. -@inline function store_tile!( - destination::Tile, acc::NTuple{NV, Vec{W, T}}, - alpha::T, beta::T, kernel::SIMDKernel{MR, NR, T, W} - ) where {MR, NR, T, W, NV} - m, n = store_prologue!(destination, alpha, beta) - (m == 0 || n == 0) && return destination - - if _vector_store_eligible(destination, T) - return _store_tile_vector!(destination, acc, alpha, beta, kernel, m, n) - end - - return _store_tile_scattered!(destination, acc, alpha, beta, kernel, m, n) -end diff --git a/src/microkernels/types.jl b/src/microkernels/types.jl new file mode 100644 index 0000000..5900991 --- /dev/null +++ b/src/microkernels/types.jl @@ -0,0 +1,214 @@ +# The vector kernels' types, their accumulator layouts and the generator-time +# pieces of the stores, ahead of every implementation: a `@generated` +# function's generator only sees methods defined before it. + +using SIMD: Vec + +""" + VectorKernel{MR, NR, T, W} + +An explicit-SIMD microkernel over `SIMD.Vec{W, real(T)}` lanes. Its +accumulator is an immutable tuple of vectors in one of three layouts +([`accumulator_layout`](@ref)), which fixes the stores; each kernel supplies +its K step, `accumulate_step`, and the generic `add_tile` loops it +`k_steps(kernel, k_block_length)` times. +""" +abstract type VectorKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} end + +""" + SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a +multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). +""" +struct SIMDKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::RealDescriptor{MR, NR, T} +end + +""" + PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Split-complex microkernel for `T = ComplexF32/ComplexF64` over +`SIMD.Vec{W,real(T)}` lanes. `MR`/`NR` are complex extents; `MR` must be a +multiple of `W`, which counts reals. +""" +struct PlanarKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} +end + +""" + OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Van Zee's induced 1m complex microkernel: 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. Used only when named in `plan_contract(...; kernel = ...)`. +""" +struct OneMKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} +end + +""" + FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over +`SIMD.Vec{W,real(T)}` lanes with `InterleavedFormat` A and `PlanarFormat` B. +`2MR` must be a multiple of `W`, and `W` must be even. The default for small-M +complex contractions on AVX-512. +""" +struct FMAddSubKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} +end + +""" + 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} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} +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} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} +end + +complex_method(::Type{<:SIMDKernel}) = RealMethod() +complex_method(::Type{<:PlanarKernel}) = PlanarMethod() +complex_method(::Type{<:OneMKernel}) = OneMMethod() +complex_method(::Type{<:FMAddSubKernel}) = FMAddSubMethod() +complex_method(::Type{<:ComplexRealKernel}) = ComplexRealMethod() +complex_method(::Type{<:RealComplexKernel}) = RealComplexMethod() +complex_method(kernel::VectorKernel) = complex_method(typeof(kernel)) + +""" + accumulator_layout(kernel::VectorKernel) + +How the accumulator holds the tile, as `(MR ÷ rows) * NR` vectors of `rows` +consecutive tile rows each, column-major (`rows_per_vector`): + + - `RealLayout()`: one real per lane (`SIMDKernel`). + - `SplitLayout()`: separate re and im vectors (`PlanarKernel`: a real and an + imaginary plane in one flat tuple; `RealComplexKernel`: the real kernel's + columns `2j-1` and `2j`). + - `LanePairLayout()`: lanes `2u-1` and `2u` are re and im of one row + (1m, fmaddsub, `ComplexRealKernel`). +""" +function accumulator_layout end + +abstract type AccumulatorLayout end +struct RealLayout <: AccumulatorLayout end +struct SplitLayout <: AccumulatorLayout end +struct LanePairLayout <: AccumulatorLayout end + +accumulator_layout(::Type{<:SIMDKernel}) = RealLayout() +accumulator_layout(::Type{<:Union{PlanarKernel, RealComplexKernel}}) = SplitLayout() +accumulator_layout(::Type{<:Union{OneMKernel, FMAddSubKernel, ComplexRealKernel}}) = LanePairLayout() +accumulator_layout(kernel::VectorKernel) = accumulator_layout(typeof(kernel)) + +rows_per_vector(::AccumulatorLayout, W::Int) = W +rows_per_vector(::LanePairLayout, W::Int) = W ÷ 2 +accumulator_length(::AccumulatorLayout, MV::Int, NR::Int) = MV * NR +accumulator_length(::SplitLayout, MV::Int, NR::Int) = 2 * MV * NR + +function (::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {K <: VectorKernel, MR, NR, T, W} + method = complex_method(K) + method === RealMethod() || T <: Complex || + throw(ArgumentError("$(nameof(K)) requires a complex element type, got $T")) + pairs = accumulator_layout(K) isa LanePairLayout + check_vector_shape(nameof(K), pairs ? 2 * MR : MR, W, pairs) + return K{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T, pack_formats(method)...)) +end +(::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}) where {K <: VectorKernel, MR, NR, T} = + K(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) + +lanewidth(::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} = W + +# Real and split stores are `@inline`: out of line, the whole accumulator is +# spilled to the stack and reloaded on every micro-tile. The lane-pair store is +# not: inlining it cost time (code growth). The scalar stores never are: they +# are scalar anyway, and inlining them bloats the tile function. +inline_store(::AccumulatorLayout) = true +inline_store(::LanePairLayout) = false + +# Generator-time pieces of the stores. `acc_bindings` names the accumulator +# vector(s) of row block `v` (zero-based) of column `j`, `lane_value` is the +# element at row `lane` of the block, `block_store` stores a full block whose +# first row is at zero-based storage index `first`. +acc_index(MV::Int, v::Int, j::Int) = v + MV * (j - 1) + 1 +split_index(::Type{<:PlanarKernel}, MV::Int, NR::Int, v::Int, j::Int) = + (acc_index(MV, v, j), MV * NR + acc_index(MV, v, j)) +split_index(::Type{<:RealComplexKernel}, MV::Int, NR::Int, v::Int, j::Int) = + (v + MV * (2j - 2) + 1, v + MV * (2j - 1) + 1) + +acc_bindings(::AccumulatorLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) = + :(vec = acc[$(acc_index(MV, v, j))]) +function acc_bindings(::SplitLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) + ire, iim = split_index(kernel, MV, NR, v, j) + return quote + revec = acc[$ire] + imvec = acc[$iim] + end +end + +lane_value(::RealLayout) = :(vec[lane]) +lane_value(::SplitLayout) = :(Complex(revec[lane], imvec[lane])) +lane_value(::LanePairLayout) = :(Complex(vec[2 * lane - 1], vec[2 * lane])) + +function block_store(::RealLayout, first, W::Int, R::Type, RC::Type) + old = :(convert(Vec{$W, $R}, vload(Vec{$W, $RC}, storage, at))) + return quote + at = $first + 1 + vstore( + convert( + Vec{$W, $RC}, + iszero(beta) ? alpha * vec : + isone(beta) ? muladd(alpha, vec, $old) : + muladd(alpha, vec, beta * $old) + ), + storage, at + ) + end +end +block_store(::SplitLayout, first, W::Int, R::Type, RC::Type) = + :(split_store_block!(sp, 2 * $first, revec, imvec, ar, ai, br, bi, beta, Val($W))) +block_store(::LanePairLayout, first, W::Int, R::Type, RC::Type) = + :(lanepair_store_block!(sp, 2 * $first, vec, ar, ai, br, bi, beta)) + +# The real lane type of the vector store's storage `S`. The complex layouts +# reinterpret the storage as reals, only sound on dense rank-1 complex storage. +store_lanetype(::RealLayout, S::Type, T::Type) = eltype(S) +function store_lanetype(::AccumulatorLayout, S::Type, T::Type) + S <: DenseVector && lane_convertible(eltype(S), T) || + throw(ArgumentError("vector_store!: storage $S is not a dense vector convertible to $T")) + return real(eltype(S)) +end + +# The vector store's body around its blocks: the complex layouts broadcast +# `alpha`/`beta` once and store through a raw pointer, only dereferenced +# inside `GC.@preserve`. +store_body(::RealLayout, W::Int, R::Type, RC::Type, blocks) = :( + @inbounds begin + $(blocks...) + end +) +store_body(::AccumulatorLayout, W::Int, R::Type, RC::Type, blocks) = quote + ar = Vec{$W, $R}(real(alpha)) + ai = Vec{$W, $R}(imag(alpha)) + br = Vec{$W, $R}(real(beta)) + bi = Vec{$W, $R}(imag(beta)) + GC.@preserve storage begin + sp = reinterpret(Ptr{$RC}, pointer(storage)) + @inbounds begin + $(blocks...) + end + end +end diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 78d6b26..3faa361 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -26,20 +26,20 @@ const _QSF = QuasiStrided @test_throws ArgumentError FMAddSubKernel(Val(8), Val(4), Float64, Val(4)) end - @testset "_fmaddsub is lane-exact x86 fmaddsub; _swap_pairs swaps pairs" begin + @testset "fmaddsub is lane-exact x86 fmaddsub; swap_pairs swaps pairs" begin rng = MersenneTwister(0xADD5) for R in (Float64, Float32), N in (2, 4, 8, 16), _ in 1:10 x, y, c = (Vec{N, R}(ntuple(_ -> randn(rng, R), N)) for _ in 1:3) # One fused rounding per lane, with an exactly negated addend in even lanes. - @test Tuple(_QSF._fmaddsub(x, y, c)) === + @test Tuple(_QSF.fmaddsub(x, y, c)) === ntuple(l -> isodd(l) ? fma(x[l], y[l], -c[l]) : fma(x[l], y[l], c[l]), N) - @test Tuple(_QSF._swap_pairs(x)) === ntuple(l -> isodd(l) ? x[l + 1] : x[l - 1], N) + @test Tuple(_QSF.swap_pairs(x)) === ntuple(l -> isodd(l) ? x[l + 1] : x[l - 1], N) end x = Vec{4, Float64}((0.0, -0.0, Inf, 1.0)) y = Vec{4, Float64}((1.0, 1.0, 1.0, NaN)) c = Vec{4, Float64}((0.0, 0.0, 1.0, 1.0)) @test isequal( - Tuple(_QSF._fmaddsub(x, y, c)), + Tuple(_QSF.fmaddsub(x, y, c)), (fma(0.0, 1.0, -0.0), fma(-0.0, 1.0, 0.0), fma(Inf, 1.0, -1.0), fma(1.0, NaN, 1.0)) ) end @@ -50,9 +50,9 @@ const _QSF = QuasiStrided b = Complex{R}(7, -2) c = Vec{2, R}((R(11), R(13))) br, bi = Vec{2, R}(real(b)), Vec{2, R}(imag(b)) - right = _QSF._fmaddsub(a, br, _QSF._fmaddsub(_QSF._swap_pairs(a), bi, c)) + right = _QSF.fmaddsub(a, br, _QSF.fmaddsub(_QSF.swap_pairs(a), bi, c)) @test Complex(right[1], right[2]) == Complex{R}(11, 13) + Complex{R}(3, 5) * b - wrong = _QSF._fmaddsub(_QSF._swap_pairs(a), bi, _QSF._fmaddsub(a, br, c)) + wrong = _QSF.fmaddsub(_QSF.swap_pairs(a), bi, _QSF.fmaddsub(a, br, c)) @test wrong[1] == R(11) + R(5) * imag(b) - R(3) * real(b) end end diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index 3dbb5bb..ac78c9f 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -33,8 +33,8 @@ 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 QuasiStrided.inner(cr) isa SIMDKernel{16, 4, Float32, 8} + @test QuasiStrided.inner(rc) isa SIMDKernel{8, 8, Float64, 4} @test (sliver_width(cr)..., sliver_width(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)) diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index fcd7e3c..c22455a 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -1,5 +1,5 @@ -using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEFormat, - PlanarFormat, complex_method, lanewidth, kernel_shapes, +using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, + complex_method, lanewidth, kernel_shapes, _kernel_from_shape, _default_method @testset "OneMKernel" begin @@ -13,8 +13,8 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEForm @testset "construction" begin k = OneMKernel(Val(12), Val(8), ComplexF64, Val(8)) - @test k.inner isa SIMDKernel{24, 8, Float64, 8} # the real kernel, at 2MR rows - @test zero_accumulator(k) === zero_accumulator(k.inner) + @test QuasiStrided.inner(k) isa SIMDKernel{24, 8, Float64, 8} # the real kernel, at 2MR rows + @test zero_accumulator(k) === zero_accumulator(QuasiStrided.inner(k)) @test complex_method(k) === OneMMethod() @test lanewidth(OneMKernel(Val(8), Val(4), ComplexF32)) == 8 # 2MR, not MR, must divide by W; W must be even even where 2MR divides. @@ -23,9 +23,6 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, OneEForm @test_throws ArgumentError OneMKernel(Val(8), Val(4), ComplexF64, Val(0)) @test_throws ArgumentError OneMKernel(Val(8), Val(4), Float64) @test_throws ArgumentError OneMKernel(Val(8), Val(4), Float64, Val(4)) - wrong = SIMDKernel(Val(12), Val(8), Float64, Val(4)) - d = Descriptor(Val(12), Val(8), ComplexF64, OneEFormat(), PlanarFormat()) - @test_throws ArgumentError OneMKernel{12, 8, ComplexF64, 8, typeof(wrong)}(d, wrong) # The error reports the logical k_block_length, not the doubled real one. err = try add_tile(k, zero_accumulator(k), Float64[], Float64[], -3) diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 0a041f9..286171f 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -1,5 +1,5 @@ -# The vectorized complex store fast paths (planar here, fmaddsub in -# test_fmaddsub_store_fastpath.jl, which reuses `mk_store_fastpath`). +# The vectorized complex store fast paths: split (planar) and lane pairs +# (1m, fmaddsub). # # Values are pinned two ways. On the elements the fast path vectorizes (full # row blocks) it must match, bitwise, an independent transcription of Base's @@ -11,7 +11,7 @@ # The fast path ships for AVX-512 only; every expectation derives from the live # profile, so `test/forced_isa_runner.jl` checks the other ISAs too. -using QuasiStrided: PlanarKernel, FMAddSubKernel, KERNEL_SHAPES_C64_FMADDSUB, KERNEL_SHAPES_C32_FMADDSUB, TargetProfile, CacheLevel, +using QuasiStrided: PlanarKernel, FMAddSubKernel, OneMKernel, KERNEL_SHAPES_C64_ONEM, KERNEL_SHAPES_C64_FMADDSUB, KERNEL_SHAPES_C32_FMADDSUB, TargetProfile, CacheLevel, target_profile, unknown_target, KERNEL_SHAPES_C64_PLANAR, KERNEL_SHAPES_C32_PLANAR const STORE_FASTPATH_ON = QuasiStrided.complex_fastpath_isa_eligible() @@ -65,7 +65,7 @@ function mk_store_fastpath(K, T, shapes; beta0_exact = false, S = T) for (alpha, beta) in mk_store_ab(T) fast = mk_dense(cold) dfast = Tile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) - @test QuasiStrided._complex_vector_eligible(dfast, T) == STORE_FASTPATH_ON + @test QuasiStrided.vector_store_eligible(QuasiStrided.accumulator_layout(k), dfast, T) == STORE_FASTPATH_ON store_tile!(dfast, acc, alpha, beta, k) scal = copy(cold) # scattered rows: always the scalar store store_tile!(Tile(scal, 0, view(collect(0:(m - 1)), 1:m), AffineAxis(0, m, n)), acc, alpha, beta, k) @@ -92,10 +92,10 @@ end m, n = 16, 4 storage = zeros(T, 4 * m * n) eligible(s, rows, cols = AffineAxis(0, m, n)) = - QuasiStrided._complex_vector_eligible(Tile(s, 0, rows, cols), T) + QuasiStrided.vector_store_eligible(QuasiStrided.SplitLayout(), Tile(s, 0, rows, cols), T) @test eligible(storage, AffineAxis(0, 1, m)) == STORE_FASTPATH_ON @test !eligible(storage, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)) - @test !QuasiStrided._complex_vector_eligible(Tile(storage, m - 1, AffineAxis(0, -1, m), AffineAxis(0, m, n)), T) + @test !QuasiStrided.vector_store_eligible(QuasiStrided.SplitLayout(), Tile(storage, m - 1, AffineAxis(0, -1, m), AffineAxis(0, m, n)), T) @test !eligible(storage, view(collect(0:(m - 1)), 1:m)) @test !eligible(view(storage, 1:(m * n)), AffineAxis(0, 1, m)) @test !eligible(reshape(storage, 4m, n), AffineAxis(0, 1, m)) @@ -105,7 +105,8 @@ end end -@testset "fmaddsub store fast path" begin +@testset "lane-pair store fast path" begin + mk_store_fastpath(OneMKernel, ComplexF64, KERNEL_SHAPES_C64_ONEM[[1, end]]; beta0_exact = true) 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) diff --git a/test/microkernels/test_simd_kernel.jl b/test/microkernels/test_simd_kernel.jl index 9fbf142..ea3696f 100644 --- a/test/microkernels/test_simd_kernel.jl +++ b/test/microkernels/test_simd_kernel.jl @@ -1,4 +1,4 @@ -using QuasiStrided: SIMDKernel, lanewidth, _vector_store_eligible, kernel_shapes +using QuasiStrided: SIMDKernel, lanewidth, vector_store_eligible, RealLayout, kernel_shapes using StridedViews: StridedView @testset "SIMDKernel" begin @@ -18,19 +18,19 @@ 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 real storage" 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) mem = parent(StridedView(zeros(T, m * n))) # `Memory{T}` on Julia >= 1.11 - @test _vector_store_eligible(Tile(mem, 0, rows, cols), T) - @test _vector_store_eligible(Tile(zeros(T, m * n), 0, rows, cols), T) - @test !_vector_store_eligible(Tile(view(zeros(T, m * n + 4), 3:(m * n + 2)), 0, rows, cols), T) - @test !_vector_store_eligible(Tile(zeros(T, m, n), 0, rows, cols), T) - @test !_vector_store_eligible(Tile(zeros(T, 2m * n), 0, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)), T) - @test !_vector_store_eligible(Tile(mem, 0, view(collect(0:(m - 1)), 1:m), cols), T) - @test _vector_store_eligible(Tile(mem, 0, rows, cols), T === Float64 ? Float32 : Float64) - @test !_vector_store_eligible(Tile(zeros(complex(T), m * n), 0, rows, cols), T) + @test vector_store_eligible(RealLayout(), Tile(mem, 0, rows, cols), T) + @test vector_store_eligible(RealLayout(), Tile(zeros(T, m * n), 0, rows, cols), T) + @test !vector_store_eligible(RealLayout(), Tile(view(zeros(T, m * n + 4), 3:(m * n + 2)), 0, rows, cols), T) + @test !vector_store_eligible(RealLayout(), Tile(zeros(T, m, n), 0, rows, cols), T) + @test !vector_store_eligible(RealLayout(), Tile(zeros(T, 2m * n), 0, AffineAxis(0, 2, m), AffineAxis(0, 2m, n)), T) + @test !vector_store_eligible(RealLayout(), Tile(mem, 0, view(collect(0:(m - 1)), 1:m), cols), T) + @test vector_store_eligible(RealLayout(), Tile(mem, 0, rows, cols), T === Float64 ? Float32 : Float64) + @test !vector_store_eligible(RealLayout(), Tile(zeros(complex(T), m * n), 0, rows, cols), T) end end @@ -42,7 +42,7 @@ using StridedViews: StridedView cold = S.(2 .* rand(rng, m * n) .- 1) fast = mk_dense(cold) dfast = Tile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) - @test _vector_store_eligible(dfast, T) + @test vector_store_eligible(RealLayout(), dfast, T) store_tile!(dfast, acc, alpha, beta, k) scal = copy(cold) store_tile!(Tile(scal, 0, view(collect(0:(m - 1)), 1:m), AffineAxis(0, m, n)), acc, alpha, beta, k) From cc78fd1d853775d6e2ab032b0231e8bfd17ed080 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 1 Oct 2026 19:17:49 -0400 Subject: [PATCH 16/60] Hoist the beta branch out of the vector stores vector_store! generates its block loop once per beta case (iszero, isone, general) and branches once per tile; the full-block expressions are the specialised ones (real: alpha*vec, muladd(alpha, vec, old), muladd(alpha, vec, beta*old); split_store_block!/lanepair_store_block! take the case as a Val, same arithmetic). Tails keep axpby_at!: inside each arm LLVM folds its repeated beta test. Codegen (AVX-512, default shapes): K loops unchanged. vector_store! instructions before -> after: SIMD 1953 -> 2726, Planar 1494 -> 2004, FMAddSub 3619 -> 4757, ComplexReal/1m 3612 -> 4485, RealComplex 1900 -> 2628. Beta compares in the vector store: SIMD 96 -> 2, Planar 62 -> 13, FMAddSub/ComplexReal 167 -> 28, RealComplex 83 -> 16. Code per executed variant (beta 0/1/general): SIMD 1040/1112/1023, Planar 492/635/707, lane pairs ~1003/1291/1723, RealComplex 637/828/924. A/B vs the previous commit, medians of 5 interleaved runs: small-K (M = N = 200) Float64 K = 1: -12%, K = 2..8: -19..-27%; ComplexF64 -3..-10%; 256^3 matmul/scattered/named-kernel cases -4..+1%. Allocations 0 under --check-bounds=yes. Cost: TTFX first call Float64 3.1 -> 3.7 s, ComplexF64 2.0 -> 2.4 s (2 runs each, consistent). Full suite 31967 pass, 28m02 (previous commit 18m07; single runs on a shared workstation, but in line with the TTFX increase). Co-Authored-By: Claude Opus 5.5 --- src/microkernels/fmaddsub.jl | 8 ++--- src/microkernels/planar.jl | 11 ++++--- src/microkernels/simd.jl | 61 ++++++++++++++++++++++-------------- src/microkernels/types.jl | 36 +++++++-------------- 4 files changed, 59 insertions(+), 57 deletions(-) diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl index 8712b88..6535358 100644 --- a/src/microkernels/fmaddsub.jl +++ b/src/microkernels/fmaddsub.jl @@ -130,14 +130,14 @@ end @inline function lanepair_store_block!( 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 {RC, R, W} + ::Val{B} + ) where {RC, R, W, B} s = swap_pairs(r) - if iszero(beta) + if B === :zero new = addsub(ar * r, ai * s) else old = convert(Vec{W, R}, vload(Vec{W, RC}, sp + sizeof(RC) * at)) - x = isone(beta) ? old : addsub(br * old, bi * swap_pairs(old)) + x = B === :one ? old : addsub(br * old, bi * swap_pairs(old)) new = fmaddsub(ar, r, fmaddsub(ai, s, x)) end vstore(convert(Vec{W, RC}, new), sp + sizeof(RC) * at) diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl index 17c969d..cdd3c1f 100644 --- a/src/microkernels/planar.jl +++ b/src/microkernels/planar.jl @@ -95,7 +95,8 @@ end return :(shufflevector(re, im, Val($idx))) end -# One full `W`-row block; `at` is the zero-based index of its first real. +# One full `W`-row block for `beta` case `B`; `at` is the zero-based index of +# its first real. # The arithmetic transcribes Base's `Complex` expression trees (the ones # `axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, # `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, @@ -105,16 +106,16 @@ end @inline function split_store_block!( 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 {RC, R, W} - if iszero(beta) + ::Val{W}, ::Val{B} + ) where {RC, R, W, B} + if B === :zero newre = ar * rev - ai * imv newim = ar * imv + ai * rev else 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) + if B === :one xr, xi = orv, oiv else xr = br * orv - bi * oiv diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl index 093ec4c..0f81467 100644 --- a/src/microkernels/simd.jl +++ b/src/microkernels/simd.jl @@ -142,7 +142,8 @@ end end # Whole row blocks are one vector load/store; a block straddling `m` is stored -# lane by lane, so nothing outside the valid rectangle is touched. +# lane by lane, so nothing outside the valid rectangle is touched. The blocks +# are generated once per `beta` case, so `beta` is tested once per tile. # `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. @generated function vector_store!( destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, @@ -155,40 +156,52 @@ end check_acc(:vector_store!, R, T, NA, accumulator_length(layout, MV, NR)) RC = store_lanetype(layout, S, T) - blocks = Any[] - for j in 1:NR - vblocks = Any[] - for v in 0:(MV - 1) - push!( - vblocks, quote - $(acc_bindings(layout, kernel, MV, NR, v, j)) - if $((v + 1) * rows) <= m - $(block_store(layout, :(colbase + $(v * rows)), W, R, RC)) - elseif $(v * rows) < m - for lane in 1:$rows - i = $(v * rows) + lane - i <= m || break - axpby_at!(storage, colbase + i, alpha, $(lane_value(layout)), beta) + function blocks(B) + out = Any[] + for j in 1:NR + vblocks = Any[] + for v in 0:(MV - 1) + push!( + vblocks, quote + $(acc_bindings(layout, kernel, MV, NR, v, j)) + if $((v + 1) * rows) <= m + $(block_store(layout, :(colbase + $(v * rows)), W, R, RC, B)) + elseif $(v * rows) < m + for lane in 1:$rows + i = $(v * rows) + lane + i <= m || break + axpby_at!(storage, colbase + i, alpha, $(lane_value(layout)), beta) + end end end + ) + end + push!( + out, quote + if $j <= n + colbase = rowbase0 + cols[$j] # the address of (1, j) + $(vblocks...) + end end ) end - push!( - blocks, quote - if $j <= n - colbase = rowbase0 + cols[$j] # the address of (1, j) - $(vblocks...) - end - end - ) + return out + end + by_beta = quote + if iszero(beta) + $(blocks(:zero)...) + elseif isone(beta) + $(blocks(:one)...) + else + $(blocks(:general)...) + end end return quote $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) storage = destination.storage cols = destination.cols rowbase0 = @inbounds destination.base + destination.rows[1] - $(store_body(layout, W, R, RC, blocks)) + $(store_body(layout, W, R, RC, by_beta)) return destination end end diff --git a/src/microkernels/types.jl b/src/microkernels/types.jl index 5900991..099a580 100644 --- a/src/microkernels/types.jl +++ b/src/microkernels/types.jl @@ -142,7 +142,8 @@ inline_store(::LanePairLayout) = false # Generator-time pieces of the stores. `acc_bindings` names the accumulator # vector(s) of row block `v` (zero-based) of column `j`, `lane_value` is the # element at row `lane` of the block, `block_store` stores a full block whose -# first row is at zero-based storage index `first`. +# first row is at zero-based storage index `first`, for `beta` case `B` +# (`:zero`, `:one` or `:general`). acc_index(MV::Int, v::Int, j::Int) = v + MV * (j - 1) + 1 split_index(::Type{<:PlanarKernel}, MV::Int, NR::Int, v::Int, j::Int) = (acc_index(MV, v, j), MV * NR + acc_index(MV, v, j)) @@ -163,25 +164,18 @@ lane_value(::RealLayout) = :(vec[lane]) lane_value(::SplitLayout) = :(Complex(revec[lane], imvec[lane])) lane_value(::LanePairLayout) = :(Complex(vec[2 * lane - 1], vec[2 * lane])) -function block_store(::RealLayout, first, W::Int, R::Type, RC::Type) +function block_store(::RealLayout, first, W::Int, R::Type, RC::Type, B::Symbol) old = :(convert(Vec{$W, $R}, vload(Vec{$W, $RC}, storage, at))) + new = B === :zero ? :(alpha * vec) : B === :one ? :(muladd(alpha, vec, $old)) : :(muladd(alpha, vec, beta * $old)) return quote at = $first + 1 - vstore( - convert( - Vec{$W, $RC}, - iszero(beta) ? alpha * vec : - isone(beta) ? muladd(alpha, vec, $old) : - muladd(alpha, vec, beta * $old) - ), - storage, at - ) + vstore(convert(Vec{$W, $RC}, $new), storage, at) end end -block_store(::SplitLayout, first, W::Int, R::Type, RC::Type) = - :(split_store_block!(sp, 2 * $first, revec, imvec, ar, ai, br, bi, beta, Val($W))) -block_store(::LanePairLayout, first, W::Int, R::Type, RC::Type) = - :(lanepair_store_block!(sp, 2 * $first, vec, ar, ai, br, bi, beta)) +block_store(::SplitLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = + :(split_store_block!(sp, 2 * $first, revec, imvec, ar, ai, br, bi, Val($W), Val($(QuoteNode(B))))) +block_store(::LanePairLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = + :(lanepair_store_block!(sp, 2 * $first, vec, ar, ai, br, bi, Val($(QuoteNode(B))))) # The real lane type of the vector store's storage `S`. The complex layouts # reinterpret the storage as reals, only sound on dense rank-1 complex storage. @@ -195,20 +189,14 @@ end # The vector store's body around its blocks: the complex layouts broadcast # `alpha`/`beta` once and store through a raw pointer, only dereferenced # inside `GC.@preserve`. -store_body(::RealLayout, W::Int, R::Type, RC::Type, blocks) = :( - @inbounds begin - $(blocks...) - end -) -store_body(::AccumulatorLayout, W::Int, R::Type, RC::Type, blocks) = quote +store_body(::RealLayout, W::Int, R::Type, RC::Type, body) = :(@inbounds $body) +store_body(::AccumulatorLayout, W::Int, R::Type, RC::Type, body) = quote ar = Vec{$W, $R}(real(alpha)) ai = Vec{$W, $R}(imag(alpha)) br = Vec{$W, $R}(real(beta)) bi = Vec{$W, $R}(imag(beta)) GC.@preserve storage begin sp = reinterpret(Ptr{$RC}, pointer(storage)) - @inbounds begin - $(blocks...) - end + @inbounds $body end end From 2c6f810e0502d6c2800b316f6a2b0970b56da35d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 10:10:43 -0400 Subject: [PATCH 17/60] Microkernels: organize by concern Five files instead of eight, included in this order: - interface.jl: the contract (Microkernel, KernelMethod and pack_formats, the add_tile/store_tile!/zero_accumulator docstrings, execute_tile!, scale_tile!, store_prologue!, axpby helpers, shape checks). - kernels.jl: the seven kernel types and their type-level facts (method and layout traits, rows_per_vector, acc_index, inner, k_steps, lanewidth, the generic constructor). The 1m and mixed-domain header comments are folded into the kernels' docstrings. - vecops.jl: fmaddsub, swap_pairs, addsub, deinterleave_re/_im, interleave_planes. - steps.jl: the K steps (scalar, real, planar, fmaddsub), the B-element hooks, the generic add_tile and zero_accumulator. - stores.jl: the layouts' generator-time pieces, the full-block stores and the generic store_tile!/vector_store!/scalar_store!, plus ScalarKernel's. Removed: types.jl, scalar.jl, simd.jl, planar.jl, onem.jl, fmaddsub.jl, mixed.jl. Holy traits: accumulator_layout(kernel) becomes AccumulatorLayout(kernel) and complex_method(kernel) becomes KernelMethod(kernel), each defined on the kernel type, with the docstring on the abstract type. The complex_method(::Any) fallback is replaced by KernelMethod for ScalarKernel (its only user). Every vector kernel's inner constructor now runs the vector-shape check (check_vector_kernel), so K{MR, NR, T, W}(descriptor) is validated and the generic constructor only builds the descriptor. A real T for a complex kernel is rejected by the descriptor (its message instead of the kernel's). Behaviour-neutral: tests 31967/31967 pass (27m59s wall); K-loop and store instruction counts identical to cc78fd1 for all six vector kernels; TTFX unchanged (2 interleaved runs each, before/after: load 0.5/0.5-0.6 s, first Float64 call 3.7-3.8/3.7 s, ComplexF64 2.4/2.4 s, K=1 outer 0.6/0.6 s). Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_blocking_model.jl | 4 +- src/QuasiStrided.jl | 11 +- src/execution/execute.jl | 2 +- src/microkernels/fmaddsub.jl | 145 --------- src/microkernels/interface.jl | 6 +- src/microkernels/kernels.jl | 195 ++++++++++++ src/microkernels/mixed.jl | 11 - src/microkernels/onem.jl | 10 - src/microkernels/planar.jl | 129 -------- src/microkernels/scalar.jl | 60 ---- src/microkernels/simd.jl | 207 ------------- src/microkernels/steps.jl | 251 ++++++++++++++++ src/microkernels/stores.jl | 279 ++++++++++++++++++ src/microkernels/types.jl | 202 ------------- src/microkernels/vecops.jl | 71 +++++ src/planning/defaults.jl | 2 +- test/microkernels/test_fmaddsub_kernel.jl | 4 +- test/microkernels/test_onem_kernel.jl | 4 +- test/microkernels/test_planar_kernel.jl | 4 +- .../test_planar_store_fastpath.jl | 2 +- test/planning/test_kernel_selection.jl | 10 +- test/runtests.jl | 2 +- 22 files changed, 819 insertions(+), 792 deletions(-) delete mode 100644 src/microkernels/fmaddsub.jl create mode 100644 src/microkernels/kernels.jl delete mode 100644 src/microkernels/mixed.jl delete mode 100644 src/microkernels/onem.jl delete mode 100644 src/microkernels/planar.jl delete mode 100644 src/microkernels/scalar.jl delete mode 100644 src/microkernels/simd.jl create mode 100644 src/microkernels/steps.jl create mode 100644 src/microkernels/stores.jl delete mode 100644 src/microkernels/types.jl create mode 100644 src/microkernels/vecops.jl diff --git a/benchmark/bench_blocking_model.jl b/benchmark/bench_blocking_model.jl index 9dd1656..aea7b2d 100644 --- a/benchmark/bench_blocking_model.jl +++ b/benchmark/bench_blocking_model.jl @@ -14,7 +14,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: Blocking, target_profile, _default_kernel, _modelled_blocking, - _fallback_blocking, _scale_blocking, complex_method + _fallback_blocking, _scale_blocking, KernelMethod const SMOKE = hasflag("smoke") const REPS = SMOKE ? 1 : argopt("reps", 21) @@ -132,7 +132,7 @@ for T in DTYPES end for T in CDTYPES kernel = _default_kernel(T) - rows = map(v -> v === nothing ? nothing : _scale_blocking(v, complex_method(kernel)), named_rows(real(T))) + rows = map(v -> v === nothing ? nothing : _scale_blocking(v, KernelMethod(kernel)), named_rows(real(T))) println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) for spec in thin(MAIN_SHAPES) sweep_shape!(raw, kernel, T, spec, named_points(rows)) diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index da194d8..bf5e869 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -26,13 +26,10 @@ include("packing/pack_contiguous.jl") # --- Microkernels: add one packed K block to a register tile, store into C --- include("microkernels/interface.jl") -include("microkernels/types.jl") -include("microkernels/scalar.jl") -include("microkernels/simd.jl") -include("microkernels/planar.jl") -include("microkernels/onem.jl") -include("microkernels/fmaddsub.jl") -include("microkernels/mixed.jl") +include("microkernels/kernels.jl") +include("microkernels/vecops.jl") +include("microkernels/steps.jl") +include("microkernels/stores.jl") # --- Planning: labels, conjugation, kernel and blocking choice, the plan --- include("planning/labels.jl") diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 33f1ed7..a9697a4 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -180,7 +180,7 @@ end # default `k_block`; 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.k_block) === (path isa _PanelPath) && - _unpacked_b_method_eligible(complex_method(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) + _unpacked_b_method_eligible(KernelMethod(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 diff --git a/src/microkernels/fmaddsub.jl b/src/microkernels/fmaddsub.jl deleted file mode 100644 index 6535358..0000000 --- a/src/microkernels/fmaddsub.jl +++ /dev/null @@ -1,145 +0,0 @@ -# FMAddSub (interleaved-accumulator) complex microkernel. One accumulator plane -# in 1m's layout (lanes `(2u-1, 2u)` are `(re, im)` of one complex row), updated -# per (A-vector, B-column) pair by two chained x86 `vfmaddsub` ops: -# -# acc = fmaddsub(a, br, fmaddsub(swap(a), bi, acc)) -# -# where `fmaddsub(x, y, c)` is `x*y - c` in even (0-based) lanes and `x*y + c` -# in odd lanes, `a = [ar0, ai0, ...]`, `swap(a) = [ai0, ar0, ...]`. The two -# sign flips compose back to an accumulation: -# -# even: ar*br - (ai*bi - c_re) = c_re + ar*br - ai*bi -# odd: ai*br + (ar*bi + c_im) = c_im + ai*br + ar*bi -# -# The lane-parity sign needs re/im of one element in adjacent lanes, hence -# `InterleavedFormat` A; B is read as broadcast scalars, so it stays planar. -# Same FMA count as planar and 1m, plus one pair-swap per A vector per K step. - -using SIMD: Vec, shufflevector - -# `x*y - c` in even lanes, `x*y + c` in odd lanes, each one fused rounding. -# Generic IR (`fneg` + two `llvm.fma` + a blend), which the X86 backend folds -# into one `vfmaddsub` wherever FMA3 exists (and elsewhere stays correct). -# GUARDRAIL: the SIMD.jl spelling `shufflevector(muladd(x, y, -c), muladd(x, y, -# c), ...)` also selects `vfmaddsub` but leaves a dead stack store of the -# accumulator in the K loop, every step. -@generated function fmaddsub(x::Vec{N, R}, y::Vec{N, R}, c::Vec{N, R}) where {N, R} - R === Float64 || R === Float32 || - return :(throw(ArgumentError("fmaddsub: unsupported lane type $R"))) - ty = "<$N x $(R === Float64 ? "double" : "float")>" - fn = "llvm.fma.v$(N)$(R === Float64 ? "f64" : "f32")" - # lane k: even -> `s` (x*y - c); odd -> `d` (x*y + c), index N + k. - mask = join(("i32 $(iseven(k) ? k : N + k)" for k in 0:(N - 1)), ", ") - ir = """ - declare $ty @$fn($ty, $ty, $ty) - define $ty @entry($ty %x, $ty %y, $ty %c) #0 { - top: - %nc = fneg $ty %c - %s = call $ty @$fn($ty %x, $ty %y, $ty %nc) - %d = call $ty @$fn($ty %x, $ty %y, $ty %c) - %r = shufflevector $ty %s, $ty %d, <$N x i32> <$mask> - ret $ty %r - } - attributes #0 = { alwaysinline } - """ - VT = NTuple{N, VecElement{R}} - return quote - Base.@_inline_meta - Vec{$N, $R}( - Base.llvmcall(($ir, "entry"), $VT, Tuple{$VT, $VT, $VT}, x.data, y.data, c.data) - ) - end -end - -# `[x1, x0, x3, x2, ...]`: one in-lane `vshufpd`/`vpermilps`. -@generated function swap_pairs(x::Vec{N, R}) where {N, R} - iseven(N) || return :(throw(ArgumentError("swap_pairs: expected even N, got $N"))) - idx = ntuple(k -> isodd(k) ? k : k - 2, N) - return :(Base.@_inline_meta; shufflevector(x, Val($idx))) -end - -# The INNER op must be the `swap(a) * bi` one: only the inner product is -# subtracted in the real lanes. Reversed, the real part is `ai*bi - ar*br`. -@generated function accumulate_step( - kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, p::Int - ) where {MR, NR, T, W, R, NA, PA, PB} - MV = (2 * MR) ÷ W - check_acc(:accumulate_step, R, T, NA, MV * NR) - - av = [Symbol(:a, v) for v in 0:(MV - 1)] - sv = [Symbol(:s, v) for v in 0:(MV - 1)] - brv = [Symbol(:br, j) for j in 1:NR] - biv = [Symbol(:bi, j) for j in 1:NR] - - load_a = Any[] - for v in 0:(MV - 1) - push!( - load_a, - :( - $(av[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) - ) - ) - ) - push!(load_a, :($(sv[v + 1]) = swap_pairs($(av[v + 1])))) - end - - load_b = Any[] - for j in 1:NR - push!(load_b, :((br_s, bi_s) = b_complex(packed_b, kernel, $j, p))) - push!(load_b, :($(brv[j]) = Vec{$W, $R}(br_s))) - push!(load_b, :($(biv[j]) = Vec{$W, $R}(bi_s))) - end - - acc_exprs = Vector{Any}(undef, NA) - for j in 1:NR, v in 0:(MV - 1) - idx = acc_index(MV, v, j) - acc_exprs[idx] = :( - fmaddsub( - $(av[v + 1]), $(brv[j]), - fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) - ) - ) - end - - return quote - Base.@_inline_meta - @inbounds begin - $(load_a...) - $(load_b...) - return $(Expr(:tuple, acc_exprs...)) - end - end -end - -# `p - q` in even lanes, `p + q` in odd lanes: Base's unfused complex `*` on -# interleaved data. -@generated function addsub(p::Vec{N, R}, q::Vec{N, R}) where {N, R} - iseven(N) || return :(throw(ArgumentError("addsub: expected even N, got $N"))) - idx = ntuple(k -> iseven(k - 1) ? k - 1 : N + k - 1, N) - return :(Base.@_inline_meta; shufflevector(p - q, p + q, Val($idx))) -end - -# One full `W÷2`-row block, already in `Complex`'s memory order, so no -# interleave shuffle. Base's `Complex` expression trees in lanes, as planar's -# `split_store_block!`: -# beta == 0: addsub(ar*r, ai*swap(r)) -# beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) -# otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) -@inline function lanepair_store_block!( - 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}, - ::Val{B} - ) where {RC, R, W, B} - s = swap_pairs(r) - if B === :zero - new = addsub(ar * r, ai * s) - else - old = convert(Vec{W, R}, vload(Vec{W, RC}, sp + sizeof(RC) * at)) - x = B === :one ? old : addsub(br * old, bi * swap_pairs(old)) - new = fmaddsub(ar, r, fmaddsub(ai, s, x)) - end - vstore(convert(Vec{W, RC}, new), sp + sizeof(RC) * at) - return nothing -end diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 5dc6488..029ad38 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -1,5 +1,5 @@ # The microkernel contract, the complex-arithmetic methods, and the validation -# and store helpers the kernels share. +# and store helpers every kernel shares. """ Microkernel{MR, NR, T} @@ -15,7 +15,7 @@ never bitwise. abstract type Microkernel{MR, NR, T} end """ - KernelMethod + KernelMethod(kernel) How a kernel maps its multiply-adds onto real arithmetic, as singletons so blocking and the shape menus dispatch on it: @@ -51,8 +51,6 @@ pack_formats(::RealComplexMethod) = (RealFormat(), InterleavedFormat()) accumulator_planes(::KernelMethod) = 1 accumulator_planes(::PlanarMethod) = 2 -complex_method(::Any) = RealMethod() - realtype(k::Microkernel) = realtype(k.descriptor) sliver_width(k::Microkernel) = sliver_width(k.descriptor) a_format(k::Microkernel) = a_format(k.descriptor) diff --git a/src/microkernels/kernels.jl b/src/microkernels/kernels.jl new file mode 100644 index 0000000..f3c759b --- /dev/null +++ b/src/microkernels/kernels.jl @@ -0,0 +1,195 @@ +# The kernel types and their type-level facts, ahead of every implementation: a +# `@generated` function's generator only sees methods defined before it. + +""" + ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) + +Scalar reference microkernel for register tile `(MR, NR)` and element type `T`. +Its accumulator is an `NTuple{MR * NR, T}` holding the tile column-major. +""" +struct ScalarKernel{MR, NR, T} <: Microkernel{MR, NR, T} + descriptor::RealDescriptor{MR, NR, T} +end + +ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = + ScalarKernel(Descriptor(Val(MR), Val(NR), T)) + +""" + VectorKernel{MR, NR, T, W} + +An explicit-SIMD microkernel over `SIMD.Vec{W, real(T)}` lanes. Its +accumulator is an immutable tuple of vectors in one of three layouts +([`AccumulatorLayout`](@ref)), which fixes the stores; each kernel supplies +its K step, `accumulate_step`, and the generic `add_tile` loops it +`k_steps(kernel, k_block_length)` times. +""" +abstract type VectorKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} end + +""" + AccumulatorLayout(kernel::VectorKernel) + +How the accumulator holds the tile, as `(MR ÷ rows) * NR` vectors of `rows` +consecutive tile rows each, column-major (`rows_per_vector`): + + - `RealLayout()`: one real per lane (`SIMDKernel`). + - `SplitLayout()`: separate re and im vectors (`PlanarKernel`: a real and an + imaginary plane in one flat tuple; `RealComplexKernel`: the real kernel's + columns `2j-1` and `2j`). + - `LanePairLayout()`: lanes `2u-1` and `2u` are re and im of one row + (1m, fmaddsub, `ComplexRealKernel`). +""" +abstract type AccumulatorLayout end +struct RealLayout <: AccumulatorLayout end +struct SplitLayout <: AccumulatorLayout end +struct LanePairLayout <: AccumulatorLayout end + +# The checks of every vector kernel's constructor; the descriptor has checked +# the tile size and the element type. +function check_vector_kernel(::Type{K}, MR, W) where {K <: VectorKernel} + pairs = AccumulatorLayout(K) isa LanePairLayout + return check_vector_shape(nameof(K), pairs ? 2 * MR : MR, W, pairs) +end + +""" + SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a +multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). +""" +struct SIMDKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::RealDescriptor{MR, NR, T} + function SIMDKernel{MR, NR, T, W}(descriptor::RealDescriptor{MR, NR, T}) where {MR, NR, T, W} + check_vector_kernel(SIMDKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +""" + PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Split-complex microkernel for `T = ComplexF32/ComplexF64` over +`SIMD.Vec{W,real(T)}` lanes. `MR`/`NR` are complex extents; `MR` must be a +multiple of `W`, which counts reals. +""" +struct PlanarKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} + function PlanarKernel{MR, NR, T, W}( + descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} + ) where {MR, NR, T, W} + check_vector_kernel(PlanarKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +""" + OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Van Zee's induced 1m complex microkernel: the real `SIMDKernel` of `2MR x NR` +over `SIMD.Vec{W,real(T)}` lanes, run verbatim over `2 * k_block_length` real +K steps against 1e-packed A and planar B. At logical K step `p`, `OneEFormat` +A holds `(re_0, im_0, re_1, im_1, ...)` then `(-im_0, re_0, -im_1, re_1, ...)`, +and planar B `re_0..` then `im_0..`, so accumulator real row `2i-1` is the +real and `2i` the imaginary part of complex row `i`. `2MR` must be a multiple +of `W`, and `W` must be even. Used only when named in +`plan_contract(...; kernel = ...)`. +""" +struct OneMKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} + function OneMKernel{MR, NR, T, W}( + descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} + ) where {MR, NR, T, W} + check_vector_kernel(OneMKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +""" + FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over +`SIMD.Vec{W,real(T)}` lanes with `InterleavedFormat` A and `PlanarFormat` B. +`2MR` must be a multiple of `W`, and `W` must be even. The default for small-M +complex contractions on AVX-512. +""" +struct FMAddSubKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} + function FMAddSubKernel{MR, NR, T, W}( + descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} + ) where {MR, NR, T, W} + check_vector_kernel(FMAddSubKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +""" + ComplexRealKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Mixed-domain microkernel for complex `A` times real `B`, computing in complex +`T`. Interleaved A is a real `2MR x k_block_length` panel, so this is the real +`SIMDKernel` of `2MR x NR` over `SIMD.Vec{W,real(T)}` lanes, run verbatim at 2 +real FMAs per complex MAC; the accumulator is 1m's lane-pair layout. `2MR` +must be a multiple of `W`, and `W` must be even. +""" +struct ComplexRealKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} + function ComplexRealKernel{MR, NR, T, W}( + descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} + ) where {MR, NR, T, W} + check_vector_kernel(ComplexRealKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +""" + RealComplexKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) + +Mixed-domain microkernel for real `A` times complex `B`, computing in complex +`T`. Interleaved B is a real `k_block_length x 2NR` panel, so this is the real +`SIMDKernel` of `MR x 2NR` over `SIMD.Vec{W,real(T)}` lanes, run verbatim at 2 +real FMAs per complex MAC; accumulator column `2j-1` holds the real and `2j` +the imaginary part of column `j`. `MR` must be a multiple of `W`. +""" +struct RealComplexKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} + descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} + function RealComplexKernel{MR, NR, T, W}( + descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} + ) where {MR, NR, T, W} + check_vector_kernel(RealComplexKernel, MR, W) + return new{MR, NR, T, W}(descriptor) + end +end + +KernelMethod(::Type{<:Union{ScalarKernel, SIMDKernel}}) = RealMethod() +KernelMethod(::Type{<:PlanarKernel}) = PlanarMethod() +KernelMethod(::Type{<:OneMKernel}) = OneMMethod() +KernelMethod(::Type{<:FMAddSubKernel}) = FMAddSubMethod() +KernelMethod(::Type{<:ComplexRealKernel}) = ComplexRealMethod() +KernelMethod(::Type{<:RealComplexKernel}) = RealComplexMethod() +KernelMethod(kernel::Microkernel) = KernelMethod(typeof(kernel)) + +AccumulatorLayout(::Type{<:SIMDKernel}) = RealLayout() +AccumulatorLayout(::Type{<:Union{PlanarKernel, RealComplexKernel}}) = SplitLayout() +AccumulatorLayout(::Type{<:Union{OneMKernel, FMAddSubKernel, ComplexRealKernel}}) = LanePairLayout() +AccumulatorLayout(kernel::VectorKernel) = AccumulatorLayout(typeof(kernel)) + +rows_per_vector(::AccumulatorLayout, W::Int) = W +rows_per_vector(::LanePairLayout, W::Int) = W ÷ 2 +accumulator_length(::AccumulatorLayout, MV::Int, NR::Int) = MV * NR +accumulator_length(::SplitLayout, MV::Int, NR::Int) = 2 * MV * NR +# The accumulator vector of row block `v` (zero-based) of column `j`. +acc_index(MV::Int, v::Int, j::Int) = v + MV * (j - 1) + 1 + +(::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {K <: VectorKernel, MR, NR, T, W} = + K{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T, pack_formats(KernelMethod(K))...)) +(::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}) where {K <: VectorKernel, MR, NR, T} = + K(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) + +lanewidth(::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} = W + +k_steps(::VectorKernel, k_block_length::Int) = k_block_length +k_steps(::OneMKernel, k_block_length::Int) = 2 * k_block_length + +# The real kernel whose K step a kernel without one of its own runs. +inner(::OneMKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) +inner(::ComplexRealKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) +inner(::RealComplexKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(MR), Val(2 * NR), real(T), Val(W)) diff --git a/src/microkernels/mixed.jl b/src/microkernels/mixed.jl deleted file mode 100644 index 138f52a..0000000 --- a/src/microkernels/mixed.jl +++ /dev/null @@ -1,11 +0,0 @@ -# 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 k_block_length` -# panel; the accumulator is 1m/fmaddsub's interleaved layout. -# * real A x complex B: interleaved B is a real `k_block_length x 2NR` panel; -# accumulator column `2j-1` holds the real and `2j` the imaginary part of -# column `j`. - -inner(::ComplexRealKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) -inner(::RealComplexKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(MR), Val(2 * NR), real(T), Val(W)) diff --git a/src/microkernels/onem.jl b/src/microkernels/onem.jl deleted file mode 100644 index ee9d9b8..0000000 --- a/src/microkernels/onem.jl +++ /dev/null @@ -1,10 +0,0 @@ -# 1m (Van Zee's induced method): the real `SIMDKernel` of `2MR x NR`, run over -# `2*k_block_length` real K steps against 1e-packed A and planar B. At logical K -# step `p`, `OneEFormat` A holds `(re_0, im_0, re_1, im_1, ...)` then `(-im_0, -# re_0, -im_1, re_1, ...)`, and planar B `re_0..` then `im_0..`, so accumulator -# real row `2i-1` is the real and `2i` the imaginary part of complex row `i`. -# There is deliberately no FMA loop here: 1m reuses the real kernel body -# verbatim. - -inner(::OneMKernel{MR, NR, T, W}) where {MR, NR, T, W} = SIMDKernel(Val(2 * MR), Val(NR), real(T), Val(W)) -k_steps(::OneMKernel, k_block_length::Int) = 2 * k_block_length diff --git a/src/microkernels/planar.jl b/src/microkernels/planar.jl deleted file mode 100644 index cdd3c1f..0000000 --- a/src/microkernels/planar.jl +++ /dev/null @@ -1,129 +0,0 @@ -# Planar (split-complex, BLIS "1r") microkernel, the default complex method. -# Both panels are `PlanarFormat` (`[re_0..re_{n-1} | im_0..im_{n-1}]` per K -# step), so the data is already in the right lanes: four real FMAs per -# (A-vector, B-scalar) pair, no shuffles. Everything works in `real(T)`. - -using SIMD: Vec, vload, vstore, shufflevector - -# GUARDRAIL: the real part is `muladd(-ai, bi, muladd(ar, br, c))`. `c - ai*bi` -# does NOT fuse (Julia sets no LLVM `contract` flag): two instructions, and -# different rounding. `-ai` is hoisted out of the `j` loop and folds into -# `vfnmadd`, so the negation is free. -@generated function accumulate_step( - kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, p::Int - ) where {MR, NR, T, W, R, NA, PA, PB} - MV = MR ÷ W - NV = MV * NR - check_acc(:accumulate_step, R, T, NA, 2NV) - - arv = [Symbol(:ar, v) for v in 0:(MV - 1)] - aiv = [Symbol(:ai, v) for v in 0:(MV - 1)] - naiv = [Symbol(:nai, v) for v in 0:(MV - 1)] - brv = [Symbol(:br, j) for j in 1:NR] - biv = [Symbol(:bi, j) for j in 1:NR] - - load_a = Any[] - for v in 0:(MV - 1) - push!( - load_a, - :( - $(arv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) - ) - ) - ) - push!( - load_a, - :( - $(aiv[v + 1]) = panel_vload( - Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p, 1) - ) - ) - ) - push!(load_a, :($(naiv[v + 1]) = -$(aiv[v + 1]))) - end - - load_b = Any[] - for j in 1:NR - push!(load_b, :(($(brv[j]), $(biv[j])) = b_complex(packed_b, kernel, $j, p))) - end - - acc_exprs = Vector{Any}(undef, NA) - for j in 1:NR, v in 0:(MV - 1) - idx = acc_index(MV, v, j) - acc_exprs[idx] = :( # re: ar*br - ai*bi - muladd( - $(naiv[v + 1]), $(biv[j]), - muladd($(arv[v + 1]), $(brv[j]), acc[$idx]) - ) - ) - acc_exprs[NV + idx] = :( # im: ar*bi + ai*br - muladd( - $(aiv[v + 1]), $(brv[j]), - muladd($(arv[v + 1]), $(biv[j]), acc[$(NV + idx)]) - ) - ) - end - - return quote - Base.@_inline_meta - @inbounds begin - $(load_a...) - $(load_b...) - return $(Expr(:tuple, acc_exprs...)) - end - end -end - -# Shuffle patterns built from `W` at specialization time, never hardcoded to -# one ISA. `v` holds `W` complex values `[re_0, im_0, ..., re_{W-1}, im_{W-1}]`. -@generated function deinterleave_re(v::Vec{N, R}, ::Val{W}) where {N, R, W} - N == 2 * W || return :(throw(ArgumentError("deinterleave_re: expected N == 2W"))) - idx = ntuple(k -> 2 * (k - 1), W) - return :(shufflevector(v, Val($idx))) -end - -@generated function deinterleave_im(v::Vec{N, R}, ::Val{W}) where {N, R, W} - N == 2 * W || return :(throw(ArgumentError("deinterleave_im: expected N == 2W"))) - idx = ntuple(k -> 2 * (k - 1) + 1, W) - return :(shufflevector(v, Val($idx))) -end - -@generated function interleave_planes(re::Vec{W, R}, im::Vec{W, R}, ::Val{W}) where {W, R} - idx = ntuple(k -> isodd(k) ? (k - 1) ÷ 2 : W + (k - 1) ÷ 2, 2 * W) - return :(shufflevector(re, im, Val($idx))) -end - -# One full `W`-row block for `beta` case `B`; `at` is the zero-based index of -# its first real. -# The arithmetic transcribes Base's `Complex` expression trees (the ones -# `axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, -# `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, -# muladd(zi, wr, xi)))`. Against the scalar fallback it is exact at -# `beta == 0/1` and ~1 ULP otherwise, because LLVM contracts Base's scalar -# complex `muladd` depending on inlining context. -@inline function split_store_block!( - 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}, - ::Val{W}, ::Val{B} - ) where {RC, R, W, B} - if B === :zero - newre = ar * rev - ai * imv - newim = ar * imv + ai * rev - else - 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 B === :one - xr, xi = orv, oiv - else - xr = br * orv - bi * oiv - xi = br * oiv + bi * orv - end - newre = muladd(ar, rev, -muladd(ai, imv, -xr)) - newim = muladd(ar, imv, muladd(ai, rev, xi)) - end - vstore(convert(Vec{2 * W, RC}, interleave_planes(newre, newim, Val(W))), sp + sizeof(RC) * at) - return nothing -end diff --git a/src/microkernels/scalar.jl b/src/microkernels/scalar.jl deleted file mode 100644 index 634efdf..0000000 --- a/src/microkernels/scalar.jl +++ /dev/null @@ -1,60 +0,0 @@ -""" - ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) - -Scalar reference microkernel for register tile `(MR, NR)` and element type `T`. -Its accumulator is an `NTuple{MR * NR, T}` holding the tile column-major. -""" -struct ScalarKernel{MR, NR, T} <: Microkernel{MR, NR, T} - descriptor::RealDescriptor{MR, NR, T} -end - -ScalarKernel(::Val{MR}, ::Val{NR}, ::Type{T}) where {MR, NR, T} = - ScalarKernel(Descriptor(Val(MR), Val(NR), T)) - -zero_accumulator(::ScalarKernel{MR, NR, T}) where {MR, NR, T} = ntuple(_ -> zero(T), Val(MR * NR)) - -function add_tile( - kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {MR, NR, T, N, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - for p in 1:k_block_length - acc = add_step(kernel, acc, packed_a, packed_b, p) - end - return acc -end - -# One K step: element `(i, j)` of the tile is `acc[i + MR * (j - 1)]`. -@inline function add_step( - kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, packed_a, packed_b, p::Int - ) where {MR, NR, T, N} - return ntuple(Val(N)) do q - i, j = mod1(q, MR), cld(q, MR) - a = panel_load(packed_a, packed_a_offset(kernel, i, p)) - b = panel_load(packed_b, packed_b_offset(kernel, j, p)) - muladd(a, b, acc[q]) - end -end - -function store_tile!( - destination::Tile, acc::NTuple{N, T}, - alpha::T, beta::T, kernel::ScalarKernel{MR} - ) where {MR, N, T} - m, n = store_prologue!(destination, alpha, beta) - (m == 0 || n == 0) && return destination - if iszero(beta) - @inbounds for j in 1:n, i in 1:m - destination[i, j] = alpha * acc[i + MR * (j - 1)] - end - elseif isone(beta) - @inbounds for j in 1:n, i in 1:m - destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], convert(T, destination[i, j])) - end - else - @inbounds for j in 1:n, i in 1:m - destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], beta * convert(T, destination[i, j])) - end - end - return destination -end diff --git a/src/microkernels/simd.jl b/src/microkernels/simd.jl deleted file mode 100644 index 0f81467..0000000 --- a/src/microkernels/simd.jl +++ /dev/null @@ -1,207 +0,0 @@ -# The vector kernels' shared `add_tile` and stores, and the explicit-SIMD real -# K step. The accumulator is an immutable tuple of `Vec`s so it stays in -# registers, and every K step and store is `@generated` straight-line code. -# -# GUARDRAIL (all kernels): every `acc[...]` must be a *literal* tuple index. -# Indexing an `NTuple` dynamically forces it to memory, and above 16 vectors -# the compiler heap-allocates it on every call. Only the lane index inside one -# `Vec` may be a runtime value. - -using SIMD: Vec, vload, vstore - -function zero_accumulator(kernel::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} - layout = accumulator_layout(kernel) - z = zero(Vec{W, real(T)}) - return ntuple(_ -> z, Val(accumulator_length(layout, MR ÷ rows_per_vector(layout, W), NR))) -end - -k_steps(::VectorKernel, k_block_length::Int) = k_block_length - -@inline function add_tile( - kernel::VectorKernel, acc::NTuple{NA, Vec{W, R}}, - packed_a::PA, packed_b::PB, k_block_length::Int - ) where {NA, W, R, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - @inbounds for p in 1:k_steps(kernel, k_block_length) - acc = accumulate_step(kernel, acc, packed_a, packed_b, p) - end - return acc -end - -# B column `j` at K step `p`, as a real or as `(re, im)`; `UnpackedBView` -# (src/execution/unpackedb.jl) overrides both to read B in place, resolved at -# compile time. -@inline b_scalar(packed_b::PB, kernel, j::Int, p::Int) where {PB} = - panel_load(packed_b, packed_b_offset(kernel, j, p)) -@inline b_complex(packed_b::PB, kernel, j::Int, p::Int) where {PB} = ( - panel_load(packed_b, packed_b_offset(kernel, j, p)), - panel_load(packed_b, packed_b_offset(kernel, j, p, 1)), -) - -@generated function accumulate_step( - kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, - packed_a::PA, packed_b::PB, p::Int - ) where {MR, NR, T, W, NV, PA, PB} - NVECA = MR ÷ W - check_acc(:accumulate_step, T, T, NV, NVECA * NR) - - avars = [Symbol(:a, v) for v in 0:(NVECA - 1)] - bvars = [Symbol(:b, j) for j in 1:NR] - - load_a = [ - :($(avars[v + 1]) = panel_vload(Vec{$W, $T}, packed_a, packed_a_offset(kernel, $(v * W + 1), p))) - for v in 0:(NVECA - 1) - ] - load_b = [ - :($(bvars[j]) = b_scalar(packed_b, kernel, $j, p)) - for j in 1:NR - ] - - acc_exprs = Vector{Any}(undef, NV) - for j in 1:NR, v in 0:(NVECA - 1) - idx = acc_index(NVECA, v, j) - acc_exprs[idx] = :(muladd($(avars[v + 1]), $(bvars[j]), acc[$idx])) - end - - return quote - Base.@_inline_meta - @inbounds begin - $(load_a...) - $(load_b...) - return $(Expr(:tuple, acc_exprs...)) - end - end -end - -# A kernel without a K step of its own runs its `inner(kernel)`'s. -@inline accumulate_step(kernel::VectorKernel, acc::NTuple, packed_a::PA, packed_b::PB, p::Int) where {PA, PB} = - accumulate_step(inner(kernel), acc, packed_a, packed_b, p) - -# Vector store eligibility: unit-stride rows into rank-1 dense 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. The complex -# layouts reinterpret `W` rows as `2W` consecutive reals, on an ISA the complex -# fast paths ship for (shared with the complex pack fast path). -@inline vector_store_eligible(::RealLayout, tile::Tile, ::Type{T}) where {T} = - is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) -@inline vector_store_eligible(::AccumulatorLayout, tile::Tile, ::Type{T}) where {T} = - is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) && - complex_fastpath_isa_eligible() - -@generated function store_tile!( - destination::Tile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W} - ) where {MR, NR, T, W, R, NA} - layout = accumulator_layout(kernel) - return quote - $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) - m, n = store_prologue!(destination, alpha, beta) - (m == 0 || n == 0) && return destination - - if vector_store_eligible($layout, destination, T) - return vector_store!(destination, acc, alpha, beta, kernel, m, n) - end - - return scalar_store!(destination, acc, alpha, beta, kernel, m, n) - end -end - -# Element by element, for every destination the vector store cannot take. -@generated function scalar_store!( - destination::Tile, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, - m::Int, n::Int - ) where {MR, NR, T, W, R, NA} - layout = accumulator_layout(kernel) - rows = rows_per_vector(layout, W) - MV = MR ÷ rows - check_acc(:scalar_store!, R, T, NA, accumulator_length(layout, MV, NR)) - - blocks = Any[] - for j in 1:NR, v in 0:(MV - 1) - push!( - blocks, quote - if $j <= n - $(acc_bindings(layout, kernel, MV, NR, v, j)) - for lane in 1:$rows - i = $(v * rows) + lane - i <= m || break - axpby_tile!(destination, i, $j, alpha, $(lane_value(layout)), beta) - end - end - end - ) - end - return quote - @inbounds begin - $(blocks...) - end - return destination - end -end - -# Whole row blocks are one vector load/store; a block straddling `m` is stored -# lane by lane, so nothing outside the valid rectangle is touched. The blocks -# are generated once per `beta` case, so `beta` is tested once per tile. -# `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. -@generated function vector_store!( - destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, - alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, - m::Int, n::Int - ) where {S, MR, NR, T, W, R, NA} - layout = accumulator_layout(kernel) - rows = rows_per_vector(layout, W) - MV = MR ÷ rows - check_acc(:vector_store!, R, T, NA, accumulator_length(layout, MV, NR)) - RC = store_lanetype(layout, S, T) - - function blocks(B) - out = Any[] - for j in 1:NR - vblocks = Any[] - for v in 0:(MV - 1) - push!( - vblocks, quote - $(acc_bindings(layout, kernel, MV, NR, v, j)) - if $((v + 1) * rows) <= m - $(block_store(layout, :(colbase + $(v * rows)), W, R, RC, B)) - elseif $(v * rows) < m - for lane in 1:$rows - i = $(v * rows) + lane - i <= m || break - axpby_at!(storage, colbase + i, alpha, $(lane_value(layout)), beta) - end - end - end - ) - end - push!( - out, quote - if $j <= n - colbase = rowbase0 + cols[$j] # the address of (1, j) - $(vblocks...) - end - end - ) - end - return out - end - by_beta = quote - if iszero(beta) - $(blocks(:zero)...) - elseif isone(beta) - $(blocks(:one)...) - else - $(blocks(:general)...) - end - end - return quote - $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) - storage = destination.storage - cols = destination.cols - rowbase0 = @inbounds destination.base + destination.rows[1] - $(store_body(layout, W, R, RC, by_beta)) - return destination - end -end diff --git a/src/microkernels/steps.jl b/src/microkernels/steps.jl new file mode 100644 index 0000000..4e5334a --- /dev/null +++ b/src/microkernels/steps.jl @@ -0,0 +1,251 @@ +# The K steps and the generic `add_tile`. A vector kernel's accumulator is an +# immutable tuple of `Vec`s so it stays in registers, and every K step and store +# is `@generated` straight-line code. +# +# GUARDRAIL (all kernels): every `acc[...]` must be a *literal* tuple index. +# Indexing an `NTuple` dynamically forces it to memory, and above 16 vectors +# the compiler heap-allocates it on every call. Only the lane index inside one +# `Vec` may be a runtime value. + +using SIMD: Vec + +zero_accumulator(::ScalarKernel{MR, NR, T}) where {MR, NR, T} = ntuple(_ -> zero(T), Val(MR * NR)) + +function add_tile( + kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, + packed_a::PA, packed_b::PB, k_block_length::Int + ) where {MR, NR, T, N, PA <: PackedPanel, PB} + k_block_length == 0 && return acc + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) + for p in 1:k_block_length + acc = add_step(kernel, acc, packed_a, packed_b, p) + end + return acc +end + +# One K step: element `(i, j)` of the tile is `acc[i + MR * (j - 1)]`. +@inline function add_step( + kernel::ScalarKernel{MR, NR, T}, acc::NTuple{N, T}, packed_a, packed_b, p::Int + ) where {MR, NR, T, N} + return ntuple(Val(N)) do q + i, j = mod1(q, MR), cld(q, MR) + a = panel_load(packed_a, packed_a_offset(kernel, i, p)) + b = panel_load(packed_b, packed_b_offset(kernel, j, p)) + muladd(a, b, acc[q]) + end +end + +function zero_accumulator(kernel::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} + layout = AccumulatorLayout(kernel) + z = zero(Vec{W, real(T)}) + return ntuple(_ -> z, Val(accumulator_length(layout, MR ÷ rows_per_vector(layout, W), NR))) +end + +@inline function add_tile( + kernel::VectorKernel, acc::NTuple{NA, Vec{W, R}}, + packed_a::PA, packed_b::PB, k_block_length::Int + ) where {NA, W, R, PA <: PackedPanel, PB} + k_block_length == 0 && return acc + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) + @inbounds for p in 1:k_steps(kernel, k_block_length) + acc = accumulate_step(kernel, acc, packed_a, packed_b, p) + end + return acc +end + +# B column `j` at K step `p`, as a real or as `(re, im)`; `UnpackedBView` +# (src/execution/unpackedb.jl) overrides both to read B in place, resolved at +# compile time. +@inline b_scalar(packed_b::PB, kernel, j::Int, p::Int) where {PB} = + panel_load(packed_b, packed_b_offset(kernel, j, p)) +@inline b_complex(packed_b::PB, kernel, j::Int, p::Int) where {PB} = ( + panel_load(packed_b, packed_b_offset(kernel, j, p)), + panel_load(packed_b, packed_b_offset(kernel, j, p, 1)), +) + +@generated function accumulate_step( + kernel::SIMDKernel{MR, NR, T, W}, acc::NTuple{NV, Vec{W, T}}, + packed_a::PA, packed_b::PB, p::Int + ) where {MR, NR, T, W, NV, PA, PB} + NVECA = MR ÷ W + check_acc(:accumulate_step, T, T, NV, NVECA * NR) + + avars = [Symbol(:a, v) for v in 0:(NVECA - 1)] + bvars = [Symbol(:b, j) for j in 1:NR] + + load_a = [ + :($(avars[v + 1]) = panel_vload(Vec{$W, $T}, packed_a, packed_a_offset(kernel, $(v * W + 1), p))) + for v in 0:(NVECA - 1) + ] + load_b = [ + :($(bvars[j]) = b_scalar(packed_b, kernel, $j, p)) + for j in 1:NR + ] + + acc_exprs = Vector{Any}(undef, NV) + for j in 1:NR, v in 0:(NVECA - 1) + idx = acc_index(NVECA, v, j) + acc_exprs[idx] = :(muladd($(avars[v + 1]), $(bvars[j]), acc[$idx])) + end + + return quote + Base.@_inline_meta + @inbounds begin + $(load_a...) + $(load_b...) + return $(Expr(:tuple, acc_exprs...)) + end + end +end + +# A kernel without a K step of its own runs its `inner(kernel)`'s. +@inline accumulate_step(kernel::VectorKernel, acc::NTuple, packed_a::PA, packed_b::PB, p::Int) where {PA, PB} = + accumulate_step(inner(kernel), acc, packed_a, packed_b, p) + +# Planar (split-complex, BLIS "1r"), the default complex method. Both panels +# are `PlanarFormat` (`[re_0..re_{n-1} | im_0..im_{n-1}]` per K step), so the +# data is already in the right lanes: four real FMAs per (A-vector, B-scalar) +# pair, no shuffles. Everything works in `real(T)`. +# +# GUARDRAIL: the real part is `muladd(-ai, bi, muladd(ar, br, c))`. `c - ai*bi` +# does NOT fuse (Julia sets no LLVM `contract` flag): two instructions, and +# different rounding. `-ai` is hoisted out of the `j` loop and folds into +# `vfnmadd`, so the negation is free. +@generated function accumulate_step( + kernel::PlanarKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, + packed_a::PA, packed_b::PB, p::Int + ) where {MR, NR, T, W, R, NA, PA, PB} + MV = MR ÷ W + NV = MV * NR + check_acc(:accumulate_step, R, T, NA, 2NV) + + arv = [Symbol(:ar, v) for v in 0:(MV - 1)] + aiv = [Symbol(:ai, v) for v in 0:(MV - 1)] + naiv = [Symbol(:nai, v) for v in 0:(MV - 1)] + brv = [Symbol(:br, j) for j in 1:NR] + biv = [Symbol(:bi, j) for j in 1:NR] + + load_a = Any[] + for v in 0:(MV - 1) + push!( + load_a, + :( + $(arv[v + 1]) = panel_vload( + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) + ) + ) + ) + push!( + load_a, + :( + $(aiv[v + 1]) = panel_vload( + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p, 1) + ) + ) + ) + push!(load_a, :($(naiv[v + 1]) = -$(aiv[v + 1]))) + end + + load_b = Any[] + for j in 1:NR + push!(load_b, :(($(brv[j]), $(biv[j])) = b_complex(packed_b, kernel, $j, p))) + end + + acc_exprs = Vector{Any}(undef, NA) + for j in 1:NR, v in 0:(MV - 1) + idx = acc_index(MV, v, j) + acc_exprs[idx] = :( # re: ar*br - ai*bi + muladd( + $(naiv[v + 1]), $(biv[j]), + muladd($(arv[v + 1]), $(brv[j]), acc[$idx]) + ) + ) + acc_exprs[NV + idx] = :( # im: ar*bi + ai*br + muladd( + $(aiv[v + 1]), $(brv[j]), + muladd($(arv[v + 1]), $(biv[j]), acc[$(NV + idx)]) + ) + ) + end + + return quote + Base.@_inline_meta + @inbounds begin + $(load_a...) + $(load_b...) + return $(Expr(:tuple, acc_exprs...)) + end + end +end + +# FMAddSub (interleaved accumulator): one accumulator plane in 1m's layout +# (lanes `(2u-1, 2u)` are `(re, im)` of one complex row), updated per +# (A-vector, B-column) pair by two chained x86 `vfmaddsub` ops: +# +# acc = fmaddsub(a, br, fmaddsub(swap(a), bi, acc)) +# +# where `fmaddsub(x, y, c)` is `x*y - c` in even (0-based) lanes and `x*y + c` +# in odd lanes, `a = [ar0, ai0, ...]`, `swap(a) = [ai0, ar0, ...]`. The two +# sign flips compose back to an accumulation: +# +# even: ar*br - (ai*bi - c_re) = c_re + ar*br - ai*bi +# odd: ai*br + (ar*bi + c_im) = c_im + ai*br + ar*bi +# +# The lane-parity sign needs re/im of one element in adjacent lanes, hence +# `InterleavedFormat` A; B is read as broadcast scalars, so it stays planar. +# Same FMA count as planar and 1m, plus one pair-swap per A vector per K step. +# +# The INNER op must be the `swap(a) * bi` one: only the inner product is +# subtracted in the real lanes. Reversed, the real part is `ai*bi - ar*br`. +@generated function accumulate_step( + kernel::FMAddSubKernel{MR, NR, T, W}, acc::NTuple{NA, Vec{W, R}}, + packed_a::PA, packed_b::PB, p::Int + ) where {MR, NR, T, W, R, NA, PA, PB} + MV = (2 * MR) ÷ W + check_acc(:accumulate_step, R, T, NA, MV * NR) + + av = [Symbol(:a, v) for v in 0:(MV - 1)] + sv = [Symbol(:s, v) for v in 0:(MV - 1)] + brv = [Symbol(:br, j) for j in 1:NR] + biv = [Symbol(:bi, j) for j in 1:NR] + + load_a = Any[] + for v in 0:(MV - 1) + push!( + load_a, + :( + $(av[v + 1]) = panel_vload( + Vec{$W, $R}, packed_a, packed_a_offset(kernel, $(v * W + 1), p) + ) + ) + ) + push!(load_a, :($(sv[v + 1]) = swap_pairs($(av[v + 1])))) + end + + load_b = Any[] + for j in 1:NR + push!(load_b, :((br_s, bi_s) = b_complex(packed_b, kernel, $j, p))) + push!(load_b, :($(brv[j]) = Vec{$W, $R}(br_s))) + push!(load_b, :($(biv[j]) = Vec{$W, $R}(bi_s))) + end + + acc_exprs = Vector{Any}(undef, NA) + for j in 1:NR, v in 0:(MV - 1) + idx = acc_index(MV, v, j) + acc_exprs[idx] = :( + fmaddsub( + $(av[v + 1]), $(brv[j]), + fmaddsub($(sv[v + 1]), $(biv[j]), acc[$idx]) + ) + ) + end + + return quote + Base.@_inline_meta + @inbounds begin + $(load_a...) + $(load_b...) + return $(Expr(:tuple, acc_exprs...)) + end + end +end diff --git a/src/microkernels/stores.jl b/src/microkernels/stores.jl new file mode 100644 index 0000000..c740792 --- /dev/null +++ b/src/microkernels/stores.jl @@ -0,0 +1,279 @@ +# The stores: one generic `store_tile!` per kernel, picking the vector or the +# scalar store; each accumulator layout supplies its full-block store and lane +# value. + +using SIMD: Vec, vload, vstore + +# Real and split stores are `@inline`: out of line, the whole accumulator is +# spilled to the stack and reloaded on every micro-tile. The lane-pair store is +# not: inlining it cost time (code growth). The scalar stores never are: they +# are scalar anyway, and inlining them bloats the tile function. +inline_store(::AccumulatorLayout) = true +inline_store(::LanePairLayout) = false + +# Generator-time pieces of the stores. `acc_bindings` names the accumulator +# vector(s) of row block `v` (zero-based) of column `j`, `lane_value` is the +# element at row `lane` of the block, `block_store` stores a full block whose +# first row is at zero-based storage index `first`, for `beta` case `B` +# (`:zero`, `:one` or `:general`). +split_index(::Type{<:PlanarKernel}, MV::Int, NR::Int, v::Int, j::Int) = + (acc_index(MV, v, j), MV * NR + acc_index(MV, v, j)) +split_index(::Type{<:RealComplexKernel}, MV::Int, NR::Int, v::Int, j::Int) = + (v + MV * (2j - 2) + 1, v + MV * (2j - 1) + 1) + +acc_bindings(::AccumulatorLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) = + :(vec = acc[$(acc_index(MV, v, j))]) +function acc_bindings(::SplitLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) + ire, iim = split_index(kernel, MV, NR, v, j) + return quote + revec = acc[$ire] + imvec = acc[$iim] + end +end + +lane_value(::RealLayout) = :(vec[lane]) +lane_value(::SplitLayout) = :(Complex(revec[lane], imvec[lane])) +lane_value(::LanePairLayout) = :(Complex(vec[2 * lane - 1], vec[2 * lane])) + +function block_store(::RealLayout, first, W::Int, R::Type, RC::Type, B::Symbol) + old = :(convert(Vec{$W, $R}, vload(Vec{$W, $RC}, storage, at))) + new = B === :zero ? :(alpha * vec) : B === :one ? :(muladd(alpha, vec, $old)) : :(muladd(alpha, vec, beta * $old)) + return quote + at = $first + 1 + vstore(convert(Vec{$W, $RC}, $new), storage, at) + end +end +block_store(::SplitLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = + :(split_store_block!(sp, 2 * $first, revec, imvec, ar, ai, br, bi, Val($W), Val($(QuoteNode(B))))) +block_store(::LanePairLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = + :(lanepair_store_block!(sp, 2 * $first, vec, ar, ai, br, bi, Val($(QuoteNode(B))))) + +# The real lane type of the vector store's storage `S`. The complex layouts +# reinterpret the storage as reals, only sound on dense rank-1 complex storage. +store_lanetype(::RealLayout, S::Type, T::Type) = eltype(S) +function store_lanetype(::AccumulatorLayout, S::Type, T::Type) + S <: DenseVector && lane_convertible(eltype(S), T) || + throw(ArgumentError("vector_store!: storage $S is not a dense vector convertible to $T")) + return real(eltype(S)) +end + +# The vector store's body around its blocks: the complex layouts broadcast +# `alpha`/`beta` once and store through a raw pointer, only dereferenced +# inside `GC.@preserve`. +store_body(::RealLayout, W::Int, R::Type, RC::Type, body) = :(@inbounds $body) +store_body(::AccumulatorLayout, W::Int, R::Type, RC::Type, body) = quote + ar = Vec{$W, $R}(real(alpha)) + ai = Vec{$W, $R}(imag(alpha)) + br = Vec{$W, $R}(real(beta)) + bi = Vec{$W, $R}(imag(beta)) + GC.@preserve storage begin + sp = reinterpret(Ptr{$RC}, pointer(storage)) + @inbounds $body + end +end + +# Vector store eligibility: unit-stride rows into rank-1 dense 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. The complex +# layouts reinterpret `W` rows as `2W` consecutive reals, on an ISA the complex +# fast paths ship for (shared with the complex pack fast path). +@inline vector_store_eligible(::RealLayout, tile::Tile, ::Type{T}) where {T} = + is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) +@inline vector_store_eligible(::AccumulatorLayout, tile::Tile, ::Type{T}) where {T} = + is_unit_stride(tile.rows) && dense_lanes(tile.storage, T) && + complex_fastpath_isa_eligible() + +# One full `W`-row block for `beta` case `B`; `at` is the zero-based index of +# its first real. +# The arithmetic transcribes Base's `Complex` expression trees (the ones +# `axpby_tile!` reaches), so full blocks match them bitwise: `*` is unfused, +# `muladd(z, w, x) = (muladd(zr, wr, -muladd(zi, wi, -xr)), muladd(zr, wi, +# muladd(zi, wr, xi)))`. Against the scalar fallback it is exact at +# `beta == 0/1` and ~1 ULP otherwise, because LLVM contracts Base's scalar +# complex `muladd` depending on inlining context. +@inline function split_store_block!( + 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}, + ::Val{W}, ::Val{B} + ) where {RC, R, W, B} + if B === :zero + newre = ar * rev - ai * imv + newim = ar * imv + ai * rev + else + 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 B === :one + xr, xi = orv, oiv + else + xr = br * orv - bi * oiv + xi = br * oiv + bi * orv + end + newre = muladd(ar, rev, -muladd(ai, imv, -xr)) + newim = muladd(ar, imv, muladd(ai, rev, xi)) + end + vstore(convert(Vec{2 * W, RC}, interleave_planes(newre, newim, Val(W))), sp + sizeof(RC) * at) + return nothing +end + +# One full `W÷2`-row block, already in `Complex`'s memory order, so no +# interleave shuffle. Base's `Complex` expression trees in lanes, as planar's +# `split_store_block!`: +# beta == 0: addsub(ar*r, ai*swap(r)) +# beta == 1: fmaddsub(ar, r, fmaddsub(ai, swap(r), C)) +# otherwise: as beta == 1 with C := addsub(br*C, bi*swap(C)) +@inline function lanepair_store_block!( + 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}, + ::Val{B} + ) where {RC, R, W, B} + s = swap_pairs(r) + if B === :zero + new = addsub(ar * r, ai * s) + else + old = convert(Vec{W, R}, vload(Vec{W, RC}, sp + sizeof(RC) * at)) + x = B === :one ? old : addsub(br * old, bi * swap_pairs(old)) + new = fmaddsub(ar, r, fmaddsub(ai, s, x)) + end + vstore(convert(Vec{W, RC}, new), sp + sizeof(RC) * at) + return nothing +end + +@generated function store_tile!( + destination::Tile, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W} + ) where {MR, NR, T, W, R, NA} + layout = AccumulatorLayout(kernel) + return quote + $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) + m, n = store_prologue!(destination, alpha, beta) + (m == 0 || n == 0) && return destination + + if vector_store_eligible($layout, destination, T) + return vector_store!(destination, acc, alpha, beta, kernel, m, n) + end + + return scalar_store!(destination, acc, alpha, beta, kernel, m, n) + end +end + +# Element by element, for every destination the vector store cannot take. +@generated function scalar_store!( + destination::Tile, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, + m::Int, n::Int + ) where {MR, NR, T, W, R, NA} + layout = AccumulatorLayout(kernel) + rows = rows_per_vector(layout, W) + MV = MR ÷ rows + check_acc(:scalar_store!, R, T, NA, accumulator_length(layout, MV, NR)) + + blocks = Any[] + for j in 1:NR, v in 0:(MV - 1) + push!( + blocks, quote + if $j <= n + $(acc_bindings(layout, kernel, MV, NR, v, j)) + for lane in 1:$rows + i = $(v * rows) + lane + i <= m || break + axpby_tile!(destination, i, $j, alpha, $(lane_value(layout)), beta) + end + end + end + ) + end + return quote + @inbounds begin + $(blocks...) + end + return destination + end +end + +# Whole row blocks are one vector load/store; a block straddling `m` is stored +# lane by lane, so nothing outside the valid rectangle is touched. The blocks +# are generated once per `beta` case, so `beta` is tested once per tile. +# `rows::AffineAxis` in the signature: an ineligible tile is a MethodError. +@generated function vector_store!( + destination::Tile{S, <:AffineAxis}, acc::NTuple{NA, Vec{W, R}}, + alpha::T, beta::T, kernel::VectorKernel{MR, NR, T, W}, + m::Int, n::Int + ) where {S, MR, NR, T, W, R, NA} + layout = AccumulatorLayout(kernel) + rows = rows_per_vector(layout, W) + MV = MR ÷ rows + check_acc(:vector_store!, R, T, NA, accumulator_length(layout, MV, NR)) + RC = store_lanetype(layout, S, T) + + function blocks(B) + out = Any[] + for j in 1:NR + vblocks = Any[] + for v in 0:(MV - 1) + push!( + vblocks, quote + $(acc_bindings(layout, kernel, MV, NR, v, j)) + if $((v + 1) * rows) <= m + $(block_store(layout, :(colbase + $(v * rows)), W, R, RC, B)) + elseif $(v * rows) < m + for lane in 1:$rows + i = $(v * rows) + lane + i <= m || break + axpby_at!(storage, colbase + i, alpha, $(lane_value(layout)), beta) + end + end + end + ) + end + push!( + out, quote + if $j <= n + colbase = rowbase0 + cols[$j] # the address of (1, j) + $(vblocks...) + end + end + ) + end + return out + end + by_beta = quote + if iszero(beta) + $(blocks(:zero)...) + elseif isone(beta) + $(blocks(:one)...) + else + $(blocks(:general)...) + end + end + return quote + $(inline_store(layout) ? :(Base.@_inline_meta) : nothing) + storage = destination.storage + cols = destination.cols + rowbase0 = @inbounds destination.base + destination.rows[1] + $(store_body(layout, W, R, RC, by_beta)) + return destination + end +end + +function store_tile!( + destination::Tile, acc::NTuple{N, T}, + alpha::T, beta::T, kernel::ScalarKernel{MR} + ) where {MR, N, T} + m, n = store_prologue!(destination, alpha, beta) + (m == 0 || n == 0) && return destination + if iszero(beta) + @inbounds for j in 1:n, i in 1:m + destination[i, j] = alpha * acc[i + MR * (j - 1)] + end + elseif isone(beta) + @inbounds for j in 1:n, i in 1:m + destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], convert(T, destination[i, j])) + end + else + @inbounds for j in 1:n, i in 1:m + destination[i, j] = muladd(alpha, acc[i + MR * (j - 1)], beta * convert(T, destination[i, j])) + end + end + return destination +end diff --git a/src/microkernels/types.jl b/src/microkernels/types.jl deleted file mode 100644 index 099a580..0000000 --- a/src/microkernels/types.jl +++ /dev/null @@ -1,202 +0,0 @@ -# The vector kernels' types, their accumulator layouts and the generator-time -# pieces of the stores, ahead of every implementation: a `@generated` -# function's generator only sees methods defined before it. - -using SIMD: Vec - -""" - VectorKernel{MR, NR, T, W} - -An explicit-SIMD microkernel over `SIMD.Vec{W, real(T)}` lanes. Its -accumulator is an immutable tuple of vectors in one of three layouts -([`accumulator_layout`](@ref)), which fixes the stores; each kernel supplies -its K step, `accumulate_step`, and the generic `add_tile` loops it -`k_steps(kernel, k_block_length)` times. -""" -abstract type VectorKernel{MR, NR, T, W} <: Microkernel{MR, NR, T} end - -""" - SIMDKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Explicit-SIMD real microkernel over `SIMD.Vec{W,T}` lanes. `MR` must be a -multiple of `W` (default: one 256-bit register, 4 for `Float64`, 8 for `Float32`). -""" -struct SIMDKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} - descriptor::RealDescriptor{MR, NR, T} -end - -""" - PlanarKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Split-complex microkernel for `T = ComplexF32/ComplexF64` over -`SIMD.Vec{W,real(T)}` lanes. `MR`/`NR` are complex extents; `MR` must be a -multiple of `W`, which counts reals. -""" -struct PlanarKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} - descriptor::Descriptor{MR, NR, T, PlanarFormat, PlanarFormat} -end - -""" - OneMKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Van Zee's induced 1m complex microkernel: 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. Used only when named in `plan_contract(...; kernel = ...)`. -""" -struct OneMKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} - descriptor::Descriptor{MR, NR, T, OneEFormat, PlanarFormat} -end - -""" - FMAddSubKernel(::Val{MR}, ::Val{NR}, ::Type{T}[, ::Val{W}]) - -Interleaved-accumulator complex microkernel using x86 `vfmaddsub`, over -`SIMD.Vec{W,real(T)}` lanes with `InterleavedFormat` A and `PlanarFormat` B. -`2MR` must be a multiple of `W`, and `W` must be even. The default for small-M -complex contractions on AVX-512. -""" -struct FMAddSubKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} - descriptor::Descriptor{MR, NR, T, InterleavedFormat, PlanarFormat} -end - -""" - 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} <: VectorKernel{MR, NR, T, W} - descriptor::Descriptor{MR, NR, T, InterleavedFormat, RealFormat} -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} <: VectorKernel{MR, NR, T, W} - descriptor::Descriptor{MR, NR, T, RealFormat, InterleavedFormat} -end - -complex_method(::Type{<:SIMDKernel}) = RealMethod() -complex_method(::Type{<:PlanarKernel}) = PlanarMethod() -complex_method(::Type{<:OneMKernel}) = OneMMethod() -complex_method(::Type{<:FMAddSubKernel}) = FMAddSubMethod() -complex_method(::Type{<:ComplexRealKernel}) = ComplexRealMethod() -complex_method(::Type{<:RealComplexKernel}) = RealComplexMethod() -complex_method(kernel::VectorKernel) = complex_method(typeof(kernel)) - -""" - accumulator_layout(kernel::VectorKernel) - -How the accumulator holds the tile, as `(MR ÷ rows) * NR` vectors of `rows` -consecutive tile rows each, column-major (`rows_per_vector`): - - - `RealLayout()`: one real per lane (`SIMDKernel`). - - `SplitLayout()`: separate re and im vectors (`PlanarKernel`: a real and an - imaginary plane in one flat tuple; `RealComplexKernel`: the real kernel's - columns `2j-1` and `2j`). - - `LanePairLayout()`: lanes `2u-1` and `2u` are re and im of one row - (1m, fmaddsub, `ComplexRealKernel`). -""" -function accumulator_layout end - -abstract type AccumulatorLayout end -struct RealLayout <: AccumulatorLayout end -struct SplitLayout <: AccumulatorLayout end -struct LanePairLayout <: AccumulatorLayout end - -accumulator_layout(::Type{<:SIMDKernel}) = RealLayout() -accumulator_layout(::Type{<:Union{PlanarKernel, RealComplexKernel}}) = SplitLayout() -accumulator_layout(::Type{<:Union{OneMKernel, FMAddSubKernel, ComplexRealKernel}}) = LanePairLayout() -accumulator_layout(kernel::VectorKernel) = accumulator_layout(typeof(kernel)) - -rows_per_vector(::AccumulatorLayout, W::Int) = W -rows_per_vector(::LanePairLayout, W::Int) = W ÷ 2 -accumulator_length(::AccumulatorLayout, MV::Int, NR::Int) = MV * NR -accumulator_length(::SplitLayout, MV::Int, NR::Int) = 2 * MV * NR - -function (::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {K <: VectorKernel, MR, NR, T, W} - method = complex_method(K) - method === RealMethod() || T <: Complex || - throw(ArgumentError("$(nameof(K)) requires a complex element type, got $T")) - pairs = accumulator_layout(K) isa LanePairLayout - check_vector_shape(nameof(K), pairs ? 2 * MR : MR, W, pairs) - return K{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T, pack_formats(method)...)) -end -(::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}) where {K <: VectorKernel, MR, NR, T} = - K(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) - -lanewidth(::VectorKernel{MR, NR, T, W}) where {MR, NR, T, W} = W - -# Real and split stores are `@inline`: out of line, the whole accumulator is -# spilled to the stack and reloaded on every micro-tile. The lane-pair store is -# not: inlining it cost time (code growth). The scalar stores never are: they -# are scalar anyway, and inlining them bloats the tile function. -inline_store(::AccumulatorLayout) = true -inline_store(::LanePairLayout) = false - -# Generator-time pieces of the stores. `acc_bindings` names the accumulator -# vector(s) of row block `v` (zero-based) of column `j`, `lane_value` is the -# element at row `lane` of the block, `block_store` stores a full block whose -# first row is at zero-based storage index `first`, for `beta` case `B` -# (`:zero`, `:one` or `:general`). -acc_index(MV::Int, v::Int, j::Int) = v + MV * (j - 1) + 1 -split_index(::Type{<:PlanarKernel}, MV::Int, NR::Int, v::Int, j::Int) = - (acc_index(MV, v, j), MV * NR + acc_index(MV, v, j)) -split_index(::Type{<:RealComplexKernel}, MV::Int, NR::Int, v::Int, j::Int) = - (v + MV * (2j - 2) + 1, v + MV * (2j - 1) + 1) - -acc_bindings(::AccumulatorLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) = - :(vec = acc[$(acc_index(MV, v, j))]) -function acc_bindings(::SplitLayout, kernel::Type, MV::Int, NR::Int, v::Int, j::Int) - ire, iim = split_index(kernel, MV, NR, v, j) - return quote - revec = acc[$ire] - imvec = acc[$iim] - end -end - -lane_value(::RealLayout) = :(vec[lane]) -lane_value(::SplitLayout) = :(Complex(revec[lane], imvec[lane])) -lane_value(::LanePairLayout) = :(Complex(vec[2 * lane - 1], vec[2 * lane])) - -function block_store(::RealLayout, first, W::Int, R::Type, RC::Type, B::Symbol) - old = :(convert(Vec{$W, $R}, vload(Vec{$W, $RC}, storage, at))) - new = B === :zero ? :(alpha * vec) : B === :one ? :(muladd(alpha, vec, $old)) : :(muladd(alpha, vec, beta * $old)) - return quote - at = $first + 1 - vstore(convert(Vec{$W, $RC}, $new), storage, at) - end -end -block_store(::SplitLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = - :(split_store_block!(sp, 2 * $first, revec, imvec, ar, ai, br, bi, Val($W), Val($(QuoteNode(B))))) -block_store(::LanePairLayout, first, W::Int, R::Type, RC::Type, B::Symbol) = - :(lanepair_store_block!(sp, 2 * $first, vec, ar, ai, br, bi, Val($(QuoteNode(B))))) - -# The real lane type of the vector store's storage `S`. The complex layouts -# reinterpret the storage as reals, only sound on dense rank-1 complex storage. -store_lanetype(::RealLayout, S::Type, T::Type) = eltype(S) -function store_lanetype(::AccumulatorLayout, S::Type, T::Type) - S <: DenseVector && lane_convertible(eltype(S), T) || - throw(ArgumentError("vector_store!: storage $S is not a dense vector convertible to $T")) - return real(eltype(S)) -end - -# The vector store's body around its blocks: the complex layouts broadcast -# `alpha`/`beta` once and store through a raw pointer, only dereferenced -# inside `GC.@preserve`. -store_body(::RealLayout, W::Int, R::Type, RC::Type, body) = :(@inbounds $body) -store_body(::AccumulatorLayout, W::Int, R::Type, RC::Type, body) = quote - ar = Vec{$W, $R}(real(alpha)) - ai = Vec{$W, $R}(imag(alpha)) - br = Vec{$W, $R}(real(beta)) - bi = Vec{$W, $R}(imag(beta)) - GC.@preserve storage begin - sp = reinterpret(Ptr{$RC}, pointer(storage)) - @inbounds $body - end -end diff --git a/src/microkernels/vecops.jl b/src/microkernels/vecops.jl new file mode 100644 index 0000000..f60175a --- /dev/null +++ b/src/microkernels/vecops.jl @@ -0,0 +1,71 @@ +# SIMD building blocks shared by the K steps and the stores. Shuffle patterns are +# built from the vector width at specialization time, never hardcoded to one ISA. + +using SIMD: Vec, shufflevector + +# `x*y - c` in even lanes, `x*y + c` in odd lanes, each one fused rounding. +# Generic IR (`fneg` + two `llvm.fma` + a blend), which the X86 backend folds +# into one `vfmaddsub` wherever FMA3 exists (and elsewhere stays correct). +# GUARDRAIL: the SIMD.jl spelling `shufflevector(muladd(x, y, -c), muladd(x, y, +# c), ...)` also selects `vfmaddsub` but leaves a dead stack store of the +# accumulator in the K loop, every step. +@generated function fmaddsub(x::Vec{N, R}, y::Vec{N, R}, c::Vec{N, R}) where {N, R} + R === Float64 || R === Float32 || + return :(throw(ArgumentError("fmaddsub: unsupported lane type $R"))) + ty = "<$N x $(R === Float64 ? "double" : "float")>" + fn = "llvm.fma.v$(N)$(R === Float64 ? "f64" : "f32")" + # lane k: even -> `s` (x*y - c); odd -> `d` (x*y + c), index N + k. + mask = join(("i32 $(iseven(k) ? k : N + k)" for k in 0:(N - 1)), ", ") + ir = """ + declare $ty @$fn($ty, $ty, $ty) + define $ty @entry($ty %x, $ty %y, $ty %c) #0 { + top: + %nc = fneg $ty %c + %s = call $ty @$fn($ty %x, $ty %y, $ty %nc) + %d = call $ty @$fn($ty %x, $ty %y, $ty %c) + %r = shufflevector $ty %s, $ty %d, <$N x i32> <$mask> + ret $ty %r + } + attributes #0 = { alwaysinline } + """ + VT = NTuple{N, VecElement{R}} + return quote + Base.@_inline_meta + Vec{$N, $R}( + Base.llvmcall(($ir, "entry"), $VT, Tuple{$VT, $VT, $VT}, x.data, y.data, c.data) + ) + end +end + +# `[x1, x0, x3, x2, ...]`: one in-lane `vshufpd`/`vpermilps`. +@generated function swap_pairs(x::Vec{N, R}) where {N, R} + iseven(N) || return :(throw(ArgumentError("swap_pairs: expected even N, got $N"))) + idx = ntuple(k -> isodd(k) ? k : k - 2, N) + return :(Base.@_inline_meta; shufflevector(x, Val($idx))) +end + +# `p - q` in even lanes, `p + q` in odd lanes: Base's unfused complex `*` on +# interleaved data. +@generated function addsub(p::Vec{N, R}, q::Vec{N, R}) where {N, R} + iseven(N) || return :(throw(ArgumentError("addsub: expected even N, got $N"))) + idx = ntuple(k -> iseven(k - 1) ? k - 1 : N + k - 1, N) + return :(Base.@_inline_meta; shufflevector(p - q, p + q, Val($idx))) +end + +# `v` holds `W` complex values `[re_0, im_0, ..., re_{W-1}, im_{W-1}]`. +@generated function deinterleave_re(v::Vec{N, R}, ::Val{W}) where {N, R, W} + N == 2 * W || return :(throw(ArgumentError("deinterleave_re: expected N == 2W"))) + idx = ntuple(k -> 2 * (k - 1), W) + return :(shufflevector(v, Val($idx))) +end + +@generated function deinterleave_im(v::Vec{N, R}, ::Val{W}) where {N, R, W} + N == 2 * W || return :(throw(ArgumentError("deinterleave_im: expected N == 2W"))) + idx = ntuple(k -> 2 * (k - 1) + 1, W) + return :(shufflevector(v, Val($idx))) +end + +@generated function interleave_planes(re::Vec{W, R}, im::Vec{W, R}, ::Val{W}) where {W, R} + idx = ntuple(k -> isodd(k) ? (k - 1) ÷ 2 : W + (k - 1) ÷ 2, 2 * W) + return :(shufflevector(re, im, Val($idx))) +end diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index c869f5e..2adb13f 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -104,5 +104,5 @@ size and clamps all three to the contraction's extents. """ function default_blocking(kernel) T = scalartype(kernel) - return _scale_blocking(_resolved_defaults(T).real_row, complex_method(kernel)) + return _scale_blocking(_resolved_defaults(T).real_row, KernelMethod(kernel)) end diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 3faa361..a9b1a9b 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -1,5 +1,5 @@ using QuasiStrided: FMAddSubKernel, FMAddSubMethod, PlanarKernel, PlanarMethod, OneMKernel, - InterleavedFormat, complex_method, lanewidth, kernel_shapes, packed_panel, + InterleavedFormat, KernelMethod, lanewidth, kernel_shapes, packed_panel, _kernel_from_shape, _default_method, accumulator_planes, target_profile, PackedPanel using SIMD: Vec using InteractiveUtils: code_native @@ -16,7 +16,7 @@ const _QSF = QuasiStrided end @testset "construction" begin - @test complex_method(FMAddSubKernel(Val(8), Val(4), ComplexF64)) === FMAddSubMethod() + @test KernelMethod(FMAddSubKernel(Val(8), Val(4), ComplexF64)) === FMAddSubMethod() @test accumulator_planes(FMAddSubMethod()) == 1 @test lanewidth(FMAddSubKernel(Val(8), Val(4), ComplexF32)) == 8 @test_throws ArgumentError FMAddSubKernel(Val(12), Val(8), ComplexF64, Val(16)) diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index c22455a..9927162 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -1,5 +1,5 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, - complex_method, lanewidth, kernel_shapes, + KernelMethod, lanewidth, kernel_shapes, _kernel_from_shape, _default_method @testset "OneMKernel" begin @@ -15,7 +15,7 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, k = OneMKernel(Val(12), Val(8), ComplexF64, Val(8)) @test QuasiStrided.inner(k) isa SIMDKernel{24, 8, Float64, 8} # the real kernel, at 2MR rows @test zero_accumulator(k) === zero_accumulator(QuasiStrided.inner(k)) - @test complex_method(k) === OneMMethod() + @test KernelMethod(k) === OneMMethod() @test lanewidth(OneMKernel(Val(8), Val(4), ComplexF32)) == 8 # 2MR, not MR, must divide by W; W must be even even where 2MR divides. @test_throws ArgumentError OneMKernel(Val(12), Val(8), ComplexF64, Val(16)) diff --git a/test/microkernels/test_planar_kernel.jl b/test/microkernels/test_planar_kernel.jl index 5035fe5..4fc147e 100644 --- a/test/microkernels/test_planar_kernel.jl +++ b/test/microkernels/test_planar_kernel.jl @@ -1,4 +1,4 @@ -using QuasiStrided: PlanarKernel, PlanarMethod, complex_method, lanewidth, +using QuasiStrided: PlanarKernel, PlanarMethod, KernelMethod, lanewidth, KERNEL_SHAPES_C64_PLANAR, KERNEL_SHAPES_C32_PLANAR @testset "PlanarKernel" begin @@ -11,7 +11,7 @@ using QuasiStrided: PlanarKernel, PlanarMethod, complex_method, lanewidth, end k = PlanarKernel(Val(16), Val(6), ComplexF64, Val(8)) - @test complex_method(k) === PlanarMethod() + @test KernelMethod(k) === PlanarMethod() @test lanewidth(PlanarKernel(Val(8), Val(4), ComplexF64)) == 4 # from the REAL type @test lanewidth(PlanarKernel(Val(8), Val(4), ComplexF32)) == 8 @test_throws ArgumentError PlanarKernel(Val(6), Val(4), ComplexF64, Val(4)) diff --git a/test/microkernels/test_planar_store_fastpath.jl b/test/microkernels/test_planar_store_fastpath.jl index 286171f..0d0dd61 100644 --- a/test/microkernels/test_planar_store_fastpath.jl +++ b/test/microkernels/test_planar_store_fastpath.jl @@ -65,7 +65,7 @@ function mk_store_fastpath(K, T, shapes; beta0_exact = false, S = T) for (alpha, beta) in mk_store_ab(T) fast = mk_dense(cold) dfast = Tile(fast, 0, AffineAxis(0, 1, m), AffineAxis(0, m, n)) - @test QuasiStrided.vector_store_eligible(QuasiStrided.accumulator_layout(k), dfast, T) == STORE_FASTPATH_ON + @test QuasiStrided.vector_store_eligible(QuasiStrided.AccumulatorLayout(k), dfast, T) == STORE_FASTPATH_ON store_tile!(dfast, acc, alpha, beta, k) scal = copy(cold) # scattered rows: always the scalar store store_tile!(Tile(scal, 0, view(collect(0:(m - 1)), 1:m), AffineAxis(0, m, n)), acc, alpha, beta, k) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index 6c46416..cdde59a 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -57,7 +57,7 @@ using StridedViews: StridedView @test (MR ÷ W) * NR + (MR ÷ W) <= 32 @test sliver_width(k) === (MR, NR) @test realtype(k) === T === scalartype(k) - @test complex_method(k) === RealMethod() + @test KernelMethod(k) === RealMethod() @test packed_a_length(k, 7) === MR * 7 && packed_b_length(k, 7) === NR * 7 end for T in (Float64, Float32) @@ -236,7 +236,7 @@ end @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 && (tile_size(k)..., lanewidth(k)) === shape + @test KernelMethod(k) === m && (tile_size(k)..., lanewidth(k)) === shape end end @test dmethod(T, T, R) === dmethod(T, ComplexF32, Float64) === CR @@ -321,7 +321,7 @@ end QS.RealComplexKernel(Val(4), Val(2), ComplexF64), ) for k in kernels - @test pack_formats(complex_method(k)) === (QS.a_format(k), QS.b_format(k)) + @test pack_formats(KernelMethod(k)) === (QS.a_format(k), QS.b_format(k)) end end @@ -350,7 +350,7 @@ end kernel = _default_kernel(T) @test kernel isa QuasiStrided.PlanarKernel @test scalartype(kernel) === T && realtype(kernel) === real(T) - @test complex_method(kernel) === PlanarMethod() === QuasiStrided._default_method(T) + @test KernelMethod(kernel) === PlanarMethod() === QuasiStrided._default_method(T) shape = (tile_size(kernel)..., lanewidth(kernel)) profile = target_profile() @test _planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) @@ -362,7 +362,7 @@ end shape1m = first(kernel_shapes(T, OneMMethod())) k1m = QuasiStrided._kernel_from_shape(shape1m, T, OneMMethod()) - @test k1m isa QuasiStrided.OneMKernel && complex_method(k1m) === OneMMethod() + @test k1m isa QuasiStrided.OneMKernel && KernelMethod(k1m) === OneMMethod() @test (tile_size(k1m)..., lanewidth(k1m)) === shape1m # A method with no kernel for `T` throws, naming itself. err = try diff --git a/test/runtests.jl b/test/runtests.jl index e8c9411..fe16056 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,7 +12,7 @@ using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, _rule_applies, isa_nregisters, sliver_width, realtype, - complex_method, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, pack_formats, + KernelMethod, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, pack_formats, reals_per_element, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, _planar_pressure, _pack_split, _NestPath # plan_contract, execute! and ContractPlan are bound in helpers.jl instead. From 3be4a87bf70a35bc07b3b82bfff48c437f843794 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 10:11:06 -0400 Subject: [PATCH 18/60] Decision log: chunks 7 and 8 done Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index cbd33ed..1fed2fb 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -18,9 +18,9 @@ TTFX when codegen is touched). | 3 | Tiles | `layout/tiles.jl` | done | | 4 | Packing formats | `packing/format.jl`, `panel.jl` (`transposed.jl` → chunk 11) | done | | 5 | Packers | `packing/pack.jl`, `pack_contiguous.jl` | done | -| 6 | Kernel interface | `microkernels/interface.jl`, `scalar.jl` | done | -| 7 | SIMD kernels | `simd.jl`, `planar.jl` | in progress (with 8) | -| 8 | Complex/mixed kernels | `onem.jl`, `fmaddsub.jl`, `mixed.jl` | in progress (with 7) | +| 6 | Kernel interface | `microkernels/interface.jl` (+ ScalarKernel) | done | +| 7 | Vector kernels | now `kernels.jl`, `vecops.jl`, `steps.jl`, `stores.jl` | done | +| 8 | Complex/mixed kernels | (merged into chunk 7's files) | done | | 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | | 10 | Kernel selection | `planning/kernel_selection.jl` | | | 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | @@ -157,7 +157,7 @@ it: each executed variant is smaller and branch-free (less code fetched per store), at the cost of compile time and code size; measure both. `b_step_load*` hooks get descriptive names. -### D12. Kernels = K step + accumulator layout (chunks 7–8) +### D12. Kernels = K step + accumulator layout (chunks 7–8, applied) Three K steps (real, planar, fmaddsub) and three accumulator layouts (real, split re/im, lane pairs). `accumulator_layout(kernel)` selects one generic @@ -166,7 +166,13 @@ layout supplies its full-block store and lane value); one generic `add_tile` over `accumulate_step` with `k_steps` (1m: 2k). `inner(k)` is computed, not a field (no `KI` parameter). One generic constructor without `W`; generic `complex_method`/`lanewidth`. 1m gets the lane-pair vector store. The `beta` -hoist is a separate commit on top, measured on its own. +hoist (cc78fd1) is kept: small K 19–27% faster (Float64) and 3–10% +(ComplexF64), large cases unchanged, vector stores ~35–40% larger, first call +Float64 3.1 → 3.7 s and ComplexF64 2.0 → 2.4 s, full suite ~18 → ~28 min. +Files, by concern: `interface.jl` (contract), `kernels.jl` (all kernel types +and their traits; `AccumulatorLayout`/`KernelMethod` are Holy traits), +`vecops.jl` (shared SIMD ops), `steps.jl` (K steps, `add_tile`), `stores.jl` +(layouts' stores). ## Possible improvements From 232b69d9c9ec077f53504afe724db4c96c22f805 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 11:00:59 -0400 Subject: [PATCH 19/60] Chunk 9: labels - `_sort_perm` (hand-written insertion sort) replaced by `TupleTools.sortperm`/`getindices`. TupleTools sorts by a recursive merge that takes from the right half only on a strict `lt`, so ties keep their original order, as before; allocation-free and inferred for NTuples. `import TupleTools` moves to the module file. - `_order_contract_labels` and `_choose_k_order` merged into `order_contract_labels` with an optional trailing `l2bytes` (default `nothing`: the core's L2 share, looked up only when a cost is needed). - `_l2_core_bytes` (both methods) moved to defaults.jl as `l2_core_bytes`. - conjugation.jl deleted: `_op_conjugates`/`_qs_isconj` now live in plan.jl as `op_conjugates`/`isconj(view, flag)`, both GUARDRAIL comments kept. - Names defined in labels.jl lose their leading underscore (`classify_labels`, `order_free_labels`, `leading_unit_run`, `prefer_swap`, `KOperand`, `k_operand`, `k_order_cost`, ...). The `_K_*` constants keep theirs, matching the package's other internal constants. A test helper `_macro_op_conjugates` became `_macro_conj_op`. Behaviour-neutral: K-order model, free-label order and swap rule unchanged. Per-call floor (probe_call_floor.jl, 4 interleaved runs each, medians, base 3be4a87 -> this commit, ns): F64 8x8x8 360 -> 369, 16^3 470 -> 472, d4 outer 352 -> 365; C64 8x8x8 417 -> 417, 16^3 933 -> 931, d4 outer 576 -> 577. Sample ranges overlap (F64 d4 outer 347-389 vs 347-373); 0 B/call throughout. Planning micro-benchmark with 3 K labels: K order 39 -> 18 ns, 0 B. Tests: 31967/31967 pass, 28m52s. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 4 +- src/integrations/tensoroperations.jl | 3 +- src/planning/conjugation.jl | 22 ----- src/planning/defaults.jl | 12 ++- src/planning/labels.jl | 113 ++++++++---------------- src/planning/pack_split.jl | 2 +- src/planning/plan.jl | 43 ++++++--- test/execution/test_macro_blocking.jl | 6 +- test/planning/test_per_call_overhead.jl | 14 +-- test/planning/test_plan_contract.jl | 60 ++++++------- 10 files changed, 127 insertions(+), 152 deletions(-) delete mode 100644 src/planning/conjugation.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index bf5e869..f365607 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -1,6 +1,7 @@ module QuasiStrided using StridedViews: StridedView, offset +import TupleTools using Base.Checked: checked_abs, checked_add, checked_mul # TensorOperations names are always qualified: a bare `using` collides on `scalartype`. @@ -31,9 +32,8 @@ include("microkernels/vecops.jl") include("microkernels/steps.jl") include("microkernels/stores.jl") -# --- Planning: labels, conjugation, kernel and blocking choice, the plan --- +# --- Planning: labels, kernel and blocking choice, the plan --- include("planning/labels.jl") -include("planning/conjugation.jl") include("planning/kernel_selection.jl") include("planning/blocking.jl") include("planning/defaults.jl") diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index a10e588..40af38b 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -1,7 +1,6 @@ # TensorOperations backend. The engine itself knows nothing about TO. import TensorOperations as TO using TensorOperations: Index2Tuple, linearize -import TupleTools using StridedViews: StridedView, isstrided """ @@ -100,7 +99,7 @@ end (Base.mightalias(Cv, Av) || Base.mightalias(Cv, Bv)) && _qs_throw( "output tensor must not be aliased with an input tensor in $(TO.tensorcontract!)" ) - _qs_isconj(Cv, false) && _qs_throw( + isconj(Cv, false) && _qs_throw( "output tensor of $(TO.tensorcontract!) must not be a conjugated view: " * "QuasiStrided writes through to the parent array and does not apply " * "`StridedView.op` on store, so a conjugated `C` would be silently wrong" diff --git a/src/planning/conjugation.jl b/src/planning/conjugation.jl deleted file mode 100644 index ef0bd46..0000000 --- a/src/planning/conjugation.jl +++ /dev/null @@ -1,22 +0,0 @@ -# Conjugation: folding `conjA`/`conjB` with each view's `.op`. The engine never -# indexes through a `StridedView`, so an unfolded `.op` would be silently dropped. - -# GUARDRAIL: a TOTAL table with a throwing fallback, not an `op === conj` test, -# which classifies a directly constructed `adjoint` view as unconjugated. -_op_conjugates(::typeof(identity)) = false -_op_conjugates(::typeof(conj)) = true -# Elementwise on a `Number`, these are identity/conj; the axes are already resolved. -_op_conjugates(::typeof(transpose)) = false -_op_conjugates(::typeof(adjoint)) = true -@noinline _op_conjugates(f) = throw( - ArgumentError( - "unsupported StridedView.op $f: QuasiStrided folds a view's `op` into the " * - "packing transform and recognizes only identity/conj/transpose/adjoint" - ) -) - -# GUARDRAIL: `⊻`, not `||`: the flag and the view's `op` are independent -# conjugations and `conj` is involutive. Always `false` for real `T`, so the -# real path never gets a `conj` specialization. -_qs_isconj(v::StridedView{T}, flag::Bool) where {T} = - (T <: Complex) && (flag ⊻ _op_conjugates(v.op)) diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index 2adb13f..3e5f433 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -15,6 +15,16 @@ struct ResolvedDefaults l2_core::Int end +# The core's private L2 share, or 1 MB when undetected. +function l2_core_bytes(profile::TargetProfile) + profile.l2.bytes > 0 || return 1 << 20 + return core_bytes(profile, profile.l2) +end + +# Through the per-eltype cache; an eltype without a slot fails later in planning. +@inline l2_core_bytes(::Type{T}) where {T} = + _defaults_slot(T) === nothing ? l2_core_bytes(target_profile()) : _resolved_defaults(T).l2_core + function _resolve_defaults(profile::TargetProfile, ::Type{T}) where {T} method = _default_method(T) # First, so an element type with no menu throws from here. @@ -22,7 +32,7 @@ function _resolve_defaults(profile::TargetProfile, ::Type{T}) where {T} fitted = _fitted_shape(profile, T, method) small_m = _small_m_candidates(Val(profile.isa), profile, T) real_row = _real_blocking_row(profile, real(T)) - l2_core = _l2_core_bytes(profile) + l2_core = l2_core_bytes(profile) return ResolvedDefaults(profile, shape, fitted, small_m, real_row, l2_core) end diff --git a/src/planning/labels.jl b/src/planning/labels.jl index ffd73e5..e272b48 100644 --- a/src/planning/labels.jl +++ b/src/planning/labels.jl @@ -3,15 +3,15 @@ # `NTuple` whose length follows from the label tuples' lengths alone, so # planning allocates nothing and every `AxisGroup` is concretely typed. -@inline _label_in(lbl::Int, t::NTuple{N, Int}) where {N} = any(==(lbl), t) +@inline label_in(lbl::Int, t::NTuple{N, Int}) where {N} = any(==(lbl), t) # The M/N/K ranks: with A = M ∪ K, B = N ∪ K, C = M ∪ N, solve # NA = |M| + |K|, NB = |N| + |K|, NC = |M| + |N|. Meaningful only for input that -# `_classify_labels` accepts, which asserts them against the counts it finds. -@inline _group_ranks(NA::Int, NB::Int, NC::Int) = +# `classify_labels` accepts, which asserts them against the counts it finds. +@inline group_ranks(NA::Int, NB::Int, NC::Int) = ((NA + NC - NB) ÷ 2, (NB + NC - NA) ÷ 2, (NA + NB - NC) ÷ 2) -@noinline _throw_rank_mismatch(which::Symbol, found::Int, derived::Int) = throw( +@noinline throw_rank_mismatch(which::Symbol, found::Int, derived::Int) = throw( ArgumentError( "internal error: $which label count $found does not match the rank $derived " * "derived from the label tuple lengths" @@ -19,13 +19,13 @@ ) # Set element `n` of `out`, ignoring an `n` past the end (input about to be rejected). -@inline function _push_label(out::NTuple{D, Int}, n::Int, x::Int) where {D} +@inline function push_label(out::NTuple{D, Int}, n::Int, x::Int) where {D} return n <= D ? Base.setindex(out, x, n) : out end # (mlabels, nlabels, klabels) in indA/indB order. Per (inA, inB, inC): # (T,F,T) -> M, (F,T,T) -> N, (T,T,F) -> K, anything else is an ArgumentError. -function _classify_labels( +function classify_labels( indA::NTuple{NA, Int}, indB::NTuple{NB, Int}, indC::NTuple{NC, Int} ) where {NA, NB, NC} @@ -36,15 +36,15 @@ function _classify_labels( allunique(indC) || throw(ArgumentError("indC has a repeated label (diagonal), not supported: $indC")) - DM, DN, DK = _group_ranks(NA, NB, NC) + DM, DN, DK = group_ranks(NA, NB, NC) mlabels = ntuple(_ -> 0, Val(max(DM, 0))) nlabels = ntuple(_ -> 0, Val(max(DN, 0))) klabels = ntuple(_ -> 0, Val(max(DK, 0))) nm = 0 nk = 0 for lbl in indA - inB = _label_in(lbl, indB) - inC = _label_in(lbl, indC) + inB = label_in(lbl, indB) + inC = label_in(lbl, indC) if inB && inC throw( ArgumentError( @@ -54,10 +54,10 @@ function _classify_labels( ) elseif inB && !inC nk += 1 - klabels = _push_label(klabels, nk, lbl) + klabels = push_label(klabels, nk, lbl) elseif !inB && inC nm += 1 - mlabels = _push_label(mlabels, nm, lbl) + mlabels = push_label(mlabels, nm, lbl) else throw( ArgumentError( @@ -70,13 +70,13 @@ function _classify_labels( nn = 0 for lbl in indB - inA = _label_in(lbl, indA) - inC = _label_in(lbl, indC) + inA = label_in(lbl, indA) + inC = label_in(lbl, indC) if inA continue # K, or already rejected while scanning indA elseif inC nn += 1 - nlabels = _push_label(nlabels, nn, lbl) + nlabels = push_label(nlabels, nn, lbl) else throw( ArgumentError( @@ -88,45 +88,30 @@ function _classify_labels( end for lbl in indC - inA = _label_in(lbl, indA) - inB = _label_in(lbl, indB) + inA = label_in(lbl, indA) + inB = label_in(lbl, indB) (inA || inB) || throw(ArgumentError("label $lbl appears in indC but not in indA or indB")) end # Unreachable after the checks above; guards against a silently truncated list. - nm == DM || _throw_rank_mismatch(:M, nm, DM) - nn == DN || _throw_rank_mismatch(:N, nn, DN) - nk == DK || _throw_rank_mismatch(:K, nk, DK) + nm == DM || throw_rank_mismatch(:M, nm, DM) + nn == DN || throw_rank_mismatch(:N, nn, DN) + nk == DK || throw_rank_mismatch(:K, nk, DK) return mlabels, nlabels, klabels end -# Stable insertion sort of the permutation `perm` by `key[perm[j]]`, ascending -# (n <= ndims, and unlike `sortperm` it allocates nothing). -@inline function _sort_perm(perm::NTuple{D, Int}, key::NTuple{D, Int}) where {D} - out = perm - @inbounds for i in 2:D - x = out[i] - kx = key[x] - j = i - 1 - while j >= 1 && key[out[j]] > kx - out = Base.setindex(out, out[j], j + 1) - j -= 1 - end - out = Base.setindex(out, x, j + 1) - end - return out -end - # `labels` stably sorted by `abs(stride)` of their axis in `C` (every label must # be in `indC`): a composite enumerates its first label fastest, so this walks # C's fastest axis fastest. -function _order_free_labels( +# `TupleTools.sortperm` is a merge sort that keeps ties in order (relied on here +# and in `order_contract_labels`). +function order_free_labels( labels::NTuple{D, Int}, indC::NTuple{NC, Int}, C::StridedView ) where {D, NC} st = Base.strides(C) key = map(l -> abs(st[findfirst(==(l), indC)::Int]), labels) - return map(i -> @inbounds(labels[i]), _sort_perm(ntuple(identity, Val(D)), key)) + return TupleTools.getindices(labels, TupleTools.sortperm(key)) end # ---------------------------------------------------------------------------- @@ -152,22 +137,12 @@ const _K_WALK_FAR_BYTES = 4096 const _K_WALK_FAR_PENALTY = 3 const _K_LINE_BYTES = 64 -# The core's private L2 share, or 1 MB when undetected. -function _l2_core_bytes(profile::TargetProfile) - profile.l2.bytes > 0 || return 1 << 20 - return core_bytes(profile, profile.l2) -end - -# Through the per-eltype cache; an eltype without a slot fails later in planning. -@inline _l2_core_bytes(::Type{T}) where {T} = - _defaults_slot(T) === nothing ? _l2_core_bytes(target_profile()) : _resolved_defaults(T).l2_core - # What the cost model reads of one operand, gathered once per plan in `klabels` # order: each K label's extent and |stride|, the element count and size, and # whether its register slivers are whole lines (the free composite's first # non-singleton axis is unit-stride and at least a line long). A candidate # order is then a permutation of `1:DK`. -struct _KOperand{DK} +struct KOperand{DK} len::NTuple{DK, Int} st::NTuple{DK, Int} n::Int @@ -175,7 +150,7 @@ struct _KOperand{DK} wholeline::Bool end -@inline function _k_operand( +@inline function k_operand( klabels::NTuple{DK, Int}, ind::NTuple{N, Int}, v::StridedView, free::NTuple{DF, Int} ) where {DK, N, DF} @@ -194,12 +169,12 @@ end wholeline = st[p] == 1 && sz[p] >= line_elems break end - return _KOperand{DK}(len, kst, length(v), S, wholeline) + return KOperand{DK}(len, kst, length(v), S, wholeline) end # Cost of packing one operand under the K order `perm`; `Qfree` is its free extent. -function _k_order_cost( - perm::NTuple{DK, Int}, op::_KOperand{DK}, Qfree::Int, l2bytes::Int +function k_order_cost( + perm::NTuple{DK, Int}, op::KOperand{DK}, Qfree::Int, l2bytes::Int ) where {DK} len, st, S = op.len, op.st, op.S kfast = 0 @@ -240,32 +215,22 @@ function _k_order_cost( end # Zero or one K label returns at a compile-time branch, before the model. -@inline function _order_contract_labels( - klabels::NTuple{DK, Int}, - indA::NTuple{NA, Int}, A::StridedView, morder::NTuple{DM, Int}, - indB::NTuple{NB, Int}, B::StridedView, norder::NTuple{DN, Int}, - m_length::Int, n_length::Int - ) where {DK, NA, NB, DM, DN} - DK <= 1 && return klabels - return _choose_k_order(klabels, indA, A, morder, indB, B, norder, m_length, n_length, nothing) -end - # `l2bytes = nothing`: the core's L2 share, looked up only once a cost is needed. -function _choose_k_order( +@inline function order_contract_labels( klabels::NTuple{DK, Int}, indA::NTuple{NA, Int}, A::StridedView, morder::NTuple{DM, Int}, indB::NTuple{NB, Int}, B::StridedView, norder::NTuple{DN, Int}, - m_length::Int, n_length::Int, l2bytes::Union{Int, Nothing} + m_length::Int, n_length::Int, l2bytes::Union{Int, Nothing} = nothing ) where {DK, NA, NB, DM, DN} DK <= 1 && return klabels - opA = _k_operand(klabels, indA, A, morder) - opB = _k_operand(klabels, indB, B, norder) + opA = k_operand(klabels, indA, A, morder) + opB = k_operand(klabels, indB, B, norder) id = ntuple(identity, Val(DK)) - permA = _sort_perm(id, opA.st) - permB = _sort_perm(id, opB.st) + permA = TupleTools.sortperm(opA.st) + permB = TupleTools.sortperm(opB.st) permA == id && permB == id && return klabels - l2 = l2bytes === nothing ? _l2_core_bytes(eltype(A)) : l2bytes - cost(perm) = _k_order_cost(perm, opA, m_length, l2) + _k_order_cost(perm, opB, n_length, l2) + l2 = l2bytes === nothing ? l2_core_bytes(eltype(A)) : l2bytes + cost(perm) = k_order_cost(perm, opA, m_length, l2) + k_order_cost(perm, opB, n_length, l2) best = id bestcost = cost(id) for cand in (permA, permB) @@ -276,14 +241,14 @@ function _choose_k_order( bestcost = c end end - return map(i -> @inbounds(klabels[i]), best) + return TupleTools.getindices(klabels, best) end # Length of the leading unit-stride run in C when `labels` is enumerated # first-label-fastest: the first non-singleton label needs C-stride exactly +1, # and each following one extends the run only if its stride equals the run so # far. 1 when no run starts, 0 if an empty axis is met first. -function _leading_unit_run( +function leading_unit_run( labels::NTuple{D, Int}, indC::NTuple{NC, Int}, C::StridedView ) where {D, NC} st = Base.strides(C) @@ -303,5 +268,5 @@ end # sliver of the kernel it would run and the swapped one can: the vectorized # store needs `m_tile` consecutive M coordinates unit-stride in C, so a shorter # run buys nothing. -_prefer_swap(run_m::Int, run_n::Int, m_tile_asis::Int, m_tile_swapped::Int = m_tile_asis) = +prefer_swap(run_m::Int, run_n::Int, m_tile_asis::Int, m_tile_swapped::Int = m_tile_asis) = run_m < m_tile_asis && run_n >= m_tile_swapped diff --git a/src/planning/pack_split.jl b/src/planning/pack_split.jl index f48fef2..2ef0a36 100644 --- a/src/planning/pack_split.jl +++ b/src/planning/pack_split.jl @@ -90,7 +90,7 @@ end # 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) + l2 = l2_core_bytes(profile) l3 = profile.l3 (aliased || l3.bytes <= 0) && return l2 return l2 + core_bytes(profile, l3) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 80fd0b4..30cec9f 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -150,6 +150,29 @@ end ) end +# Conjugation: folding `conjA`/`conjB` with each view's `.op`. The engine never +# indexes through a `StridedView`, so an unfolded `.op` would be silently dropped. + +# GUARDRAIL: a TOTAL table with a throwing fallback, not an `op === conj` test, +# which classifies a directly constructed `adjoint` view as unconjugated. +op_conjugates(::typeof(identity)) = false +op_conjugates(::typeof(conj)) = true +# Elementwise on a `Number`, these are identity/conj; the axes are already resolved. +op_conjugates(::typeof(transpose)) = false +op_conjugates(::typeof(adjoint)) = true +@noinline op_conjugates(f) = throw( + ArgumentError( + "unsupported StridedView.op $f: QuasiStrided folds a view's `op` into the " * + "packing transform and recognizes only identity/conj/transpose/adjoint" + ) +) + +# GUARDRAIL: `⊻`, not `||`: the flag and the view's `op` are independent +# conjugations and `conj` is involutive. Always `false` for real `T`, so the +# real path never gets a `conj` specialization. +isconj(v::StridedView{T}, flag::Bool) where {T} = + (T <: Complex) && (flag ⊻ op_conjugates(v.op)) + # `plan_contract`'s body, positional, with a continuation `f` applied to the # plan inside the barrier, where its type is concrete (the TensorOperations # adapter passes an executor). @@ -167,23 +190,23 @@ function _planned( # GUARDRAIL: a conjugated `C` is rejected; the engine writes through to # its parent, so there is nowhere to absorb its `op`. - _qs_isconj(C, false) && throw( + isconj(C, false) && throw( ArgumentError( "plan_contract: cannot write into a conjugated view (C has op $(C.op)); " * "writing a conjugated output is not supported" ) ) - atransform = _qs_isconj(A, conjA) ? conj : identity - btransform = _qs_isconj(B, conjB) ? conj : identity + atransform = isconj(A, conjA) ? conj : identity + btransform = isconj(B, conjB) ? conj : identity - mlabels, nlabels, klabels = _classify_labels(indA, indB, indC) + mlabels, nlabels, klabels = classify_labels(indA, indB, indC) - morder = _order_free_labels(mlabels, indC, C) - norder = _order_free_labels(nlabels, indC, C) + morder = order_free_labels(mlabels, indC, C) + norder = order_free_labels(nlabels, indC, C) # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. - run_m = T <: Real || method isa _MixedMethod ? _leading_unit_run(morder, indC, C) : 0 - run_n = T <: Real ? _leading_unit_run(norder, indC, C) : 0 + run_m = T <: Real || method isa _MixedMethod ? leading_unit_run(morder, indC, C) : 0 + run_n = T <: Real ? leading_unit_run(norder, indC, C) : 0 mgroup = AxisGroup(morder, (indA, A), (indC, C)) # maps: (A, C) ngroup = AxisGroup(norder, (indB, B), (indC, C)) # maps: (B, C) @@ -191,14 +214,14 @@ function _planned( m_length = axis_length(mgroup) n_length = axis_length(ngroup) - korder = _order_contract_labels(klabels, indA, A, morder, indB, B, norder, m_length, n_length) + korder = order_contract_labels(klabels, indA, A, morder, indB, B, norder, m_length, n_length) kgroup = AxisGroup(korder, (indA, A), (indB, B)) # maps: (A, B) k_length = axis_length(kgroup) m_tile_asis, m_tile_swapped = _candidate_m_tiles(T, method, kernel, m_length, n_length, run_m, run_n) - if T <: Real && _prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) + if T <: Real && prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) # B takes the M role: groups, K maps, storages, run and transforms move # together; the sum is unchanged. kgroup_swapped = AxisGroup(korder, (indB, B), (indA, A)) # maps: (B, A) diff --git a/test/execution/test_macro_blocking.jl b/test/execution/test_macro_blocking.jl index 616d8a2..72d1034 100644 --- a/test/execution/test_macro_blocking.jl +++ b/test/execution/test_macro_blocking.jl @@ -19,8 +19,8 @@ end # The engine's conjugation rule re-derived independently: the flag and the # view's `op` compose by XOR; real eltypes never conjugate. -_macro_op_conjugates(op) = op === conj || op === adjoint -_macro_conjugated(::Type{T}, flag::Bool, op) where {T} = (T <: Complex) && (flag ⊻ _macro_op_conjugates(op)) +_macro_conj_op(op) = op === conj || op === adjoint +_macro_conjugated(::Type{T}, flag::Bool, op) where {T} = (T <: Complex) && (flag ⊻ _macro_conj_op(op)) const _MACRO_OPS = (identity, conj, adjoint, transpose) const _MACRO_SHAPES = ((Val(4), Val(3)), (Val(8), Val(6))) @@ -160,7 +160,7 @@ end _macro_op_view(C, op), _macro_op_view(Amat, op), (1, 2), _macro_op_view(Bmat, op), (2, 3), (1, 3); kernel = kernel, conjA = true, conjB = true, m_block = 4, k_block = 3, n_block = 4 ) - if T <: Complex && _macro_op_conjugates(op) + if T <: Complex && _macro_conj_op(op) @test_throws ArgumentError mkplan() else plan = mkplan() diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 713f8fa..fe3f07f 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -9,7 +9,7 @@ const _pcf_exec = QuasiStrided.execute! const _pcf_exec_tw = QuasiStrided.execute_tilewise! const _pcf_Plan = QuasiStrided.ContractPlan -@testset "per-call floor: _classify_labels" begin +@testset "per-call floor: classify_labels" begin for (indA, indB, indC, want) in ( ((1, 2), (2, 3), (1, 3), ((1,), (3,), (2,))), # GEMM ((1, 2), (2, 3), (3, 1), ((1,), (3,), (2,))), # transposed C @@ -19,14 +19,14 @@ const _pcf_Plan = QuasiStrided.ContractPlan ((1,), (1, 2), (2,), ((), (2,), (1,))), # no M ((1, 2), (3, 4), (1, 2, 3, 4), ((1, 2), (3, 4), ())), # outer product ) - got = QS._classify_labels(indA, indB, indC) + got = QS.classify_labels(indA, indB, indC) @test got === want - @test map(length, got) == QS._group_ranks(length(indA), length(indB), length(indC)) + @test map(length, got) == QS.group_ranks(length(indA), length(indB), length(indC)) end - @test_throws ArgumentError QS._classify_labels((1, 1), (1, 2), (1, 2)) - @test_throws ArgumentError QS._classify_labels((1, 2), (2, 3), (1, 2, 3)) # all three - @test_throws ArgumentError QS._classify_labels((1, 2), (3, 4), (1, 3)) # dangling in A - @test_throws ArgumentError QS._classify_labels((1, 2), (3, 4), (1, 2, 3, 9)) + @test_throws ArgumentError QS.classify_labels((1, 1), (1, 2), (1, 2)) + @test_throws ArgumentError QS.classify_labels((1, 2), (2, 3), (1, 2, 3)) # all three + @test_throws ArgumentError QS.classify_labels((1, 2), (3, 4), (1, 3)) # dangling in A + @test_throws ArgumentError QS.classify_labels((1, 2), (3, 4), (1, 2, 3, 9)) end @testset "per-call floor: AxisGroup from labels matches a direct construction" begin diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 0b7c2fa..6d13ac4 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -17,7 +17,7 @@ @test_throws ArgumentError contract!(Cv, 1.0, Av, (1, 2), Bv, (2, 3), 0.0, (1, 9)) # Dangling label: only in indA, then only in indB. Distinct code paths - # (the A loop vs. the B loop in _classify_labels). + # (the A loop vs. the B loop in classify_labels). @test_throws ArgumentError contract!(Cv, 1.0, Av, (1, 2), Bv, (5, 3), 0.0, (1, 3)) @test_throws ArgumentError contract!(Cv, 1.0, Av, (1, 2), Bv, (2, 9), 0.0, (1, 3)) @@ -76,23 +76,23 @@ end end -@testset "_op_conjugates is a total table with a throwing fallback" begin - @test QuasiStrided._op_conjugates(identity) === false - @test QuasiStrided._op_conjugates(conj) === true - @test QuasiStrided._op_conjugates(transpose) === false # elementwise identity - @test QuasiStrided._op_conjugates(adjoint) === true - @test_throws ArgumentError QuasiStrided._op_conjugates(sin) +@testset "op_conjugates is a total table with a throwing fallback" begin + @test QuasiStrided.op_conjugates(identity) === false + @test QuasiStrided.op_conjugates(conj) === true + @test QuasiStrided.op_conjugates(transpose) === false # elementwise identity + @test QuasiStrided.op_conjugates(adjoint) === true + @test_throws ArgumentError QuasiStrided.op_conjugates(sin) # A real element type is conjugated by nothing; a complex one composes # the flag and the op with XOR. for op in (identity, conj, transpose, adjoint), flag in (false, true) v = StridedView(randn(16), (4, 4), (1, 4), 0, op) - @test QuasiStrided._qs_isconj(v, flag) === false + @test QuasiStrided.isconj(v, flag) === false end for (op, oc) in ((identity, false), (conj, true), (transpose, false), (adjoint, true)), flag in (false, true) v = StridedView(randn(ComplexF64, 16), (4, 4), (1, 4), 0, op) - @test QuasiStrided._qs_isconj(v, flag) === (flag ⊻ oc) + @test QuasiStrided.isconj(v, flag) === (flag ⊻ oc) end end @@ -153,7 +153,7 @@ end @test err isa ArgumentError @test occursin("conjugated", err.msg) end - # StridedViews bounds `op` to exactly the four functions `_op_conjugates` + # StridedViews bounds `op` to exactly the four functions `op_conjugates` # tabulates; fail here if that ever widens. @test_throws TypeError StridedView(Cc, (6, 4), (1, 6), 0, sin) Fbound = fieldtype(typeof(StridedView(Cc, (6, 4), (1, 6), 0, conj)), :op) @@ -165,9 +165,9 @@ end # Label order within M/N and the M/N orientation swap. -const _lo_order = QuasiStrided._order_free_labels -const _lo_run = QuasiStrided._leading_unit_run -_lo_swap(morder, norder, indC, C, m_tile_asis, m_tile_swapped = m_tile_asis) = QuasiStrided._prefer_swap( +const _lo_order = QuasiStrided.order_free_labels +const _lo_run = QuasiStrided.leading_unit_run +_lo_swap(morder, norder, indC, C, m_tile_asis, m_tile_swapped = m_tile_asis) = QuasiStrided.prefer_swap( _lo_run(morder, indC, C), _lo_run(norder, indC, C), m_tile_asis, m_tile_swapped ) @@ -224,7 +224,7 @@ function _lo_reference_loop!( return acc end -@testset "label order: _order_free_labels sorts by |C-stride|, stably" begin +@testset "label order: order_free_labels sorts by |C-stride|, stably" begin Cv = _lo_view((5, 3, 2, 4), (12, 1, 60, 3)) indC = (10, 20, 30, 40) labels = (10, 20, 30, 40) @@ -259,7 +259,7 @@ end end end -@testset "label order: _leading_unit_run / _prefer_swap" begin +@testset "label order: leading_unit_run / prefer_swap" begin d = 4 C6 = StridedView(zeros(Float64, d, d, d, d, d, d)) # strides 1, d, d^2, ... indC = (1, 2, 3, 4, 5, 6) # a,b,c,i,j,k @@ -298,7 +298,7 @@ end @test !_lo_swap((b, i, j), (), indC, C6, 1) end -@testset "label order: _leading_unit_run's full-coverage condition is a single-map affine_ramp on C's own strides" begin +@testset "label order: leading_unit_run's full-coverage condition is a single-map affine_ramp on C's own strides" begin # The whole composite is one run iff C's own single-map group is a unit # ramp. Empty and all-singleton composites are excluded: `affine_ramp` # reports step 0 there while the run is trivially 1. @@ -316,7 +316,7 @@ end m_length == 1 && continue ntested += 1 - run = QuasiStrided._leading_unit_run(order, indCr, Cr) + run = QuasiStrided.leading_unit_run(order, indCr, Cr) clens = ntuple(d -> lens[order[d]], nlabels) cstrides = ntuple(d -> strides[order[d]], nlabels) @@ -337,7 +337,7 @@ end Av, Bv, Cv = StridedView(A), StridedView(B), StridedView(C) cst = Base.strides(Cv) cstride(l) = cst[findfirst(==(l), indC)] - mlab, nlab, klab = QuasiStrided._classify_labels(indA, indB, indC) + mlab, nlab, klab = QuasiStrided.classify_labels(indA, indB, indC) @test length(klab) == 1 kA = Base.strides(Av)[findfirst(==(klab[1]), indA)] kB = Base.strides(Bv)[findfirst(==(klab[1]), indB)] @@ -405,7 +405,7 @@ end C = zeros(T, d, extra, extra, extra, extra, extra) # (a, b, c, i, j, k) Av, Bv, Cv = StridedView(A), StridedView(B), StridedView(C) - mlab, nlab, klab = QuasiStrided._classify_labels(indA, indB, indC) + mlab, nlab, klab = QuasiStrided.classify_labels(indA, indB, indC) msorted = _lo_order(mlab, indC, Cv) run = _lo_run(msorted, indC, Cv) cpos(l) = findfirst(==(l), indC)::Int @@ -494,7 +494,7 @@ end Cv2, Av2, Bv2 = _run_demote_fixture(T, a, m_shallow) plan_shallow = plan_contract(Cv2, Av2, indA, Bv2, indB, indC) @test plan_shallow.Astorage === parent(Av2) - mlab, = QuasiStrided._classify_labels(indA, indB, indC) + mlab, = QuasiStrided.classify_labels(indA, indB, indC) msorted = _lo_order(mlab, indC, Cv2) run = _lo_run(msorted, indC, Cv2) if m_length == run || run % tile_size(default_kernel, 1) == 0 @@ -640,7 +640,7 @@ end # `_KO_EXT` fits any L2; `_KO_BIG` makes B's lines touched before `d` advances # (b*c*e*64 B) exceed THIS host's L2 share, with e >= 40, so the model reorders # K wherever the suite runs. -function _ko_big(l2 = QuasiStrided._l2_core_bytes(Float64)) +function _ko_big(l2 = QuasiStrided.l2_core_bytes(Float64)) e = max(40, fld(l2, 40 * 40 * 64) + 1) return Dict(1 => 9, 2 => 40, 3 => 40, 4 => 5, 5 => e) end @@ -650,11 +650,11 @@ _ko_array(::Type{T}, ind, ext = _KO_EXT) where {T} = randn(T, Tuple(ext[l] for l _ko_stride(v, ind, l) = Base.strides(v)[findfirst(==(l), ind)] function _ko_order(Av, indA, Bv, indB, Cv, indC; l2bytes) - mlabels, nlabels, klabels = QuasiStrided._classify_labels(indA, indB, indC) - morder = QuasiStrided._order_free_labels(mlabels, indC, Cv) - norder = QuasiStrided._order_free_labels(nlabels, indC, Cv) + mlabels, nlabels, klabels = QuasiStrided.classify_labels(indA, indB, indC) + morder = QuasiStrided.order_free_labels(mlabels, indC, Cv) + norder = QuasiStrided.order_free_labels(nlabels, indC, Cv) ext(ls) = prod((size(Cv, findfirst(==(l), indC)) for l in ls); init = 1) - return QuasiStrided._choose_k_order( + return QuasiStrided.order_contract_labels( klabels, indA, Av, morder, indB, Bv, norder, ext(mlabels), ext(nlabels), l2bytes ) end @@ -701,15 +701,15 @@ end Bs = StridedView(_ko_array(T, (4, 3, 2, 5), ext)) Cv = StridedView(zeros(T, ext[1], ext[5])) klabels, morder, norder = (2, 3, 4), (1,), (5,) - f(l2) = QuasiStrided._choose_k_order(klabels, (1, 2, 3, 4), Av, morder, (4, 3, 2, 5), Bs, norder, ext[1], ext[5], l2) - g() = QuasiStrided._order_contract_labels(klabels, (1, 2, 3, 4), Av, morder, (4, 3, 2, 5), Bs, norder, ext[1], ext[5]) + f(l2) = QuasiStrided.order_contract_labels(klabels, (1, 2, 3, 4), Av, morder, (4, 3, 2, 5), Bs, norder, ext[1], ext[5], l2) + g() = QuasiStrided.order_contract_labels(klabels, (1, 2, 3, 4), Av, morder, (4, 3, 2, 5), Bs, norder, ext[1], ext[5]) @test @inferred(f(0)) === (4, 3, 2) @test @inferred(g()) isa NTuple{3, Int} f(0); g() @test (@allocated f(0)) == 0 skip = (VERSION < v"1.11") @test (@allocated g()) == 0 skip = (VERSION < v"1.11") - h1() = QuasiStrided._order_contract_labels((2,), (1, 2), Av, (1,), (2, 5), Bs, (5,), 9, 40) - h0() = QuasiStrided._order_contract_labels((), (1,), Av, (1,), (5,), Bs, (5,), 9, 40) + h1() = QuasiStrided.order_contract_labels((2,), (1, 2), Av, (1,), (2, 5), Bs, (5,), 9, 40) + h0() = QuasiStrided.order_contract_labels((), (1,), Av, (1,), (5,), Bs, (5,), 9, 40) @test @inferred(h1()) === (2,) @test @inferred(h0()) === () end @@ -717,7 +717,7 @@ end @testset "K order: plan_contract flips contract_scrambled at an L2-exceeding size" begin T = Float64 ext = _KO_BIG - @test ext[2] * ext[3] * ext[5] * 64 > QuasiStrided._l2_core_bytes(T) + @test ext[2] * ext[3] * ext[5] * 64 > QuasiStrided.l2_core_bytes(T) Av = StridedView(_ko_array(T, (1, 2, 3, 4), ext)) Bs = StridedView(_ko_array(T, (4, 3, 2, 5), ext)) # B[d,c,b,e] Cv = StridedView(zeros(T, ext[1], ext[5])) From 7032aed390c319c03691e19b0a3b9647a4453f04 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 11:01:46 -0400 Subject: [PATCH 20/60] Decision log: chunk 9 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/CLEANUP.md b/CLEANUP.md index 1fed2fb..246fc16 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -21,7 +21,7 @@ TTFX when codegen is touched). | 6 | Kernel interface | `microkernels/interface.jl` (+ ScalarKernel) | done | | 7 | Vector kernels | now `kernels.jl`, `vecops.jl`, `steps.jl`, `stores.jl` | done | | 8 | Complex/mixed kernels | (merged into chunk 7's files) | done | -| 9 | Labels | `planning/labels.jl`, `conjugation.jl` | | +| 9 | Labels | `planning/labels.jl` (`conjugation.jl` → plan.jl) | done | | 10 | Kernel selection | `planning/kernel_selection.jl` | | | 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | @@ -174,6 +174,14 @@ and their traits; `AccumulatorLayout`/`KernelMethod` are Holy traits), `vecops.jl` (shared SIMD ops), `steps.jl` (K steps, `add_tile`), `stores.jl` (layouts' stores). +### D13. Labels (chunk 9, applied) + +`TupleTools.sortperm`/`getindices` (its merge sort is stable; relied on and +tested) replace the insertion sort; one K-order function; `l2_core_bytes` +moves to defaults.jl; `isconj`/`op_conjugates` move into plan.jl. The K-order +model, free-label order and swap rule are unchanged; labels.jl stays one +file. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: @@ -184,8 +192,16 @@ and their traits; `AccumulatorLayout`/`KernelMethod` are Holy traits), - `PackedPanel` without a raw pointer: https://github.com/lkdvos/QuasiStrided.jl/issues/15. +- M and N are both ordered by their C strides. Only the M side matters for the + vector store; for N the order mainly decides ramp detection and C locality, + while B's pack (and the unpacked-B path) would prefer B's strides. Measure on + suite layouts where B's and C's N strides disagree. + ## Open items (to revisit in their chunk) +- Internal constants still use `_UPPER` names; drop the underscores in one pass + at the end (D2). + - Detected `l1d.line` is unused: `_K_LINE_BYTES = 64` in `labels.jl` (K-order cost model) and `pack_split.jl`; thread the detected line size through planning (chunks 9/11). From e7261a45e329919436972f98ac3e93880f094d9f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 12:32:51 -0400 Subject: [PATCH 21/60] Chunk 10: one shape-selection pipeline `select_shape(T, method, m_length, k_length, run)` in kernel_selection.jl runs every automatic selection step in order (host shape, extent step-down, small M, `fit_to_run`), previously split between `_default_shape` in defaults.jl and `_plan_with_kernel` in plan.jl; one docstring lists the steps. defaults.jl keeps only the per-eltype cache. The mixed-method mapping now wraps the whole pipeline once. The swap decision calls it with an unbounded `k_length`, so, as before, it judges orientations without the K-dependent divisor step. `fit_to_run` merges `_store_shape` and `_demote_shape_for_run`: the MV = 4 -> MV = 2 step at any K, then the divisor step at small K, with their conditions unchanged. It now runs after the small-M step; the store step only fires on the tall shape, which the small-M step never leaves, so the order is equivalent. Also: the NEON planar override rows go (the register fit selects the same shapes; a test pins them), the dead `vb % sizeof(real(T)) == 0` conditions go, `_MixedMethod` -> `MixedMethod`, and names defined in kernel_selection.jl lose their leading underscore. The uncached `_kernel_for` and the unused `n_length` of the old pipeline are dropped. Behaviour-neutral: the automatic kernel type, shape, blocking and orientation are identical to 7032aed for Float32/64, ComplexF32/64 and both mixed pairs, m, n, k in {1,2,3,5,8,13,17,31,48,64,100,257}, four C layouts, under :avx512 (Intel and znver4), :avx2, :neon and :unknown profiles (138270 plans), and on a dense (m <= 300, every run dividing m, 7 K values) grid of the shape pipeline itself (579040 cases). Per-call floor (medians of 3 interleaved runs, ns, before -> after): Float64 8^3 372 -> 368, 16^3 484 -> 479, d4 outer 344 -> 340; ComplexF64 8^3 424 -> 415, 16^3 966 -> 963, d4 outer 580 -> 589; 0 B. TTFX: first call Float64 3.8 -> 3.8-3.9 s, ComplexF64 2.2-2.5 -> 2.4-2.5 s. Tests: 32198 pass (was 31967; the mixed-selection grid gained a K dimension), 48 min on a loaded workstation. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_kernels.jl | 6 +- benchmark/bench_mixed.jl | 6 +- src/microkernels/interface.jl | 2 +- src/planning/blocking.jl | 2 +- src/planning/defaults.jl | 42 +--- src/planning/kernel_selection.jl | 272 +++++++++++----------- src/planning/plan.jl | 21 +- test/microkernels/test_fmaddsub_kernel.jl | 6 +- test/microkernels/test_onem_kernel.jl | 8 +- test/planning/test_kernel_selection.jl | 204 ++++++++-------- test/runtests.jl | 6 +- 11 files changed, 276 insertions(+), 299 deletions(-) diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index beb58d1..5c6d1c0 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -22,7 +22,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: RealMethod, PlanarMethod, OneMMethod, FMAddSubMethod, kernel_shapes, - _kernel_from_shape, default_blocking, packed_panel, execute_tile!, + kernel_from_shape, default_blocking, packed_panel, execute_tile!, Tile, AffineAxis, target_profile const RUN_DTYPES = parse_dtypes(argopt("dtypes", "Float64,Float32,ComplexF64,ComplexF32")) @@ -69,7 +69,7 @@ function kernel_hot!(kernel, C, apack, bpack, k_block, reps) end for T in filter(t -> t <: Real, RUN_DTYPES), sh in kernel_shapes(T) - kernel = _kernel_from_shape(sh, T) + kernel = kernel_from_shape(sh, T) MR, NR = sh k_block = default_blocking(kernel).k_block apack = rand(T, MR * k_block); bpack = rand(T, NR * k_block); C = zeros(T, MR * NR) @@ -92,7 +92,7 @@ for spec in SHAPES, T in RUN_DTYPES @assert isapprox(fx.Cmat, fx.Amat * fx.Bmat; rtol = sqrt(eps(real(T)))) record!("engine", T, spec, "default", tag(pd.kernel), pd.blocking, td, tb) for m in menu_methods(T), sh in kernel_shapes(T, m) - k = _kernel_from_shape(sh, T, m) + k = kernel_from_shape(sh, T, m) b = default_blocking(k) t, _ = time_plan(fx, T; kernel = k, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block) record!("engine", T, spec, lowercase(replace(string(nameof(typeof(m))), "Method" => "")), tag(k), b, t, tb) diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index 9e3a604..04f07c0 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -10,7 +10,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: ComplexRealKernel, RealComplexKernel, PlanarMethod, target_profile, - _derived_shape, _kernel_from_shape + derived_shape, kernel_from_shape const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) const SIZES = parse_ints(argopt("sizes", "512,2048")) @@ -29,8 +29,8 @@ function run(io) 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()) + 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)), diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index 029ad38..e06fde0 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -29,7 +29,7 @@ blocking and the shape menus dispatch on it: No method ranking is hardcoded: planar is the default, 1m and fmaddsub are used only when named (plus fmaddsub for the AVX-512 small-M demotion, -`_small_m_shape` in src/planning/kernel_selection.jl). +`small_m_shape` in src/planning/kernel_selection.jl). """ abstract type KernelMethod end struct RealMethod <: KernelMethod end diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index 509595f..dfdb631 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -61,7 +61,7 @@ function _modelled_blocking(profile::TargetProfile, ::Type{T}, MR::Int, NR::Int) end function _modelled_blocking(profile::TargetProfile, ::Type{T}) where {T <: Real} - MR, NR, _ = _derived_shape(profile, T) + MR, NR, _ = derived_shape(profile, T) return _modelled_blocking(profile, T, MR, NR) end diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index 3e5f433..ae6af77 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -26,11 +26,11 @@ end _defaults_slot(T) === nothing ? l2_core_bytes(target_profile()) : _resolved_defaults(T).l2_core function _resolve_defaults(profile::TargetProfile, ::Type{T}) where {T} - method = _default_method(T) + method = default_method(T) # First, so an element type with no menu throws from here. - shape = _derived_shape(profile, T, method) - fitted = _fitted_shape(profile, T, method) - small_m = _small_m_candidates(Val(profile.isa), profile, T) + shape = derived_shape(profile, T, method) + fitted = fitted_shape(profile, T, method) + small_m = small_m_candidates(Val(profile.isa), profile, T) real_row = _real_blocking_row(profile, real(T)) l2_core = l2_core_bytes(profile) return ResolvedDefaults(profile, shape, fitted, small_m, real_row, l2_core) @@ -67,40 +67,12 @@ end end _default_kernel(::Type{T}) where {T} = - _kernel_from_shape(_resolved_defaults(T).shape, T, _default_method(T)) - -# The automatic `(shape, method)` for extents `m_length`/`n_length` and C's -# unit-stride run along M (`run = m_length`: no layout known). Returned as plain -# values so `plan_contract` never holds a menu-wide kernel Union. The extent and -# store step-downs apply first; then an `m_length` that cannot fill one tile -# demotes to the fitted shape, or on AVX-512 complex to FMAddSub. -@inline function _default_shape(::Type{T}, m_length::Int, n_length::Int, run::Int = m_length) where {T} - d = _resolved_defaults(T) - method = _default_method(T) - shape = _store_shape(_extent_shape(d.shape, T, method, m_length), T, method, m_length, run) - (m_length > 0 && m_length < shape[1]) || return (shape, method) - # Static, so a real `T`'s method stays a concrete `RealMethod`. - T <: Complex || return (d.fitted, method) - small = _small_m_shape(d.small_m, m_length) - small === nothing || return (small, FMAddSubMethod()) - return (d.fitted, method) -end - -# The automatic `(shape, method)` under `method`, `_default_method(T, TA, TB)`. -@inline _default_shape(::Type{T}, ::KernelMethod, m_length::Int, n_length::Int, run::Int) where {T} = - _default_shape(T, m_length, n_length, 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, m_length::Int, n_length::Int, run::Int) where {T} - shape, _ = _default_shape(real(T), _real_problem(method, m_length, n_length, run)...) - return (_mixed_shape(method, shape), method) -end + kernel_from_shape(_resolved_defaults(T).shape, T, default_method(T)) # `@noinline`: the return type is the Union of `T`'s menu kernels. @noinline function _default_kernel(::Type{T}, m_length::Int, n_length::Int) where {T} - shape, method = _default_shape(T, m_length, n_length) - return _kernel_from_shape(shape, T, method) + shape, method = select_shape(T, default_method(T), m_length, 0, m_length) + return kernel_from_shape(shape, T, method) end """ diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index 0204cc4..bdf5792 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -4,30 +4,30 @@ # the `Val(profile.isa)` dispatch below is dynamic. # `OneMMethod`/`FMAddSubMethod` are chosen only by naming the kernel (FMAddSub -# also by the AVX-512 small-M demotion, `_small_m_shape`). -_default_method(::Type{<:Real}) = RealMethod() -_default_method(::Type{<:Complex}) = PlanarMethod() +# also by the AVX-512 small-M demotion, `small_m_shape`). +default_method(::Type{<:Real}) = RealMethod() +default_method(::Type{<:Complex}) = PlanarMethod() -const _MixedMethod = Union{ComplexRealMethod, RealComplexMethod} +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} = +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} = +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. -_fallback_shape(::Type{T}) where {T} = (8, NR_DEFAULT, default_lanewidth(real(T))) +fallback_shape(::Type{T}) where {T} = (8, NR_DEFAULT, default_lanewidth(real(T))) # Closed menus of `(MR, NR, W)` shapes, so the compiled specializations stay # bounded. Complex `MR` counts complex rows and `W` real lanes; 1m runs a real # kernel of `2MR` rows, so only `2MR` must be a multiple of `W`. The trailing -# `MV = 1` planar entries guarantee `_fitted_shape` a fit on any register file +# `MV = 1` planar entries guarantee `fitted_shape` a fit on any register file # of >= 16 registers. const KERNEL_SHAPES_F64 = ((8, 6, 4), (16, 6, 8), (32, 6, 8)) const KERNEL_SHAPES_F32 = ((8, 6, 8), (32, 6, 16), (16, 6, 8), (64, 6, 16)) @@ -45,12 +45,12 @@ const KERNEL_SHAPES_C64_FMADDSUB = ((12, 8, 8), (8, 8, 8), (4, 6, 4), (4, 5, 4)) const KERNEL_SHAPES_C32_FMADDSUB = ((24, 8, 16), (16, 8, 16), (8, 6, 8), (8, 5, 8)) """ - kernel_shapes(T, method::KernelMethod = _default_method(T)) -> NTuple{<:Any,NTuple{3,Int}} + kernel_shapes(T, method::KernelMethod = default_method(T)) -> NTuple{<:Any,NTuple{3,Int}} The closed menu of `(MR, NR, W)` register shapes the engine may build for element type `T` under `method`. """ -kernel_shapes(::Type{T}) where {T} = kernel_shapes(T, _default_method(T)) +kernel_shapes(::Type{T}) where {T} = kernel_shapes(T, default_method(T)) kernel_shapes(::Type{Float64}, ::RealMethod) = KERNEL_SHAPES_F64 kernel_shapes(::Type{Float32}, ::RealMethod) = KERNEL_SHAPES_F32 kernel_shapes(::Type{ComplexF64}, ::PlanarMethod) = KERNEL_SHAPES_C64_PLANAR @@ -59,38 +59,36 @@ 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())) +# 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) +# A real inner shape to the mixed kernel's `(MR, NR, W)`: 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) -# The real problem the inner kernel sees: extents and C's unit-stride M run in reals. -_real_problem(::ComplexRealMethod, m_length::Int, n_length::Int, run::Int) = (2 * m_length, n_length, 2 * run) -_real_problem(::RealComplexMethod, m_length::Int, n_length::Int, run::Int) = (m_length, 2 * n_length, run) +# The real problem the inner kernel sees: `m_length` and C's unit-stride M run in reals. +real_problem(::ComplexRealMethod, m_length::Int, run::Int) = (2 * m_length, 2 * run) +real_problem(::RealComplexMethod, m_length::Int, run::Int) = (m_length, run) # The kernel type implementing `method` for `T`, or `nothing` (exactly the # pairs with a menu). -_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 +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 # are built from literals, so each is concrete; a shape outside the menu throws. # Each branch builds the kernel, or `Val(shape)` when `valonly`. -function _menu_ladder(::Type{T}, method, valonly::Bool) where {T} - K = _kernel_type(method, T) - K === nothing && return :(_throw_no_kernel(shape, T, method)) - ex = :(_throw_shape_not_in_menu(shape, T, method)) +function menu_ladder(::Type{T}, method, valonly::Bool) where {T} + K = kernel_type(method, T) + K === nothing && return :(throw_no_kernel(shape, T, method)) + ex = :(throw_shape_not_in_menu(shape, T, method)) for (MR, NR, W) in reverse(kernel_shapes(T, method)) arm = valonly ? :(Val(($MR, $NR, $W))) : :($K(Val($MR), Val($NR), T, Val($W))) ex = :(shape === ($MR, $NR, $W) ? $arm : $ex) @@ -99,17 +97,17 @@ function _menu_ladder(::Type{T}, method, valonly::Bool) where {T} end # The `method` kernel for `T` at a menu `shape` (one dynamic dispatch per plan). -_kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = - _kernel_from_shape(shape, T, _default_method(T)) -@generated _kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = - _menu_ladder(T, M.instance, false) +kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = + kernel_from_shape(shape, T, default_method(T)) +@generated kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = + menu_ladder(T, M.instance, false) # `Val(shape)` for a menu `shape`: a singleton, so it crosses `plan_contract`'s # kernel barrier without boxing. -@generated _menu_val(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = - _menu_ladder(T, M.instance, true) +@generated menu_val(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = + menu_ladder(T, M.instance, true) -@noinline _throw_no_kernel(shape, ::Type{T}, method) where {T} = throw( +@noinline throw_no_kernel(shape, ::Type{T}, method) where {T} = throw( ArgumentError( "no microkernel is available for $T at shape $shape under $(method). " * "Pass an explicit `kernel = ...` to plan_contract to use a kernel this " * @@ -117,7 +115,7 @@ _kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = ) ) -@noinline _throw_shape_not_in_menu(shape, ::Type{T}, method) where {T} = throw( +@noinline throw_shape_not_in_menu(shape, ::Type{T}, method) where {T} = throw( ArgumentError( "shape $shape is not in the $(method) menu for $T, " * "$(kernel_shapes(T, method)); only menu shapes are compiled" @@ -129,88 +127,85 @@ _kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = # ---------------------------------------------------------------------------- # Per-ISA shapes, consulted first. None for real types: the rule below is the optimum. -_shape_override(key::Val, ::Type{T}) where {T} = _shape_override(key, T, _default_method(T)) -_shape_override(::Val, ::Type, ::KernelMethod) = nothing +shape_override(key::Val, ::Type{T}) where {T} = shape_override(key, T, default_method(T)) +shape_override(::Val, ::Type, ::KernelMethod) = nothing # On AVX-512 the rule's planar `MR = 2W, NR = 6` spills; `NR = 3` tiles do not. -_shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::PlanarMethod) = (24, 3, 8) -_shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::PlanarMethod) = (48, 3, 16) -# What the register fit selects on NEON anyway; pinned against menu edits. -_shape_override(::Val{:neon}, ::Type{ComplexF64}, ::PlanarMethod) = (4, 6, 2) -_shape_override(::Val{:neon}, ::Type{ComplexF32}, ::PlanarMethod) = (8, 6, 4) +shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::PlanarMethod) = (24, 3, 8) +shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::PlanarMethod) = (48, 3, 16) # AVX2 has 16 registers: at `MV = 1` planar `NR = 6` needs all 16, `NR = 5` needs 14. -_shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::PlanarMethod) = (4, 5, 4) -_shape_override(::Val{:avx2}, ::Type{ComplexF32}, ::PlanarMethod) = (8, 5, 8) +shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::PlanarMethod) = (4, 5, 4) +shape_override(::Val{:avx2}, ::Type{ComplexF32}, ::PlanarMethod) = (8, 5, 8) # The AVX2-sized 1m shape, for a caller naming `OneMMethod()` (the fit would # hand it an AVX-512 shape). -_shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::OneMMethod) = (4, 6, 4) +shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::OneMMethod) = (4, 6, 4) # Where the `MR = MV*W` rule applies. Not on NEON (2W on 128-bit lanes is no # better than the fallback), nor complex off AVX-512 (planar's two accumulator # planes leave AVX2's 16 registers nothing spare). -_rule_applies(::Val{:avx512}, ::RealMethod) = true -_rule_applies(::Val{:avx512}, ::Union{PlanarMethod, OneMMethod}) = true -_rule_applies(::Val{:avx2}, ::RealMethod) = true -_rule_applies(::Val, ::KernelMethod) = false +rule_applies(::Val{:avx512}, ::RealMethod) = true +rule_applies(::Val{:avx512}, ::Union{PlanarMethod, OneMMethod}) = true +rule_applies(::Val{:avx2}, ::RealMethod) = true +rule_applies(::Val, ::KernelMethod) = false # `W` real lanes per register, `MV` A vectors per column. -_rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = +rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = (mv * (vb ÷ sizeof(real(T))), NR_DEFAULT, vb ÷ sizeof(real(T))) # MV = 4 for the real kernel on AVX-512: the MV = 2 tile is front-end bound on # cores with 2 FMA ports, and 4*6 accumulators + 4 A vectors + 2 still fit 32 # registers. Complex methods keep MV = 2 (planar already spills there). -_rule_mv(::Val{:avx512}, ::RealMethod) = 4 -_rule_mv(::Val, ::KernelMethod) = 2 +rule_mv(::Val{:avx512}, ::RealMethod) = 4 +rule_mv(::Val, ::KernelMethod) = 2 # AMD's AVX-512 cores double-pump 512-bit FMAs, so their MV = 2 tile is not # front-end bound and the taller tile only adds edge and store cost. const _MV4_UNPROFITABLE_CPUS = ("znver4", "znver5") -_profile_mv(profile::TargetProfile, method) = _rule_mv(Val(profile.isa), method) -function _profile_mv(profile::TargetProfile, method::RealMethod) - mv = _rule_mv(Val(profile.isa), method) +profile_mv(profile::TargetProfile, method) = rule_mv(Val(profile.isa), method) +function profile_mv(profile::TargetProfile, method::RealMethod) + mv = rule_mv(Val(profile.isa), method) return (mv == 4 && profile.cpu_name in _MV4_UNPROFITABLE_CPUS) ? 2 : mv end # The register shape for `T` under `method` on `profile`: an override row, else -# the rule where it applies (and lands in the menu), else `_fitted_shape`. -_derived_shape(profile::TargetProfile, ::Type{T}) where {T} = - _derived_shape(profile, T, _default_method(T)) +# the rule where it applies (and lands in the menu), else `fitted_shape`. +derived_shape(profile::TargetProfile, ::Type{T}) where {T} = + derived_shape(profile, T, default_method(T)) -function _derived_shape(profile::TargetProfile, ::Type{T}, method) where {T} +function derived_shape(profile::TargetProfile, ::Type{T}, method) where {T} key = Val(profile.isa) - ovr = _shape_override(key, T, method) + ovr = shape_override(key, T, method) ovr === nothing || return ovr vb = profile.vector_bytes - if _rule_applies(key, method) && vb > 0 && vb % sizeof(real(T)) == 0 - shape = _rule_shape(vb, T, _profile_mv(profile, method)) + if rule_applies(key, method) && vb > 0 + shape = rule_shape(vb, T, profile_mv(profile, method)) shape in kernel_shapes(T, method) && return shape end - return _fitted_shape(profile, T, method) + return fitted_shape(profile, T, method) end # Vector registers a planar kernel holds live per K step (two planes of # accumulators and A vectors, plus 2 broadcasts). Spilling is not monotone in # it, so it only excludes shapes, never ranks them. -_planar_pressure(MR::Int, NR::Int, W::Int) = 2 * (MR ÷ W) * NR + 2 * (MR ÷ W) + 2 +planar_pressure(MR::Int, NR::Int, W::Int) = 2 * (MR ÷ W) * NR + 2 * (MR ÷ W) + 2 # The shape where no rule applies, and the small-M demotion target. Planar: the # largest menu shape with `W <= lanes` whose pressure fits the register file # (16 when unknown). 1m/FMAddSub: the menu head. -_fitted_shape(::TargetProfile, ::Type{T}, ::RealMethod) where {T} = _fallback_shape(T) -_fitted_shape(::TargetProfile, ::Type{T}, method::Union{OneMMethod, FMAddSubMethod}) where {T} = +fitted_shape(::TargetProfile, ::Type{T}, ::RealMethod) where {T} = fallback_shape(T) +fitted_shape(::TargetProfile, ::Type{T}, method::Union{OneMMethod, FMAddSubMethod}) where {T} = first(kernel_shapes(T, method)) -function _fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) where {T} +function fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) where {T} R = real(T) vb = profile.vector_bytes - lanes = (vb > 0 && vb % sizeof(R) == 0) ? vb ÷ sizeof(R) : default_lanewidth(R) + lanes = vb > 0 ? vb ÷ sizeof(R) : default_lanewidth(R) budget = profile.nregisters > 0 ? profile.nregisters : 16 best = nothing for shape in kernel_shapes(T, method) MR, NR, W = shape (W <= lanes && MR % W == 0) || continue - _planar_pressure(MR, NR, W) <= budget || continue + planar_pressure(MR, NR, W) <= budget || continue # Largest logical tile wins; ties by the wider vector. if best === nothing || (MR * NR, W) > (best[1] * best[2], best[3]) best = shape @@ -221,53 +216,74 @@ function _fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) end # ---------------------------------------------------------------------------- -# The default kernel for a profile, and the plan-time demotions +# The automatic shape for a contraction # ---------------------------------------------------------------------------- -# The default kernel for `T` on `profile`, uncached. -@noinline _kernel_for(profile::TargetProfile, ::Type{T}) where {T} = - _kernel_from_shape(_derived_shape(profile, T), T, _default_method(T)) +""" + select_shape(T, method, m_length, k_length, run) -> (shape, method) + +The automatic register shape `(MR, NR, W)`, and the method it runs under, for +compute type `T` under `method = default_method(T, TA, TB)`, extents +`m_length`/`k_length` and C's unit-stride run along M (`run`; `m_length` when +no layout is known). The steps, in order: + + 1. The host's shape for `T` (`derived_shape`, cached in `ResolvedDefaults`). + 2. Extent (real): the MV = 4 shape steps down to MV = 2 where `m_length` + pads less (`extent_shape`). + 3. Small M: an `m_length` below one tile takes the fitted shape, or for + complex `T` on AVX-512 the FMAddSub shape that pads least (`small_m_shape`). + 4. Run (real): the shape shrinks until C's run fills its slivers (`fit_to_run`). + +A mixed method runs the steps on the real problem of `real(T)` and maps the +shape. Plain values, so `plan_contract` never holds a menu-wide kernel Union. +""" +@inline function select_shape(::Type{T}, method, m_length::Int, k_length::Int, run::Int) where {T} + d = _resolved_defaults(T) + shape = extent_shape(d.shape, T, method, m_length) + if m_length > 0 && m_length < shape[1] + # Static, so a real `T`'s method stays concrete. + if T <: Complex + small = small_m_shape(d.small_m, m_length) + small === nothing || return (small, FMAddSubMethod()) + end + shape = d.fitted + end + return (fit_to_run(shape, T, method, m_length, k_length, run), method) +end + +@inline function select_shape(::Type{T}, method::MixedMethod, m_length::Int, k_length::Int, run::Int) where {T} + m, r = real_problem(method, m_length, run) + shape, _ = select_shape(real(T), RealMethod(), m, k_length, r) + return (mixed_shape(method, shape), method) +end # The real MV = 4 shape steps down to its MV = 2 sibling when `m_length` is # below one tall tile, or below two and the half tile pads to fewer rows. Above # `2MR` the tall tile's padding excess is under 1.2x, its speed advantage. -@inline _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int) where {T} = shape +@inline extent_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int) where {T} = shape -@inline function _extent_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int) where {T} +@inline function extent_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int) where {T} MR, NR, W = shape (m_length > 0 && MR == 4 * W && m_length < 2 * MR) || return shape half = (2 * W, NR, W) half in kernel_shapes(T, method) || return shape # Below one tall tile, always: the tall shape would fall to the fitted - # shape (two steps down) in `_default_shape`. + # shape (two steps down) at the small-M step. m_length < MR && return half return cld(m_length, half[1]) * half[1] < cld(m_length, MR) * MR ? half : shape end -# The real MV = 4 shape steps down to MV = 2 when C's unit-stride run along M -# (`run`) can't fill the tall tile's slivers but fills a half tile: a sliver -# that is not unit-stride in C takes the scattered store, which costs more than -# the tall tile gains. Below one half tile both scatter, and the tall one stays. -@inline _store_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int, run::Int) where {T} = shape - -@inline function _store_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int, run::Int) where {T} - MR, NR, W = shape - (MR == 4 * W && m_length != run && run % MR != 0 && run >= 2 * W) || return shape - half = (2 * W, NR, W) - return half in kernel_shapes(T, method) ? half : shape -end - # Small-M demotion for complex `T` on AVX-512, where the planar fitted shape is # the spilling `MR = 2W` tile: the native-width FMAddSub shape that pads -# `m_length` least, ties by the larger tile. `_small_m_candidates` is the +# `m_length` least, ties by the larger tile. `small_m_candidates` is the # cached, host-dependent half (empty where the rule does not apply). -_small_m_candidates(::Val, ::TargetProfile, ::Type) = NTuple{3, Int}[] -function _small_m_candidates(::Val{:avx512}, profile::TargetProfile, ::Type{T}) where {T <: Complex} +small_m_candidates(::Val, ::TargetProfile, ::Type) = NTuple{3, Int}[] +function small_m_candidates(::Val{:avx512}, profile::TargetProfile, ::Type{T}) where {T <: Complex} lanes = profile.vector_bytes ÷ sizeof(real(T)) return [shape for shape in kernel_shapes(T, FMAddSubMethod()) if shape[3] == lanes] end -function _small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) +function small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) best = nothing for shape in candidates MR, NR, _ = shape @@ -279,42 +295,32 @@ function _small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) return best === nothing ? nothing : best[2] end -# Run-length demotion (real `T` only): the vectorized store needs every -# register sliver unit-stride in C, i.e. `m_length == run || run % m_tile == 0` -# (not `m_tile <= run`). Otherwise demote to the largest menu shape whose -# `m_tile` divides `run`, unless `k_length` is deep enough that the smaller -# kernel's cost dominates. -@inline function _demote_shape_for_run( - ::Type{T}, shape::NTuple{3, Int}, method, run::Int, m_length::Int, k_length::Int - ) where {T} - target = _run_demotion_target(T, shape[1], method, run, m_length, k_length) - return target === nothing ? shape : target -end - -@inline function _demote_shape_for_run( - ::Type{T}, shape::NTuple{3, Int}, method::_MixedMethod, run::Int, m_length::Int, k_length::Int - ) where {T} - m, _, r = _real_problem(method, m_length, 0, run) - real_shape = _demote_shape_for_run(real(T), _real_shape(method, shape), RealMethod(), r, m, k_length) - return _mixed_shape(method, real_shape) -end - -function _run_demotion_target(::Type{T}, m_tile::Int, method, run::Int, m_length::Int, k_length::Int) where {T} - T <: Real || return nothing +# A register sliver takes the vectorized store only if it is unit-stride in C: +# `m_length == run || run % MR == 0` (not `MR <= run`); otherwise it takes the +# scattered one. Real `T` only, two steps in this order: +# - the MV = 4 shape steps down to MV = 2 when the run fills a half tile but +# not the tall tile's slivers (below one half tile both scatter, and the tall +# one stays); +# - at a shallow `k_length`, where the smaller kernel's cost does not dominate, +# the largest menu shape whose `MR` divides the run. +@inline fit_to_run(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int, k_length::Int, run::Int) where {T} = shape + +@inline function fit_to_run(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int, k_length::Int, run::Int) where {T} + m_length == run && return shape + MR, NR, W = shape + half = (2 * W, NR, W) + if MR == 4 * W && run % MR != 0 && run >= 2 * W && half in kernel_shapes(T, method) + shape = half + end kmax = T === Float64 ? _RUN_DEMOTE_KMAX_F64 : _RUN_DEMOTE_KMAX_F32 - k_length > kmax && return nothing - m_length == run && return nothing - run % m_tile == 0 && return nothing + (k_length > kmax || run % shape[1] == 0) && return shape best = nothing - for shape in kernel_shapes(T, method) - m = shape[1] - if run % m == 0 && (best === nothing || m > best[1]) - best = shape - end + for s in kernel_shapes(T, method) + run % s[1] == 0 && (best === nothing || s[1] > best[1]) && (best = s) end - return best + return something(best, shape) end -# Deepest `k_length` at which run-length demotion still wins. +# Deepest `k_length` at which the divisor step still wins. const _RUN_DEMOTE_KMAX_F64 = 32 const _RUN_DEMOTE_KMAX_F32 = 64 diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 30cec9f..5ee747c 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -185,7 +185,7 @@ function _planned( accumulator::AC ) where {F, NA, NB, NC, AC} T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) - method = _default_method(T, eltype(A), eltype(B)) + 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 @@ -205,7 +205,7 @@ function _planned( norder = order_free_labels(nlabels, indC, C) # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. - run_m = T <: Real || method isa _MixedMethod ? leading_unit_run(morder, indC, C) : 0 + run_m = T <: Real || method isa MixedMethod ? leading_unit_run(morder, indC, C) : 0 run_n = T <: Real ? leading_unit_run(norder, indC, C) : 0 mgroup = AxisGroup(morder, (indA, A), (indC, C)) # maps: (A, C) @@ -274,13 +274,14 @@ end ) # `m_tile` 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. +# named kernel either way, else `select_shape`'s pick at that orientation's +# extent and C run. The swap is judged before `fit_to_run`'s K-dependent +# divisor step, which an unbounded `k_length` skips. @inline _candidate_m_tiles(::Type{T}, method, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = (tile_size(kernel, 1), tile_size(kernel, 1)) @inline function _candidate_m_tiles(::Type{T}, method, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} - m_tile_asis = _default_shape(T, method, m_length, n_length, run_m)[1][1] - m_tile_swapped = T <: Real ? _default_shape(T, method, n_length, m_length, run_n)[1][1] : m_tile_asis + m_tile_asis = select_shape(T, method, m_length, typemax(Int), run_m)[1][1] + m_tile_swapped = T <: Real ? select_shape(T, method, n_length, typemax(Int), run_n)[1][1] : m_tile_asis return m_tile_asis, m_tile_swapped end @@ -293,10 +294,8 @@ end @inline _plan_with_kernel(kernel, method, atransform, btransform, req::_PlanRequest) = _plan_contract(kernel, atransform, btransform, req, nothing) @inline function _plan_with_kernel(::Nothing, method, atransform, btransform, req::_PlanRequest{T}) where {T} - m_length = axis_length(req.mgroup) - shape, method = _default_shape(T, method, m_length, axis_length(req.ngroup), req.run) - shape = _demote_shape_for_run(T, shape, method, req.run, m_length, axis_length(req.kgroup)) - vshape = _menu_val(shape, T, method) + shape, method = select_shape(T, method, axis_length(req.mgroup), axis_length(req.kgroup), req.run) + vshape = menu_val(shape, T, method) # The execution path, predicted so the callee is specialised on it. Only a # `Bool` crosses: a call union-split on `method` is emitted out of line and # boxes `req`. @@ -315,7 +314,7 @@ function _plan_resolved( slot::Base.RefValue{R}, Astorage::SA, Bstorage::SB, Cstorage::SC ) where {S, M, TA, TB, H, T, R <: _PlanRequest{T}, SA, SB, SC} req = _with_storage(slot[], Astorage, Bstorage, Cstorage) - return _plan_contract(_kernel_from_shape(S, T, method), atransform, btransform, req, hint) + return _plan_contract(kernel_from_shape(S, T, method), atransform, btransform, req, hint) end # The storages cross the barrier as arguments, so no slot retains a user array. diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index a9b1a9b..2b9310c 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -1,6 +1,6 @@ using QuasiStrided: FMAddSubKernel, FMAddSubMethod, PlanarKernel, PlanarMethod, OneMKernel, InterleavedFormat, KernelMethod, lanewidth, kernel_shapes, packed_panel, - _kernel_from_shape, _default_method, accumulator_planes, target_profile, PackedPanel + kernel_from_shape, default_method, accumulator_planes, target_profile, PackedPanel using SIMD: Vec using InteractiveUtils: code_native @@ -138,9 +138,9 @@ const _QSF = QuasiStrided bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) @test (bf.m_block, bf.k_block, bf.n_block) == (bp.m_block, bp.k_block, bp.n_block) # planar's packed reals for s in kernel_shapes(T, FMAddSubMethod()) - @test _kernel_from_shape(s, T, FMAddSubMethod()) isa FMAddSubKernel{s[1], s[2], T, s[3]} + @test kernel_from_shape(s, T, FMAddSubMethod()) isa FMAddSubKernel{s[1], s[2], T, s[3]} end - @test_throws ArgumentError _kernel_from_shape((7, 7, 7), T, FMAddSubMethod()) + @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, FMAddSubMethod()) @test !(_QSF._default_kernel(T, 1024, 1024) isa FMAddSubKernel) end end diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 9927162..9227e0d 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -1,6 +1,6 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, KernelMethod, lanewidth, kernel_shapes, - _kernel_from_shape, _default_method + kernel_from_shape, default_method @testset "OneMKernel" begin full = ((ComplexF64, (12, 8, 8)), (ComplexF32, (8, 6, 8))) @@ -38,11 +38,11 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) @test (bm.m_block, bm.k_block, bm.n_block) == (bp.m_block ÷ 2, bp.k_block, bp.n_block) # twice the packed A reals for s in kernel_shapes(T, OneMMethod()) - @test _kernel_from_shape(s, T, OneMMethod()) isa OneMKernel{s[1], s[2], T, s[3]} + @test kernel_from_shape(s, T, OneMMethod()) isa OneMKernel{s[1], s[2], T, s[3]} end - @test_throws ArgumentError _kernel_from_shape((7, 7, 7), T, OneMMethod()) + @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, OneMMethod()) # Never the default: method ranking does not transfer between machines. - @test _default_method(T) === PlanarMethod() + @test default_method(T) === PlanarMethod() @test QuasiStrided._default_kernel(T, 1024, 1024) isa PlanarKernel end end diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index cdde59a..e49c930 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -6,10 +6,10 @@ using StridedViews: StridedView @testset "unknown target: fallback shape and blocking" begin u = unknown_target() for T in (Float64, Float32) - @test _derived_shape(u, T) === _fallback_shape(T) - k = _kernel_for(u, T) + @test derived_shape(u, T) === fallback_shape(T) + k = kernel_from_shape(derived_shape(u, T), T) @test k isa SIMDKernel - @test (tile_size(k)..., lanewidth(k)) === _fallback_shape(T) + @test (tile_size(k)..., lanewidth(k)) === fallback_shape(T) @test _real_blocking_row(u, T) === _fallback_blocking(T) end @test _fallback_blocking(Float64) === Blocking(128, 256, 768) @@ -22,31 +22,31 @@ using StridedViews: StridedView @testset "rule: MR = MV*W, NR = 6; MV = 4 on Intel :avx512, 2 on :avx2 and AMD" begin for (isakey, vb, MV) in ((:avx512, 64, 4), (:avx2, 32, 2)), T in (Float64, Float32) W = vb ÷ sizeof(T) - @test _derived_shape(synthetic(isakey), T) === (MV * W, NR_DEFAULT, W) - @test QuasiStrided._rule_mv(Val(isakey), RealMethod()) == MV - @test QuasiStrided._rule_shape(vb, T, MV) === (MV * W, NR_DEFAULT, W) + @test derived_shape(synthetic(isakey), T) === (MV * W, NR_DEFAULT, W) + @test QuasiStrided.rule_mv(Val(isakey), RealMethod()) == MV + @test QuasiStrided.rule_shape(vb, T, MV) === (MV * W, NR_DEFAULT, W) end for m in (PlanarMethod(), OneMMethod(), FMAddSubMethod()), key in (:avx512, :avx2) - @test QuasiStrided._rule_mv(Val(key), m) == 2 + @test QuasiStrided.rule_mv(Val(key), m) == 2 end for (cpu, MV) in (("znver4", 2), ("znver5", 2), ("icelake-server", 4), ("cascadelake", 4)), T in (Float64, Float32) p = TargetProfile(:avx512, cpu, CacheLevel(), CacheLevel(), CacheLevel()) W = 64 ÷ sizeof(T) - @test _derived_shape(p, T) === (MV * W, NR_DEFAULT, W) - @test _derived_shape(p, ComplexF64) === _derived_shape(synthetic(:avx512), ComplexF64) + @test derived_shape(p, T) === (MV * W, NR_DEFAULT, W) + @test derived_shape(p, ComplexF64) === derived_shape(synthetic(:avx512), ComplexF64) end - @test _derived_shape(synthetic(:avx2), Float64) === _fallback_shape(Float64) + @test derived_shape(synthetic(:avx2), Float64) === fallback_shape(Float64) # ISAs without a rule get the fallback shape, whatever their width. for T in (Float64, Float32) - @test _derived_shape(synthetic(:neon), T) === _fallback_shape(T) - @test _derived_shape(synthetic(:unknown), T) === _fallback_shape(T) + @test derived_shape(synthetic(:neon), T) === fallback_shape(T) + @test derived_shape(synthetic(:unknown), T) === fallback_shape(T) end # No real override row on any ISA: the rule is the optimum. for T in (Float64, Float32), key in VALID_ISAS - @test _shape_override(Val(key), T) === nothing + @test shape_override(Val(key), T) === nothing end - @test _rule_applies(Val(:avx2), RealMethod()) && _rule_applies(Val(:avx512), RealMethod()) + @test rule_applies(Val(:avx2), RealMethod()) && rule_applies(Val(:avx512), RealMethod()) end @testset "every real menu shape is constructible and fits the register file" begin @@ -85,15 +85,15 @@ using StridedViews: StridedView MR = tile_size(_default_kernel(T), 1) @test tile_size(_default_kernel(T, 4 * MR, 256), 1) == MR small = _default_kernel(T, 1, 256) - @test (tile_size(small)..., lanewidth(small)) === _fallback_shape(T) + @test (tile_size(small)..., lanewidth(small)) === fallback_shape(T) @test tile_size(_default_kernel(T, 0, 256), 1) == MR # empty must not demote end end - @testset "_extent_shape: MV = 4 steps down to MV = 2 where it pads less" begin - ext(shape, T, m_length) = QuasiStrided._extent_shape(shape, T, RealMethod(), m_length) + @testset "extent_shape: MV = 4 steps down to MV = 2 where it pads less" begin + ext(shape, T, m_length) = QuasiStrided.extent_shape(shape, T, RealMethod(), m_length) for T in (Float64, Float32) - tall = _derived_shape(synthetic(:avx512), T) + tall = derived_shape(synthetic(:avx512), T) MR, NR, W = tall half = (2 * W, NR, W) @test ext(tall, T, 0) === tall @@ -104,54 +104,72 @@ using StridedViews: StridedView @test ext(tall, T, m_length) === tall end # Never below MV = 2, and complex methods are untouched. - avx2 = _derived_shape(synthetic(:avx2), T) + avx2 = derived_shape(synthetic(:avx2), T) @test ext(avx2, T, 3) === avx2 - cs = _derived_shape(synthetic(:avx512), ComplexF64) - @test QuasiStrided._extent_shape(cs, ComplexF64, PlanarMethod(), 3) === cs + cs = derived_shape(synthetic(:avx512), ComplexF64) + @test QuasiStrided.extent_shape(cs, ComplexF64, PlanarMethod(), 3) === cs # Through `_default_kernel`, where the host's shape is the tall one. - if _derived_shape(target_profile(), T) === tall + if derived_shape(target_profile(), T) === tall @test tile_size(_default_kernel(T, MR + W, 256)) == (2 * W, NR) @test tile_size(_default_kernel(T, 2 * W, 256), 1) == 2 * W - @test tile_size(_default_kernel(T, 2 * W - 1, 256), 1) == _fallback_shape(T)[1] + @test tile_size(_default_kernel(T, 2 * W - 1, 256), 1) == fallback_shape(T)[1] @test tile_size(_default_kernel(T, 2 * MR, 256), 1) == MR end end end - @testset "_store_shape: MV = 4 steps down where C's run breaks its slivers" begin - st = QuasiStrided._store_shape + @testset "fit_to_run: C's run along M decides the real shape" begin + fit(shape, T, m_length, k_length, run) = + QuasiStrided.fit_to_run(shape, T, RealMethod(), m_length, k_length, run) + largest_divisor(T, run) = argmax(s -> (run % s[1] == 0, s[1]), kernel_shapes(T)) + deep = 10^6 for T in (Float64, Float32) - tall = _derived_shape(synthetic(:avx512), T) + tall = derived_shape(synthetic(:avx512), T) MR, NR, W = tall half = (2 * W, NR, W) m_length = 64 * MR + # MV = 4 -> MV = 2 at any K. for run in (m_length, MR, 2 * MR, 5 * MR) # every tall sliver whole - @test st(tall, T, RealMethod(), m_length, run) === tall + @test fit(tall, T, m_length, deep, run) === tall end - @test st(tall, T, RealMethod(), 3 * W, 3 * W) === tall + @test fit(tall, T, 3 * W, deep, 3 * W) === tall for run in (2 * W, 3 * W, 6 * W, MR + 2 * W, m_length - 1) - @test st(tall, T, RealMethod(), m_length, run) === half + @test fit(tall, T, m_length, deep, run) === half end for run in (1, W, 2 * W - 1) # below one half tile - @test st(tall, T, RealMethod(), m_length, run) === tall + @test fit(tall, T, m_length, deep, run) === tall end - @test st(half, T, RealMethod(), m_length, 3 * W) === half - @test st(_fallback_shape(T), T, RealMethod(), m_length, 3 * W) === _fallback_shape(T) - cs = _derived_shape(synthetic(:avx512), ComplexF64) - @test st(cs, ComplexF64, PlanarMethod(), m_length, 3 * W) === cs - if _derived_shape(target_profile(), T) === tall - @test QuasiStrided._default_shape(T, m_length, 256, 2 * W)[1] === half - @test QuasiStrided._default_shape(T, m_length, 256, MR)[1] === tall - @test QuasiStrided._default_shape(T, m_length, 256)[1] === tall # no run given + @test fit(half, T, m_length, deep, 3 * W) === half + # Then, at shallow K, the largest menu shape whose `MR` divides the run. + @test fit(tall, T, m_length, 1, 3 * W) === largest_divisor(T, 3 * W) + kmax = T === Float64 ? QuasiStrided._RUN_DEMOTE_KMAX_F64 : QuasiStrided._RUN_DEMOTE_KMAX_F32 + for shape in kernel_shapes(T) + S = shape[1] + @test fit(shape, T, 1000, 1, 3 * S) === shape + @test fit(shape, T, 7, 1, 7) === shape + @test fit(shape, T, 1000, kmax + 1, 4) === shape + for run in (1, 4, 8, 12, 16, 20, 24, 48) + (run % S == 0 || (S == 4 * shape[3] && run >= 2 * shape[3])) && continue + want = any(s -> run % s[1] == 0, kernel_shapes(T)) ? largest_divisor(T, run) : shape + @test fit(shape, T, 1000, kmax, run) === want + end + end + cs = derived_shape(synthetic(:avx512), ComplexF64) + @test QuasiStrided.fit_to_run(cs, ComplexF64, PlanarMethod(), m_length, 1, 1) === cs + if derived_shape(target_profile(), T) === tall + select(run) = QuasiStrided.select_shape(T, RealMethod(), m_length, deep, run)[1] + @test select(2 * W) === half + @test select(MR) === tall + @test select(m_length) === tall end end end - @testset "_store_shape feeds the swap: ccsd_t_2 / ao2mo_2 at dim 16" begin + @testset "fit_to_run feeds the swap: ccsd_t_2 / ao2mo_2 at dim 16" begin # Regression: C's only unit run is 16 long, on the N side. Judged # against m_tile = 32 the swap was declined and every tile stored scattered. T = Float64 - if _derived_shape(target_profile(), T)[1] == 32 + if derived_shape(target_profile(), T)[1] == 32 d = 16 a, b, c, i, j, k, m = 1, 2, 3, 4, 5, 6, -1 Bv = StridedView(randn(T, d, d, d, d)) @@ -177,35 +195,38 @@ end @testset "complex shape selection" begin for (T, W, swept) in ((ComplexF64, 8, (24, 3, 8)), (ComplexF32, 16, (48, 3, 16))) # W counts REAL lanes; the AVX-512 override replaces the spilling rule shape. - @test QuasiStrided._rule_shape(64, T, 2) === (2 * W, NR_DEFAULT, W) - @test _shape_override(Val(:avx512), T) === swept - @test _derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarMethod())) + @test QuasiStrided.rule_shape(64, T, 2) === (2 * W, NR_DEFAULT, W) + @test shape_override(Val(:avx512), T) === swept + @test derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarMethod())) @test Set(kernel_shapes(T, PlanarMethod())) == Set( ( - (2 * W, NR_DEFAULT, W), swept, (W, 8, W), _shape_override(Val(:avx2), T), - _shape_override(Val(:neon), T), (W ÷ 4, NR_DEFAULT, W ÷ 4), + (2 * W, NR_DEFAULT, W), swept, (W, 8, W), shape_override(Val(:avx2), T), + (W ÷ 2, NR_DEFAULT, W ÷ 4), (W ÷ 4, NR_DEFAULT, W ÷ 4), ) ) # Every override row fits its ISA's register file with room to spare. - for (key, nreg) in ((:avx512, 32), (:avx2, 16), (:neon, 32)) - ovr = _shape_override(Val(key), T) - @test _planar_pressure(ovr...) < nreg + for (key, nreg) in ((:avx512, 32), (:avx2, 16)) + ovr = shape_override(Val(key), T) + @test planar_pressure(ovr...) < nreg @test ovr in kernel_shapes(T, PlanarMethod()) end - @test _fallback_shape(T) === (8, NR_DEFAULT, _fallback_shape(real(T))[3]) + @test fallback_shape(T) === (8, NR_DEFAULT, fallback_shape(real(T))[3]) # Off :avx512 an override row wins whatever the width; without a row # the shape is fitted to the register budget. - for key in (:avx2, :neon) - @test !_rule_applies(Val(key), PlanarMethod()) - @test _derived_shape(synthetic(key), T) === _shape_override(Val(key), T) + @test derived_shape(synthetic(:avx2), T) === shape_override(Val(:avx2), T) + for key in (:avx2, :neon, :unknown) + @test !rule_applies(Val(key), PlanarMethod()) + end + for key in (:neon, :unknown) + @test shape_override(Val(key), T) === nothing end - @test !_rule_applies(Val(:unknown), PlanarMethod()) - @test _shape_override(Val(:unknown), T) === nothing - @test _planar_pressure(_derived_shape(synthetic(:unknown), T)...) <= 16 + @test planar_pressure(derived_shape(synthetic(:unknown), T)...) <= 16 end - @test _rule_applies(Val(:avx512), PlanarMethod()) - @test _planar_pressure(_fallback_shape(ComplexF64)...) == 30 - @test _planar_pressure(_fallback_shape(ComplexF32)...) == 16 + @test rule_applies(Val(:avx512), PlanarMethod()) + @test derived_shape(synthetic(:neon), ComplexF64) === (4, 6, 2) + @test derived_shape(synthetic(:neon), ComplexF32) === (8, 6, 4) + @test planar_pressure(fallback_shape(ComplexF64)...) == 30 + @test planar_pressure(fallback_shape(ComplexF32)...) == 16 # Complex menus fit AVX-512 and stay bounded. nreg = isa_nregisters(:avx512) @@ -228,14 +249,14 @@ 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 + 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) + k = QuasiStrided.kernel_from_shape(shape, T, m) @test KernelMethod(k) === m && (tile_size(k)..., lanewidth(k)) === shape end end @@ -244,15 +265,16 @@ end @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. + # Every real step carries over. saved = QuasiStrided.TARGET[] try for key in (:avx512, :avx2, :neon) QuasiStrided.TARGET[] = synthetic(key) - for T in (ComplexF64, ComplexF32), m in (CR, RC), m_length in (1, 3, 17, 40, 4096), run in (0, 1, 8, m_length) - real_problem = m === CR ? (2 * m_length, 50, 2 * run) : (m_length, 100, run) - @test QuasiStrided._default_shape(T, m, m_length, 50, run) === - (mapped(m, QuasiStrided._default_shape(real(T), real_problem...)[1]), m) + for T in (ComplexF64, ComplexF32), m in (CR, RC), m_length in (1, 3, 17, 40, 4096), + run in (0, 1, 8, m_length), k_length in (1, 1000) + real_m, real_run = m === CR ? (2 * m_length, 2 * run) : (m_length, run) + real_shape = QuasiStrided.select_shape(real(T), RealMethod(), real_m, k_length, real_run)[1] + @test QuasiStrided.select_shape(T, m, m_length, k_length, run) === (mapped(m, real_shape), m) end end finally @@ -260,15 +282,15 @@ end end end -@testset "_derived_shape is always a constructible member of the method's menu" begin +@testset "derived_shape is always a constructible member of the method's menu" begin methods = ( (Float64, RealMethod()), (Float32, RealMethod()), (ComplexF64, PlanarMethod()), (ComplexF32, PlanarMethod()), (ComplexF64, OneMMethod()), (ComplexF32, OneMMethod()), ) for (T, m) in methods, isakey in VALID_ISAS - shape = _derived_shape(synthetic(isakey), T, m) + shape = derived_shape(synthetic(isakey), T, m) @test shape in kernel_shapes(T, m) - k = QuasiStrided._kernel_from_shape(shape, T, m) + k = QuasiStrided.kernel_from_shape(shape, T, m) @test (tile_size(k)..., lanewidth(k)) === shape end end @@ -290,7 +312,7 @@ end @test _modelled_blocking(nosmt, Float64) === _modelled_blocking(nosmt, Float64, 8, 6) for p in (nosmt, smt2), T in (Float64, Float32) b = _modelled_blocking(p, T) - MR, NR, _ = _derived_shape(p, T) + MR, NR, _ = derived_shape(p, T) @test b.m_block % MR == 0 && b.n_block % NR == 0 @test NR * b.k_block * sizeof(T) <= p.l1d.bytes ÷ 2 @test _real_blocking_row(p, T) === b @@ -350,10 +372,10 @@ end kernel = _default_kernel(T) @test kernel isa QuasiStrided.PlanarKernel @test scalartype(kernel) === T && realtype(kernel) === real(T) - @test KernelMethod(kernel) === PlanarMethod() === QuasiStrided._default_method(T) + @test KernelMethod(kernel) === PlanarMethod() === QuasiStrided.default_method(T) shape = (tile_size(kernel)..., lanewidth(kernel)) profile = target_profile() - @test _planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) + @test planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarMethod())) @test _default_kernel(T, 1024, 1024) isa QuasiStrided.PlanarKernel # The small-M demotion: FMAddSub on AVX-512, planar elsewhere. @@ -361,12 +383,12 @@ end @test small isa (profile.isa === :avx512 ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) shape1m = first(kernel_shapes(T, OneMMethod())) - k1m = QuasiStrided._kernel_from_shape(shape1m, T, OneMMethod()) + k1m = QuasiStrided.kernel_from_shape(shape1m, T, OneMMethod()) @test k1m isa QuasiStrided.OneMKernel && KernelMethod(k1m) === OneMMethod() @test (tile_size(k1m)..., lanewidth(k1m)) === shape1m # A method with no kernel for `T` throws, naming itself. err = try - QuasiStrided._kernel_from_shape(shape1m, T, RealMethod()) + QuasiStrided.kernel_from_shape(shape1m, T, RealMethod()) catch e e end @@ -374,17 +396,17 @@ end @test occursin("RealMethod", err.msg) && occursin(string(T), err.msg) end for T in (Float16, Int, ComplexF16) - @test_throws ArgumentError QuasiStrided._kernel_from_shape((8, 6, 4), T) + @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 4), T) end - @test_throws ArgumentError QuasiStrided._kernel_from_shape((8, 6, 4), Float64, OneMMethod()) - @test_throws ArgumentError QuasiStrided._kernel_from_shape((8, 6, 5), Float64) - @test_throws ArgumentError QuasiStrided._menu_val((8, 6, 5), Float64, RealMethod()) - @test QuasiStrided._menu_val((8, 6, 4), Float64, RealMethod()) === Val((8, 6, 4)) + @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 4), Float64, OneMMethod()) + @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 5), Float64) + @test_throws ArgumentError QuasiStrided.menu_val((8, 6, 5), Float64, RealMethod()) + @test QuasiStrided.menu_val((8, 6, 4), Float64, RealMethod()) === Val((8, 6, 4)) end -@testset "_small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin +@testset "small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin sms(T, m_length, p = synthetic(:avx512)) = - QuasiStrided._small_m_shape(QuasiStrided._small_m_candidates(Val(p.isa), p, T), m_length) + QuasiStrided.small_m_shape(QuasiStrided.small_m_candidates(Val(p.isa), p, T), m_length) # Least padded rows, then the larger tile. @test sms(ComplexF64, 12) === (12, 8, 8) @test sms(ComplexF64, 16) === (8, 8, 8) @@ -401,25 +423,3 @@ end @test sms(T, 2, synthetic(isakey)) === nothing end end - -@testset "_demote_shape_for_run: largest menu shape whose m_tile divides the run" begin - demote = QuasiStrided._demote_shape_for_run - for T in (Float64, Float32), shape in kernel_shapes(T) - MR = shape[1] - @test demote(T, shape, RealMethod(), 3 * MR, 1000, 1) === shape - @test demote(T, shape, RealMethod(), 7, 7, 1) === shape - kmax = T === Float64 ? QuasiStrided._RUN_DEMOTE_KMAX_F64 : QuasiStrided._RUN_DEMOTE_KMAX_F32 - @test demote(T, shape, RealMethod(), 1, 1000, kmax + 1) === shape - for run in (1, 4, 8, 12, 16, 20, 24, 48) - run % MR == 0 && continue - fits = [s for s in kernel_shapes(T) if run % s[1] == 0] - want = isempty(fits) ? shape : fits[argmax(first.(fits))] - @test demote(T, shape, RealMethod(), run, 1000, 1) === want - end - end - # Complex kernels are never demoted. - for T in (ComplexF64, ComplexF32) - shape = first(kernel_shapes(T, PlanarMethod())) - @test demote(T, shape, PlanarMethod(), 1, 1000, 1) === shape - end -end diff --git a/test/runtests.jl b/test/runtests.jl index fe16056..8dea1d5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,12 +9,12 @@ using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, zero_accumulator, add_tile, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, target_profile, cache_topology, unknown_target, detect_isa, detect_target, - _derived_shape, _fallback_shape, _shape_override, _kernel_for, _default_kernel, + derived_shape, fallback_shape, shape_override, _default_kernel, kernel_from_shape, _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, - _rule_applies, isa_nregisters, sliver_width, realtype, + rule_applies, isa_nregisters, sliver_width, realtype, KernelMethod, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, pack_formats, reals_per_element, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, - _planar_pressure, _pack_split, _NestPath + 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. From e85ea7ac2f5ed1dbf0906342e6fb2f820ceb1af3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 13:28:06 -0400 Subject: [PATCH 22/60] Kernel types replace KernelMethod The unparameterised kernel type (`SIMDKernel`, `PlanarKernel`, ...) is now the shape-free name of a complex-arithmetic scheme. `KernelMethod`, its six singletons, `KernelMethod(kernel)` and `_kernel_type` are gone; `default_kernel_type(T[, TA, TB])` replaces `default_method`. Menus, pack formats, accumulator planes, blocking scaling, the shape rules, the mixed mapping (`MixedKernel`), unpacked-B eligibility and the menu ladder dispatch on `::Type{<:PlanarKernel}` etc.; a pair without a menu has the empty menu. The per-scheme summary moves to the `pack_formats` docstring in kernels.jl. ScalarKernel shares SIMDKernel's traits, so `default_blocking` works on it. With eligibility predicted from the kernel type itself, `_hint_holds` loses its eligibility clause. Two places keep a type out of a value position, since a plain `Type` regressed the per-call floor: - the plan barrier receives `Val{}()` (from `menu_val`) instead of `Val(shape)` plus the method: a `Type` argument misses the dynamic dispatch's fast cache (+350-430 ns per call measured with the bare type); - `select_shape` returns the kernel type as a `Val`: inference widens a type in a returned tuple to `UnionAll`, which made `menu_val` a dynamic call that boxed the shape (32 B per complex plan); recovering the Union by an `isa` assertion or `===` on types cost a subtype check or structural egal (+30-100 ns). Behaviour-neutral: the automatic kernel type, shape, blocking and orientation are identical to 7032aed on the same 138270-plan grid as the previous commit, and the shape pipeline on the dense 579040-case grid. Per-call floor (medians of 3 interleaved runs vs 7032aed, ns): Float64 8^3 372 -> 369, 16^3 484 -> 474, d4 outer 352 -> 341; ComplexF64 8^3 426 -> 422, 16^3 957 -> 954, d4 outer 586 -> 575; 0 B. TTFX: first call Float64 3.7-3.8 -> 3.7 s, ComplexF64 2.4 -> 2.3 s. Tests: 32190 pass (redundant `KernelMethod(k) === ...` checks removed, one ScalarKernel blocking check added), 34 min. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_blocking_model.jl | 4 +- benchmark/bench_kernels.jl | 10 +- benchmark/bench_mixed.jl | 4 +- src/execution/execute.jl | 6 +- src/execution/unpackedb.jl | 15 +- src/microkernels/interface.jl | 41 +--- src/microkernels/kernels.jl | 39 ++- src/planning/blocking.jl | 6 +- src/planning/defaults.jl | 16 +- src/planning/kernel_selection.jl | 235 +++++++++---------- src/planning/plan.jl | 55 +++-- test/microkernels/test_fmaddsub_kernel.jl | 19 +- test/microkernels/test_onem_kernel.jl | 16 +- test/microkernels/test_planar_kernel.jl | 3 +- test/packing/test_pack_complex_contiguous.jl | 16 +- test/planning/test_kernel_selection.jl | 106 ++++----- test/runtests.jl | 4 +- 17 files changed, 286 insertions(+), 309 deletions(-) diff --git a/benchmark/bench_blocking_model.jl b/benchmark/bench_blocking_model.jl index aea7b2d..6ee6ec0 100644 --- a/benchmark/bench_blocking_model.jl +++ b/benchmark/bench_blocking_model.jl @@ -14,7 +14,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: Blocking, target_profile, _default_kernel, _modelled_blocking, - _fallback_blocking, _scale_blocking, KernelMethod + _fallback_blocking, _scale_blocking const SMOKE = hasflag("smoke") const REPS = SMOKE ? 1 : argopt("reps", 21) @@ -132,7 +132,7 @@ for T in DTYPES end for T in CDTYPES kernel = _default_kernel(T) - rows = map(v -> v === nothing ? nothing : _scale_blocking(v, KernelMethod(kernel)), named_rows(real(T))) + rows = map(v -> v === nothing ? nothing : _scale_blocking(v, typeof(kernel)), named_rows(real(T))) println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) for spec in thin(MAIN_SHAPES) sweep_shape!(raw, kernel, T, spec, named_points(rows)) diff --git a/benchmark/bench_kernels.jl b/benchmark/bench_kernels.jl index 5c6d1c0..9ee0833 100644 --- a/benchmark/bench_kernels.jl +++ b/benchmark/bench_kernels.jl @@ -21,7 +21,7 @@ include(joinpath(@__DIR__, "harness.jl")) -using QuasiStrided: RealMethod, PlanarMethod, OneMMethod, FMAddSubMethod, kernel_shapes, +using QuasiStrided: SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel, kernel_shapes, kernel_from_shape, default_blocking, packed_panel, execute_tile!, Tile, AffineAxis, target_profile @@ -36,7 +36,7 @@ const CSV_PATH = joinpath(OUTDIR, "bench_kernels.csv") const SUMMARY_PATH = joinpath(OUTDIR, "summary_kernels.txt") tag(k) = "$(tile_size(k, 1))x$(tile_size(k, 2))/W$(lanewidth(k))" -menu_methods(::Type{T}) where {T} = T <: Complex ? (PlanarMethod(), OneMMethod(), FMAddSubMethod()) : (RealMethod(),) +menu_kernels(::Type{T}) where {T} = T <: Complex ? (PlanarKernel, OneMKernel, FMAddSubKernel) : (SIMDKernel,) csv = open(CSV_PATH, "w") println(csv, "arm,dtype,shape,M,K,N,config,kernel,m_block,k_block,n_block,seconds,gflops,frac_openblas") @@ -91,11 +91,11 @@ for spec in SHAPES, T in RUN_DTYPES td, pd = time_plan(fx, T) @assert isapprox(fx.Cmat, fx.Amat * fx.Bmat; rtol = sqrt(eps(real(T)))) record!("engine", T, spec, "default", tag(pd.kernel), pd.blocking, td, tb) - for m in menu_methods(T), sh in kernel_shapes(T, m) - k = kernel_from_shape(sh, T, m) + for K in menu_kernels(T), sh in kernel_shapes(T, K) + k = kernel_from_shape(sh, T, K) b = default_blocking(k) t, _ = time_plan(fx, T; kernel = k, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block) - record!("engine", T, spec, lowercase(replace(string(nameof(typeof(m))), "Method" => "")), tag(k), b, t, tb) + record!("engine", T, spec, lowercase(replace(string(nameof(K)), "Kernel" => "")), tag(k), b, t, tb) end println(" ", rpad(string(T), 11), rpad(spec.name, 22), @sprintf("default %7.2f GF/s %.2fx OpenBLAS", gflops(T, spec.Ma, spec.Ka, spec.Na, td), tb / td)) end diff --git a/benchmark/bench_mixed.jl b/benchmark/bench_mixed.jl index 04f07c0..a3c4845 100644 --- a/benchmark/bench_mixed.jl +++ b/benchmark/bench_mixed.jl @@ -9,7 +9,7 @@ include(joinpath(@__DIR__, "harness.jl")) -using QuasiStrided: ComplexRealKernel, RealComplexKernel, PlanarMethod, target_profile, +using QuasiStrided: ComplexRealKernel, RealComplexKernel, PlanarKernel, target_profile, derived_shape, kernel_from_shape const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) @@ -30,7 +30,7 @@ function run(io) for T in RUN_DTYPES, n in SIZES R = real(T) MRr, NRr, W = derived_shape(target_profile(), R) - planar = kernel_from_shape(derived_shape(target_profile(), T), T, PlanarMethod()) + planar = kernel_from_shape(derived_shape(target_profile(), T), T, PlanarKernel) cases = ( ( "CR", ComplexRealKernel(Val(MRr ÷ 2), Val(NRr), T, Val(W)), diff --git a/src/execution/execute.jl b/src/execution/execute.jl index a9697a4..bbcf4d9 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -174,13 +174,11 @@ end U, plan.mgroup, plan.ngroup, plan.kgroup, _is_split(plan.mpack), _is_split(plan.npack), true ) -# The prediction differs from `_select_path(plan)` in three inputs only: the -# kernel's unpacked-B eligibility (predicted from its method; folds), the +# The prediction differs from `_select_path(plan)` in two inputs only: the # dot path's workspace capacity (assumed) and the panel decision (made at the # default `k_block`; 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.k_block) === (path isa _PanelPath) && - _unpacked_b_method_eligible(KernelMethod(plan.kernel)) === _unpacked_b_kernel_eligible(plan.kernel) + _c_panel_needed(T, plan.Cstorage, axis_length(plan.kgroup), plan.blocking.k_block) === (path isa _PanelPath) @inline _hint_holds(plan::ContractPlan, ::_DotPath) = _dot_capacity_ok(plan) @inline _hint_holds(plan::ContractPlan, ::_OuterPath) = true diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 5f5536d..e7588f7 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -36,16 +36,11 @@ end end # The kernels whose K step reads B through `b_scalar`/`b_complex` -# (not 1m, whose inner kernel walks the planar panel, nor the scalar one). -@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, ComplexRealMethod}) = true -@inline _unpacked_b_method_eligible(::Any) = false +# (not 1m, whose inner kernel walks the planar panel, nor the scalar one); by +# type, to predict the path before the kernel exists. +@inline _unpacked_b_kernel_eligible(::Type{<:Union{SIMDKernel, PlanarKernel, FMAddSubKernel, ComplexRealKernel}}) = true +@inline _unpacked_b_kernel_eligible(::Type) = false +@inline _unpacked_b_kernel_eligible(kernel::Microkernel) = _unpacked_b_kernel_eligible(typeof(kernel)) # `:always`/`:never` override the rule below, for benchmarks and tests. const _UNPACKED_B_MODE = Ref{Symbol}(:auto) diff --git a/src/microkernels/interface.jl b/src/microkernels/interface.jl index e06fde0..647d2b4 100644 --- a/src/microkernels/interface.jl +++ b/src/microkernels/interface.jl @@ -1,5 +1,5 @@ -# The microkernel contract, the complex-arithmetic methods, and the validation -# and store helpers every kernel shares. +# The microkernel contract, and the validation and store helpers every kernel +# shares. """ Microkernel{MR, NR, T} @@ -14,43 +14,6 @@ never bitwise. """ abstract type Microkernel{MR, NR, T} end -""" - KernelMethod(kernel) - -How a kernel maps its multiply-adds onto real arithmetic, as singletons so -blocking and the shape menus dispatch on it: - - - `RealMethod`: a real kernel. - - `PlanarMethod`: split re/im planes, 4 real FMAs per complex MAC. - - `OneMMethod`: Van Zee's 1m, a real `2MR x NR` kernel. - - `FMAddSubMethod`: interleaved A, x86 `vfmaddsub`. - - `ComplexRealMethod`: complex A, real B; the real kernel on `2MR` rows. - - `RealComplexMethod`: real A, complex B; the real kernel on `2NR` columns. - -No method ranking is hardcoded: planar is the default, 1m and fmaddsub are used -only when named (plus fmaddsub for the AVX-512 small-M demotion, -`small_m_shape` in src/planning/kernel_selection.jl). -""" -abstract type KernelMethod end -struct RealMethod <: KernelMethod end -struct PlanarMethod <: KernelMethod end -struct OneMMethod <: KernelMethod end -struct FMAddSubMethod <: KernelMethod end -struct ComplexRealMethod <: KernelMethod end -struct RealComplexMethod <: KernelMethod end - -# The (A, B) packing formats of every kernel of a method. -pack_formats(::RealMethod) = (RealFormat(), RealFormat()) -pack_formats(::PlanarMethod) = (PlanarFormat(), PlanarFormat()) -pack_formats(::OneMMethod) = (OneEFormat(), PlanarFormat()) -pack_formats(::FMAddSubMethod) = (InterleavedFormat(), PlanarFormat()) -pack_formats(::ComplexRealMethod) = (InterleavedFormat(), RealFormat()) -pack_formats(::RealComplexMethod) = (RealFormat(), InterleavedFormat()) - -# Accumulator planes held live: planar keeps separate re/im planes. -accumulator_planes(::KernelMethod) = 1 -accumulator_planes(::PlanarMethod) = 2 - realtype(k::Microkernel) = realtype(k.descriptor) sliver_width(k::Microkernel) = sliver_width(k.descriptor) a_format(k::Microkernel) = a_format(k.descriptor) diff --git a/src/microkernels/kernels.jl b/src/microkernels/kernels.jl index f3c759b..6d030a4 100644 --- a/src/microkernels/kernels.jl +++ b/src/microkernels/kernels.jl @@ -159,13 +159,34 @@ struct RealComplexKernel{MR, NR, T, W} <: VectorKernel{MR, NR, T, W} end end -KernelMethod(::Type{<:Union{ScalarKernel, SIMDKernel}}) = RealMethod() -KernelMethod(::Type{<:PlanarKernel}) = PlanarMethod() -KernelMethod(::Type{<:OneMKernel}) = OneMMethod() -KernelMethod(::Type{<:FMAddSubKernel}) = FMAddSubMethod() -KernelMethod(::Type{<:ComplexRealKernel}) = ComplexRealMethod() -KernelMethod(::Type{<:RealComplexKernel}) = RealComplexMethod() -KernelMethod(kernel::Microkernel) = KernelMethod(typeof(kernel)) +""" + pack_formats(K::Type{<:Microkernel}) -> (a_format, b_format) + +The packing formats of every kernel of type `K`. A kernel type is also the +shape-free name of its complex-arithmetic scheme, on which blocking and the +shape menus dispatch: + + - `SIMDKernel` (and `ScalarKernel`): real. + - `PlanarKernel`: split re/im planes, 4 real FMAs per complex MAC. + - `OneMKernel`: Van Zee's 1m, a real `2MR x NR` kernel. + - `FMAddSubKernel`: interleaved A, x86 `vfmaddsub`. + - `ComplexRealKernel`: complex A, real B; the real kernel on `2MR` rows. + - `RealComplexKernel`: real A, complex B; the real kernel on `2NR` columns. + +No ranking is hardcoded: planar is the complex default, 1m and fmaddsub are +used only when named (plus fmaddsub for the AVX-512 small-M demotion, +`small_m_shape` in src/planning/kernel_selection.jl). +""" +pack_formats(::Type{<:Union{ScalarKernel, SIMDKernel}}) = (RealFormat(), RealFormat()) +pack_formats(::Type{<:PlanarKernel}) = (PlanarFormat(), PlanarFormat()) +pack_formats(::Type{<:OneMKernel}) = (OneEFormat(), PlanarFormat()) +pack_formats(::Type{<:FMAddSubKernel}) = (InterleavedFormat(), PlanarFormat()) +pack_formats(::Type{<:ComplexRealKernel}) = (InterleavedFormat(), RealFormat()) +pack_formats(::Type{<:RealComplexKernel}) = (RealFormat(), InterleavedFormat()) + +# Accumulator planes held live: planar keeps separate re/im planes. +accumulator_planes(::Type{<:Microkernel}) = 1 +accumulator_planes(::Type{<:PlanarKernel}) = 2 AccumulatorLayout(::Type{<:SIMDKernel}) = RealLayout() AccumulatorLayout(::Type{<:Union{PlanarKernel, RealComplexKernel}}) = SplitLayout() @@ -180,7 +201,9 @@ accumulator_length(::SplitLayout, MV::Int, NR::Int) = 2 * MV * NR acc_index(MV::Int, v::Int, j::Int) = v + MV * (j - 1) + 1 (::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}, ::Val{W}) where {K <: VectorKernel, MR, NR, T, W} = - K{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T, pack_formats(KernelMethod(K))...)) + K{MR, NR, T, W}(Descriptor(Val(MR), Val(NR), T, pack_formats(K)...)) +(::Type{K})() where {MR, NR, T, W, K <: VectorKernel{MR, NR, T, W}} = + K(Descriptor(Val(MR), Val(NR), T, pack_formats(K)...)) (::Type{K})(::Val{MR}, ::Val{NR}, ::Type{T}) where {K <: VectorKernel, MR, NR, T} = K(Val(MR), Val(NR), T, Val(default_lanewidth(real(T)))) diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index dfdb631..d7a1661 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -20,10 +20,10 @@ struct Blocking end # Complex blocking is the real row divided by the packed reals per element of -# each operand, so every method gets the same packed BYTE budget (1m's +# each operand, so every kernel type gets the same packed BYTE budget (1m's # `m_block` is half planar's). -@inline function _scale_blocking(base::Blocking, m::KernelMethod) - a_reals, b_reals = map(reals_per_element, pack_formats(m)) +@inline function _scale_blocking(base::Blocking, ::Type{K}) where {K <: Microkernel} + a_reals, b_reals = map(reals_per_element, pack_formats(K)) return Blocking(max(1, base.m_block ÷ a_reals), base.k_block, max(1, base.n_block ÷ b_reals)) end diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl index ae6af77..5407e7a 100644 --- a/src/planning/defaults.jl +++ b/src/planning/defaults.jl @@ -26,10 +26,10 @@ end _defaults_slot(T) === nothing ? l2_core_bytes(target_profile()) : _resolved_defaults(T).l2_core function _resolve_defaults(profile::TargetProfile, ::Type{T}) where {T} - method = default_method(T) + K = default_kernel_type(T) # First, so an element type with no menu throws from here. - shape = derived_shape(profile, T, method) - fitted = fitted_shape(profile, T, method) + shape = derived_shape(profile, T, K) + fitted = fitted_shape(profile, T, K) small_m = small_m_candidates(Val(profile.isa), profile, T) real_row = _real_blocking_row(profile, real(T)) l2_core = l2_core_bytes(profile) @@ -67,12 +67,12 @@ end end _default_kernel(::Type{T}) where {T} = - kernel_from_shape(_resolved_defaults(T).shape, T, default_method(T)) + kernel_from_shape(_resolved_defaults(T).shape, T, default_kernel_type(T)) # `@noinline`: the return type is the Union of `T`'s menu kernels. @noinline function _default_kernel(::Type{T}, m_length::Int, n_length::Int) where {T} - shape, method = select_shape(T, default_method(T), m_length, 0, m_length) - return kernel_from_shape(shape, T, method) + shape, kernel_type = select_shape(T, default_kernel_type(T), m_length, 0, m_length) + return kernel_from_shape(shape, T, kernel_type) end """ @@ -80,11 +80,11 @@ end Cache-blocking factors for `kernel` on this host: an analytical model of the detected cache geometry ([`target_profile`](@ref)), or fixed constants where -L1d or L2 is undetected, scaled for complex methods by packed reals per +L1d or L2 is undetected, scaled for complex kernels by packed reals per element. `plan_contract` rounds `m_block`/`n_block` to multiples of the tile size and clamps all three to the contraction's extents. """ function default_blocking(kernel) T = scalartype(kernel) - return _scale_blocking(_resolved_defaults(T).real_row, KernelMethod(kernel)) + return _scale_blocking(_resolved_defaults(T).real_row, typeof(kernel)) end diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index bdf5792..81f5674 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -3,21 +3,22 @@ # src/planning/defaults.jl caches the host-dependent results per eltype, since # the `Val(profile.isa)` dispatch below is dynamic. -# `OneMMethod`/`FMAddSubMethod` are chosen only by naming the kernel (FMAddSub +# Kernel types, unparameterised, name the complex-arithmetic schemes here. +# `OneMKernel`/`FMAddSubKernel` are chosen only by naming the kernel (FMAddSub # also by the AVX-512 small-M demotion, `small_m_shape`). -default_method(::Type{<:Real}) = RealMethod() -default_method(::Type{<:Complex}) = PlanarMethod() +default_kernel_type(::Type{<:Real}) = SIMDKernel +default_kernel_type(::Type{<:Complex}) = PlanarKernel -const MixedMethod = Union{ComplexRealMethod, RealComplexMethod} +const MixedKernel = Union{Type{<:ComplexRealKernel}, Type{<:RealComplexKernel}} -# 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() +# The kernel type for compute type `T` and (oriented) operand eltypes: a +# mixed-domain kernel when exactly one operand is real. +default_kernel_type(::Type{T}, ::Type{TA}, ::Type{TB}) where {T, TA, TB} = + default_kernel_type(T) +default_kernel_type(::Type{T}, ::Type{<:Complex}, ::Type{<:Real}) where {T <: Complex} = + ComplexRealKernel +default_kernel_type(::Type{T}, ::Type{<:Real}, ::Type{<:Complex}) where {T <: Complex} = + RealComplexKernel const NR_DEFAULT = 6 @@ -45,80 +46,74 @@ const KERNEL_SHAPES_C64_FMADDSUB = ((12, 8, 8), (8, 8, 8), (4, 6, 4), (4, 5, 4)) const KERNEL_SHAPES_C32_FMADDSUB = ((24, 8, 16), (16, 8, 16), (8, 6, 8), (8, 5, 8)) """ - kernel_shapes(T, method::KernelMethod = default_method(T)) -> NTuple{<:Any,NTuple{3,Int}} + kernel_shapes(T, K = default_kernel_type(T)) -> NTuple{<:Any,NTuple{3,Int}} The closed menu of `(MR, NR, W)` register shapes the engine may build for -element type `T` under `method`. +element type `T` with kernel type `K` (empty where it builds none). """ -kernel_shapes(::Type{T}) where {T} = kernel_shapes(T, default_method(T)) -kernel_shapes(::Type{Float64}, ::RealMethod) = KERNEL_SHAPES_F64 -kernel_shapes(::Type{Float32}, ::RealMethod) = KERNEL_SHAPES_F32 -kernel_shapes(::Type{ComplexF64}, ::PlanarMethod) = KERNEL_SHAPES_C64_PLANAR -kernel_shapes(::Type{ComplexF64}, ::OneMMethod) = KERNEL_SHAPES_C64_ONEM -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 +kernel_shapes(::Type{T}) where {T} = kernel_shapes(T, default_kernel_type(T)) +kernel_shapes(::Type, ::Type) = () +kernel_shapes(::Type{Float64}, ::Type{<:SIMDKernel}) = KERNEL_SHAPES_F64 +kernel_shapes(::Type{Float32}, ::Type{<:SIMDKernel}) = KERNEL_SHAPES_F32 +kernel_shapes(::Type{ComplexF64}, ::Type{<:PlanarKernel}) = KERNEL_SHAPES_C64_PLANAR +kernel_shapes(::Type{ComplexF64}, ::Type{<:OneMKernel}) = KERNEL_SHAPES_C64_ONEM +kernel_shapes(::Type{ComplexF32}, ::Type{<:PlanarKernel}) = KERNEL_SHAPES_C32_PLANAR +kernel_shapes(::Type{ComplexF32}, ::Type{<:OneMKernel}) = KERNEL_SHAPES_C32_ONEM +kernel_shapes(::Type{ComplexF64}, ::Type{<:FMAddSubKernel}) = KERNEL_SHAPES_C64_FMADDSUB +kernel_shapes(::Type{ComplexF32}, ::Type{<:FMAddSubKernel}) = 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())) +kernel_shapes(::Type{T}, K::MixedKernel) where {T <: Union{ComplexF32, ComplexF64}} = + map(s -> mixed_shape(K, s), kernel_shapes(real(T), SIMDKernel)) # A real inner shape to the mixed kernel's `(MR, NR, W)`: 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) +mixed_shape(::Type{<:ComplexRealKernel}, (MR, NR, W)::NTuple{3, Int}) = (MR ÷ 2, NR, W) +mixed_shape(::Type{<:RealComplexKernel}, (MR, NR, W)::NTuple{3, Int}) = (MR, NR ÷ 2, W) # The real problem the inner kernel sees: `m_length` and C's unit-stride M run in reals. -real_problem(::ComplexRealMethod, m_length::Int, run::Int) = (2 * m_length, 2 * run) -real_problem(::RealComplexMethod, m_length::Int, run::Int) = (m_length, run) - -# The kernel type implementing `method` for `T`, or `nothing` (exactly the -# pairs with a menu). -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 +real_problem(::Type{<:ComplexRealKernel}, m_length::Int, run::Int) = (2 * m_length, 2 * run) +real_problem(::Type{<:RealComplexKernel}, m_length::Int, run::Int) = (m_length, run) # An unrolled `shape === menu[i] ? ... :` ladder over the menu whose branches # are built from literals, so each is concrete; a shape outside the menu throws. -# Each branch builds the kernel, or `Val(shape)` when `valonly`. -function menu_ladder(::Type{T}, method, valonly::Bool) where {T} - K = kernel_type(method, T) - K === nothing && return :(throw_no_kernel(shape, T, method)) - ex = :(throw_shape_not_in_menu(shape, T, method)) - for (MR, NR, W) in reverse(kernel_shapes(T, method)) - arm = valonly ? :(Val(($MR, $NR, $W))) : :($K(Val($MR), Val($NR), T, Val($W))) +# Each branch builds the kernel, or a `Val` of its concrete type when `valonly`. +function menu_ladder(::Type{T}, ::Type{K}, valonly::Bool) where {T, K} + menu = kernel_shapes(T, K) + isempty(menu) && return :(throw_no_kernel(shape, T, K)) + ex = :(throw_shape_not_in_menu(shape, T, K)) + for (MR, NR, W) in reverse(menu) + arm = valonly ? :(Val{$(K{MR, NR, T, W})}()) : :(K(Val($MR), Val($NR), T, Val($W))) ex = :(shape === ($MR, $NR, $W) ? $arm : $ex) end return ex end -# The `method` kernel for `T` at a menu `shape` (one dynamic dispatch per plan). +# The kernel of type `K` for `T` at a menu `shape` (one dynamic dispatch per plan). kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = - kernel_from_shape(shape, T, default_method(T)) -@generated kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = - menu_ladder(T, M.instance, false) - -# `Val(shape)` for a menu `shape`: a singleton, so it crosses `plan_contract`'s -# kernel barrier without boxing. -@generated menu_val(shape::Tuple{Int, Int, Int}, ::Type{T}, method::M) where {T, M} = - menu_ladder(T, M.instance, true) - -@noinline throw_no_kernel(shape, ::Type{T}, method) where {T} = throw( + kernel_from_shape(shape, T, default_kernel_type(T)) +@generated kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, ::Type{K}) where {T, K} = + menu_ladder(T, K, false) +kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, ::Val{K}) where {T, K} = + kernel_from_shape(shape, T, K) + +# The concrete kernel type at a menu `shape`, as a `Val`: a singleton crosses +# `plan_contract`'s kernel barrier without boxing, and dispatch on it hits the +# fast method cache, which a bare `Type` argument misses. +@generated menu_val(shape::Tuple{Int, Int, Int}, ::Type{T}, ::Type{K}) where {T, K} = + menu_ladder(T, K, true) + +@noinline throw_no_kernel(shape, ::Type{T}, K) where {T} = throw( ArgumentError( - "no microkernel is available for $T at shape $shape under $(method). " * + "no microkernel is available for $T at shape $shape with $(K). " * "Pass an explicit `kernel = ...` to plan_contract to use a kernel this " * "engine does not pick itself." ) ) -@noinline throw_shape_not_in_menu(shape, ::Type{T}, method) where {T} = throw( +@noinline throw_shape_not_in_menu(shape, ::Type{T}, K) where {T} = throw( ArgumentError( - "shape $shape is not in the $(method) menu for $T, " * - "$(kernel_shapes(T, method)); only menu shapes are compiled" + "shape $shape is not in the $(K) menu for $T, " * + "$(kernel_shapes(T, K)); only menu shapes are compiled" ) ) @@ -127,25 +122,25 @@ kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}) where {T} = # ---------------------------------------------------------------------------- # Per-ISA shapes, consulted first. None for real types: the rule below is the optimum. -shape_override(key::Val, ::Type{T}) where {T} = shape_override(key, T, default_method(T)) -shape_override(::Val, ::Type, ::KernelMethod) = nothing +shape_override(key::Val, ::Type{T}) where {T} = shape_override(key, T, default_kernel_type(T)) +shape_override(::Val, ::Type, ::Type) = nothing # On AVX-512 the rule's planar `MR = 2W, NR = 6` spills; `NR = 3` tiles do not. -shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::PlanarMethod) = (24, 3, 8) -shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::PlanarMethod) = (48, 3, 16) +shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::Type{<:PlanarKernel}) = (24, 3, 8) +shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::Type{<:PlanarKernel}) = (48, 3, 16) # AVX2 has 16 registers: at `MV = 1` planar `NR = 6` needs all 16, `NR = 5` needs 14. -shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::PlanarMethod) = (4, 5, 4) -shape_override(::Val{:avx2}, ::Type{ComplexF32}, ::PlanarMethod) = (8, 5, 8) -# The AVX2-sized 1m shape, for a caller naming `OneMMethod()` (the fit would -# hand it an AVX-512 shape). -shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::OneMMethod) = (4, 6, 4) +shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::Type{<:PlanarKernel}) = (4, 5, 4) +shape_override(::Val{:avx2}, ::Type{ComplexF32}, ::Type{<:PlanarKernel}) = (8, 5, 8) +# The AVX2-sized 1m shape, for a caller asking for `OneMKernel`'s shape (the +# fit would hand it an AVX-512 shape). +shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::Type{<:OneMKernel}) = (4, 6, 4) # Where the `MR = MV*W` rule applies. Not on NEON (2W on 128-bit lanes is no # better than the fallback), nor complex off AVX-512 (planar's two accumulator # planes leave AVX2's 16 registers nothing spare). -rule_applies(::Val{:avx512}, ::RealMethod) = true -rule_applies(::Val{:avx512}, ::Union{PlanarMethod, OneMMethod}) = true -rule_applies(::Val{:avx2}, ::RealMethod) = true -rule_applies(::Val, ::KernelMethod) = false +rule_applies(::Val{:avx512}, ::Type{<:SIMDKernel}) = true +rule_applies(::Val{:avx512}, ::Type{<:Union{PlanarKernel, OneMKernel}}) = true +rule_applies(::Val{:avx2}, ::Type{<:SIMDKernel}) = true +rule_applies(::Val, ::Type) = false # `W` real lanes per register, `MV` A vectors per column. rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = @@ -153,35 +148,35 @@ rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = # MV = 4 for the real kernel on AVX-512: the MV = 2 tile is front-end bound on # cores with 2 FMA ports, and 4*6 accumulators + 4 A vectors + 2 still fit 32 -# registers. Complex methods keep MV = 2 (planar already spills there). -rule_mv(::Val{:avx512}, ::RealMethod) = 4 -rule_mv(::Val, ::KernelMethod) = 2 +# registers. Complex kernels keep MV = 2 (planar already spills there). +rule_mv(::Val{:avx512}, ::Type{<:SIMDKernel}) = 4 +rule_mv(::Val, ::Type) = 2 # AMD's AVX-512 cores double-pump 512-bit FMAs, so their MV = 2 tile is not # front-end bound and the taller tile only adds edge and store cost. const _MV4_UNPROFITABLE_CPUS = ("znver4", "znver5") -profile_mv(profile::TargetProfile, method) = rule_mv(Val(profile.isa), method) -function profile_mv(profile::TargetProfile, method::RealMethod) - mv = rule_mv(Val(profile.isa), method) +profile_mv(profile::TargetProfile, ::Type{K}) where {K} = rule_mv(Val(profile.isa), K) +function profile_mv(profile::TargetProfile, ::Type{K}) where {K <: SIMDKernel} + mv = rule_mv(Val(profile.isa), K) return (mv == 4 && profile.cpu_name in _MV4_UNPROFITABLE_CPUS) ? 2 : mv end -# The register shape for `T` under `method` on `profile`: an override row, else -# the rule where it applies (and lands in the menu), else `fitted_shape`. +# The register shape for `T` with kernel type `K` on `profile`: an override +# row, else the rule where it applies (and lands in the menu), else `fitted_shape`. derived_shape(profile::TargetProfile, ::Type{T}) where {T} = - derived_shape(profile, T, default_method(T)) + derived_shape(profile, T, default_kernel_type(T)) -function derived_shape(profile::TargetProfile, ::Type{T}, method) where {T} +function derived_shape(profile::TargetProfile, ::Type{T}, ::Type{K}) where {T, K} key = Val(profile.isa) - ovr = shape_override(key, T, method) + ovr = shape_override(key, T, K) ovr === nothing || return ovr vb = profile.vector_bytes - if rule_applies(key, method) && vb > 0 - shape = rule_shape(vb, T, profile_mv(profile, method)) - shape in kernel_shapes(T, method) && return shape + if rule_applies(key, K) && vb > 0 + shape = rule_shape(vb, T, profile_mv(profile, K)) + shape in kernel_shapes(T, K) && return shape end - return fitted_shape(profile, T, method) + return fitted_shape(profile, T, K) end # Vector registers a planar kernel holds live per K step (two planes of @@ -192,17 +187,17 @@ planar_pressure(MR::Int, NR::Int, W::Int) = 2 * (MR ÷ W) * NR + 2 * (MR ÷ W) + # The shape where no rule applies, and the small-M demotion target. Planar: the # largest menu shape with `W <= lanes` whose pressure fits the register file # (16 when unknown). 1m/FMAddSub: the menu head. -fitted_shape(::TargetProfile, ::Type{T}, ::RealMethod) where {T} = fallback_shape(T) -fitted_shape(::TargetProfile, ::Type{T}, method::Union{OneMMethod, FMAddSubMethod}) where {T} = - first(kernel_shapes(T, method)) +fitted_shape(::TargetProfile, ::Type{T}, ::Type{<:SIMDKernel}) where {T} = fallback_shape(T) +fitted_shape(::TargetProfile, ::Type{T}, ::Type{K}) where {T, K <: Union{OneMKernel, FMAddSubKernel}} = + first(kernel_shapes(T, K)) -function fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) where {T} +function fitted_shape(profile::TargetProfile, ::Type{T}, ::Type{K}) where {T, K <: PlanarKernel} R = real(T) vb = profile.vector_bytes lanes = vb > 0 ? vb ÷ sizeof(R) : default_lanewidth(R) budget = profile.nregisters > 0 ? profile.nregisters : 16 best = nothing - for shape in kernel_shapes(T, method) + for shape in kernel_shapes(T, K) MR, NR, W = shape (W <= lanes && MR % W == 0) || continue planar_pressure(MR, NR, W) <= budget || continue @@ -212,7 +207,7 @@ function fitted_shape(profile::TargetProfile, ::Type{T}, method::PlanarMethod) w end end # Unreachable for any budget >= 16; a slow kernel beats throwing. - return best === nothing ? last(kernel_shapes(T, method)) : best + return best === nothing ? last(kernel_shapes(T, K)) : best end # ---------------------------------------------------------------------------- @@ -220,12 +215,12 @@ end # ---------------------------------------------------------------------------- """ - select_shape(T, method, m_length, k_length, run) -> (shape, method) + select_shape(T, K, m_length, k_length, run) -> (shape, Val(kernel_type)) -The automatic register shape `(MR, NR, W)`, and the method it runs under, for -compute type `T` under `method = default_method(T, TA, TB)`, extents -`m_length`/`k_length` and C's unit-stride run along M (`run`; `m_length` when -no layout is known). The steps, in order: +The automatic register shape `(MR, NR, W)`, and the kernel type it is built +with, for compute type `T` and kernel type `K = default_kernel_type(T, TA, TB)`, +extents `m_length`/`k_length` and C's unit-stride run along M (`run`; +`m_length` when no layout is known). The steps, in order: 1. The host's shape for `T` (`derived_shape`, cached in `ResolvedDefaults`). 2. Extent (real): the MV = 4 shape steps down to MV = 2 where `m_length` @@ -234,39 +229,41 @@ no layout is known). The steps, in order: complex `T` on AVX-512 the FMAddSub shape that pads least (`small_m_shape`). 4. Run (real): the shape shrinks until C's run fills its slivers (`fit_to_run`). -A mixed method runs the steps on the real problem of `real(T)` and maps the -shape. Plain values, so `plan_contract` never holds a menu-wide kernel Union. +A mixed kernel runs the steps on the real problem of `real(T)` and maps the +shape. Plain values, so `plan_contract` never holds a menu-wide kernel Union; +the kernel type as a `Val`, since inference widens a type inside a returned +tuple to `UnionAll`. """ -@inline function select_shape(::Type{T}, method, m_length::Int, k_length::Int, run::Int) where {T} +@inline function select_shape(::Type{T}, ::Type{K}, m_length::Int, k_length::Int, run::Int) where {T, K} d = _resolved_defaults(T) - shape = extent_shape(d.shape, T, method, m_length) + shape = extent_shape(d.shape, T, K, m_length) if m_length > 0 && m_length < shape[1] - # Static, so a real `T`'s method stays concrete. + # Static, so a real `T`'s kernel type stays concrete. if T <: Complex small = small_m_shape(d.small_m, m_length) - small === nothing || return (small, FMAddSubMethod()) + small === nothing || return (small, Val(FMAddSubKernel)) end shape = d.fitted end - return (fit_to_run(shape, T, method, m_length, k_length, run), method) + return (fit_to_run(shape, T, K, m_length, k_length, run), Val(K)) end -@inline function select_shape(::Type{T}, method::MixedMethod, m_length::Int, k_length::Int, run::Int) where {T} - m, r = real_problem(method, m_length, run) - shape, _ = select_shape(real(T), RealMethod(), m, k_length, r) - return (mixed_shape(method, shape), method) +@inline function select_shape(::Type{T}, K::MixedKernel, m_length::Int, k_length::Int, run::Int) where {T} + m, r = real_problem(K, m_length, run) + shape, _ = select_shape(real(T), SIMDKernel, m, k_length, r) + return (mixed_shape(K, shape), Val(K)) end # The real MV = 4 shape steps down to its MV = 2 sibling when `m_length` is # below one tall tile, or below two and the half tile pads to fewer rows. Above # `2MR` the tall tile's padding excess is under 1.2x, its speed advantage. -@inline extent_shape(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int) where {T} = shape +@inline extent_shape(shape::NTuple{3, Int}, ::Type{T}, ::Type, m_length::Int) where {T} = shape -@inline function extent_shape(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int) where {T} +@inline function extent_shape(shape::NTuple{3, Int}, ::Type{T}, ::Type{K}, m_length::Int) where {T, K <: SIMDKernel} MR, NR, W = shape (m_length > 0 && MR == 4 * W && m_length < 2 * MR) || return shape half = (2 * W, NR, W) - half in kernel_shapes(T, method) || return shape + half in kernel_shapes(T, K) || return shape # Below one tall tile, always: the tall shape would fall to the fitted # shape (two steps down) at the small-M step. m_length < MR && return half @@ -280,7 +277,7 @@ end small_m_candidates(::Val, ::TargetProfile, ::Type) = NTuple{3, Int}[] function small_m_candidates(::Val{:avx512}, profile::TargetProfile, ::Type{T}) where {T <: Complex} lanes = profile.vector_bytes ÷ sizeof(real(T)) - return [shape for shape in kernel_shapes(T, FMAddSubMethod()) if shape[3] == lanes] + return [shape for shape in kernel_shapes(T, FMAddSubKernel) if shape[3] == lanes] end function small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) @@ -303,19 +300,19 @@ end # one stays); # - at a shallow `k_length`, where the smaller kernel's cost does not dominate, # the largest menu shape whose `MR` divides the run. -@inline fit_to_run(shape::NTuple{3, Int}, ::Type{T}, method, m_length::Int, k_length::Int, run::Int) where {T} = shape +@inline fit_to_run(shape::NTuple{3, Int}, ::Type{T}, ::Type, m_length::Int, k_length::Int, run::Int) where {T} = shape -@inline function fit_to_run(shape::NTuple{3, Int}, ::Type{T}, method::RealMethod, m_length::Int, k_length::Int, run::Int) where {T} +@inline function fit_to_run(shape::NTuple{3, Int}, ::Type{T}, ::Type{K}, m_length::Int, k_length::Int, run::Int) where {T, K <: SIMDKernel} m_length == run && return shape MR, NR, W = shape half = (2 * W, NR, W) - if MR == 4 * W && run % MR != 0 && run >= 2 * W && half in kernel_shapes(T, method) + if MR == 4 * W && run % MR != 0 && run >= 2 * W && half in kernel_shapes(T, K) shape = half end kmax = T === Float64 ? _RUN_DEMOTE_KMAX_F64 : _RUN_DEMOTE_KMAX_F32 (k_length > kmax || run % shape[1] == 0) && return shape best = nothing - for s in kernel_shapes(T, method) + for s in kernel_shapes(T, K) run % s[1] == 0 && (best === nothing || s[1] > best[1]) && (best = s) end return something(best, shape) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 5ee747c..a3f28db 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -185,7 +185,7 @@ function _planned( accumulator::AC ) where {F, NA, NB, NC, AC} T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) - method = default_method(T, eltype(A), eltype(B)) + K = default_kernel_type(T, eltype(A), eltype(B)) _check_kernel_domain(kernel, eltype(A), eltype(B)) # GUARDRAIL: a conjugated `C` is rejected; the engine writes through to @@ -205,7 +205,7 @@ function _planned( norder = order_free_labels(nlabels, indC, C) # Only real `T` swaps and only real kernels run-demote; `0` is a placeholder. - run_m = T <: Real || method isa MixedMethod ? leading_unit_run(morder, indC, C) : 0 + run_m = T <: Real || K isa MixedKernel ? leading_unit_run(morder, indC, C) : 0 run_n = T <: Real ? leading_unit_run(norder, indC, C) : 0 mgroup = AxisGroup(morder, (indA, A), (indC, C)) # maps: (A, C) @@ -219,7 +219,7 @@ function _planned( k_length = axis_length(kgroup) - m_tile_asis, m_tile_swapped = _candidate_m_tiles(T, method, kernel, m_length, n_length, run_m, run_n) + m_tile_asis, m_tile_swapped = _candidate_m_tiles(T, K, kernel, m_length, n_length, run_m, run_n) if T <: Real && prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) # B takes the M role: groups, K maps, storages, run and transforms move @@ -230,14 +230,14 @@ function _planned( parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), run_n, m_block, k_block, n_block, workspace, allocator, oracle ) - return _plan_with_kernel(kernel, method, btransform, atransform, req_swapped) + return _plan_with_kernel(kernel, K, 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, m_block, k_block, n_block, workspace, allocator, oracle ) - return _plan_with_kernel(kernel, method, atransform, btransform, req) + return _plan_with_kernel(kernel, K, atransform, btransform, req) end const _QS_ELTYPES = (Float32, Float64, ComplexF32, ComplexF64) @@ -277,44 +277,49 @@ end # named kernel either way, else `select_shape`'s pick at that orientation's # extent and C run. The swap is judged before `fit_to_run`'s K-dependent # divisor step, which an unbounded `k_length` skips. -@inline _candidate_m_tiles(::Type{T}, method, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = +@inline _candidate_m_tiles(::Type{T}, ::Type, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = (tile_size(kernel, 1), tile_size(kernel, 1)) -@inline function _candidate_m_tiles(::Type{T}, method, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} - m_tile_asis = select_shape(T, method, m_length, typemax(Int), run_m)[1][1] - m_tile_swapped = T <: Real ? select_shape(T, method, n_length, typemax(Int), run_n)[1][1] : m_tile_asis +@inline function _candidate_m_tiles(::Type{T}, ::Type{K}, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T, K} + m_tile_asis = select_shape(T, K, m_length, typemax(Int), run_m)[1][1] + m_tile_swapped = T <: Real ? select_shape(T, K, n_length, typemax(Int), run_n)[1][1] : m_tile_asis return m_tile_asis, m_tile_swapped end # Kernel resolution. A named kernel goes straight through, never demoted. An -# automatic one is chosen as a `(shape, method)` value and the plan is built -# across a dispatch barrier specialised on that one shape, so only the chosen -# 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, method, atransform, btransform, req::_PlanRequest) = +# automatic one is chosen as a `(shape, kernel type)` value and the plan is +# built across a dispatch barrier specialised on that one concrete kernel type, +# so only the chosen 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, ::Type, atransform, btransform, req::_PlanRequest) = _plan_contract(kernel, atransform, btransform, req, nothing) -@inline function _plan_with_kernel(::Nothing, method, atransform, btransform, req::_PlanRequest{T}) where {T} - shape, method = select_shape(T, method, axis_length(req.mgroup), axis_length(req.kgroup), req.run) - vshape = menu_val(shape, T, method) +@inline function _plan_with_kernel(::Nothing, ::Type{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} + shape, kernel_type = select_shape(T, K, axis_length(req.mgroup), axis_length(req.kgroup), req.run) + return _plan_at_shape(shape, kernel_type, atransform, btransform, req) +end + +@inline function _plan_at_shape(shape, ::Val{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} + vkernel = menu_val(shape, T, K) # The execution path, predicted so the callee is specialised on it. Only a - # `Bool` crosses: a call union-split on `method` is emitted out of line and - # boxes `req`. - hint = _path_hint(req.f, req, _unpacked_b_method_eligible(method)) + # `Bool` crosses: a call union-split on the kernel type is emitted out of + # line and boxes `req`. + hint = _path_hint(req.f, req, _unpacked_b_kernel_eligible(K)) core = _strip_storage(req) slot = _barrier_slot!(req.workspace, typeof(core)) slot[] = core return Base.inferencebarrier(_plan_resolved)( - vshape, method, atransform, btransform, hint, slot, + vkernel, atransform, btransform, hint, slot, req.Astorage, req.Bstorage, req.Cstorage ) end function _plan_resolved( - ::Val{S}, method::M, atransform::TA, btransform::TB, hint::H, + ::Val{Kern}, atransform::TA, btransform::TB, hint::H, slot::Base.RefValue{R}, Astorage::SA, Bstorage::SB, Cstorage::SC - ) where {S, M, TA, TB, H, T, R <: _PlanRequest{T}, SA, SB, SC} + ) where {Kern, TA, TB, H, T, R <: _PlanRequest{T}, SA, SB, SC} req = _with_storage(slot[], Astorage, Bstorage, Cstorage) - return _plan_contract(kernel_from_shape(S, T, method), atransform, btransform, req, hint) + return _plan_contract(Kern(), atransform, btransform, req, hint) end # The storages cross the barrier as arguments, so no slot retains a user array. diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 2b9310c..555ac17 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -1,6 +1,6 @@ -using QuasiStrided: FMAddSubKernel, FMAddSubMethod, PlanarKernel, PlanarMethod, OneMKernel, - InterleavedFormat, KernelMethod, lanewidth, kernel_shapes, packed_panel, - kernel_from_shape, default_method, accumulator_planes, target_profile, PackedPanel +using QuasiStrided: FMAddSubKernel, PlanarKernel, OneMKernel, + InterleavedFormat, lanewidth, kernel_shapes, packed_panel, + kernel_from_shape, default_kernel_type, accumulator_planes, target_profile, PackedPanel using SIMD: Vec using InteractiveUtils: code_native @@ -11,13 +11,12 @@ const _QSF = QuasiStrided for (T, (MR, NR, W)) in full mk_contract(FMAddSubKernel(Val(MR), Val(NR), T, Val(W))) end - for T in (ComplexF64, ComplexF32), s in (kernel_shapes(T, FMAddSubMethod())..., (2, 3, 4)) + for T in (ComplexF64, ComplexF32), s in (kernel_shapes(T, FMAddSubKernel)..., (2, 3, 4)) (T, s) in full || mk_contract(FMAddSubKernel(Val(s[1]), Val(s[2]), T, Val(s[3])); full = false) end @testset "construction" begin - @test KernelMethod(FMAddSubKernel(Val(8), Val(4), ComplexF64)) === FMAddSubMethod() - @test accumulator_planes(FMAddSubMethod()) == 1 + @test accumulator_planes(FMAddSubKernel) == 1 @test lanewidth(FMAddSubKernel(Val(8), Val(4), ComplexF32)) == 8 @test_throws ArgumentError FMAddSubKernel(Val(12), Val(8), ComplexF64, Val(16)) @test_throws ArgumentError FMAddSubKernel(Val(3), Val(4), ComplexF64, Val(3)) # odd W @@ -110,7 +109,7 @@ const _QSF = QuasiStrided @testset "pack! InterleavedFormat A: scalar loop and fast path vs local layout" begin # Bitwise (`isequal` separates -0.0): packing is a copy, or a sign flip under conj. - for T in (ComplexF64, ComplexF32), (MR, NR, W) in (kernel_shapes(T, FMAddSubMethod())..., (3, 2, 2)) + for T in (ComplexF64, ComplexF32), (MR, NR, W) in (kernel_shapes(T, FMAddSubKernel)..., (3, 2, 2)) R = real(T) k = FMAddSubKernel(Val(MR), Val(NR), T, Val(W)) k_block_length, lda, base = 7, MR + 3, 2 @@ -137,10 +136,10 @@ const _QSF = QuasiStrided bf = default_blocking(FMAddSubKernel(Val(8), Val(8), T, Val(8))) bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) @test (bf.m_block, bf.k_block, bf.n_block) == (bp.m_block, bp.k_block, bp.n_block) # planar's packed reals - for s in kernel_shapes(T, FMAddSubMethod()) - @test kernel_from_shape(s, T, FMAddSubMethod()) isa FMAddSubKernel{s[1], s[2], T, s[3]} + for s in kernel_shapes(T, FMAddSubKernel) + @test kernel_from_shape(s, T, FMAddSubKernel) isa FMAddSubKernel{s[1], s[2], T, s[3]} end - @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, FMAddSubMethod()) + @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, FMAddSubKernel) @test !(_QSF._default_kernel(T, 1024, 1024) isa FMAddSubKernel) end end diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 9227e0d..5549b9b 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -1,13 +1,12 @@ -using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, - KernelMethod, lanewidth, kernel_shapes, - kernel_from_shape, default_method +using QuasiStrided: OneMKernel, PlanarKernel, lanewidth, kernel_shapes, + kernel_from_shape, default_kernel_type @testset "OneMKernel" begin full = ((ComplexF64, (12, 8, 8)), (ComplexF32, (8, 6, 8))) for (T, (MR, NR, W)) in full mk_contract(OneMKernel(Val(MR), Val(NR), T, Val(W))) end - for T in (ComplexF64, ComplexF32), s in (kernel_shapes(T, OneMMethod())..., (8, 4, 4)) + for T in (ComplexF64, ComplexF32), s in (kernel_shapes(T, OneMKernel)..., (8, 4, 4)) (T, s) in full || mk_contract(OneMKernel(Val(s[1]), Val(s[2]), T, Val(s[3])); full = false) end @@ -15,7 +14,6 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, k = OneMKernel(Val(12), Val(8), ComplexF64, Val(8)) @test QuasiStrided.inner(k) isa SIMDKernel{24, 8, Float64, 8} # the real kernel, at 2MR rows @test zero_accumulator(k) === zero_accumulator(QuasiStrided.inner(k)) - @test KernelMethod(k) === OneMMethod() @test lanewidth(OneMKernel(Val(8), Val(4), ComplexF32)) == 8 # 2MR, not MR, must divide by W; W must be even even where 2MR divides. @test_throws ArgumentError OneMKernel(Val(12), Val(8), ComplexF64, Val(16)) @@ -37,12 +35,12 @@ using QuasiStrided: OneMKernel, OneMMethod, PlanarKernel, PlanarMethod, bm = default_blocking(OneMKernel(Val(8), Val(8), T, Val(8))) bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) @test (bm.m_block, bm.k_block, bm.n_block) == (bp.m_block ÷ 2, bp.k_block, bp.n_block) # twice the packed A reals - for s in kernel_shapes(T, OneMMethod()) - @test kernel_from_shape(s, T, OneMMethod()) isa OneMKernel{s[1], s[2], T, s[3]} + for s in kernel_shapes(T, OneMKernel) + @test kernel_from_shape(s, T, OneMKernel) isa OneMKernel{s[1], s[2], T, s[3]} end - @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, OneMMethod()) + @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, OneMKernel) # Never the default: method ranking does not transfer between machines. - @test default_method(T) === PlanarMethod() + @test default_kernel_type(T) === PlanarKernel @test QuasiStrided._default_kernel(T, 1024, 1024) isa PlanarKernel end end diff --git a/test/microkernels/test_planar_kernel.jl b/test/microkernels/test_planar_kernel.jl index 4fc147e..f1c9b17 100644 --- a/test/microkernels/test_planar_kernel.jl +++ b/test/microkernels/test_planar_kernel.jl @@ -1,4 +1,4 @@ -using QuasiStrided: PlanarKernel, PlanarMethod, KernelMethod, lanewidth, +using QuasiStrided: PlanarKernel, lanewidth, KERNEL_SHAPES_C64_PLANAR, KERNEL_SHAPES_C32_PLANAR @testset "PlanarKernel" begin @@ -11,7 +11,6 @@ using QuasiStrided: PlanarKernel, PlanarMethod, KernelMethod, lanewidth, end k = PlanarKernel(Val(16), Val(6), ComplexF64, Val(8)) - @test KernelMethod(k) === PlanarMethod() @test lanewidth(PlanarKernel(Val(8), Val(4), ComplexF64)) == 4 # from the REAL type @test lanewidth(PlanarKernel(Val(8), Val(4), ComplexF32)) == 8 @test_throws ArgumentError PlanarKernel(Val(6), Val(4), ComplexF64, Val(4)) diff --git a/test/packing/test_pack_complex_contiguous.jl b/test/packing/test_pack_complex_contiguous.jl index e379bed..2939449 100644 --- a/test/packing/test_pack_complex_contiguous.jl +++ b/test/packing/test_pack_complex_contiguous.jl @@ -4,21 +4,21 @@ # checks the same outputs with the fast path off. using QuasiStrided: target_profile, unknown_target, TargetProfile, CacheLevel, - kernel_shapes, PlanarMethod, OneMMethod, FMAddSubMethod, SliverSpec, complex_contiguous_eligible, + kernel_shapes, PlanarKernel, OneMKernel, FMAddSubKernel, SliverSpec, complex_contiguous_eligible, complex_fastpath_isa_eligible const FASTPATH_ON = complex_fastpath_isa_eligible() -# The extremes of each method's menu plus extents that are not multiples of any -# lane width (interleaved shares 1e's shuffle); B is planar under every method. -_menu_extremes(T, method, i) = extrema(s[i] for s in kernel_shapes(T, method)) +# The extremes of each kernel's menu plus extents that are not multiples of any +# lane width (interleaved shares 1e's shuffle); B is planar under every kernel. +_menu_extremes(T, K, i) = extrema(s[i] for s in kernel_shapes(T, K)) const FAST_A_CASES = [ (T, fa, MR) for T in (ComplexF64, ComplexF32) - for (fa, method, extra) in ( - (PlanarFormat(), PlanarMethod(), (1, 7)), (OneEFormat(), OneMMethod(), (1, 7)), - (InterleavedFormat(), FMAddSubMethod(), ()), + for (fa, K, extra) in ( + (PlanarFormat(), PlanarKernel, (1, 7)), (OneEFormat(), OneMKernel, (1, 7)), + (InterleavedFormat(), FMAddSubKernel, ()), ) - for MR in unique((_menu_extremes(T, method, 1)..., extra...)) + for MR in unique((_menu_extremes(T, K, 1)..., extra...)) ] const FAST_B_NRS = (1, 3, 7, 8) diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index e49c930..a9c7b63 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -23,11 +23,11 @@ using StridedViews: StridedView for (isakey, vb, MV) in ((:avx512, 64, 4), (:avx2, 32, 2)), T in (Float64, Float32) W = vb ÷ sizeof(T) @test derived_shape(synthetic(isakey), T) === (MV * W, NR_DEFAULT, W) - @test QuasiStrided.rule_mv(Val(isakey), RealMethod()) == MV + @test QuasiStrided.rule_mv(Val(isakey), SIMDKernel) == MV @test QuasiStrided.rule_shape(vb, T, MV) === (MV * W, NR_DEFAULT, W) end - for m in (PlanarMethod(), OneMMethod(), FMAddSubMethod()), key in (:avx512, :avx2) - @test QuasiStrided.rule_mv(Val(key), m) == 2 + for K in (PlanarKernel, OneMKernel, FMAddSubKernel), key in (:avx512, :avx2) + @test QuasiStrided.rule_mv(Val(key), K) == 2 end for (cpu, MV) in (("znver4", 2), ("znver5", 2), ("icelake-server", 4), ("cascadelake", 4)), T in (Float64, Float32) @@ -46,7 +46,7 @@ using StridedViews: StridedView for T in (Float64, Float32), key in VALID_ISAS @test shape_override(Val(key), T) === nothing end - @test rule_applies(Val(:avx2), RealMethod()) && rule_applies(Val(:avx512), RealMethod()) + @test rule_applies(Val(:avx2), SIMDKernel) && rule_applies(Val(:avx512), SIMDKernel) end @testset "every real menu shape is constructible and fits the register file" begin @@ -57,14 +57,13 @@ using StridedViews: StridedView @test (MR ÷ W) * NR + (MR ÷ W) <= 32 @test sliver_width(k) === (MR, NR) @test realtype(k) === T === scalartype(k) - @test KernelMethod(k) === RealMethod() @test packed_a_length(k, 7) === MR * 7 && packed_b_length(k, 7) === NR * 7 end for T in (Float64, Float32) W = 64 ÷ sizeof(T) # The MV = 2 sibling stays in the menu as the step-down target. @test (2 * W, NR_DEFAULT, W) in kernel_shapes(T) - @test kernel_shapes(T, RealMethod()) === kernel_shapes(T) + @test kernel_shapes(T, SIMDKernel) === kernel_shapes(T) k = _default_kernel(T) @test k isa SIMDKernel && scalartype(k) === T @test (tile_size(k)..., lanewidth(k)) in kernel_shapes(T) @@ -91,7 +90,7 @@ using StridedViews: StridedView end @testset "extent_shape: MV = 4 steps down to MV = 2 where it pads less" begin - ext(shape, T, m_length) = QuasiStrided.extent_shape(shape, T, RealMethod(), m_length) + ext(shape, T, m_length) = QuasiStrided.extent_shape(shape, T, SIMDKernel, m_length) for T in (Float64, Float32) tall = derived_shape(synthetic(:avx512), T) MR, NR, W = tall @@ -103,11 +102,11 @@ using StridedViews: StridedView for m_length in (MR, MR + 2 * W + 1, 2 * MR - 1, 2 * MR, 2 * MR + 1, 3 * MR + W, 100 * MR) @test ext(tall, T, m_length) === tall end - # Never below MV = 2, and complex methods are untouched. + # Never below MV = 2, and complex kernels are untouched. avx2 = derived_shape(synthetic(:avx2), T) @test ext(avx2, T, 3) === avx2 cs = derived_shape(synthetic(:avx512), ComplexF64) - @test QuasiStrided.extent_shape(cs, ComplexF64, PlanarMethod(), 3) === cs + @test QuasiStrided.extent_shape(cs, ComplexF64, PlanarKernel, 3) === cs # Through `_default_kernel`, where the host's shape is the tall one. if derived_shape(target_profile(), T) === tall @test tile_size(_default_kernel(T, MR + W, 256)) == (2 * W, NR) @@ -120,7 +119,7 @@ using StridedViews: StridedView @testset "fit_to_run: C's run along M decides the real shape" begin fit(shape, T, m_length, k_length, run) = - QuasiStrided.fit_to_run(shape, T, RealMethod(), m_length, k_length, run) + QuasiStrided.fit_to_run(shape, T, SIMDKernel, m_length, k_length, run) largest_divisor(T, run) = argmax(s -> (run % s[1] == 0, s[1]), kernel_shapes(T)) deep = 10^6 for T in (Float64, Float32) @@ -155,9 +154,9 @@ using StridedViews: StridedView end end cs = derived_shape(synthetic(:avx512), ComplexF64) - @test QuasiStrided.fit_to_run(cs, ComplexF64, PlanarMethod(), m_length, 1, 1) === cs + @test QuasiStrided.fit_to_run(cs, ComplexF64, PlanarKernel, m_length, 1, 1) === cs if derived_shape(target_profile(), T) === tall - select(run) = QuasiStrided.select_shape(T, RealMethod(), m_length, deep, run)[1] + select(run) = QuasiStrided.select_shape(T, SIMDKernel, m_length, deep, run)[1] @test select(2 * W) === half @test select(MR) === tall @test select(m_length) === tall @@ -197,8 +196,8 @@ end # W counts REAL lanes; the AVX-512 override replaces the spilling rule shape. @test QuasiStrided.rule_shape(64, T, 2) === (2 * W, NR_DEFAULT, W) @test shape_override(Val(:avx512), T) === swept - @test derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarMethod())) - @test Set(kernel_shapes(T, PlanarMethod())) == Set( + @test derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarKernel)) + @test Set(kernel_shapes(T, PlanarKernel)) == Set( ( (2 * W, NR_DEFAULT, W), swept, (W, 8, W), shape_override(Val(:avx2), T), (W ÷ 2, NR_DEFAULT, W ÷ 4), (W ÷ 4, NR_DEFAULT, W ÷ 4), @@ -208,21 +207,21 @@ end for (key, nreg) in ((:avx512, 32), (:avx2, 16)) ovr = shape_override(Val(key), T) @test planar_pressure(ovr...) < nreg - @test ovr in kernel_shapes(T, PlanarMethod()) + @test ovr in kernel_shapes(T, PlanarKernel) end @test fallback_shape(T) === (8, NR_DEFAULT, fallback_shape(real(T))[3]) # Off :avx512 an override row wins whatever the width; without a row # the shape is fitted to the register budget. @test derived_shape(synthetic(:avx2), T) === shape_override(Val(:avx2), T) for key in (:avx2, :neon, :unknown) - @test !rule_applies(Val(key), PlanarMethod()) + @test !rule_applies(Val(key), PlanarKernel) end for key in (:neon, :unknown) @test shape_override(Val(key), T) === nothing end @test planar_pressure(derived_shape(synthetic(:unknown), T)...) <= 16 end - @test rule_applies(Val(:avx512), PlanarMethod()) + @test rule_applies(Val(:avx512), PlanarKernel) @test derived_shape(synthetic(:neon), ComplexF64) === (4, 6, 2) @test derived_shape(synthetic(:neon), ComplexF32) === (8, 6, 4) @test planar_pressure(fallback_shape(ComplexF64)...) == 30 @@ -231,9 +230,9 @@ end # Complex menus fit AVX-512 and stay bounded. nreg = isa_nregisters(:avx512) target_profile().isa === :avx512 && @test target_profile().nregisters == nreg - for T in (ComplexF64, ComplexF32), m in (PlanarMethod(), OneMMethod()) - planes = accumulator_planes(m) - for (MR, NR, W) in kernel_shapes(T, m) + for T in (ComplexF64, ComplexF32), K in (PlanarKernel, OneMKernel) + planes = accumulator_planes(K) + for (MR, NR, W) in kernel_shapes(T, K) rows = (2 * MR) ÷ planes @test rows % W == 0 mv = rows ÷ W @@ -241,15 +240,15 @@ end end end for T in (ComplexF64, ComplexF32) - @test length(kernel_shapes(T, PlanarMethod())) <= 6 - @test length(kernel_shapes(T, OneMMethod())) <= 4 + @test length(kernel_shapes(T, PlanarKernel)) <= 6 + @test length(kernel_shapes(T, OneMKernel)) <= 4 end end @testset "mixed-domain selection: the real default of real(T), mapped" begin - CR, RC = QuasiStrided.ComplexRealMethod(), QuasiStrided.RealComplexMethod() + CR, RC = QuasiStrided.ComplexRealKernel, QuasiStrided.RealComplexKernel mapped(m, (MR, NR, W)) = m === CR ? (MR ÷ 2, NR, W) : (MR, NR ÷ 2, W) - dmethod = QuasiStrided.default_method + dkernel = QuasiStrided.default_kernel_type for T in (ComplexF64, ComplexF32) R = real(T) @test all(((MR, NR, W),) -> iseven(NR) && iseven(W), kernel_shapes(R)) @@ -257,13 +256,13 @@ end @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 KernelMethod(k) === m && (tile_size(k)..., lanewidth(k)) === shape + @test k isa m && (tile_size(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() + @test dkernel(T, T, R) === dkernel(T, ComplexF32, Float64) === CR + @test dkernel(T, R, T) === dkernel(T, Float32, ComplexF64) === RC + @test dkernel(T, R, R) === dkernel(T, T, T) === dkernel(T, ComplexF32, T) === PlanarKernel + @test dkernel(R, Float32, Float64) === SIMDKernel end # Every real step carries over. saved = QuasiStrided.TARGET[] @@ -273,8 +272,8 @@ end for T in (ComplexF64, ComplexF32), m in (CR, RC), m_length in (1, 3, 17, 40, 4096), run in (0, 1, 8, m_length), k_length in (1, 1000) real_m, real_run = m === CR ? (2 * m_length, 2 * run) : (m_length, run) - real_shape = QuasiStrided.select_shape(real(T), RealMethod(), real_m, k_length, real_run)[1] - @test QuasiStrided.select_shape(T, m, m_length, k_length, run) === (mapped(m, real_shape), m) + real_shape = QuasiStrided.select_shape(real(T), SIMDKernel, real_m, k_length, real_run)[1] + @test QuasiStrided.select_shape(T, m, m_length, k_length, run) === (mapped(m, real_shape), Val(m)) end end finally @@ -282,12 +281,12 @@ end end end -@testset "derived_shape is always a constructible member of the method's menu" begin - methods = ( - (Float64, RealMethod()), (Float32, RealMethod()), (ComplexF64, PlanarMethod()), - (ComplexF32, PlanarMethod()), (ComplexF64, OneMMethod()), (ComplexF32, OneMMethod()), +@testset "derived_shape is always a constructible member of the kernel's menu" begin + pairs = ( + (Float64, SIMDKernel), (Float32, SIMDKernel), (ComplexF64, PlanarKernel), + (ComplexF32, PlanarKernel), (ComplexF64, OneMKernel), (ComplexF32, OneMKernel), ) - for (T, m) in methods, isakey in VALID_ISAS + for (T, m) in pairs, isakey in VALID_ISAS shape = derived_shape(synthetic(isakey), T, m) @test shape in kernel_shapes(T, m) k = QuasiStrided.kernel_from_shape(shape, T, m) @@ -343,27 +342,28 @@ end QS.RealComplexKernel(Val(4), Val(2), ComplexF64), ) for k in kernels - @test pack_formats(KernelMethod(k)) === (QS.a_format(k), QS.b_format(k)) + @test pack_formats(typeof(k)) === (QS.a_format(k), QS.b_format(k)) end end @testset "complex blocking is the real row scaled by packed reals" begin for base in (_fallback_blocking(Float64), _fallback_blocking(Float32), Blocking(1, 5, 1)) - for m in (PlanarMethod(), OneMMethod(), FMAddSubMethod()) + for m in (PlanarKernel, OneMKernel, FMAddSubKernel) b = _scale_blocking(base, m) @test b.k_block === base.k_block @test b.m_block === max(1, base.m_block ÷ reals_per_element(pack_formats(m)[1])) @test b.n_block === max(1, base.n_block ÷ reals_per_element(pack_formats(m)[2])) end end - @test _scale_blocking(_fallback_blocking(Float64), PlanarMethod()) === Blocking(64, 256, 384) - @test _scale_blocking(_fallback_blocking(Float64), OneMMethod()) === Blocking(32, 256, 384) + @test _scale_blocking(_fallback_blocking(Float64), PlanarKernel) === Blocking(64, 256, 384) + @test _scale_blocking(_fallback_blocking(Float64), OneMKernel) === Blocking(32, 256, 384) for T in (ComplexF64, ComplexF32) row = _real_blocking_row(target_profile(), real(T)) - @test default_blocking(_default_kernel(T)) === _scale_blocking(row, PlanarMethod()) + @test default_blocking(_default_kernel(T)) === _scale_blocking(row, PlanarKernel) end for T in (Float64, Float32) @test default_blocking(_default_kernel(T)) === _real_blocking_row(target_profile(), T) + @test default_blocking(ScalarKernel(Val(4), Val(2), T)) === _real_blocking_row(target_profile(), T) end end @@ -372,36 +372,36 @@ end kernel = _default_kernel(T) @test kernel isa QuasiStrided.PlanarKernel @test scalartype(kernel) === T && realtype(kernel) === real(T) - @test KernelMethod(kernel) === PlanarMethod() === QuasiStrided.default_method(T) + @test QuasiStrided.default_kernel_type(T) === PlanarKernel shape = (tile_size(kernel)..., lanewidth(kernel)) profile = target_profile() @test planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) - profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarMethod())) + profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarKernel)) @test _default_kernel(T, 1024, 1024) isa QuasiStrided.PlanarKernel # The small-M demotion: FMAddSub on AVX-512, planar elsewhere. small = _default_kernel(T, 1, 1) @test small isa (profile.isa === :avx512 ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) - shape1m = first(kernel_shapes(T, OneMMethod())) - k1m = QuasiStrided.kernel_from_shape(shape1m, T, OneMMethod()) - @test k1m isa QuasiStrided.OneMKernel && KernelMethod(k1m) === OneMMethod() + shape1m = first(kernel_shapes(T, OneMKernel)) + k1m = QuasiStrided.kernel_from_shape(shape1m, T, OneMKernel) + @test k1m isa OneMKernel @test (tile_size(k1m)..., lanewidth(k1m)) === shape1m - # A method with no kernel for `T` throws, naming itself. + # A kernel type with no menu for `T` throws, naming itself. err = try - QuasiStrided.kernel_from_shape(shape1m, T, RealMethod()) + QuasiStrided.kernel_from_shape(shape1m, T, SIMDKernel) catch e e end @test err isa ArgumentError - @test occursin("RealMethod", err.msg) && occursin(string(T), err.msg) + @test occursin("SIMDKernel", err.msg) && occursin(string(T), err.msg) end for T in (Float16, Int, ComplexF16) @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 4), T) end - @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 4), Float64, OneMMethod()) + @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 4), Float64, OneMKernel) @test_throws ArgumentError QuasiStrided.kernel_from_shape((8, 6, 5), Float64) - @test_throws ArgumentError QuasiStrided.menu_val((8, 6, 5), Float64, RealMethod()) - @test QuasiStrided.menu_val((8, 6, 4), Float64, RealMethod()) === Val((8, 6, 4)) + @test_throws ArgumentError QuasiStrided.menu_val((8, 6, 5), Float64, SIMDKernel) + @test QuasiStrided.menu_val((8, 6, 4), Float64, SIMDKernel) === Val(SIMDKernel{8, 6, Float64, 4}) end @testset "small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin @@ -415,7 +415,7 @@ end @test sms(ComplexF32, 16) === (16, 8, 16) for T in (ComplexF64, ComplexF32), m_length in 1:47 shape = sms(T, m_length) - @test shape in kernel_shapes(T, FMAddSubMethod()) + @test shape in kernel_shapes(T, FMAddSubKernel) @test shape[3] == 64 ÷ sizeof(real(T)) end @test sms(Float64, 4) === nothing diff --git a/test/runtests.jl b/test/runtests.jl index 8dea1d5..286a0a6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,8 +12,8 @@ using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, derived_shape, fallback_shape, shape_override, _default_kernel, kernel_from_shape, _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, rule_applies, isa_nregisters, sliver_width, realtype, - KernelMethod, RealMethod, PlanarMethod, OneMMethod, accumulator_planes, pack_formats, - reals_per_element, FMAddSubMethod, _modelled_blocking, _scale_blocking, _real_blocking_row, + PlanarKernel, OneMKernel, FMAddSubKernel, accumulator_planes, pack_formats, + reals_per_element, _modelled_blocking, _scale_blocking, _real_blocking_row, planar_pressure, _pack_split, _NestPath # plan_contract, execute! and ContractPlan are bound in helpers.jl instead. From 135ac03010670b590917dc740b796c550c8184d3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 13:29:07 -0400 Subject: [PATCH 23/60] Decision log: chunk 10 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/CLEANUP.md b/CLEANUP.md index 246fc16..9b33028 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -22,7 +22,7 @@ TTFX when codegen is touched). | 7 | Vector kernels | now `kernels.jl`, `vecops.jl`, `steps.jl`, `stores.jl` | done | | 8 | Complex/mixed kernels | (merged into chunk 7's files) | done | | 9 | Labels | `planning/labels.jl` (`conjugation.jl` → plan.jl) | done | -| 10 | Kernel selection | `planning/kernel_selection.jl` | | +| 10 | Kernel selection | `planning/kernel_selection.jl` | done | | 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | @@ -182,6 +182,15 @@ moves to defaults.jl; `isconj`/`op_conjugates` move into plan.jl. The K-order model, free-label order and swap rule are unchanged; labels.jl stays one file. +### D14. Kernel selection (chunk 10, applied) + +One `select_shape` pipeline (host shape → extent → small M → `fit_to_run`, +which merges the two C-run rules); NEON shapes pinned by a test instead of +override rows. The kernel type (unparameterised, e.g. `PlanarKernel`) replaces +`KernelMethod` as the shape-free identifier of a scheme; it crosses the plan +barrier as `Val{K}()` (a bare `Type` argument misses the dispatch fast path, ++350–430 ns per call). + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: @@ -192,6 +201,8 @@ file. - `PackedPanel` without a raw pointer: https://github.com/lkdvos/QuasiStrided.jl/issues/15. +- Analytical shape selection over a bounded search space instead of the menu: + https://github.com/lkdvos/QuasiStrided.jl/issues/17. - M and N are both ordered by their C strides. Only the M side matters for the vector store; for N the order mainly decides ramp detection and C locality, while B's pack (and the unpacked-B path) would prefer B's strides. Measure on From 4c6f17ef61f0ca7b53ee1ff5f806e0fc5145cd9f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 15:23:09 -0400 Subject: [PATCH 24/60] Chunk 11: blocking, host defaults and line packing Blocking: `default_blocking(kernel, profile = target_profile())` evaluates the analytical model with the kernel's own tile and packed bytes per element (B sliver = NR * packed reals * sizeof(real(T))), so the "same packed byte budget" falls out of the model; undetected caches take the fallback rows scaled by packed reals as before. `_scale_blocking`, `_real_blocking_row`, the default-shape `_modelled_blocking` method and `default_blocking(::Type)` are gone. Names in blocking.jl lose their underscores (`modelled_blocking`, `fallback_blocking`, `kernel_blocking`, `roundup`). No host-defaults cache: `ResolvedDefaults`, the per-eltype slots and defaults.jl are removed. The `Val(profile.isa)` tables become functions branching on the `isa` Symbol; the small-M FMAddSub candidates are a filter over the static menu. select_shape, blocking, the K-order model, the split decision and the TensorOperations pool read `target_profile()` directly. The direct version first cost +25-57 ns per Float64 call: the `cpu_name` string compare in `profile_mv` (7 ns, three times per plan) and the divisions of `core_bytes` in the blocking model. Passing the profile down cannot help (`target_profile()` is a plain load), so instead of an eltype-keyed cache the profile itself carries the three derived values, computed by its constructor: `l2_share`, `l3_share` (the core's cache shares) and `double_pumped` (znver4/znver5). The path hint computes a kernel's default `k_block` only under a C narrower than T. Line size: `line_bytes(profile)` (detected L1d line, else 64) replaces `_K_LINE_BYTES` in the K-order model and the split decision; `_K_WALK_FAR_BYTES`/`_K_WALK_FAR_PENALTY` move to target.jl. Line packing: src/packing/line_packing.jl holds `PackSplit`, the split decision, the split enumeration and the block packer, now `pack_block_by_lines!(..., split::PackSplit)`. `pack_split(g, kgroup, spec, S, k_block, eff, rounded, k_block_requested; l2bytes = nothing)` reads the K map, tile and complexity from the operand's `SliverSpec`. src/planning/pack_split.jl is deleted. Test-only API removed: `_default_kernel(T, m, n)` and `_default_kernel(T)`; tests use `select_shape` + `kernel_from_shape` (helpers `auto_kernel`, `host_kernel`), the pool builds the host-shape kernel. Behaviour changes (equivalence grid over 7 profiles, eltypes, extents, C layouts, every menu kernel's blocking, line splits and K orders; the dense select_shape grid is identical): - Real kernels: no blocking changed on any profile. Mixed kernels: none. - Complex planar/1m/FMAddSub kernels on detected caches block by their own tile. On this host (cascadelake): planar 24x3 unchanged (96,341,786); planar 16x6 192,170,1578; 1m 12x8 120,128,2096 (was 48,341,786); FMAddSub 12x8 252,128,2096 (was 96,341,786); the AVX2 planar 4x5 160,204,1315 (was 96,341,786). Undetected caches: unchanged. - 128-byte lines (an M1-like :neon profile, and :avx512 with line 128): K orders flip on 44 grid cases, and line splits change on 22 (new splits where a line is 8 instead of 4 elements; larger L chunks). 64-byte profiles: no split or K-order decision changed. Per-call floor (tensorcontract!, 6 interleaved rounds, medians, ns, base -> new, 0 B both): Float64 8^3 362 -> 363, 16^3 473 -> 459, d4 outer 362 -> 334; ComplexF64 8^3 412 -> 405, 16^3 967 -> 948, d4 outer 570 -> 551. TTFX (3 rounds): Float64 3.8 -> 3.7 s, ComplexF64 2.4 -> 2.3 s, K=1 outer 0.6 -> 0.6 s. Tests: 32300 pass (was 32190; the per-kernel blocking test adds the rest), 54 min wall on the loaded workstation. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_blocking_model.jl | 26 ++-- src/QuasiStrided.jl | 3 +- src/execution/execute.jl | 39 +++--- src/hardware/target.jl | 33 ++++- src/integrations/tensoroperations.jl | 2 +- .../pack_split.jl => packing/line_packing.jl} | 118 ++++++++++++----- src/packing/panel.jl | 55 -------- src/planning/blocking.jl | 68 +++++----- src/planning/defaults.jl | 90 ------------- src/planning/kernel_selection.jl | 86 ++++++------- src/planning/labels.jl | 26 ++-- src/planning/plan.jl | 30 ++--- test/execution/test_execute.jl | 15 ++- test/execution/test_mixed.jl | 2 +- test/execution/test_workspace.jl | 2 +- test/helpers.jl | 8 ++ test/microkernels/test_fmaddsub_kernel.jl | 2 +- test/microkernels/test_onem_kernel.jl | 5 +- test/planning/test_kernel_selection.jl | 121 ++++++++---------- test/planning/test_plan_contract.jl | 10 +- test/runtests.jl | 8 +- 21 files changed, 340 insertions(+), 409 deletions(-) rename src/{planning/pack_split.jl => packing/line_packing.jl} (54%) delete mode 100644 src/planning/defaults.jl diff --git a/benchmark/bench_blocking_model.jl b/benchmark/bench_blocking_model.jl index 6ee6ec0..8fb4d7a 100644 --- a/benchmark/bench_blocking_model.jl +++ b/benchmark/bench_blocking_model.jl @@ -3,18 +3,18 @@ # # julia --project=benchmark benchmark/bench_blocking_model.jl [--reps 21] [--smoke] [--outdir DIR] # -# Per real dtype: the named rows (`model` = `_modelled_blocking`, `fallback` = -# `_fallback_blocking`) and the `wide` grid on every harness shape, plus 1-D -# `n_block` and `m_block` slices through the model row on a 2048x256x2048 shape -# (the only one large enough for `n_block` to bind). Per complex dtype: the -# named rows scaled by `_scale_blocking`. Every configuration of a shape is timed back to back. +# Per real dtype: the named rows (`model` = `default_blocking` on this host, +# `fallback` = on an undetected one) and the `wide` grid on every harness +# shape, plus 1-D `n_block` and `m_block` slices through the model row on a +# 2048x256x2048 shape (the only one large enough for `n_block` to bind). Per +# complex dtype: the named rows. Every configuration of a shape is timed back to back. # The score of a configuration is the geomean over shapes of its time over the # best time at that shape. `--smoke` runs 1 rep on a few points. include(joinpath(@__DIR__, "harness.jl")) -using QuasiStrided: Blocking, target_profile, _default_kernel, _modelled_blocking, - _fallback_blocking, _scale_blocking +using QuasiStrided: Blocking, target_profile, unknown_target, default_blocking, + kernel_from_shape, derived_shape const SMOKE = hasflag("smoke") const REPS = SMOKE ? 1 : argopt("reps", 21) @@ -36,7 +36,9 @@ const SUMMARY_PATH = joinpath(OUTDIR, "summary_blocking_model.txt") tup(b::Blocking) = (b.m_block, b.k_block, b.n_block) named_points(rows) = [(string(k), tup(v)) for (k, v) in pairs(rows) if v !== nothing] -named_rows(::Type{T}) where {T <: Real} = (model = _modelled_blocking(PROFILE, T), fallback = _fallback_blocking(T)) +detected = PROFILE.l1d.bytes > 0 && PROFILE.l2.bytes > 0 +named_rows(kernel) = (model = detected ? default_blocking(kernel, PROFILE) : nothing, fallback = default_blocking(kernel, unknown_target())) +host_kernel(::Type{T}) where {T} = kernel_from_shape(derived_shape(PROFILE, T), T) csv = open(CSV_PATH, "w") println(csv, "set,dtype,shape,Ma,Ka,Na,m_block_requested,k_block_requested,n_block_requested,m_block,k_block,n_block,reps,median_seconds,gflops") @@ -114,8 +116,8 @@ raw = NamedTuple[] crng = MersenneTwister(0x0CA9A121) canaries = [run_canary(crng, "start")] for T in DTYPES - kernel = _default_kernel(T) - rows = named_rows(T) + kernel = host_kernel(T) + rows = named_rows(kernel) println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) rows.model === nothing && @warn "no cache geometry detected: model undefined" grid = [("wide", c) for c in thin(WIDE[T])] @@ -131,8 +133,8 @@ for T in DTYPES push!(canaries, run_canary(crng, "after-$T")) end for T in CDTYPES - kernel = _default_kernel(T) - rows = map(v -> v === nothing ? nothing : _scale_blocking(v, typeof(kernel)), named_rows(real(T))) + kernel = host_kernel(T) + rows = named_rows(kernel) println("\n$T kernel $(tile_size(kernel, 1))x$(tile_size(kernel, 2))/W$(lanewidth(kernel)) ", rows) for spec in thin(MAIN_SHAPES) sweep_shape!(raw, kernel, T, spec, named_points(rows)) diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index f365607..f0c45f5 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/line_packing.jl") # --- Microkernels: add one packed K block to a register tile, store into C --- include("microkernels/interface.jl") @@ -36,8 +37,6 @@ include("microkernels/stores.jl") include("planning/labels.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/execute.jl b/src/execution/execute.jl index bbcf4d9..8cd3818 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -59,14 +59,14 @@ 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, - plan.blocking.k_block, _is_split(plan.mpack), _is_split(plan.npack) + plan.blocking.k_block, 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. `k_block`: the requested K block, `nothing` for -# the default. `split_a`/`split_b`: the plan packs A/B line by line. +# `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, k_block, @@ -89,11 +89,10 @@ end 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, k_length::Int, k_block) where {T} +# rather than in a C of narrower eltype. Static `false` unless C is narrower. +@inline function _c_panel_needed(::Type{T}, Cstorage, k_length::Int, k_block::Int) where {T} sizeof(real(eltype(Cstorage))) < sizeof(real(T)) || return false - return k_length > (k_block === nothing ? _resolved_defaults(T).real_row.k_block : k_block) + return k_length > k_block end # Run `path` behind a dynamic call. The plan crosses in the workspace slot @@ -138,12 +137,18 @@ end # The path predicted before the kernel barrier, so its callee is specialised # 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, +@inline _path_hint(::_Execute, req::_PlanRequest{T}, shape, ::Val{K}) where {T, K} = _select_path( + T, _unpacked_b_kernel_eligible(K), req.Astorage, req.Bstorage, req.Cstorage, req.mgroup, req.ngroup, req.kgroup, nothing, - req.k_block + req.k_block === nothing ? _default_k_block(T, req.Cstorage, shape, K) : req.k_block ) +# The kernel's default `k_block`, which decides the path only under a C +# narrower than `T`; elsewhere a static placeholder. +@inline _default_k_block(::Type{T}, Cstorage, shape, ::Type{K}) where {T, K} = + sizeof(real(eltype(Cstorage))) < sizeof(real(T)) ? + kernel_blocking(target_profile(), T, K, shape[1], shape[2]).k_block : 0 + @inline _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} = (_execute_hinted!(plan, e.alpha, e.beta, hint); nothing) @inline _continue(e::_Execute{T}, plan::ContractPlan{T}, ::Nothing) where {T} = @@ -166,12 +171,12 @@ end end @inline _splits(plan::ContractPlan, hint) = false -@inline _splits(plan::ContractPlan, ::Union{_NestPath, _PanelPath}) = _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) + 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 + 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 @@ -295,8 +300,8 @@ function _execute_nest!( btransform = plan.btransform 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 + 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. @@ -366,7 +371,7 @@ function _execute_nest!( beta_eff = firstpanel ? betaT : one(T) if split_b - _pack_block_transposed!( + pack_block_by_lines!( packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), sliver_spec(kernel, 2), plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, plan.npack @@ -405,7 +410,7 @@ function _execute_nest!( checked_span_bounds(cplan.Cbase, rng_mC, rng_nC, lenC) if split_a - _pack_block_transposed!( + pack_block_by_lines!( packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), sliver_spec(kernel, 1), plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, plan.mpack diff --git a/src/hardware/target.jl b/src/hardware/target.jl index ed1e7ae..9dd0263 100644 --- a/src/hardware/target.jl +++ b/src/hardware/target.jl @@ -19,7 +19,9 @@ CacheLevel() = CacheLevel(0, 0, 0) TargetProfile(isa, cpu_name, l1d, l2, l3) What was detected about the host CPU. `isa` is one of `:avx512`, `:avx2`, -`:neon` or `:unknown`, and fixes `vector_bytes` and `nregisters`. +`:neon` or `:unknown`, and fixes `vector_bytes` and `nregisters`. The core's +share of L2 and L3 (`l2_share`, `l3_share`, in bytes) and whether 512-bit FMAs +are double-pumped (`double_pumped`) are derived as well. """ struct TargetProfile isa::Symbol @@ -29,14 +31,28 @@ struct TargetProfile l1d::CacheLevel l2::CacheLevel l3::CacheLevel + # Derived here: planning reads them on every call. + l2_share::Int + l3_share::Int + double_pumped::Bool function TargetProfile(isa::Symbol, cpu_name::AbstractString, l1d::CacheLevel, l2::CacheLevel, l3::CacheLevel) isa in (:avx512, :avx2, :neon, :unknown) || throw(ArgumentError("TargetProfile: unknown ISA $(repr(isa))")) - return new(isa, cpu_name, isa_vector_bytes(isa), isa_nregisters(isa), l1d, l2, l3) + return new( + isa, cpu_name, isa_vector_bytes(isa), isa_nregisters(isa), l1d, l2, l3, + core_bytes(l1d, l2), core_bytes(l1d, l3), cpu_name in _DOUBLE_PUMPED_CPUS + ) end end +# This core's share of `level`: the CPUs sharing it over the SMT threads per +# core, which are the CPUs sharing L1d. +core_bytes(l1d::CacheLevel, level::CacheLevel) = level.bytes ÷ max(1, level.sharing ÷ max(1, l1d.sharing)) + +# AMD's AVX-512 cores, which double-pump 512-bit FMAs. +const _DOUBLE_PUMPED_CPUS = ("znver4", "znver5") + unknown_target() = TargetProfile(:unknown, "", CacheLevel(), CacheLevel(), CacheLevel()) # The register count is the ISA's, not the lane width's: AVX512VL gives 32 ymm registers. @@ -174,7 +190,12 @@ target_profile() = TARGET[] init_target!() = (TARGET[] = detect_target(); nothing) -# This core's share of `level`: the CPUs sharing it over the SMT threads per -# core, which are the CPUs sharing L1d. -core_bytes(profile::TargetProfile, level::CacheLevel) = - level.bytes ÷ max(1, level.sharing ÷ max(1, profile.l1d.sharing)) +# The core's private L2 share, or 1 MB when undetected. +l2_core_bytes(profile::TargetProfile) = profile.l2_share > 0 ? profile.l2_share : 1 << 20 + +line_bytes(profile::TargetProfile) = profile.l1d.line > 0 ? profile.l1d.line : 64 + +# A K step at least a page apart: a chain of demand misses no prefetcher +# follows, which the K-order model charges this factor. +const _K_WALK_FAR_BYTES = 4096 +const _K_WALK_FAR_PENALTY = 3 diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index 40af38b..3c3809a 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -52,7 +52,7 @@ end end @noinline function _qs_build_task_workspace!(pool::Dict{DataType, ContractWorkspace}, ::Type{T}) where {T} - kernel = _default_kernel(T) + kernel = kernel_from_shape(derived_shape(target_profile(), T), T) new_ws = ContractWorkspace(T, kernel, default_blocking(kernel), false, TO.DefaultAllocator()) pool[T] = new_ws return new_ws diff --git a/src/planning/pack_split.jl b/src/packing/line_packing.jl similarity index 54% rename from src/planning/pack_split.jl rename to src/packing/line_packing.jl index 2ef0a36..a2fd1b0 100644 --- a/src/planning/pack_split.jl +++ b/src/packing/line_packing.jl @@ -4,7 +4,7 @@ # 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. +# (`pack_block_by_lines!`); the packed format and the K order are unchanged. """ PackSplit(q, dj, L, kinner, Eq, E) @@ -27,20 +27,21 @@ 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 `k_block_requested` when the K extent clamps `k_block`. -# `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, k_block::Int, - eff::Int, rounded::Int, k_block_requested::Int, l2bytes::Union{Int, Nothing} = nothing - ) where {D} +is_split(s::PackSplit) = s.L != 0 + +# The split for the free group `g` of the operand packed as `spec` (A for M, B +# for N), whose storage eltype takes `S` bytes, and the block extent it needs, +# or `(eff, _NO_SPLIT)`. `eff` is the plan's extent, `rounded` the requested +# one rounded to the tile; a split block may take the budget the blocking +# reserved for `k_block_requested` when the K extent clamps `k_block`. +# `l2bytes = nothing`: `split_capacity`, doubled for complex, whose block walk +# costs more per element. +@inline function pack_split( + g::AxisGroup{D}, kg::AxisGroup, spec::SliverSpec{I, R, F}, S::Int, k_block::Int, + eff::Int, rounded::Int, k_block_requested::Int; l2bytes::Union{Int, Nothing} = nothing + ) where {D, I, R, F} D < 256 || return (eff, _NO_SPLIT) - line = _K_LINE_BYTES + line = line_bytes(target_profile()) d1 = 0 for d in 1:D g.lengths[d] > 1 && (d1 = d; break) @@ -51,16 +52,15 @@ _is_split(s::PackSplit) = s.L != 0 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, k_block, eff, rounded, k_block_requested, l2bytes) + return pack_split_window(g, kg, I, d1, dj, R, S, F !== RealFormat, line, k_block, eff, rounded, k_block_requested, l2bytes) end # Out of line: shared by every eltype, and most plans never get here. -@noinline function _pack_split_window( +@noinline function pack_split_window( g::AxisGroup, kg::AxisGroup{DK}, kmap::Int, d1::Int, dj::Int, R::Int, S::Int, complex::Bool, - k_block::Int, eff::Int, rounded::Int, k_block_requested::Int, l2bytes::Union{Int, Nothing} + line::Int, k_block::Int, eff::Int, rounded::Int, k_block_requested::Int, l2bytes::Union{Int, Nothing} ) where {DK} - line = _K_LINE_BYTES - L = _largest_divisor_upto(g.lengths[dj], max(1, line ÷ S)) + L = largest_divisor_upto(g.lengths[dj], max(1, line ÷ S)) L >= 2 || return (eff, _NO_SPLIT) ks = 0 @@ -81,26 +81,25 @@ end end klines = ks * S >= line ? k_block : cld(k_block * ks * S, line) aliased = abs(g.strides[1][d1]) * S % _K_WALK_FAR_BYTES == 0 - cap = something(l2bytes, _split_capacity(target_profile(), aliased) << complex) + 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, k_block, k_block_requested, kinner) + return pack_split_dynamic(g.lengths, d1, dj, L, psi, R, eff, rounded, k_block, k_block_requested, 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) +function split_capacity(profile::TargetProfile, aliased::Bool) l2 = l2_core_bytes(profile) - l3 = profile.l3 - (aliased || l3.bytes <= 0) && return l2 - return l2 + core_bytes(profile, l3) + (aliased || profile.l3.bytes <= 0) && return l2 + return l2 + profile.l3_share 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} +@noinline pack_split_dynamic(args...) = Base.inferencebarrier(pack_split_blocks)(args...)::Tuple{Int, PackSplit} -function _pack_split_blocks( +function pack_split_blocks( lengths::NTuple{D, Int}, d1::Int, dj::Int, L::Int, psi::Int, R::Int, eff::Int, rounded::Int, k_block::Int, k_block_requested::Int, kinner::Bool ) where {D} @@ -128,11 +127,11 @@ function _pack_split_blocks( (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(prod(lengths), R)) + neweff = min(budget ÷ blk * blk, roundup(prod(lengths), R)) return (neweff, PackSplit(q, dj, L, kinner, Eq, E)) end -@inline function _largest_divisor_upto(n::Int, cap::Int) +@inline function largest_divisor_upto(n::Int, cap::Int) for L in min(n, cap):-1:1 n % L == 0 && return L end @@ -141,15 +140,15 @@ 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) +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) +# 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) @@ -161,3 +160,56 @@ end k == dj + 2 && return p == 0 ? len(dj) ÷ L : L * str(dj) return p == 0 ? len(k - 2) : str(k - 2) end + +# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated by +# a split group, whole `E x L` groups) by `k_block_length` 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_by_lines!( + buffer::PK, spec::SliverSpec{I, R}, storage::ST, base::Int, fbuf::Vector{Int}, + kaxis::KA, transform::F, fcount::Int, k_block_length::Int, split::PackSplit + ) where {PK, I, R, ST, KA <: AbstractVector{Int}, 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!( + buffer, spec, sliver_width(spec) * k_block_length, 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 1:k_block_length + kbase = @inbounds base + 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 1:k_block_length + kbase = @inbounds base + 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) * sliver_width(spec) * k_block_length + for p in 1:k_block_length, t in (valid + 1):R + emit_padding!(buffer, spec, rb, t, p) + end + end + return nothing +end + +# Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane +# `i % R + 1`) at K step `p`. +@inline function pack_line_element!( + buffer::PK, spec::SliverSpec{I, R}, panel::Int, x, i::Int, p::Int + ) where {PK, I, R} + r, t = divrem(i, R) + emit_value!(buffer, spec, r * panel, t + 1, p, convert(element_type(spec), x)) + return nothing +end diff --git a/src/packing/panel.jl b/src/packing/panel.jl index 661673f..f3b2a3f 100644 --- a/src/packing/panel.jl +++ b/src/packing/panel.jl @@ -24,58 +24,3 @@ Base.eltype(::Type{PackedPanel{T}}) where {T} = T @inline panel_load(p::PackedPanel{T}, o::Int) where {T} = unsafe_load(p.ptr + sizeof(T) * o) @inline panel_store!(p::PackedPanel{T}, o::Int, x::T) where {T} = unsafe_store!(p.ptr + sizeof(T) * o, x) - -# --- Line-by-line packing of a split macro block --- - -# Pack a macro block of `fcount` free coordinates (offsets `fbuf`, enumerated by -# a split group, whole `E x L` groups) by `k_block_length` 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!( - buffer::PK, spec::SliverSpec{I, R}, storage::ST, base::Int, fbuf::Vector{Int}, - kaxis::KA, transform::F, fcount::Int, k_block_length::Int, split - ) where {PK, I, R, ST, KA <: AbstractVector{Int}, 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!( - buffer, spec, sliver_width(spec) * k_block_length, 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 1:k_block_length - kbase = @inbounds base + 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 1:k_block_length - kbase = @inbounds base + 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) * sliver_width(spec) * k_block_length - for p in 1:k_block_length, t in (valid + 1):R - emit_padding!(buffer, spec, rb, t, p) - end - end - return nothing -end - -# Element `x` of zero-based block coordinate `i` (sliver `i ÷ R`, lane -# `i % R + 1`) at K step `p`. -@inline function _pack_line_element!( - buffer::PK, spec::SliverSpec{I, R}, panel::Int, x, i::Int, p::Int - ) where {PK, I, R} - r, t = divrem(i, R) - emit_value!(buffer, spec, r * panel, t + 1, p, convert(element_type(spec), x)) - return nothing -end diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index d7a1661..152ad99 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -19,55 +19,57 @@ struct Blocking end end -# Complex blocking is the real row divided by the packed reals per element of -# each operand, so every kernel type gets the same packed BYTE budget (1m's -# `m_block` is half planar's). -@inline function _scale_blocking(base::Blocking, ::Type{K}) where {K <: Microkernel} +""" + default_blocking(kernel, profile = target_profile()) -> Blocking + +Cache-blocking factors for `kernel` on `profile`: an analytical model of the +cache geometry, or fixed constants where L1d or L2 is undetected. +`plan_contract` rounds `m_block`/`n_block` to multiples of the tile size and +clamps all three to the contraction's extents. +""" +default_blocking(kernel, profile::TargetProfile = target_profile()) = + kernel_blocking(profile, scalartype(kernel), typeof(kernel), tile_size(kernel)...) + +# The default blocking of a kernel of type `K` for `T` with an `MR x NR` tile. +# The fallback rows are divided by the packed reals per element of each +# operand, so every kernel type gets the same packed byte budget. +@inline function kernel_blocking(profile::TargetProfile, ::Type{T}, ::Type{K}, MR::Int, NR::Int) where {T, K} + R = real(T) a_reals, b_reals = map(reals_per_element, pack_formats(K)) + modelled = modelled_blocking(profile, MR, NR, a_reals * sizeof(R), b_reals * sizeof(R)) + modelled === nothing || return modelled + base = fallback_blocking(R) return Blocking(max(1, base.m_block ÷ a_reals), base.k_block, max(1, base.n_block ÷ b_reals)) end # For a host whose L1d or L2 size is undetected. -_fallback_blocking(::Type{Float64}) = Blocking(128, 256, 768) -_fallback_blocking(::Type{Float32}) = Blocking(96, 768, 1152) +fallback_blocking(::Type{Float64}) = Blocking(128, 256, 768) +fallback_blocking(::Type{Float32}) = Blocking(96, 768, 1152) -_real_blocking_row(profile::TargetProfile, ::Type{T}) where {T <: Union{Float32, Float64}} = - something(_modelled_blocking(profile, T), _fallback_blocking(T)) - -# Analytical real row from the cache geometry. Per (N, K) block the packed B -# panel is `n_block*k_block` elements, per (N, K, M) block the A block +# Analytical blocking from the cache geometry, with `SA`/`SB` the packed bytes +# per element of A and B. Per (N, K) block the packed B panel is +# `n_block*k_block` elements, per (N, K, M) block the A block # `m_block*k_block`, and one `NR x k_block` B sliver is reused by every A # sliver of the block: # -# NR*k_block*S <= L1/2 the reused B sliver -# m_block*k_block*S <= L2core/2 the A block -# n_block*k_block*S <= L2core + L3core the B panel +# NR*k_block*SB <= L1/2 the reused B sliver +# m_block*k_block*SA <= L2core/2 the A block +# n_block*k_block*SB <= L2core + L3core the B panel # # rounded down to MR/NR multiples, shared levels divided by the cores sharing # them. `n_block` takes the whole per-core capacity, not half the L3: an # oversized B panel only re-streams B, an undersized one repacks A once per N # block. # `nothing` when L1d or L2 is undetected. -function _modelled_blocking(profile::TargetProfile, ::Type{T}, MR::Int, NR::Int) where {T} - l1, l2, l3 = profile.l1d, profile.l2, profile.l3 - (l1.bytes > 0 && l2.bytes > 0) || return nothing - l2core = core_bytes(profile, l2) - l3core = core_bytes(profile, l3) - S = sizeof(T) - k_block = max(1, (l1.bytes ÷ 2) ÷ (NR * S)) - m_block = max(MR, ((l2core ÷ 2) ÷ (k_block * S)) ÷ MR * MR) - n_block = max(NR, ((l2core + l3core) ÷ (k_block * S)) ÷ NR * NR) +@inline function modelled_blocking(profile::TargetProfile, MR::Int, NR::Int, SA::Int, SB::Int) + l1 = profile.l1d + (l1.bytes > 0 && profile.l2.bytes > 0) || return nothing + l2core, l3core = profile.l2_share, profile.l3_share + k_block = max(1, (l1.bytes ÷ 2) ÷ (NR * SB)) + m_block = max(MR, ((l2core ÷ 2) ÷ (k_block * SA)) ÷ MR * MR) + n_block = max(NR, ((l2core + l3core) ÷ (k_block * SB)) ÷ NR * NR) return Blocking(m_block, k_block, n_block) end -function _modelled_blocking(profile::TargetProfile, ::Type{T}) where {T <: Real} - MR, NR, _ = derived_shape(profile, T) - return _modelled_blocking(profile, T, MR, NR) -end - -# For a bare scalar type: the fallback row, independent of the host. -default_blocking(::Type{Float64}) = _fallback_blocking(Float64) -default_blocking(::Type{Float32}) = _fallback_blocking(Float32) - # Smallest multiple of `n` that is >= `x` (`x >= 0`, `n >= 1`). -@inline _roundup(x::Int, n::Int) = cld(x, n) * n +@inline roundup(x::Int, n::Int) = cld(x, n) * n diff --git a/src/planning/defaults.jl b/src/planning/defaults.jl deleted file mode 100644 index 5407e7a..0000000 --- a/src/planning/defaults.jl +++ /dev/null @@ -1,90 +0,0 @@ -# The engine's per-host defaults for an element type, resolved once per -# (detected profile, eltype) and cached: every lookup in kernel_selection.jl and -# blocking.jl is keyed on `Val(profile.isa)`, a dynamic dispatch far too slow to -# repeat on every `plan_contract`. - -# The default `shape`, the `fitted` small-M demotion shape, the AVX-512 complex -# small-M FMAddSub candidates, the unscaled real blocking row and the core's L2 -# share, all derived from `profile`, which is the cache key. -struct ResolvedDefaults - profile::TargetProfile - shape::NTuple{3, Int} - fitted::NTuple{3, Int} - small_m::Vector{NTuple{3, Int}} - real_row::Blocking - l2_core::Int -end - -# The core's private L2 share, or 1 MB when undetected. -function l2_core_bytes(profile::TargetProfile) - profile.l2.bytes > 0 || return 1 << 20 - return core_bytes(profile, profile.l2) -end - -# Through the per-eltype cache; an eltype without a slot fails later in planning. -@inline l2_core_bytes(::Type{T}) where {T} = - _defaults_slot(T) === nothing ? l2_core_bytes(target_profile()) : _resolved_defaults(T).l2_core - -function _resolve_defaults(profile::TargetProfile, ::Type{T}) where {T} - K = default_kernel_type(T) - # First, so an element type with no menu throws from here. - shape = derived_shape(profile, T, K) - fitted = fitted_shape(profile, T, K) - small_m = small_m_candidates(Val(profile.isa), profile, T) - real_row = _real_blocking_row(profile, real(T)) - l2_core = l2_core_bytes(profile) - return ResolvedDefaults(profile, shape, fitted, small_m, real_row, l2_core) -end - -# One slot per supported element type (a method dispatch, not a `Dict` probe); -# other types take the uncached path and fail there. -const _DEFAULTS_F64 = Ref{Union{Nothing, ResolvedDefaults}}(nothing) -const _DEFAULTS_F32 = Ref{Union{Nothing, ResolvedDefaults}}(nothing) -const _DEFAULTS_C64 = Ref{Union{Nothing, ResolvedDefaults}}(nothing) -const _DEFAULTS_C32 = Ref{Union{Nothing, ResolvedDefaults}}(nothing) -_defaults_slot(::Type{Float64}) = _DEFAULTS_F64 -_defaults_slot(::Type{Float32}) = _DEFAULTS_F32 -_defaults_slot(::Type{ComplexF64}) = _DEFAULTS_C64 -_defaults_slot(::Type{ComplexF32}) = _DEFAULTS_C32 -_defaults_slot(::Type) = nothing - -# Refilled whenever the slot's profile is not `===` the current one, so a -# profile installed after load (as test/forced_isa_runner.jl does) takes effect -# without an invalidation hook. Unlocked: a racing refill stores the same value. -@inline function _resolved_defaults(::Type{T}) where {T} - slot = _defaults_slot(T) - profile = target_profile() - slot === nothing && return _resolve_defaults(profile, T) - cached = slot[] - (cached !== nothing && cached.profile === profile) && return cached - return _refill_defaults!(slot, profile, T) -end - -@noinline function _refill_defaults!(slot, profile::TargetProfile, ::Type{T}) where {T} - fresh = _resolve_defaults(profile, T) - slot[] = fresh - return fresh -end - -_default_kernel(::Type{T}) where {T} = - kernel_from_shape(_resolved_defaults(T).shape, T, default_kernel_type(T)) - -# `@noinline`: the return type is the Union of `T`'s menu kernels. -@noinline function _default_kernel(::Type{T}, m_length::Int, n_length::Int) where {T} - shape, kernel_type = select_shape(T, default_kernel_type(T), m_length, 0, m_length) - return kernel_from_shape(shape, T, kernel_type) -end - -""" - default_blocking(kernel) -> Blocking - -Cache-blocking factors for `kernel` on this host: an analytical model of the -detected cache geometry ([`target_profile`](@ref)), or fixed constants where -L1d or L2 is undetected, scaled for complex kernels by packed reals per -element. `plan_contract` rounds `m_block`/`n_block` to multiples of the tile -size and clamps all three to the contraction's extents. -""" -function default_blocking(kernel) - T = scalartype(kernel) - return _scale_blocking(_resolved_defaults(T).real_row, typeof(kernel)) -end diff --git a/src/planning/kernel_selection.jl b/src/planning/kernel_selection.jl index 81f5674..d021407 100644 --- a/src/planning/kernel_selection.jl +++ b/src/planning/kernel_selection.jl @@ -1,7 +1,5 @@ # Kernel selection: the microkernel `plan_contract` builds when the caller does -# not name one. Everything here is a pure function of a `TargetProfile` and `T`; -# src/planning/defaults.jl caches the host-dependent results per eltype, since -# the `Val(profile.isa)` dispatch below is dynamic. +# not name one, a pure function of a `TargetProfile` and `T`. # Kernel types, unparameterised, name the complex-arithmetic schemes here. # `OneMKernel`/`FMAddSubKernel` are chosen only by naming the kernel (FMAddSub @@ -122,25 +120,30 @@ kernel_from_shape(shape::Tuple{Int, Int, Int}, ::Type{T}, ::Val{K}) where {T, K} # ---------------------------------------------------------------------------- # Per-ISA shapes, consulted first. None for real types: the rule below is the optimum. -shape_override(key::Val, ::Type{T}) where {T} = shape_override(key, T, default_kernel_type(T)) -shape_override(::Val, ::Type, ::Type) = nothing -# On AVX-512 the rule's planar `MR = 2W, NR = 6` spills; `NR = 3` tiles do not. -shape_override(::Val{:avx512}, ::Type{ComplexF64}, ::Type{<:PlanarKernel}) = (24, 3, 8) -shape_override(::Val{:avx512}, ::Type{ComplexF32}, ::Type{<:PlanarKernel}) = (48, 3, 16) -# AVX2 has 16 registers: at `MV = 1` planar `NR = 6` needs all 16, `NR = 5` needs 14. -shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::Type{<:PlanarKernel}) = (4, 5, 4) -shape_override(::Val{:avx2}, ::Type{ComplexF32}, ::Type{<:PlanarKernel}) = (8, 5, 8) -# The AVX2-sized 1m shape, for a caller asking for `OneMKernel`'s shape (the -# fit would hand it an AVX-512 shape). -shape_override(::Val{:avx2}, ::Type{ComplexF64}, ::Type{<:OneMKernel}) = (4, 6, 4) +shape_override(isa::Symbol, ::Type{T}) where {T} = shape_override(isa, T, default_kernel_type(T)) +function shape_override(isa::Symbol, ::Type{T}, ::Type{K}) where {T, K} + if K <: PlanarKernel && T === ComplexF64 + # On AVX-512 the rule's planar `MR = 2W, NR = 6` spills; `NR = 3` tiles + # do not. AVX2 has 16 registers: at `MV = 1` planar `NR = 6` needs all + # 16, `NR = 5` needs 14. + isa === :avx512 && return (24, 3, 8) + isa === :avx2 && return (4, 5, 4) + elseif K <: PlanarKernel && T === ComplexF32 + isa === :avx512 && return (48, 3, 16) + isa === :avx2 && return (8, 5, 8) + elseif K <: OneMKernel && T === ComplexF64 + # The AVX2-sized 1m shape, for a caller asking for `OneMKernel`'s shape + # (the fit would hand it an AVX-512 shape). + isa === :avx2 && return (4, 6, 4) + end + return nothing +end # Where the `MR = MV*W` rule applies. Not on NEON (2W on 128-bit lanes is no # better than the fallback), nor complex off AVX-512 (planar's two accumulator # planes leave AVX2's 16 registers nothing spare). -rule_applies(::Val{:avx512}, ::Type{<:SIMDKernel}) = true -rule_applies(::Val{:avx512}, ::Type{<:Union{PlanarKernel, OneMKernel}}) = true -rule_applies(::Val{:avx2}, ::Type{<:SIMDKernel}) = true -rule_applies(::Val, ::Type) = false +rule_applies(isa::Symbol, ::Type{K}) where {K} = + isa === :avx512 ? K <: Union{SIMDKernel, PlanarKernel, OneMKernel} : isa === :avx2 && K <: SIMDKernel # `W` real lanes per register, `MV` A vectors per column. rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = @@ -149,17 +152,13 @@ rule_shape(vb::Int, ::Type{T}, mv::Int) where {T} = # MV = 4 for the real kernel on AVX-512: the MV = 2 tile is front-end bound on # cores with 2 FMA ports, and 4*6 accumulators + 4 A vectors + 2 still fit 32 # registers. Complex kernels keep MV = 2 (planar already spills there). -rule_mv(::Val{:avx512}, ::Type{<:SIMDKernel}) = 4 -rule_mv(::Val, ::Type) = 2 - -# AMD's AVX-512 cores double-pump 512-bit FMAs, so their MV = 2 tile is not -# front-end bound and the taller tile only adds edge and store cost. -const _MV4_UNPROFITABLE_CPUS = ("znver4", "znver5") +rule_mv(isa::Symbol, ::Type{K}) where {K} = isa === :avx512 && K <: SIMDKernel ? 4 : 2 -profile_mv(profile::TargetProfile, ::Type{K}) where {K} = rule_mv(Val(profile.isa), K) -function profile_mv(profile::TargetProfile, ::Type{K}) where {K <: SIMDKernel} - mv = rule_mv(Val(profile.isa), K) - return (mv == 4 && profile.cpu_name in _MV4_UNPROFITABLE_CPUS) ? 2 : mv +# Where 512-bit FMAs are double-pumped the MV = 2 tile is not front-end bound, +# and the taller tile only adds edge and store cost. +function profile_mv(profile::TargetProfile, ::Type{K}) where {K} + mv = rule_mv(profile.isa, K) + return (mv == 4 && profile.double_pumped) ? 2 : mv end # The register shape for `T` with kernel type `K` on `profile`: an override @@ -168,11 +167,10 @@ derived_shape(profile::TargetProfile, ::Type{T}) where {T} = derived_shape(profile, T, default_kernel_type(T)) function derived_shape(profile::TargetProfile, ::Type{T}, ::Type{K}) where {T, K} - key = Val(profile.isa) - ovr = shape_override(key, T, K) + ovr = shape_override(profile.isa, T, K) ovr === nothing || return ovr vb = profile.vector_bytes - if rule_applies(key, K) && vb > 0 + if rule_applies(profile.isa, K) && vb > 0 shape = rule_shape(vb, T, profile_mv(profile, K)) shape in kernel_shapes(T, K) && return shape end @@ -222,7 +220,7 @@ with, for compute type `T` and kernel type `K = default_kernel_type(T, TA, TB)`, extents `m_length`/`k_length` and C's unit-stride run along M (`run`; `m_length` when no layout is known). The steps, in order: - 1. The host's shape for `T` (`derived_shape`, cached in `ResolvedDefaults`). + 1. The host's shape for `T` (`derived_shape`). 2. Extent (real): the MV = 4 shape steps down to MV = 2 where `m_length` pads less (`extent_shape`). 3. Small M: an `m_length` below one tile takes the fitted shape, or for @@ -235,15 +233,15 @@ the kernel type as a `Val`, since inference widens a type inside a returned tuple to `UnionAll`. """ @inline function select_shape(::Type{T}, ::Type{K}, m_length::Int, k_length::Int, run::Int) where {T, K} - d = _resolved_defaults(T) - shape = extent_shape(d.shape, T, K, m_length) + profile = target_profile() + shape = extent_shape(derived_shape(profile, T, K), T, K, m_length) if m_length > 0 && m_length < shape[1] # Static, so a real `T`'s kernel type stays concrete. if T <: Complex - small = small_m_shape(d.small_m, m_length) + small = small_m_shape(profile, T, m_length) small === nothing || return (small, Val(FMAddSubKernel)) end - shape = d.fitted + shape = fitted_shape(profile, T, K) end return (fit_to_run(shape, T, K, m_length, k_length, run), Val(K)) end @@ -272,18 +270,14 @@ end # Small-M demotion for complex `T` on AVX-512, where the planar fitted shape is # the spilling `MR = 2W` tile: the native-width FMAddSub shape that pads -# `m_length` least, ties by the larger tile. `small_m_candidates` is the -# cached, host-dependent half (empty where the rule does not apply). -small_m_candidates(::Val, ::TargetProfile, ::Type) = NTuple{3, Int}[] -function small_m_candidates(::Val{:avx512}, profile::TargetProfile, ::Type{T}) where {T <: Complex} +# `m_length` least, ties by the larger tile. +function small_m_shape(profile::TargetProfile, ::Type{T}, m_length::Int) where {T} + profile.isa === :avx512 || return nothing lanes = profile.vector_bytes ÷ sizeof(real(T)) - return [shape for shape in kernel_shapes(T, FMAddSubKernel) if shape[3] == lanes] -end - -function small_m_shape(candidates::Vector{NTuple{3, Int}}, m_length::Int) best = nothing - for shape in candidates - MR, NR, _ = shape + for shape in kernel_shapes(T, FMAddSubKernel) + MR, NR, W = shape + W == lanes || continue key = (-(cld(m_length, MR) * MR), MR * NR) if best === nothing || key > best[1] best = (key, shape) diff --git a/src/planning/labels.jl b/src/planning/labels.jl index e272b48..c224f1e 100644 --- a/src/planning/labels.jl +++ b/src/planning/labels.jl @@ -133,10 +133,6 @@ end # `min(L_u, line_bytes / stride_u)`. # ---------------------------------------------------------------------------- -const _K_WALK_FAR_BYTES = 4096 -const _K_WALK_FAR_PENALTY = 3 -const _K_LINE_BYTES = 64 - # What the cost model reads of one operand, gathered once per plan in `klabels` # order: each K label's extent and |stride|, the element count and size, and # whether its register slivers are whole lines (the free composite's first @@ -152,7 +148,7 @@ end @inline function k_operand( klabels::NTuple{DK, Int}, ind::NTuple{N, Int}, v::StridedView, - free::NTuple{DF, Int} + free::NTuple{DF, Int}, line::Int ) where {DK, N, DF} st = Base.strides(v) sz = size(v) @@ -161,7 +157,7 @@ end ps = map(pos, klabels) len = map(p -> sz[p], ps) kst = map(p -> abs(st[p]), ps) - line_elems = max(1, _K_LINE_BYTES ÷ S) + line_elems = max(1, line ÷ S) wholeline = false for l in free p = pos(l) @@ -174,7 +170,7 @@ end # Cost of packing one operand under the K order `perm`; `Qfree` is its free extent. function k_order_cost( - perm::NTuple{DK, Int}, op::KOperand{DK}, Qfree::Int, l2bytes::Int + perm::NTuple{DK, Int}, op::KOperand{DK}, Qfree::Int, l2bytes::Int, line::Int ) where {DK} len, st, S = op.len, op.st, op.S kfast = 0 @@ -201,16 +197,16 @@ function k_order_cost( # One element of `u` per line exactly when su*S > line/2 (len[u] >= 2); # tested first so the division is only paid where its value is used. suS = su * S - (2 * suS > _K_LINE_BYTES || u == kfast) && return op.n * walk + (2 * suS > line || u == kfast) && return op.n * walk faster = 1 @inbounds for i in perm i == u && break faster *= len[i] end - footprint = Int128(faster) * Int128(Qfree) * Int128(_K_LINE_BYTES) + footprint = Int128(faster) * Int128(Qfree) * Int128(line) footprint <= l2bytes && return op.n * walk - share = min(@inbounds(len[u]), _K_LINE_BYTES ÷ suS) + share = min(@inbounds(len[u]), line ÷ suS) return op.n * walk * share end @@ -223,14 +219,16 @@ end m_length::Int, n_length::Int, l2bytes::Union{Int, Nothing} = nothing ) where {DK, NA, NB, DM, DN} DK <= 1 && return klabels - opA = k_operand(klabels, indA, A, morder) - opB = k_operand(klabels, indB, B, norder) + profile = target_profile() + line = line_bytes(profile) + opA = k_operand(klabels, indA, A, morder, line) + opB = k_operand(klabels, indB, B, norder, line) id = ntuple(identity, Val(DK)) permA = TupleTools.sortperm(opA.st) permB = TupleTools.sortperm(opB.st) permA == id && permB == id && return klabels - l2 = l2bytes === nothing ? l2_core_bytes(eltype(A)) : l2bytes - cost(perm) = k_order_cost(perm, opA, m_length, l2) + k_order_cost(perm, opB, n_length, l2) + l2 = l2bytes === nothing ? l2_core_bytes(profile) : l2bytes + cost(perm) = k_order_cost(perm, opA, m_length, l2, line) + k_order_cost(perm, opB, n_length, l2, line) best = id bestcost = cost(id) for cand in (permA, permB) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index a3f28db..62c2265 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -301,10 +301,10 @@ end @inline function _plan_at_shape(shape, ::Val{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} vkernel = menu_val(shape, T, K) - # The execution path, predicted so the callee is specialised on it. Only a - # `Bool` crosses: a call union-split on the kernel type is emitted out of - # line and boxes `req`. - hint = _path_hint(req.f, req, _unpacked_b_kernel_eligible(K)) + # The execution path, predicted so the callee is specialised on it. From + # the shape and kernel type, not the kernel: a call union-split on the + # kernel type is emitted out of line and boxes `req`. + hint = _path_hint(req.f, req, shape, Val(K)) core = _strip_storage(req) slot = _barrier_slot!(req.workspace, typeof(core)) slot[] = core @@ -335,7 +335,7 @@ end ) # `nothing` for a continuation that does not execute (see src/execution/execute.jl). -@inline _path_hint(f, req::_PlanRequest, unpack_ok::Bool) = nothing +@inline _path_hint(f, req::_PlanRequest, shape, kernel_type) = nothing # Plan construction on a concrete kernel and transform pair (a named kernel's # transform Unions die here). `req.f` runs in here, on the concrete plan type. @@ -359,10 +359,10 @@ function _plan_contract( m_tile, n_tile = tile_size(kernel) # Empty extents: the drivers never read these; the floors keep them valid. - m_block_rounded = _roundup(requested.m_block, m_tile) - n_block_rounded = _roundup(requested.n_block, n_tile) - m_block = m_length == 0 ? m_tile : min(m_block_rounded, _roundup(m_length, m_tile)) - n_block = n_length == 0 ? n_tile : min(n_block_rounded, _roundup(n_length, n_tile)) + m_block_rounded = roundup(requested.m_block, m_tile) + n_block_rounded = roundup(requested.n_block, n_tile) + m_block = m_length == 0 ? m_tile : min(m_block_rounded, roundup(m_length, m_tile)) + n_block = n_length == 0 ? n_tile : min(n_block_rounded, roundup(n_length, n_tile)) k_block = k_length == 0 ? 1 : min(requested.k_block, k_length) panel = _c_panel_needed(T, req.Cstorage, k_length, k_block) mpack = npack = _NO_SPLIT @@ -371,14 +371,14 @@ function _plan_contract( # the panel holds N blocks in C's own N order, which a split N group does # not enumerate contiguously. if m_length > 0 && n_length > 0 && k_length > 0 - m_block, mpack = _pack_split( - req.mgroup, req.kgroup, 1, m_tile, sizeof(eltype(req.Astorage)), - !(a_format(kernel) isa RealFormat), k_block, m_block, m_block_rounded, requested.k_block + m_block, mpack = pack_split( + req.mgroup, req.kgroup, sliver_spec(kernel, 1), sizeof(eltype(req.Astorage)), + k_block, m_block, m_block_rounded, requested.k_block ) if !panel - n_block, npack = _pack_split( - req.ngroup, req.kgroup, 2, n_tile, sizeof(eltype(req.Bstorage)), - !(b_format(kernel) isa RealFormat), k_block, n_block, n_block_rounded, requested.k_block + n_block, npack = pack_split( + req.ngroup, req.kgroup, sliver_spec(kernel, 2), sizeof(eltype(req.Bstorage)), + k_block, n_block, n_block_rounded, requested.k_block ) end end diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 92b9808..6712dbc 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -258,9 +258,12 @@ end function _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; kw...) p = plan_contract(Cv, Av, iA, Bv, iB, iC; kw...) k, k_block, d, T = p.kernel, p.blocking.k_block, default_blocking(p.kernel), eltype(Cv) - split(g, map, R, eff, req) = _pack_split(g, p.kgroup, map, R, sizeof(T), T <: Complex, k_block, eff, cld(req, R) * R, d.k_block, 0) - (m_block, ms) = split(p.mgroup, 1, tile_size(k, 1), p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) - (n_block, ns) = split(p.ngroup, 2, tile_size(k, 2), p.blocking.n_block, d.n_block) + function split(g, i, eff, req) + R = tile_size(k, i) + return pack_split(g, p.kgroup, sliver_spec(k, i), sizeof(T), k_block, eff, cld(req, R) * R, d.k_block; l2bytes = 0) + end + (m_block, ms) = split(p.mgroup, 1, p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) + (n_block, ns) = split(p.ngroup, 2, p.blocking.n_block, d.n_block) q = plan_contract(Cv, Av, iA, Bv, iB, iC; kw..., kernel = k, m_block, n_block) @assert q.mgroup == p.mgroup && q.ngroup == p.ngroup return ContractPlan( @@ -273,10 +276,10 @@ end 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, tile_size(k, 1), 8, false, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block, cap)[2]) + splits(l2bytes) = is_split(pack_split(p.mgroup, p.kgroup, sliver_spec(k, 1), 8, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block; l2bytes)[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 = splits(QuasiStrided.split_capacity(target_profile(), true)) + @test 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 diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index ea354e0..62299ca 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -77,7 +77,7 @@ end 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.k_block + @test K ÷ 2 > default_blocking(auto_kernel(Float64, M)).k_block backend = QuasiStrided.QuasiStridedBackend(accumulator = acc) TO.tensorcontract!(Cv, Av, ((1,), (2,)), false, Bv, ((1,), (2,)), false, ((1, 2), ()), 1, beta, backend) else diff --git a/test/execution/test_workspace.jl b/test/execution/test_workspace.jl index 6fcf798..0b1bd65 100644 --- a/test/execution/test_workspace.jl +++ b/test/execution/test_workspace.jl @@ -15,7 +15,7 @@ _ws_lengths(ws) = [ # The shape is hardware- and extent-dependent: pin the resolution. @test plan.kernel isa QuasiStrided.SIMDKernel @test QuasiStrided.scalartype(plan.kernel) === T - @test plan.kernel === QuasiStrided._default_kernel(T, size(Amat, 1), size(Bmat, 2)) + @test plan.kernel === auto_kernel(T, size(Amat, 1)) execute!(plan, one(T), zero(T)) @test Cmat ≈ Amat * Bmat diff --git a/test/helpers.jl b/test/helpers.jl index 7cf9e68..23c7100 100644 --- a/test/helpers.jl +++ b/test/helpers.jl @@ -53,6 +53,14 @@ using QuasiStrided: TargetProfile, CacheLevel const VALID_ISAS = (:avx512, :avx2, :neon, :unknown) synthetic(isakey) = TargetProfile(isakey, "synthetic", CacheLevel(), CacheLevel(), CacheLevel()) +# The kernel `plan_contract` picks for `T` on this host: at its own shape, and +# for an M extent of `m_length` with C contiguous. +host_kernel(::Type{T}) where {T} = kernel_from_shape(derived_shape(target_profile(), T), T) +function auto_kernel(::Type{T}, m_length::Int) where {T} + shape, K = QuasiStrided.select_shape(T, QuasiStrided.default_kernel_type(T), m_length, 0, m_length) + return kernel_from_shape(shape, T, K) +end + # Permuted A (with a zero-stride axis), negative-stride B, offset sliced C. function scattered_fixture(::Type{T}, a_n = 32, k_n = 32, b_n = 8, n_n = 32) where {T} A2 = randn(MersenneTwister(11), T, a_n, k_n) diff --git a/test/microkernels/test_fmaddsub_kernel.jl b/test/microkernels/test_fmaddsub_kernel.jl index 555ac17..342bc72 100644 --- a/test/microkernels/test_fmaddsub_kernel.jl +++ b/test/microkernels/test_fmaddsub_kernel.jl @@ -140,7 +140,7 @@ const _QSF = QuasiStrided @test kernel_from_shape(s, T, FMAddSubKernel) isa FMAddSubKernel{s[1], s[2], T, s[3]} end @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, FMAddSubKernel) - @test !(_QSF._default_kernel(T, 1024, 1024) isa FMAddSubKernel) + @test !(auto_kernel(T, 1024) isa FMAddSubKernel) end end diff --git a/test/microkernels/test_onem_kernel.jl b/test/microkernels/test_onem_kernel.jl index 5549b9b..adffb7d 100644 --- a/test/microkernels/test_onem_kernel.jl +++ b/test/microkernels/test_onem_kernel.jl @@ -34,14 +34,15 @@ using QuasiStrided: OneMKernel, PlanarKernel, lanewidth, kernel_shapes, for T in (ComplexF64, ComplexF32) bm = default_blocking(OneMKernel(Val(8), Val(8), T, Val(8))) bp = default_blocking(PlanarKernel(Val(8), Val(8), T, Val(8))) - @test (bm.m_block, bm.k_block, bm.n_block) == (bp.m_block ÷ 2, bp.k_block, bp.n_block) # twice the packed A reals + # Twice the packed A reals, rounded to MR = 8. + @test (bm.k_block, bm.n_block) == (bp.k_block, bp.n_block) && bp.m_block - 2 * bm.m_block in (0, 8) for s in kernel_shapes(T, OneMKernel) @test kernel_from_shape(s, T, OneMKernel) isa OneMKernel{s[1], s[2], T, s[3]} end @test_throws ArgumentError kernel_from_shape((7, 7, 7), T, OneMKernel) # Never the default: method ranking does not transfer between machines. @test default_kernel_type(T) === PlanarKernel - @test QuasiStrided._default_kernel(T, 1024, 1024) isa PlanarKernel + @test auto_kernel(T, 1024) isa PlanarKernel end end diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index a9c7b63..a4a159b 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -10,24 +10,21 @@ using StridedViews: StridedView k = kernel_from_shape(derived_shape(u, T), T) @test k isa SIMDKernel @test (tile_size(k)..., lanewidth(k)) === fallback_shape(T) - @test _real_blocking_row(u, T) === _fallback_blocking(T) + @test default_blocking(k, u) === fallback_blocking(T) end - @test _fallback_blocking(Float64) === Blocking(128, 256, 768) - @test _fallback_blocking(Float32) === Blocking(96, 768, 1152) - # The type-argument form ignores detection. - @test default_blocking(Float64) === _fallback_blocking(Float64) - @test default_blocking(Float32) === _fallback_blocking(Float32) + @test fallback_blocking(Float64) === Blocking(128, 256, 768) + @test fallback_blocking(Float32) === Blocking(96, 768, 1152) end @testset "rule: MR = MV*W, NR = 6; MV = 4 on Intel :avx512, 2 on :avx2 and AMD" begin for (isakey, vb, MV) in ((:avx512, 64, 4), (:avx2, 32, 2)), T in (Float64, Float32) W = vb ÷ sizeof(T) @test derived_shape(synthetic(isakey), T) === (MV * W, NR_DEFAULT, W) - @test QuasiStrided.rule_mv(Val(isakey), SIMDKernel) == MV + @test QuasiStrided.rule_mv(isakey, SIMDKernel) == MV @test QuasiStrided.rule_shape(vb, T, MV) === (MV * W, NR_DEFAULT, W) end for K in (PlanarKernel, OneMKernel, FMAddSubKernel), key in (:avx512, :avx2) - @test QuasiStrided.rule_mv(Val(key), K) == 2 + @test QuasiStrided.rule_mv(key, K) == 2 end for (cpu, MV) in (("znver4", 2), ("znver5", 2), ("icelake-server", 4), ("cascadelake", 4)), T in (Float64, Float32) @@ -44,9 +41,9 @@ using StridedViews: StridedView end # No real override row on any ISA: the rule is the optimum. for T in (Float64, Float32), key in VALID_ISAS - @test shape_override(Val(key), T) === nothing + @test shape_override(key, T) === nothing end - @test rule_applies(Val(:avx2), SIMDKernel) && rule_applies(Val(:avx512), SIMDKernel) + @test rule_applies(:avx2, SIMDKernel) && rule_applies(:avx512, SIMDKernel) end @testset "every real menu shape is constructible and fits the register file" begin @@ -64,7 +61,7 @@ using StridedViews: StridedView # The MV = 2 sibling stays in the menu as the step-down target. @test (2 * W, NR_DEFAULT, W) in kernel_shapes(T) @test kernel_shapes(T, SIMDKernel) === kernel_shapes(T) - k = _default_kernel(T) + k = host_kernel(T) @test k isa SIMDKernel && scalartype(k) === T @test (tile_size(k)..., lanewidth(k)) in kernel_shapes(T) end @@ -81,11 +78,11 @@ using StridedViews: StridedView @testset "extent demotion when M cannot fill a register tile" begin for T in (Float64, Float32) - MR = tile_size(_default_kernel(T), 1) - @test tile_size(_default_kernel(T, 4 * MR, 256), 1) == MR - small = _default_kernel(T, 1, 256) + MR = tile_size(host_kernel(T), 1) + @test tile_size(auto_kernel(T, 4 * MR), 1) == MR + small = auto_kernel(T, 1) @test (tile_size(small)..., lanewidth(small)) === fallback_shape(T) - @test tile_size(_default_kernel(T, 0, 256), 1) == MR # empty must not demote + @test tile_size(auto_kernel(T, 0), 1) == MR # empty must not demote end end @@ -107,12 +104,12 @@ using StridedViews: StridedView @test ext(avx2, T, 3) === avx2 cs = derived_shape(synthetic(:avx512), ComplexF64) @test QuasiStrided.extent_shape(cs, ComplexF64, PlanarKernel, 3) === cs - # Through `_default_kernel`, where the host's shape is the tall one. + # Through `auto_kernel`, where the host's shape is the tall one. if derived_shape(target_profile(), T) === tall - @test tile_size(_default_kernel(T, MR + W, 256)) == (2 * W, NR) - @test tile_size(_default_kernel(T, 2 * W, 256), 1) == 2 * W - @test tile_size(_default_kernel(T, 2 * W - 1, 256), 1) == fallback_shape(T)[1] - @test tile_size(_default_kernel(T, 2 * MR, 256), 1) == MR + @test tile_size(auto_kernel(T, MR + W)) == (2 * W, NR) + @test tile_size(auto_kernel(T, 2 * W), 1) == 2 * W + @test tile_size(auto_kernel(T, 2 * W - 1), 1) == fallback_shape(T)[1] + @test tile_size(auto_kernel(T, 2 * MR), 1) == MR end end end @@ -195,33 +192,33 @@ end for (T, W, swept) in ((ComplexF64, 8, (24, 3, 8)), (ComplexF32, 16, (48, 3, 16))) # W counts REAL lanes; the AVX-512 override replaces the spilling rule shape. @test QuasiStrided.rule_shape(64, T, 2) === (2 * W, NR_DEFAULT, W) - @test shape_override(Val(:avx512), T) === swept + @test shape_override(:avx512, T) === swept @test derived_shape(synthetic(:avx512), T) === swept === first(kernel_shapes(T, PlanarKernel)) @test Set(kernel_shapes(T, PlanarKernel)) == Set( ( - (2 * W, NR_DEFAULT, W), swept, (W, 8, W), shape_override(Val(:avx2), T), + (2 * W, NR_DEFAULT, W), swept, (W, 8, W), shape_override(:avx2, T), (W ÷ 2, NR_DEFAULT, W ÷ 4), (W ÷ 4, NR_DEFAULT, W ÷ 4), ) ) # Every override row fits its ISA's register file with room to spare. for (key, nreg) in ((:avx512, 32), (:avx2, 16)) - ovr = shape_override(Val(key), T) + ovr = shape_override(key, T) @test planar_pressure(ovr...) < nreg @test ovr in kernel_shapes(T, PlanarKernel) end @test fallback_shape(T) === (8, NR_DEFAULT, fallback_shape(real(T))[3]) # Off :avx512 an override row wins whatever the width; without a row # the shape is fitted to the register budget. - @test derived_shape(synthetic(:avx2), T) === shape_override(Val(:avx2), T) + @test derived_shape(synthetic(:avx2), T) === shape_override(:avx2, T) for key in (:avx2, :neon, :unknown) - @test !rule_applies(Val(key), PlanarKernel) + @test !rule_applies(key, PlanarKernel) end for key in (:neon, :unknown) - @test shape_override(Val(key), T) === nothing + @test shape_override(key, T) === nothing end @test planar_pressure(derived_shape(synthetic(:unknown), T)...) <= 16 end - @test rule_applies(Val(:avx512), PlanarKernel) + @test rule_applies(:avx512, PlanarKernel) @test derived_shape(synthetic(:neon), ComplexF64) === (4, 6, 2) @test derived_shape(synthetic(:neon), ComplexF32) === (8, 6, 4) @test planar_pressure(fallback_shape(ComplexF64)...) == 30 @@ -305,31 +302,35 @@ end :avx512, "cascadelake", CacheLevel(32KiB, 64, 2), CacheLevel(1MiB, 64, 2), CacheLevel(25952256, 64, 16) ) - @test _modelled_blocking(nosmt, Float64, 8, 6) === Blocking(96, 341, 1728) - @test _modelled_blocking(nosmt, Float32, 16, 6) === Blocking(96, 682, 1728) - @test _modelled_blocking(smt2, Float64, 16, 6) === Blocking(192, 341, 1572) - @test _modelled_blocking(nosmt, Float64) === _modelled_blocking(nosmt, Float64, 8, 6) - for p in (nosmt, smt2), T in (Float64, Float32) - b = _modelled_blocking(p, T) - MR, NR, _ = derived_shape(p, T) - @test b.m_block % MR == 0 && b.n_block % NR == 0 - @test NR * b.k_block * sizeof(T) <= p.l1d.bytes ÷ 2 - @test _real_blocking_row(p, T) === b + @test modelled_blocking(nosmt, 8, 6, 8, 8) === Blocking(96, 341, 1728) + @test modelled_blocking(nosmt, 16, 6, 4, 4) === Blocking(96, 682, 1728) + @test modelled_blocking(smt2, 16, 6, 8, 8) === Blocking(192, 341, 1572) + # A complex kernel blocks by its packed bytes: 1m packs 4 reals per A + # element, planar 2 per B element. + @test default_blocking(kernel_from_shape((12, 8, 8), ComplexF64, OneMKernel), nosmt) === Blocking(60, 128, 2304) + for p in (nosmt, smt2), T in (Float64, Float32, ComplexF64, ComplexF32), K in (SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel) + for shape in kernel_shapes(T, K) + k = kernel_from_shape(shape, T, K) + b = default_blocking(k, p) + MR, NR = tile_size(k) + @test b.m_block % MR == 0 && b.n_block % NR == 0 + @test sliver_width(k, 2) * b.k_block * sizeof(real(T)) <= p.l1d.bytes ÷ 2 + end end # No L3: the B panel is budgeted from the L2 alone. nol3 = TargetProfile( :neon, "", CacheLevel(64KiB, 64, 1), CacheLevel(4MiB, 64, 4), CacheLevel() ) - @test _modelled_blocking(nol3, Float64, 4, 6).n_block == (1MiB ÷ (682 * 8)) ÷ 6 * 6 + @test modelled_blocking(nol3, 4, 6, 8, 8).n_block == (1MiB ÷ (682 * 8)) ÷ 6 * 6 tiny = TargetProfile( :avx2, "", CacheLevel(64, 64, 1), CacheLevel(64, 64, 1), CacheLevel() ) - @test _modelled_blocking(tiny, Float64, 8, 6) === Blocking(8, 1, 6) - @test _modelled_blocking(unknown_target(), Float64, 8, 6) === nothing - @test _modelled_blocking( + @test modelled_blocking(tiny, 8, 6, 8, 8) === Blocking(8, 1, 6) + @test modelled_blocking(unknown_target(), 8, 6, 8, 8) === nothing + @test modelled_blocking( TargetProfile(:avx2, "", CacheLevel(32KiB, 64, 1), CacheLevel(), CacheLevel()), - Float64, 8, 6 + 8, 6, 8, 8 ) === nothing end @@ -346,30 +347,21 @@ end end end -@testset "complex blocking is the real row scaled by packed reals" begin - for base in (_fallback_blocking(Float64), _fallback_blocking(Float32), Blocking(1, 5, 1)) - for m in (PlanarKernel, OneMKernel, FMAddSubKernel) - b = _scale_blocking(base, m) - @test b.k_block === base.k_block - @test b.m_block === max(1, base.m_block ÷ reals_per_element(pack_formats(m)[1])) - @test b.n_block === max(1, base.n_block ÷ reals_per_element(pack_formats(m)[2])) - end - end - @test _scale_blocking(_fallback_blocking(Float64), PlanarKernel) === Blocking(64, 256, 384) - @test _scale_blocking(_fallback_blocking(Float64), OneMKernel) === Blocking(32, 256, 384) - for T in (ComplexF64, ComplexF32) - row = _real_blocking_row(target_profile(), real(T)) - @test default_blocking(_default_kernel(T)) === _scale_blocking(row, PlanarKernel) - end - for T in (Float64, Float32) - @test default_blocking(_default_kernel(T)) === _real_blocking_row(target_profile(), T) - @test default_blocking(ScalarKernel(Val(4), Val(2), T)) === _real_blocking_row(target_profile(), T) +@testset "undetected caches: the fallback row scaled by packed reals" begin + u = unknown_target() + for T in (ComplexF64, ComplexF32), K in (PlanarKernel, OneMKernel, FMAddSubKernel) + b = default_blocking(kernel_from_shape(first(kernel_shapes(T, K)), T, K), u) + base = fallback_blocking(real(T)) + @test b.k_block === base.k_block + @test b.m_block === base.m_block ÷ reals_per_element(pack_formats(K)[1]) + @test b.n_block === base.n_block ÷ reals_per_element(pack_formats(K)[2]) end + @test default_blocking(kernel_from_shape((24, 3, 8), ComplexF64, PlanarKernel), u) === Blocking(64, 256, 384) end @testset "kernel construction: the complex default, 1m by name, and throws" begin for T in (ComplexF64, ComplexF32) - kernel = _default_kernel(T) + kernel = host_kernel(T) @test kernel isa QuasiStrided.PlanarKernel @test scalartype(kernel) === T && realtype(kernel) === real(T) @test QuasiStrided.default_kernel_type(T) === PlanarKernel @@ -377,9 +369,9 @@ end profile = target_profile() @test planar_pressure(shape...) <= (profile.nregisters > 0 ? profile.nregisters : 16) profile.isa === :avx512 && @test shape === first(kernel_shapes(T, PlanarKernel)) - @test _default_kernel(T, 1024, 1024) isa QuasiStrided.PlanarKernel + @test auto_kernel(T, 1024) isa QuasiStrided.PlanarKernel # The small-M demotion: FMAddSub on AVX-512, planar elsewhere. - small = _default_kernel(T, 1, 1) + small = auto_kernel(T, 1) @test small isa (profile.isa === :avx512 ? QuasiStrided.FMAddSubKernel : QuasiStrided.PlanarKernel) shape1m = first(kernel_shapes(T, OneMKernel)) @@ -405,8 +397,7 @@ end end @testset "small_m_shape: AVX-512 complex small-M demotion to FMAddSub" begin - sms(T, m_length, p = synthetic(:avx512)) = - QuasiStrided.small_m_shape(QuasiStrided.small_m_candidates(Val(p.isa), p, T), m_length) + sms(T, m_length, p = synthetic(:avx512)) = QuasiStrided.small_m_shape(p, T, m_length) # Least padded rows, then the larger tile. @test sms(ComplexF64, 12) === (12, 8, 8) @test sms(ComplexF64, 16) === (8, 8, 8) diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 6d13ac4..489e8ca 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -415,7 +415,7 @@ end plan = plan_contract(Cv, Av, indA, Bv, indB, indC) @test plan.Astorage === parent(Av) - default_kernel = QuasiStrided._default_kernel(T, m_length, n_length) + default_kernel = auto_kernel(T, m_length) default_m_tile = tile_size(default_kernel, 1) if m_length == run || run % default_m_tile == 0 @test plan.kernel === default_kernel @@ -451,7 +451,7 @@ end Cmat = zeros(T, Ma, Na) Av, Bv, Cv = StridedView(Amat), StridedView(Bmat), StridedView(Cmat) plan = plan_contract(Cv, Av, (1, 2), Bv, (2, 3), (1, 3)) - @test plan.kernel === QuasiStrided._default_kernel(T, Ma, Na) + @test plan.kernel === auto_kernel(T, Ma) Cref = Amat * Bmat execute!(plan, 1.0, 0.0) @@ -480,7 +480,7 @@ end for (T, a) in ((Float64, 8), (Float32, 16)) m_length = a * 36 n_length = 216 - default_kernel = QuasiStrided._default_kernel(T, m_length, n_length) + default_kernel = auto_kernel(T, m_length) # Deep K never demotes; the fixture never swaps. m_deep = 2 * kmax_of(T) @@ -640,7 +640,7 @@ end # `_KO_EXT` fits any L2; `_KO_BIG` makes B's lines touched before `d` advances # (b*c*e*64 B) exceed THIS host's L2 share, with e >= 40, so the model reorders # K wherever the suite runs. -function _ko_big(l2 = QuasiStrided.l2_core_bytes(Float64)) +function _ko_big(l2 = QuasiStrided.l2_core_bytes(target_profile())) e = max(40, fld(l2, 40 * 40 * 64) + 1) return Dict(1 => 9, 2 => 40, 3 => 40, 4 => 5, 5 => e) end @@ -717,7 +717,7 @@ end @testset "K order: plan_contract flips contract_scrambled at an L2-exceeding size" begin T = Float64 ext = _KO_BIG - @test ext[2] * ext[3] * ext[5] * 64 > QuasiStrided.l2_core_bytes(T) + @test ext[2] * ext[3] * ext[5] * 64 > QuasiStrided.l2_core_bytes(target_profile()) Av = StridedView(_ko_array(T, (1, 2, 3, 4), ext)) Bs = StridedView(_ko_array(T, (4, 3, 2, 5), ext)) # B[d,c,b,e] Cv = StridedView(zeros(T, ext[1], ext[5])) diff --git a/test/runtests.jl b/test/runtests.jl index 286a0a6..85c20ce 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,12 +9,12 @@ using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, zero_accumulator, add_tile, store_tile!, execute_tile!, lanewidth, contract!, Blocking, default_blocking, ScalarKernel, SIMDKernel, TargetProfile, CacheLevel, target_profile, cache_topology, unknown_target, detect_isa, detect_target, - derived_shape, fallback_shape, shape_override, _default_kernel, kernel_from_shape, - _fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, + derived_shape, fallback_shape, shape_override, kernel_from_shape, + fallback_blocking, kernel_shapes, parse_size, count_cpu_list, NR_DEFAULT, rule_applies, isa_nregisters, sliver_width, realtype, PlanarKernel, OneMKernel, FMAddSubKernel, accumulator_planes, pack_formats, - reals_per_element, _modelled_blocking, _scale_blocking, _real_blocking_row, - planar_pressure, _pack_split, _NestPath + reals_per_element, modelled_blocking, kernel_blocking, + planar_pressure, pack_split, is_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. From eaa42d7f3badc18a8620d0ef1c510d98d19708d3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 15:23:35 -0400 Subject: [PATCH 25/60] Decision log: chunk 11 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/CLEANUP.md b/CLEANUP.md index 9b33028..bd521c7 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -23,7 +23,7 @@ TTFX when codegen is touched). | 8 | Complex/mixed kernels | (merged into chunk 7's files) | done | | 9 | Labels | `planning/labels.jl` (`conjugation.jl` → plan.jl) | done | | 10 | Kernel selection | `planning/kernel_selection.jl` | done | -| 11 | Blocking | `blocking.jl`, `defaults.jl`, `pack_split.jl`, the line-by-line packer | | +| 11 | Blocking | `blocking.jl`, line packing (now `packing/line_packing.jl`) | done | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | | 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | @@ -191,6 +191,18 @@ override rows. The kernel type (unparameterised, e.g. `PlanarKernel`) replaces barrier as `Val{K}()` (a bare `Type` argument misses the dispatch fast path, +350–430 ns per call). +### D15. Blocking, host defaults, line packing (chunk 11, applied) + +`default_blocking(kernel, profile)` evaluates the analytical model with the +kernel's own tile and sliver bytes (complex and named kernels now get their +own optimum; real defaults unchanged); fixed rows only when caches are +undetected. No host-defaults cache: shape rules branch on `profile.isa`, and +`TargetProfile` derives `l2_share`/`l3_share`/`double_pumped` at +construction (computing them per plan cost 25–57 ns). The detected cache-line +size is used (`line_bytes`). The line-packing feature lives in +`packing/line_packing.jl`; `pack_split` takes the `SliverSpec`; the packer is +`pack_block_by_lines!`. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: @@ -201,6 +213,9 @@ barrier as `Val{K}()` (a bare `Type` argument misses the dispatch fast path, - `PackedPanel` without a raw pointer: https://github.com/lkdvos/QuasiStrided.jl/issues/15. +- The A/B swap compares tile heights with `select_shape(...; k_length = ∞)`, + so it never sees the small-K run demotion. Benchmark whether letting it see + that step picks better orientations for small K. - Analytical shape selection over a bounded search space instead of the menu: https://github.com/lkdvos/QuasiStrided.jl/issues/17. - M and N are both ordered by their C strides. Only the M side matters for the @@ -210,6 +225,9 @@ barrier as `Val{K}()` (a bare `Type` argument misses the dispatch fast path, ## Open items (to revisit in their chunk) +- One `pack_block!(panel, block, spec, transform, order)` for per-sliver and + by-lines packing, so the nest makes one call per block (chunk 14). + - Internal constants still use `_UPPER` names; drop the underscores in one pass at the end (D2). From da7e8eb1bc35d83b66f52c6fca66bb2d3ce157b7 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 15:37:29 -0400 Subject: [PATCH 26/60] Benchmark: complex blocking A/B (own model vs scaled real row) Same-node, interleaved per case, over every planar/1m/fmaddsub menu shape that fits the host's vectors. Validates chunk 11's change of complex blocking. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_complex_blocking.jl | 85 +++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 benchmark/bench_complex_blocking.jl diff --git a/benchmark/bench_complex_blocking.jl b/benchmark/bench_complex_blocking.jl new file mode 100644 index 0000000..d2731ad --- /dev/null +++ b/benchmark/bench_complex_blocking.jl @@ -0,0 +1,85 @@ +# Same-node A/B of complex cache blocking: the kernel's own analytical blocking +# (`default_blocking`) against the real row of `real(T)`'s default kernel +# divided by packed reals per element, interleaved per case, for every menu +# shape of the planar, 1m and fmaddsub kernels that fits the host's vectors. +# +# julia --project=benchmark benchmark/bench_complex_blocking.jl +# [--dtypes ComplexF64,ComplexF32] [--rounds 3] [--reps 3] [--outdir DIR] + +include(joinpath(@__DIR__, "harness.jl")) + +using QuasiStrided: PlanarKernel, OneMKernel, FMAddSubKernel, Blocking, default_blocking, + kernel_shapes, kernel_from_shape, derived_shape, target_profile, pack_formats, + reals_per_element, scalartype + +const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) +const ROUNDS = argopt("rounds", 3) +const REPS = argopt("reps", 3) +const OUTDIR = outdir() +const CSV_PATH = joinpath(OUTDIR, "bench_complex_blocking.csv") + +const CASES = [ + ShapeSpec("512^3", 512, 512, 512), + ShapeSpec("2048^3", 2048, 2048, 2048), + ShapeSpec("smallM_12x2048x2048", 12, 2048, 2048), + ShapeSpec("smallM_24x2048x2048", 24, 2048, 2048), +] + +function scaled_real_row(kernel) + T = scalartype(kernel) + R = real(T) + row = default_blocking(kernel_from_shape(derived_shape(target_profile(), R), R, SIMDKernel)) + a_reals, b_reals = map(reals_per_element, pack_formats(typeof(kernel))) + return Blocking(max(1, row.m_block ÷ a_reals), row.k_block, max(1, row.n_block ÷ b_reals)) +end + +function time_blocking(fx, kernel, b::Blocking, ::Type{T}) where {T} + plan = plan_contract( + fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; + kernel, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block, oracle = false + ) + return median_time_s(() -> execute!(plan, one(T), zero(T)); reps = REPS) +end + +function run(io) + print_env_header(stdout, "bench_complex_blocking.jl") + println(io, "dtype,kernel,MR,NR,W,case,old_m,old_k,old_n,new_m,new_k,new_n,round,t_old,t_new") + rng = MersenneTwister(1) + lanes(T) = target_profile().vector_bytes ÷ sizeof(real(T)) + for T in RUN_DTYPES, K in (PlanarKernel, OneMKernel, FMAddSubKernel) + for (MR, NR, W) in kernel_shapes(T, K) + W <= max(lanes(T), 2) || continue + kernel = K(Val(MR), Val(NR), T, Val(W)) + new = default_blocking(kernel) + old = scaled_real_row(kernel) + if new == old + println("skip $T $(nameof(K)) ($MR,$NR,$W): blocking unchanged $new") + continue + end + for spec in CASES + fx = build_plain(T, spec, rng) + for r in 1:ROUNDS + t_old = time_blocking(fx, kernel, old, T) + t_new = time_blocking(fx, kernel, new, T) + println( + io, join( + ( + T, nameof(K), MR, NR, W, spec.name, old.m_block, old.k_block, old.n_block, + new.m_block, new.k_block, new.n_block, r, t_old, t_new, + ), "," + ) + ) + flush(io) + @printf( + "%s %s (%d,%d,%d) %s round %d: old %.4g s, new %.4g s, new/old %.3f\n", + T, nameof(K), MR, NR, W, spec.name, r, t_old, t_new, t_new / t_old + ) + end + end + end + end + return nothing +end + +open(io -> run(io), CSV_PATH, "w") +println("wrote ", CSV_PATH) From 428f16e40f7ebccc227e4739b4a3f591da9f86ce Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 16:55:06 -0400 Subject: [PATCH 27/60] Chunk 12: workspace without pooling or oracle Workspace pooling is gone: the TensorOperations task-local pool, reserve!, the reuse helpers and plan_contract's `workspace =` keyword. A plan always builds its own workspace through the given allocator, and reusing the plan is how buffers get reused. The two TO.tensorcontract! methods are now one: checkpoint, then `_planned` with an `_Execute` continuation that executes and releases the plan inside the kernel barrier (so nothing is boxed), then reset. There is no try/finally, same as TO's own blas_contract!. The oracle is gone too: execute_tilewise!, its tw_* buffers and the `oracle` keyword. The tests that compared against it already compared execute! against a dense matmul or `_lo_reference`, so those comparisons stay. The ramp-vs-buffer test now checks against `_lo_reference`. Workspace: the ten offset/descriptor fields become `m`/`n::GroupBuffers` (offsets and descriptors, each a pair) and `k::NTuple{2, Vector{Int}}`. The tile buffers are dropped; the beta-only pass uses the M/N offsets. Offsets are plain Vectors whatever the allocator. There is one ContractWorkspace constructor (tensoralloc for every allocator) plus `workspace_sizes`, and a `release!(plan, allocator)` method. SlotCache/barrier_slot! move to barrier.jl, which is now included before workspace.jl. The planning barrier passes a fresh `Ref` because no workspace exists yet. ContractPlan documents that it is used by one task at a time. Tests: plan reuse with poisoned buffers is still exact, and so is the beta-only pass. Removed the pool, reuse and reserve! tests. The TO zero-byte-steady-state test is replaced by "the TO call allocates no more than plan_contract" (the boxing guard). The plan_contract allocation ceiling now adds the workspace's own bytes. Plan-reuse execute! stays 0 B. Full suite: 32151 pass (was 32300), 20.5 min. Cost of dropping the pool (workstation, taskset 4-7, medians; before -> after): call floor, DefaultAllocator, ns (bytes after; before 0 B): Float64 8^3 336 -> 1023 (3.4 KB), 16^3 432 -> 2125 (6.9 KB), d4 outer 323 -> 754 (2.6 KB) ComplexF64 8^3 407 -> 1263 (4.1 KB), 16^3 940 -> 2917 (10.7 KB), d4 outer 545 -> 1046 (2.8 KB) @tensor n^3 matmul, DefaultAllocator, us: Float64 32: 1.26 -> 4.2, 64: 6.7 -> 16.1, 128: 46 -> 73, 256: 382 -> 500 (0.94 MB/call) ComplexF64 32: 5.5 -> 11.6, 64: 27 -> 45, 128: 199 -> 270, 256: 1420 -> 1600 (1.47 MB/call) @tensor n^3 matmul, Bumper default_buffer, us: Float64 32: 2.6 -> 2.0, 64: 10.0 -> 9.6, 128: 62 -> 61, 256: 456 -> 463 ComplexF64 32: 6.9 -> 6.7, 64: 32 -> 32, 128: 227 -> 228, 256: 1600 -> 1600 TTFX first call: Float64 3.6 -> 3.3 s, ComplexF64 2.2 -> 2.0 s. Co-Authored-By: Claude Opus 5.5 --- README.md | 10 +- benchmark/bench_complex_blocking.jl | 2 +- benchmark/probes/probe_call_floor.jl | 2 +- src/QuasiStrided.jl | 6 +- src/execution/barrier.jl | 31 ++- src/execution/dot.jl | 6 +- src/execution/execute.jl | 90 +++---- src/execution/macrokernel.jl | 12 +- src/execution/oracle.jl | 89 ------ src/execution/outer.jl | 4 +- src/execution/unpackedb.jl | 8 +- src/execution/workspace.jl | 300 ++++----------------- src/integrations/tensoroperations.jl | 73 +---- src/planning/plan.jl | 77 ++---- test/execution/test_dot.jl | 2 +- test/execution/test_execute.jl | 17 +- test/execution/test_macro_blocking.jl | 31 +-- test/execution/test_outer.jl | 2 +- test/execution/test_workspace.jl | 173 ++---------- test/helpers.jl | 1 - test/integrations/test_tensoroperations.jl | 90 ++----- test/planning/test_kernel_selection.jl | 4 +- test/planning/test_per_call_overhead.jl | 37 +-- test/planning/test_plan_contract.jl | 25 +- 24 files changed, 247 insertions(+), 845 deletions(-) delete mode 100644 src/execution/oracle.jl diff --git a/README.md b/README.md index 82c85a4..e2d4e49 100644 --- a/README.md +++ b/README.md @@ -64,8 +64,8 @@ plan = plan_contract(StridedView(C), StridedView(A), (1, 2, 3), execute!(plan, 1.0, 0.0) ``` -`plan_contract` also takes `workspace=`/`allocator=` keywords (e.g. for -Bumper-backed buffers); see its docstring. +`plan_contract` also takes an `allocator=` keyword (e.g. for Bumper-backed +buffers); see its docstring. ## API @@ -95,7 +95,7 @@ and `B` into panels and drives a microkernel over register tiles of `C`. | `src/packing/` | packed-panel formats and the `pack_a!`/`pack_b!` packers with contiguous fast paths | | `src/microkernels/` | the kernel interface and the scalar, SIMD, planar, 1m and fmaddsub kernels | | `src/planning/` | label classification, conjugation, kernel selection, cache blocking, `plan_contract` | -| `src/execution/` | the workspace, the five-loop nest, the specialised paths and the tile-by-tile oracle | +| `src/execution/` | the workspace, the barriers, the five-loop nest and the specialised paths | | `src/integrations/` | `QuasiStridedBackend` for TensorOperations.jl | `test/` mirrors this layout. @@ -127,8 +127,8 @@ Implemented: `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 + plan; `tensorcontract!` through the backend plans each call, so it allocates + the plan's workspace through the given allocator. On Julia 1.10 `SIMDKernel`'s accumulator is not kept in registers, so calls allocate; results are unaffected. Not implemented: diff --git a/benchmark/bench_complex_blocking.jl b/benchmark/bench_complex_blocking.jl index d2731ad..2cc7c25 100644 --- a/benchmark/bench_complex_blocking.jl +++ b/benchmark/bench_complex_blocking.jl @@ -36,7 +36,7 @@ end function time_blocking(fx, kernel, b::Blocking, ::Type{T}) where {T} plan = plan_contract( fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; - kernel, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block, oracle = false + kernel, m_block = b.m_block, k_block = b.k_block, n_block = b.n_block ) return median_time_s(() -> execute!(plan, one(T), zero(T)); reps = REPS) end diff --git a/benchmark/probes/probe_call_floor.jl b/benchmark/probes/probe_call_floor.jl index b67628c..a3e5bdf 100644 --- a/benchmark/probes/probe_call_floor.jl +++ b/benchmark/probes/probe_call_floor.jl @@ -1,6 +1,6 @@ # Per-call floor of `tensorcontract!` on tiny problems: QuasiStridedBackend vs # StridedBLAS, through TensorOperations' public entry point with the default -# allocator (so QuasiStrided uses its pooled task workspace). +# allocator. # # JULIA_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 taskset -c 4-7 \ # julia --project=benchmark benchmark/probes/probe_call_floor.jl [label] diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index f0c45f5..3449b52 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -37,15 +37,13 @@ include("microkernels/stores.jl") include("planning/labels.jl") include("planning/kernel_selection.jl") include("planning/blocking.jl") -include("execution/workspace.jl") include("execution/barrier.jl") +include("execution/workspace.jl") include("planning/plan.jl") -# --- Execution: the five-loop nest, the tile-by-tile oracle and the -# specialised paths --- +# --- Execution: the five-loop nest and the specialised paths --- include("execution/macrokernel.jl") include("execution/execute.jl") -include("execution/oracle.jl") include("execution/unpackedb.jl") include("execution/dot.jl") include("execution/outer.jl") diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl index 60ec9b3..036d93f 100644 --- a/src/execution/barrier.jl +++ b/src/execution/barrier.jl @@ -8,26 +8,31 @@ # or a singleton, so each barrier passes: # * static choices (kernel shape, method, transforms, path) as singletons; # * operand storages (already heap objects) as themselves; -# * everything else through a typed slot, a `Base.RefValue{P}` owned by the -# workspace's `_SlotCache`: written just before the call, read back first -# thing in the callee. -# The payload never holds an operand array, so a pooled workspace retains no -# user data between calls. A workspace belongs to one task at a time, and -# the callee copies the payload out before it could run another contraction. +# * everything else through a typed slot, a `Base.RefValue{P}` (owned by the +# workspace's `SlotCache` once there is a workspace): written just before +# the call, read back first thing in the callee. +# The payload never holds an operand array, so a slot retains no user data +# between calls. A workspace belongs to one task at a time, and the callee +# copies the payload out before it could run another contraction. + +# Typed payload slots: one `Base.RefValue{P}` per payload type, plus the last +# one used. They live as long as the workspace, which is what lets a barrier +# call allocate nothing. +mutable struct SlotCache + last::Any + const slots::IdDict{Any, Any} +end +SlotCache() = SlotCache(nothing, IdDict{Any, Any}()) # The typed slot for payload type `P`; the common case (same payload type as # the previous crossing) is one type-tag test. -@inline function _barrier_slot!(ws::ContractWorkspace, ::Type{P}) where {P} - cache = ws.slots +@inline function barrier_slot!(cache::SlotCache, ::Type{P}) where {P} slot = cache.last slot isa Base.RefValue{P} && return slot - return _barrier_slot_slow!(cache, P) + return barrier_slot_slow!(cache, P) end -# No workspace yet: a one-off slot, on a path that allocates a workspace anyway. -@inline _barrier_slot!(::Nothing, ::Type{P}) where {P} = Base.RefValue{P}() - -@noinline function _barrier_slot_slow!(cache::_SlotCache, ::Type{P}) where {P} +@noinline function barrier_slot_slow!(cache::SlotCache, ::Type{P}) where {P} slot = get!(() -> Base.RefValue{P}(), cache.slots, P)::Base.RefValue{P} cache.last = slot return slot diff --git a/src/execution/dot.jl b/src/execution/dot.jl index 3ad3230..de07e78 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -52,14 +52,14 @@ function _execute_dot!(plan::ContractPlan{T}, alphaT::T, betaT::T, matB::Bool, : _dot_nest!( plan, ws, plan.Bstorage, plan.Bbase, plan.btransform, plan.Astorage, plan.Abase, plan.atransform, - plan.ngroup, ws.n_buf_B, ws.n_buf_C, plan.blocking.n_block, ws.k_buf_A, + plan.ngroup, ws.n.offsets[1], ws.n.offsets[2], plan.blocking.n_block, ws.k[1], alphaT, betaT, Val(W) ) else _dot_nest!( plan, ws, plan.Astorage, plan.Abase, plan.atransform, plan.Bstorage, plan.Bbase, plan.btransform, - plan.mgroup, ws.m_buf_A, ws.m_buf_C, plan.blocking.m_block, ws.k_buf_B, + plan.mgroup, ws.m.offsets[1], ws.m.offsets[2], plan.blocking.m_block, ws.k[2], alphaT, betaT, Val(W) ) end @@ -108,7 +108,7 @@ function _dot_nest!( k_block_length = min(k_block, k_length - k_block_start) beta_eff = firstblock ? betaT : one(T) - fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, k_block_start, k_block_length) + fill_offsets!(ws.k, plan.kgroup, k_block_start, k_block_length) checked_span_bounds(vbase, _buffer_range(vkbuf, k_block_length), (0, 0), lenv) if vflip _dot_gather!(gptr, vstorage, vbase, vkbuf, k_block_length, conj) diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 8cd3818..6d8fa0d 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -1,11 +1,10 @@ # `C *= beta` tile by tile, never reading A or B: the -# `k_length == 0 || alpha == 0` pass of `execute!` and `execute_tilewise!`. Uses -# the tile-sized `tile_*` buffers, which is why those exist even under -# `oracle = false`. +# `k_length == 0 || alpha == 0` pass of `execute!`. Borrows the first tile of +# the M/N offset buffers. function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::Int, n_length::Int) where {T} ws = plan.workspace - m_bufs = (ws.tile_m_buf_A, ws.tile_m_buf_C) - n_bufs = (ws.tile_n_buf_B, ws.tile_n_buf_C) + m_bufs = ws.m.offsets + n_bufs = ws.n.offsets m_tile_start = 0 while m_tile_start < m_length m_tile_length = min(m_tile, m_length - m_tile_start) @@ -14,8 +13,8 @@ function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::In while n_tile_start < n_length n_tile_length = min(n_tile, n_length - n_tile_start) (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) - colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) + rowsC = _axis_of(dM_C, m_bufs[2], 0) + colsC = _axis_of(dN_C, n_bufs[2], 0) GC.@preserve ws _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) n_tile_start += n_tile_length end @@ -99,7 +98,7 @@ end # with its storages stripped; they cross as arguments. @inline function _execute_across_barrier!(plan::ContractPlan{T}, alphaT::T, betaT::T, path) where {T} core = _strip_storage(plan) - slot = _barrier_slot!(plan.workspace, Tuple{typeof(core), T, T}) + slot = barrier_slot!(plan.workspace.slots, Tuple{typeof(core), T, T}) slot[] = (core, alphaT, betaT) Base.inferencebarrier(_execute_resolved!)( path, slot, plan.Astorage, plan.Bstorage, plan.Cstorage @@ -127,13 +126,13 @@ end ) # The continuation the TensorOperations adapter hands `_planned`: `execute!` -# inside the planning barrier. A struct, not a closure, so the scalars are -# concretely typed. -struct _Execute{T} +# and `release!` inside the planning barrier. A struct, not a closure, so the +# fields are concretely typed. +struct _Execute{T, AL} alpha::T beta::T + allocator::AL end -(e::_Execute)(plan::ContractPlan) = (execute!(plan, e.alpha, e.beta); nothing) # The path predicted before the kernel barrier, so its callee is specialised # on the path as well as the kernel. @@ -149,10 +148,15 @@ end sizeof(real(eltype(Cstorage))) < sizeof(real(T)) ? kernel_blocking(target_profile(), T, K, shape[1], shape[2]).k_block : 0 -@inline _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} = - (_execute_hinted!(plan, e.alpha, e.beta, hint); nothing) -@inline _continue(e::_Execute{T}, plan::ContractPlan{T}, ::Nothing) where {T} = - (execute!(plan, e.alpha, e.beta); nothing) +@inline function _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} + if hint === nothing + execute!(plan, e.alpha, e.beta) + else + _execute_hinted!(plan, e.alpha, e.beta, hint) + end + release!(plan, e.allocator) + return nothing +end # Runs the predicted path statically when it provably equals the plan's own # `_select_path` (always, for an automatically chosen kernel); otherwise @@ -267,13 +271,15 @@ function _panel_copy!(target::ContractPlan, n_block_start::Int, n_block_length:: C = target.Cstorage m_length = axis_length(target.mgroup) m_block = target.blocking.m_block - fill_offsets!((ws.n_buf_B, ws.n_buf_C), target.ngroup, n_block_start, n_block_length) + n_bufC = ws.n.offsets[2] + fill_offsets!(ws.n.offsets, target.ngroup, n_block_start, n_block_length) m_block_start = 0 while m_block_start < m_length m_block_length = min(m_block, m_length - m_block_start) - fill_offsets!((ws.m_buf_A, ws.m_buf_C), target.mgroup, m_block_start, m_block_length) + m_bufC = ws.m.offsets[2] + fill_offsets!(ws.m.offsets, target.mgroup, m_block_start, m_block_length) for j in 1:n_block_length, i in 1:m_block_length - c = target.Cbase + ws.m_buf_C[i] + ws.n_buf_C[j] + 1 + c = target.Cbase + m_bufC[i] + n_bufC[j] + 1 q = m_block_start + i + (j - 1) * m_length if load panel[q] = C[c] @@ -325,15 +331,11 @@ function _execute_nest!( n_tiles = cld(n_block_length, n_tile) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( - ws.n_desc_B, ws.n_desc_C, n_step[1], n_step[2], n_block_start, - n_block_length, n_tile, n_tiles + ws.n, n_step[1], n_step[2], n_block_start, n_block_length, n_tile, n_tiles ) else - fill_offsets!((ws.n_buf_B, ws.n_buf_C), ngroup, n_block_start, n_block_length) - _classify_slivers!( - ws.n_desc_B, ws.n_desc_C, ws.n_buf_B, ws.n_buf_C, - n_block_length, n_tile, n_tiles - ) + fill_offsets!(ws.n.offsets, ngroup, n_block_start, n_block_length) + _classify_slivers!(ws.n, n_block_length, n_tile, n_tiles) end # --- loop over K blocks --- @@ -349,18 +351,18 @@ function _execute_nest!( _ramp_offset_range(k_step[2], k_block_start, k_block_length), ) else - fill_offsets!((ws.k_buf_A, ws.k_buf_B), plan.kgroup, k_block_start, k_block_length) - dA = describe_block(ws.k_buf_A, 0, k_block_length) - dB = describe_block(ws.k_buf_B, 0, k_block_length) + fill_offsets!(ws.k, plan.kgroup, k_block_start, k_block_length) + dA = describe_block(ws.k[1], 0, k_block_length) + dB = describe_block(ws.k[2], 0, k_block_length) ( dA, dB, - descriptor_offset_range(dA, ws.k_buf_A, 0), - descriptor_offset_range(dB, ws.k_buf_B, 0), + descriptor_offset_range(dA, ws.k[1], 0), + descriptor_offset_range(dB, ws.k[2], 0), ) end - colsA_k = _axis_of(dK_A, ws.k_buf_A, 0, aff_kA) - rowsB_k = _axis_of(dK_B, ws.k_buf_B, 0, aff_kB) + colsA_k = _axis_of(dK_A, ws.k[1], 0, aff_kA) + rowsB_k = _axis_of(dK_B, ws.k[2], 0, aff_kB) # Hoisted bounds checks (B here, A and C per M block): each # rectangle is exactly the union of the per-sliver/per-tile @@ -373,14 +375,14 @@ function _execute_nest!( if split_b pack_block_by_lines!( packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), sliver_spec(kernel, 2), - plan.Bstorage, plan.Bbase, ws.n_buf_B, rowsB_k, btransform, n_block_length, k_block_length, + plan.Bstorage, plan.Bbase, ws.n.offsets[1], rowsB_k, btransform, n_block_length, k_block_length, plan.npack ) elseif !UNPACKED_B for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - colsB = _axis_of(ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, aff_nB) + colsB = _axis_of(ws.n.descriptors[1][n_tile_index + 1], ws.n.offsets[1], n_tile_start, aff_nB) @inbounds _pack_sliver!( bpanel, plan.Bstorage, plan.Bbase, colsB, rowsB_k, sliver_spec(kernel, 2), btransform @@ -395,15 +397,11 @@ function _execute_nest!( m_tiles = cld(m_block_length, m_tile) (rng_mA, rng_mC) = if m_ramp _ramp_slivers!( - ws.m_desc_A, ws.m_desc_C, m_step[1], m_step[2], m_block_start, - m_block_length, m_tile, m_tiles + ws.m, m_step[1], m_step[2], m_block_start, m_block_length, m_tile, m_tiles ) else - fill_offsets!((ws.m_buf_A, ws.m_buf_C), mgroup, m_block_start, m_block_length) - _classify_slivers!( - ws.m_desc_A, ws.m_desc_C, ws.m_buf_A, ws.m_buf_C, - m_block_length, m_tile, m_tiles - ) + fill_offsets!(ws.m.offsets, mgroup, m_block_start, m_block_length) + _classify_slivers!(ws.m, m_block_length, m_tile, m_tiles) end checked_span_bounds(plan.Abase, rng_mA, rng_kA, lenA) @@ -412,14 +410,14 @@ function _execute_nest!( if split_a pack_block_by_lines!( packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), sliver_spec(kernel, 1), - plan.Astorage, plan.Abase, ws.m_buf_A, colsA_k, atransform, m_block_length, k_block_length, + plan.Astorage, plan.Abase, ws.m.offsets[1], colsA_k, atransform, m_block_length, k_block_length, plan.mpack ) else for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsA = _axis_of(ws.m_desc_A[m_tile_index + 1], ws.m_buf_A, m_tile_start, aff_mA) + rowsA = _axis_of(ws.m.descriptors[1][m_tile_index + 1], ws.m.offsets[1], m_tile_start, aff_mA) @inbounds _pack_sliver!( apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, sliver_spec(kernel, 1), atransform @@ -465,11 +463,11 @@ end for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) + colsC = _axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + rowsC = _axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) # Inside the caller's C check. @inbounds _execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl index 2a48f4d..dd0536a 100644 --- a/src/execution/macrokernel.jl +++ b/src/execution/macrokernel.jl @@ -57,11 +57,9 @@ end # Classify each register sliver of a just-filled macro block, and return the # two maps' offset ranges over the whole block (the slivers partition it) for # the hoisted bounds checks. -@inline function _classify_slivers!( - desc1::Vector{BlockDescriptor}, desc2::Vector{BlockDescriptor}, - buf1::Vector{Int}, buf2::Vector{Int}, - block_length::Int, tile::Int, tile_count::Int - ) +@inline function _classify_slivers!(g::GroupBuffers, block_length::Int, tile::Int, tile_count::Int) + desc1, desc2 = g.descriptors + buf1, buf2 = g.offsets lo1 = typemax(Int); hi1 = typemin(Int) lo2 = typemax(Int); hi2 = typemin(Int) for tile_index in 0:(tile_count - 1) @@ -100,10 +98,10 @@ end # `_classify_slivers!`'s closed-form twin. @inline function _ramp_slivers!( - desc1::Vector{BlockDescriptor}, desc2::Vector{BlockDescriptor}, - step1::Int, step2::Int, first::Int, + g::GroupBuffers, step1::Int, step2::Int, first::Int, block_length::Int, tile::Int, tile_count::Int ) + desc1, desc2 = g.descriptors for tile_index in 0:(tile_count - 1) tile_start = tile_index * tile tile_length = min(tile, block_length - tile_start) diff --git a/src/execution/oracle.jl b/src/execution/oracle.jl deleted file mode 100644 index cbb7303..0000000 --- a/src/execution/oracle.jl +++ /dev/null @@ -1,89 +0,0 @@ -# The tile-by-tile correctness oracle for `execute!`: per register tile, K in -# `k_block` 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`. Unlike `execute!`, rounds -# a C narrower than the compute type once per K 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`. - isempty(ws.tw_packed_a) && throw( - ArgumentError( - "execute_tilewise! needs the oracle buffers, which this plan was built " * - "without; re-plan with `oracle = true`" - ) - ) - - alphaT = convert(T, alpha) - betaT = convert(T, beta) - - kernel = plan.kernel - m_tile, n_tile = tile_size(kernel) - - m_length = axis_length(plan.mgroup) - n_length = axis_length(plan.ngroup) - k_length = axis_length(plan.kgroup) - - (m_length == 0 || n_length == 0) && return plan.Cstorage - - if k_length == 0 || iszero(alphaT) - _scale_all_of_C!(plan, betaT, m_tile, n_tile, m_length, n_length) - return plan.Cstorage - end - - m_bufs = (ws.tile_m_buf_A, ws.tile_m_buf_C) - n_bufs = (ws.tile_n_buf_B, ws.tile_n_buf_C) - k_bufs = (ws.tw_k_buf_A, ws.tw_k_buf_B) - k_block = plan.blocking.k_block - - # Scatter axes borrow pointers into `ws`. - GC.@preserve ws begin - m_tile_start = 0 - while m_tile_start < m_length - m_tile_length = min(m_tile, m_length - m_tile_start) - (dM_A, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) - rowsA = _axis_of(dM_A, ws.tile_m_buf_A, 0) - rowsC = _axis_of(dM_C, ws.tile_m_buf_C, 0) - - n_tile_start = 0 - while n_tile_start < n_length - n_tile_length = min(n_tile, n_length - n_tile_start) - (dN_B, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - colsB = _axis_of(dN_B, ws.tile_n_buf_B, 0) - colsC = _axis_of(dN_C, ws.tile_n_buf_C, 0) - - k_block_start = 0 - firstpanel = true - while k_block_start < k_length - k_block_length = min(k_block, k_length - k_block_start) - (dK_A, dK_B) = block_descriptors!(k_bufs, plan.kgroup, k_block_start, k_block_length) - colsK_A = _axis_of(dK_A, ws.tw_k_buf_A, 0) - rowsK_B = _axis_of(dK_B, ws.tw_k_buf_B, 0) - - # Whole single-sliver buffers, so no `_sliver_panel`. - _pack_sliver!( - ws.tw_packed_a, plan.Astorage, plan.Abase, rowsA, colsK_A, - sliver_spec(kernel, 1), plan.atransform - ) - _pack_sliver!( - ws.tw_packed_b, plan.Bstorage, plan.Bbase, colsB, rowsK_B, - sliver_spec(kernel, 2), plan.btransform - ) - - beta_eff = firstpanel ? betaT : one(T) - _execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - ws.tw_packed_a, ws.tw_packed_b, k_block_length, alphaT, beta_eff - ) - - firstpanel = false - k_block_start += k_block_length - end - - n_tile_start += n_tile_length - end - m_tile_start += m_tile_length - end - end - - return plan.Cstorage -end diff --git a/src/execution/outer.jl b/src/execution/outer.jl index ce3bf14..1e802fe 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -27,8 +27,8 @@ function _execute_outer!( Bstorage = plan.Bstorage Cstorage = plan.Cstorage Bbase = plan.Bbase - bufB = ws.n_buf_B - bufC = ws.n_buf_C + bufB = ws.n.offsets[1] + bufC = ws.n.offsets[2] n_block = plan.blocking.n_block lenB = length(Bstorage) lenC = length(Cstorage) diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index e7588f7..70004d7 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -63,7 +63,7 @@ end const _UNPACKED_B_MMAX = 256 # The view for one N sliver; `buf`/`n_tile_start` locate an irregular sliver's -# offsets in `ws.n_buf_B`. +# offsets in `ws.n.offsets[1]`. @inline function _unpacked_b_view( kernel::Microkernel{MR, NR, T}, storage::S, base::Int, d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F @@ -92,13 +92,13 @@ end for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile bview = _unpacked_b_view( - kernel, Bstorage, Bbase, ws.n_desc_B[n_tile_index + 1], ws.n_buf_B, n_tile_start, rowsB_k, btransform + kernel, Bstorage, Bbase, ws.n.descriptors[1][n_tile_index + 1], ws.n.offsets[1], n_tile_start, rowsB_k, btransform ) - colsC = _axis_of(ws.n_desc_C[n_tile_index + 1], ws.n_buf_C, n_tile_start, aff_nC) + colsC = _axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsC = _axis_of(ws.m_desc_C[m_tile_index + 1], ws.m_buf_C, m_tile_start, aff_mC) + rowsC = _axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) # Inside the caller's C check. @inbounds _execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 698571a..1051332 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -1,72 +1,45 @@ -# Typed payload slots for the dispatch barriers (src/execution/barrier.jl): -# one `Base.RefValue{P}` per payload type, plus the last one used. They live -# as long as the workspace, which is what lets a barrier call allocate nothing. -mutable struct _SlotCache - last::Any - const slots::IdDict{Any, Any} +# One M or N macro block's offsets in the group's two operands, and the +# per-sliver descriptors of each. +struct GroupBuffers + offsets::NTuple{2, Vector{Int}} + descriptors::NTuple{2, Vector{BlockDescriptor}} end -_SlotCache() = _SlotCache(nothing, IdDict{Any, Any}()) + +GroupBuffers(block::Int, tiles::Int) = GroupBuffers( + (Vector{Int}(undef, block), Vector{Int}(undef, block)), + (Vector{BlockDescriptor}(undef, tiles), Vector{BlockDescriptor}(undef, tiles)) +) """ ContractWorkspace{T,VT<:AbstractVector,PT<:AbstractVector} -The buffers [`execute!`](@ref) needs. Pass one back as -`plan_contract(...; workspace = ws)` to reuse (and grow) it across -contractions; one from a non-default allocator is single-use and must be -[`release!`](@ref)d. `T` is the plan's compute type; the packed panels (`VT`) -hold `real(T)`, the panel of C (`PT`) holds `T`. Field layout is an +The buffers [`execute!`](@ref) needs, built by [`plan_contract`](@ref) and +reused by every `execute!` on that plan. One from a non-default allocator must +be [`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-two GC-tracked fields. +# pointer to it instead of copying every GC-tracked field. 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} - const n_buf_B::Vector{Int} - const n_buf_C::Vector{Int} - const k_buf_A::Vector{Int} - const k_buf_B::Vector{Int} - const m_desc_A::Vector{BlockDescriptor} - const m_desc_C::Vector{BlockDescriptor} - const n_desc_B::Vector{BlockDescriptor} - const n_desc_C::Vector{BlockDescriptor} + const m::GroupBuffers + const n::GroupBuffers + # K blocks are described whole, so K has offsets only. + const k::NTuple{2, Vector{Int}} - # Packed macro panels and the panel of C: the only allocator-routed buffers. + # The only allocator-routed buffers. const packed_a::VT const packed_b::VT - - # One register tile's offsets: the beta-only pass and the oracle. - const tile_m_buf_A::Vector{Int} - const tile_m_buf_C::Vector{Int} - const tile_n_buf_B::Vector{Int} - const tile_n_buf_C::Vector{Int} - - # `execute_tilewise!`'s own buffers, empty under `oracle = false`; not - # shared, so the oracle has no state in common with what it checks. - const tw_k_buf_A::Vector{Int} - const tw_k_buf_B::Vector{Int} - const tw_packed_a::VT - const tw_packed_b::VT - - 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 + const slots::SlotCache + # GUARDRAIL: the packed panels hold `real(T)`, never `T`. - 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}, - m_desc_A::Vector{BlockDescriptor}, m_desc_C::Vector{BlockDescriptor}, - n_desc_B::Vector{BlockDescriptor}, n_desc_C::Vector{BlockDescriptor}, - packed_a::VT, packed_b::VT, - tile_m_buf_A::Vector{Int}, tile_m_buf_C::Vector{Int}, - tile_n_buf_B::Vector{Int}, tile_n_buf_C::Vector{Int}, - tw_k_buf_A::Vector{Int}, tw_k_buf_B::Vector{Int}, - tw_packed_a::VT, tw_packed_b::VT, c_panel::PT + function ContractWorkspace{T}( + m::GroupBuffers, n::GroupBuffers, k::NTuple{2, Vector{Int}}, + packed_a::VT, packed_b::VT, c_panel::PT ) where {T, VT <: AbstractVector, PT <: AbstractVector} eltype(VT) === real(T) || throw( ArgumentError( @@ -74,234 +47,61 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector, PT <: AbstractVector} "(the real type of the compute type $T), got $(eltype(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(), c_panel, - ) + return new{T, VT, PT}(m, n, k, packed_a, packed_b, c_panel, SlotCache()) end end -# Buffer lengths for `kernel` at the effective `blocking`, shared by the -# constructors and `reserve!`. GUARDRAIL: `packed_a_length` already counts -# reals, the sliver counts are in logical `m_tile`/`n_tile`; do not rescale -# either. -@inline function _workspace_sizes(kernel, blocking::Blocking) +# GUARDRAIL: `packed_a_length` already counts reals, the sliver counts are in +# logical `m_tile`/`n_tile`; do not rescale either. +@inline function workspace_sizes(kernel, blocking::Blocking) m_tile, n_tile = tile_size(kernel) - k_block = blocking.k_block - pa = packed_a_length(kernel, k_block) - pb = packed_b_length(kernel, k_block) m_tiles = cld(blocking.m_block, m_tile) n_tiles = cld(blocking.n_block, n_tile) return ( - m_block = blocking.m_block, n_block = blocking.n_block, k_block = k_block, - m_tile = m_tile, n_tile = n_tile, m_tiles = m_tiles, n_tiles = n_tiles, - packed_a = m_tiles * pa, packed_b = n_tiles * pb, - tw_packed_a = pa, tw_packed_b = pb, - ) -end - -# A temporary, never `resize!`d. -@inline function _alloc_temp(::Type{T}, n::Int, allocator) where {T} - return TO.tensoralloc(Vector{T}, (n,), Val(true), allocator) -end - -# A non-temporary, which every allocator serves as a plain `Vector{Int}`. -@inline function _alloc_offsets(n::Int, allocator) - return TO.tensoralloc(Vector{Int}, (n,), Val(false), allocator)::Vector{Int} -end - -# Not tensor buffers, so never through `tensoralloc`. -@inline _alloc_descriptors(n::Int) = Vector{BlockDescriptor}(undef, n) - -@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, c_panel::PT - ) where {T, F, VT <: AbstractVector, PT <: AbstractVector} - return ContractWorkspace{T, VT, PT}( - ints(s.m_block), ints(s.m_block), - ints(s.n_block), ints(s.n_block), - ints(s.k_block), ints(s.k_block), - _alloc_descriptors(s.m_tiles), _alloc_descriptors(s.m_tiles), - _alloc_descriptors(s.n_tiles), _alloc_descriptors(s.n_tiles), - packed_a, packed_b, - ints(s.m_tile), ints(s.m_tile), - ints(s.n_tile), ints(s.n_tile), - ints(ntw), ints(ntw), - tw_packed_a, tw_packed_b, c_panel, + packed_a = m_tiles * packed_a_length(kernel, blocking.k_block), + packed_b = n_tiles * packed_b_length(kernel, blocking.k_block), ) end """ ContractWorkspace(T, kernel, blocking::Blocking; - oracle = true, allocator = TensorOperations.DefaultAllocator(), - panel = 0) + 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). +`panel` is the length of the compute-type panel of C (see +[`plan_contract`](@ref)). The packed panels and the panel of C come from +`allocator` (`tensoralloc`) and must be handed back with [`release!`](@ref) +unless it is the default one. """ function ContractWorkspace( - ::Type{T}, kernel, blocking::Blocking, oracle::Bool, ::TO.DefaultAllocator, - panel::Int = 0 - ) where {T} - s = _workspace_sizes(kernel, blocking) - R = realtype(kernel) - return _build_workspace( - T, s, oracle ? s.k_block : 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{T}(undef, panel), - ) -end - -function ContractWorkspace( - ::Type{T}, kernel, blocking::Blocking, oracle::Bool, allocator, panel::Int = 0 + ::Type{T}, kernel, blocking::Blocking; + allocator = TO.DefaultAllocator(), panel::Int = 0 ) where {T} - s = _workspace_sizes(kernel, blocking) + s = workspace_sizes(kernel, blocking) R = realtype(kernel) - # `release!` frees in exactly the reverse order, for arena allocators. - packed_a = _alloc_temp(R, s.packed_a, allocator) - 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.k_block : 0, n -> _alloc_offsets(n, allocator), - packed_a, packed_b, tw_packed_a, tw_packed_b, c_panel + packed_a = TO.tensoralloc(Vector{R}, (s.packed_a,), Val(true), allocator) + packed_b = TO.tensoralloc(Vector{R}, (s.packed_b,), Val(true), allocator) + c_panel = TO.tensoralloc(Vector{T}, (panel,), Val(true), allocator) + k = (Vector{Int}(undef, blocking.k_block), Vector{Int}(undef, blocking.k_block)) + return ContractWorkspace{T}( + GroupBuffers(blocking.m_block, s.m_tiles), GroupBuffers(blocking.n_block, s.n_tiles), k, + packed_a, packed_b, c_panel ) end -function ContractWorkspace( - ::Type{T}, kernel, blocking::Blocking; - oracle::Bool = true, allocator = TO.DefaultAllocator(), panel::Int = 0 - ) where {T} - 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}, Vector{T}}, kernel, blocking::Blocking, - oracle::Bool, panel::Int = 0 - ) where {T, R} - s = _workspace_sizes(kernel, blocking) - - _grow!(ws.m_buf_A, s.m_block) - _grow!(ws.m_buf_C, s.m_block) - _grow!(ws.n_buf_B, s.n_block) - _grow!(ws.n_buf_C, s.n_block) - _grow!(ws.k_buf_A, s.k_block) - _grow!(ws.k_buf_B, s.k_block) - - _grow!(ws.m_desc_A, s.m_tiles) - _grow!(ws.m_desc_C, s.m_tiles) - _grow!(ws.n_desc_B, s.n_tiles) - _grow!(ws.n_desc_C, s.n_tiles) - - _grow!(ws.packed_a, s.packed_a) - _grow!(ws.packed_b, s.packed_b) - - _grow!(ws.tile_m_buf_A, s.m_tile) - _grow!(ws.tile_m_buf_C, s.m_tile) - _grow!(ws.tile_n_buf_B, s.n_tile) - _grow!(ws.tile_n_buf_C, s.n_tile) - _grow!(ws.c_panel, panel) - - if oracle - _grow!(ws.tw_k_buf_A, s.k_block) - _grow!(ws.tw_k_buf_B, s.k_block) - _grow!(ws.tw_packed_a, s.tw_packed_a) - _grow!(ws.tw_packed_b, s.tw_packed_b) - end - - return ws -end - -@inline function _grow!(v::Vector, n::Int) - length(v) < n && resize!(v, n) - return v -end - """ release!(ws::ContractWorkspace, allocator) + release!(plan::ContractPlan, 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. +Hand the packed panels and panel of C back to `allocator` +(`TensorOperations.tensorfree!`) in reverse acquisition order; the workspace +must not be used afterwards. A no-op for the default 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) TO.tensorfree!(ws.packed_a, allocator) return nothing 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, panel::Int - ) where {T} - workspace === nothing && - 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, - panel::Int - ) where {T} - workspace === nothing || throw( - ArgumentError( - "plan_contract: `workspace` cannot be combined with a non-default " * - "`allocator` ($(typeof(allocator))); an allocator-provided workspace is " * - "sized once at construction and must not be resized or reused" - ) - ) - return ContractWorkspace(T, kernel, blocking, oracle, allocator, panel) -end - -# The pool is keyed by the compute type alone, so also check the packed type -# (folds at compile time). -@inline function _reuse_workspace( - ::Type{T}, ws::ContractWorkspace{T, Vector{R}, 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, panel) -end - -@noinline function _throw_packed_eltype_mismatch(::Type{T}, ws, kernel) where {T} - throw( - ArgumentError( - "plan_contract: cannot reuse a $(typeof(ws)) whose packed panels hold " * - "$(eltype(ws.packed_a)) for a kernel packing $(realtype(kernel))" - ) - ) -end - -@noinline function _reuse_workspace( - ::Type{T}, ws::ContractWorkspace, kernel, blocking::Blocking, oracle::Bool, - panel::Int - ) where {T} - throw( - ArgumentError( - "plan_contract: cannot reuse a $(typeof(ws)) for a contraction of compute " * - "type $T on the default allocator; only a " * - "ContractWorkspace{$T,Vector{$(realtype(kernel))},Vector{$T}} -- compute " * - "type $T, packed panels of $(realtype(kernel)) -- is `reserve!`-able" - ) - ) -end diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index 3c3809a..3a00477 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -31,33 +31,6 @@ end QuasiStridedBackend(; accumulator = nothing) = QuasiStridedBackend{accumulator}() -# Task-local pool of workspaces, keyed by the compute type only: all complex -# methods pack into `Vector{real(T)}` and `reserve!` is grow-only, so they can -# share one. -const _QS_WORKSPACE_KEY = :quasistrided_contract_workspaces - -@inline function _qs_workspace_pool() - return get!(task_local_storage(), _QS_WORKSPACE_KEY) do - return Dict{DataType, ContractWorkspace}() - end::Dict{DataType, ContractWorkspace} -end - -@inline function _qs_task_workspace(::Type{T}) where {T} - pool = _qs_workspace_pool() - 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)}, 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} - kernel = kernel_from_shape(derived_shape(target_profile(), T), T) - new_ws = ContractWorkspace(T, kernel, default_blocking(kernel), false, TO.DefaultAllocator()) - pool[T] = new_ws - return new_ws -end - # TO's `pA`/`pB`/`pAB` -> one `Int` label per axis: `1:NoA` for A's open axes # (in `pA[1]` order), `NoA+1:NoA+NoB` for B's open axes, `-1:-1:-Nk` for the # contracted pairs; `indC` is then `linearize(pAB)`. For example @@ -84,9 +57,6 @@ end return nothing end -# 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) @@ -110,51 +80,24 @@ end return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end -# Default allocator: pooled task-local workspace. `_planned` builds and runs the -# plan behind the kernel dispatch barrier, so the plan is never boxed (as -# `execute!(plan_contract(...), ...)` would be). -function TO.tensorcontract!( - C::AbstractArray, - A::AbstractArray, pA::Index2Tuple, conjA::Bool, - B::AbstractArray, pB::Index2Tuple, conjB::Bool, - pAB::Index2Tuple, - α::Number, β::Number, - backend::QuasiStridedBackend{AC}, - allocator::TO.DefaultAllocator = TO.DefaultAllocator() - ) 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(typeof(α′)), allocator, false, AC # α′ has the compute type - ) - return C -end - -# Explicit allocator: a workspace scoped to this call, bracketed with -# checkpoint/reset like TO's own `blas_contract!`. This method needs the plan -# back to release it, hence `plan_contract` rather than `_planned`. +# `_planned` builds, runs and releases the plan behind the kernel dispatch +# barrier, so the plan is never boxed (as `execute!(plan_contract(...), ...)` +# would be). Bracketed with checkpoint/reset like TO's own `blas_contract!`. function TO.tensorcontract!( C::AbstractArray, A::AbstractArray, pA::Index2Tuple, conjA::Bool, B::AbstractArray, pB::Index2Tuple, conjB::Bool, pAB::Index2Tuple, α::Number, β::Number, - backend::QuasiStridedBackend{AC}, allocator + backend::QuasiStridedBackend{AC}, allocator = TO.DefaultAllocator() ) 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, accumulator = AC + _planned( + _Execute(α′, β′, allocator), Cv, Av, indA, Bv, indB, indC, + nothing, conjA, conjB, nothing, nothing, nothing, allocator, AC ) - try - execute!(plan, α′, β′) - finally - release!(plan.workspace, allocator) - TO.allocator_reset!(allocator, checkpoint) - end + TO.allocator_reset!(allocator, checkpoint) return C end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 62c2265..2b332c4 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -8,8 +8,10 @@ Reusable, concretely typed plan from [`plan_contract`](@ref): the M/N/K `AxisGroup`s, kernel, operand storages, effective [`Blocking`](@ref), packing transforms and the [`ContractWorkspace`](@ref) holding every buffer -[`execute!`](@ref) needs. Field layout is not part of the public interface; -after an M/N orientation swap the `A*` fields describe the original `B`. +[`execute!`](@ref) needs. Reusing a plan reuses its buffers. A plan, and so +its workspace, must be used by one task at a time. Field layout is not part +of the public interface; 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, @@ -47,9 +49,8 @@ end kernel = nothing, conjA = false, conjB = false, m_block = nothing, k_block = nothing, n_block = nothing, - workspace = nothing, allocator = TensorOperations.DefaultAllocator(), - oracle = true, accumulator = nothing) -> ContractPlan + 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, @@ -83,11 +84,9 @@ invalid input. is packed line by line (`PackSplit`); its `m_block` (or `n_block`) then becomes a whole number of line groups, up to `requested k_block / k_block` times the request. - * `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. + * The plan's [`ContractWorkspace`](@ref) takes its packed panels from the + TensorOperations `allocator`; for a non-default one, [`release!`](@ref) + is left to the caller. * `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` is rejected. @@ -102,14 +101,12 @@ function plan_contract( m_block::Union{Int, Nothing} = nothing, k_block::Union{Int, Nothing} = nothing, n_block::Union{Int, Nothing} = nothing, - workspace::Union{Nothing, ContractWorkspace} = nothing, allocator = TO.DefaultAllocator(), - 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, m_block, k_block, n_block, workspace, allocator, oracle, accumulator + kernel, conjA, conjB, m_block, k_block, n_block, allocator, accumulator ) end @@ -117,8 +114,7 @@ end # 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 compute type. struct _PlanRequest{ - T, F, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, - WS <: Union{Nothing, ContractWorkspace}, AL, + T, F, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, AL, } f::F mgroup::GM @@ -134,19 +130,17 @@ struct _PlanRequest{ m_block::Union{Int, Nothing} k_block::Union{Int, Nothing} n_block::Union{Int, Nothing} - workspace::WS allocator::AL - oracle::Bool end @inline function _plan_request( ::Type{T}, f::F, mgroup::GM, ngroup::GN, kgroup::GK, Astorage::SA, Abase::Int, Bstorage::SB, Bbase::Int, Cstorage::SC, Cbase::Int, - run::Int, m_block, k_block, n_block, workspace::WS, allocator::AL, oracle::Bool - ) where {T, F, GM, GN, GK, SA, SB, SC, WS, AL} - return _PlanRequest{T, F, GM, GN, GK, SA, SB, SC, WS, AL}( + run::Int, m_block, k_block, n_block, allocator::AL + ) where {T, F, GM, GN, GK, SA, SB, SC, AL} + return _PlanRequest{T, F, GM, GN, GK, SA, SB, SC, AL}( f, mgroup, ngroup, kgroup, Astorage, Abase, Bstorage, Bbase, Cstorage, Cbase, - run, m_block, k_block, n_block, workspace, allocator, oracle + run, m_block, k_block, n_block, allocator ) end @@ -181,8 +175,7 @@ function _planned( B::StridedView, indB::NTuple{NB, Int}, indC::NTuple{NC, Int}, kernel, conjA::Bool, conjB::Bool, m_block::Union{Int, Nothing}, k_block::Union{Int, Nothing}, n_block::Union{Int, Nothing}, - workspace::Union{Nothing, ContractWorkspace}, allocator, oracle::Bool, - accumulator::AC + allocator, accumulator::AC ) where {F, NA, NB, NC, AC} T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) K = default_kernel_type(T, eltype(A), eltype(B)) @@ -228,14 +221,14 @@ function _planned( req_swapped = _plan_request( T, f, ngroup, mgroup, kgroup_swapped, parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), - run_n, m_block, k_block, n_block, workspace, allocator, oracle + run_n, m_block, k_block, n_block, allocator ) return _plan_with_kernel(kernel, K, 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, m_block, k_block, n_block, workspace, allocator, oracle + run_m, m_block, k_block, n_block, allocator ) return _plan_with_kernel(kernel, K, atransform, btransform, req) end @@ -290,8 +283,8 @@ end # built across a dispatch barrier specialised on that one concrete kernel type, # so only the chosen 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. +# kernel). All barrier arguments are singletons or heap objects (the request +# in a `Ref`): one method-cache hit. @inline _plan_with_kernel(kernel, ::Type, atransform, btransform, req::_PlanRequest) = _plan_contract(kernel, atransform, btransform, req, nothing) @inline function _plan_with_kernel(::Nothing, ::Type{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} @@ -305,35 +298,17 @@ end # the shape and kernel type, not the kernel: a call union-split on the # kernel type is emitted out of line and boxes `req`. hint = _path_hint(req.f, req, shape, Val(K)) - core = _strip_storage(req) - slot = _barrier_slot!(req.workspace, typeof(core)) - slot[] = core return Base.inferencebarrier(_plan_resolved)( - vkernel, atransform, btransform, hint, slot, - req.Astorage, req.Bstorage, req.Cstorage + vkernel, atransform, btransform, hint, Base.RefValue{typeof(req)}(req) ) end function _plan_resolved( - ::Val{Kern}, atransform::TA, btransform::TB, hint::H, - slot::Base.RefValue{R}, Astorage::SA, Bstorage::SB, Cstorage::SC - ) where {Kern, TA, TB, H, T, R <: _PlanRequest{T}, SA, SB, SC} - req = _with_storage(slot[], Astorage, Bstorage, Cstorage) - return _plan_contract(Kern(), atransform, btransform, req, hint) + ::Val{Kern}, atransform::TA, btransform::TB, hint::H, slot::Base.RefValue{R} + ) where {Kern, TA, TB, H, R <: _PlanRequest} + return _plan_contract(Kern(), atransform, btransform, slot[], hint) end -# The storages cross the barrier as arguments, so no slot retains a user array. -@inline _strip_storage(req::_PlanRequest{T}) where {T} = _plan_request( - T, req.f, req.mgroup, req.ngroup, req.kgroup, - nothing, req.Abase, nothing, req.Bbase, nothing, req.Cbase, - req.run, req.m_block, req.k_block, req.n_block, req.workspace, req.allocator, req.oracle -) -@inline _with_storage(req::_PlanRequest{T}, Astorage, Bstorage, Cstorage) where {T} = _plan_request( - T, req.f, req.mgroup, req.ngroup, req.kgroup, - Astorage, req.Abase, Bstorage, req.Bbase, Cstorage, req.Cbase, - req.run, req.m_block, req.k_block, req.n_block, req.workspace, req.allocator, req.oracle -) - # `nothing` for a continuation that does not execute (see src/execution/execute.jl). @inline _path_hint(f, req::_PlanRequest, shape, kernel_type) = nothing @@ -383,8 +358,8 @@ function _plan_contract( end end blocking = Blocking(m_block, k_block, n_block) - ws = _resolve_workspace( - T, req.workspace, kernel, blocking, req.oracle, req.allocator, panel ? m_length * min(n_block, n_length) : 0 + ws = ContractWorkspace( + T, kernel, blocking; allocator = req.allocator, panel = panel ? m_length * min(n_block, n_length) : 0 ) plan = ContractPlan( @@ -397,3 +372,5 @@ end # `hint` is used only by an executing continuation (src/execution/execute.jl). @inline _continue(f::F, plan::ContractPlan, hint) where {F} = f(plan) + +release!(plan::ContractPlan, allocator) = release!(plan.workspace, allocator) diff --git a/test/execution/test_dot.jl b/test/execution/test_dot.jl index f510557..c3a0227 100644 --- a/test/execution/test_dot.jl +++ b/test/execution/test_dot.jl @@ -112,7 +112,7 @@ end @testset "dot path: allocation-free, and through the backend ($T)" for T in _DOT_TYPES for (mk, kw) in ((_dot_gemv_maker(T, 6, 4), (;)), (_dot_n1_maker(T, 5, 300, 4), (k_block = 64, m_block = 8))) - plan = plan_contract(mk()...; oracle = false, kw...) + plan = plan_contract(mk()...; kw...) @test _dot_takes(plan) execute!(plan, 1.0, 0.0) @test (@allocated execute!(plan, 1.0, 0.0)) == 0 skip = (VERSION < v"1.11") diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 6712dbc..57ab405 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -167,7 +167,6 @@ end ) @test_throws BoundsError execute!(plan, 1.0, 0.0) @test all(iszero, Cstore) - @test_throws BoundsError execute_tilewise!(plan, 1.0, 0.0) end end @@ -175,14 +174,13 @@ end Ma, Ka, Na = 19, 23, 17 Amat, Bmat = randn(MersenneTwister(1), Ma, Ka), randn(MersenneTwister(2), Ka, Na) kernel = SIMDKernel(Val(8), Val(6), Float64, Val(4)) - # Default everything; a reused workspace without the oracle (the backend's - # shape); several blocks with tail tiles in M and N (the vectorized store's - # row tail); a conj transform, which the real path never builds but which - # must still cross `_pack_sliver!` concretely typed. + # Default everything; several blocks with tail tiles in M and N (the + # vectorized store's row tail); a conj transform, which the real path never + # builds but which must still cross `_pack_sliver!` concretely typed. Cmat = zeros(Ma, Na) p = _mm_plan(Cmat, Amat, Bmat) plans = ( - p, _mm_plan(Cmat, Amat, Bmat; workspace = p.workspace, oracle = false), + p, _mm_plan(Cmat, Amat, Bmat; kernel = kernel, m_block = 16, k_block = 5, n_block = 8), ContractPlan( p.kernel, p.mgroup, p.ngroup, p.kgroup, p.blocking, @@ -196,9 +194,6 @@ end @test Cmat ≈ Amat * Bmat @test allocs == 0 skip = (VERSION < v"1.11") end - allocs_tw = _steady_allocs!(execute_tilewise!, plans[4], Cmat) - @test Cmat ≈ Amat * Bmat - @test allocs_tw == 0 skip = (VERSION < v"1.11") ps = _mm_plan(Cmat, Amat, Bmat; kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 6, n_block = 7) @test _steady_allocs!(execute!, ps, Cmat) == 0 skip = (VERSION < v"1.11") end @@ -236,9 +231,7 @@ end @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!) - @test isempty(_ws_nonconcrete_types(f, (typeof(plan), Float64, Float64))) - end + @test isempty(_ws_nonconcrete_types(execute!, (typeof(plan), Float64, Float64))) @test isconcretetype(only(Base.return_types(execute!, (typeof(plan), Float64, Float64)))) end diff --git a/test/execution/test_macro_blocking.jl b/test/execution/test_macro_blocking.jl index 72d1034..e84fd8c 100644 --- a/test/execution/test_macro_blocking.jl +++ b/test/execution/test_macro_blocking.jl @@ -1,6 +1,5 @@ -# Randomized agreement of the five-loop `execute!` with a dense matmul and with -# `execute_tilewise!`, at block sizes small enough to split every dimension, -# over real and complex kernels, alpha/beta and the conjugation flags and ops. +# Randomized agreement of the five-loop `execute!` with a dense matmul, at +# block sizes small enough to split every dimension, over real and complex kernels, alpha/beta and the conjugation flags and ops. # helpers.jl binds `plan_contract`/`execute!` unqualified in this scope, so # names here are written `QuasiStrided.`. @@ -51,7 +50,7 @@ function _macro_scalar(rng, ::Type{T}) where {T} return T <: Complex ? T(rand(rng, -3.0:0.5:3.0), rand(rng, -3.0:0.5:3.0)) : T(rand(rng, -3.0:0.5:3.0)) end -@testset "macro driver: randomized agreement vs dense matmul and the oracle ($T)" for T in (Float64, Float32, ComplexF64, ComplexF32) +@testset "macro driver: randomized agreement vs dense matmul ($T)" for T in (Float64, Float32, ComplexF64, ComplexF32) rng = MersenneTwister(0x5A17_D817) kernels = _macro_kernels(T) @test !isempty(kernels) @@ -72,9 +71,7 @@ end C = copy(Cstart) QuasiStrided.execute!(_macro_plan(C, Amat, Bmat, kernel, m_block, k_block, n_block; kw...), alpha, beta) - Ctw = copy(Cstart) - QuasiStrided.execute_tilewise!(_macro_plan(Ctw, Amat, Bmat, kernel, m_block, k_block, n_block; kw...), alpha, beta) - ok = isapprox(C, expected; rtol = _macro_rtol(T, Ka)) && isapprox(C, Ctw; rtol = _macro_rtol(T, Ka)) + ok = isapprox(C, expected; rtol = _macro_rtol(T, Ka)) ok || @error "macro driver case $i failed" kernel Ma Ka Na m_block k_block n_block kw alpha beta @test ok end @@ -105,17 +102,15 @@ end Amat, Bmat = randn(a_n, q_n, k_n), randn(k_n, n_n) Cstart = randn(a_n, q_n, n_n) Cref = _brute_ref(Amat, (1, 2, 3), Bmat, (3, 4), Cstart, (1, 2, 4), 1.7, -0.4) - for run! in (QuasiStrided.execute!, QuasiStrided.execute_tilewise!) - Cbig = randn(a_n + 3, q_n, n_n) - Csub = view(Cbig, 1:a_n, :, :) - Csub .= Cstart - plan = QuasiStrided.plan_contract( - StridedView(Csub), StridedView(Amat), (1, 2, 3), StridedView(Bmat), (3, 4), (1, 2, 4); - kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 5, n_block = 3 - ) - run!(plan, 1.7, -0.4) - @test isapprox(Array(Csub), Cref; rtol = _macro_rtol(Float64, k_n)) - end + Cbig = randn(a_n + 3, q_n, n_n) + Csub = view(Cbig, 1:a_n, :, :) + Csub .= Cstart + plan = QuasiStrided.plan_contract( + StridedView(Csub), StridedView(Amat), (1, 2, 3), StridedView(Bmat), (3, 4), (1, 2, 4); + kernel = ScalarKernel(Val(4), Val(3), Float64), m_block = 8, k_block = 5, n_block = 3 + ) + QuasiStrided.execute!(plan, 1.7, -0.4) + @test isapprox(Array(Csub), Cref; rtol = _macro_rtol(Float64, k_n)) end @testset "macro driver: several blocks at the default blocking ($(nameof(typeof(k))){$T})" for T in (ComplexF64, ComplexF32), k in _macro_kernels(T) diff --git a/test/execution/test_outer.jl b/test/execution/test_outer.jl index ee33559..b1fb921 100644 --- a/test/execution/test_outer.jl +++ b/test/execution/test_outer.jl @@ -86,7 +86,7 @@ end @test_throws BoundsError execute!(plan, 1.0, 0.0) @test all(iszero, plan.Cstorage) end - plan = plan_contract(_outer_maker(T, 64, 50, 6)()...; oracle = false, n_block = 12) + plan = plan_contract(_outer_maker(T, 64, 50, 6)()...; n_block = 12) @test _outer_takes(plan) execute!(plan, 1.0, 0.0) @test (@allocated execute!(plan, 1.0, 0.0)) == 0 skip = (VERSION < v"1.11") diff --git a/test/execution/test_workspace.jl b/test/execution/test_workspace.jl index 0b1bd65..56d3bb6 100644 --- a/test/execution/test_workspace.jl +++ b/test/execution/test_workspace.jl @@ -1,9 +1,4 @@ -# ContractWorkspace, the `workspace`/`allocator`/`oracle` keywords, and the -# SIMDKernel default. - -_ws_lengths(ws) = [ - length(getfield(ws, f)) for f in fieldnames(typeof(ws)) if getfield(ws, f) isa AbstractVector -] +# ContractWorkspace, the `allocator` keyword, and the SIMDKernel default. @testset "plan_contract: SIMDKernel is the engine-wide default kernel" begin for T in (Float64, Float32) @@ -28,131 +23,42 @@ _ws_lengths(ws) = [ end end -@testset "ContractWorkspace: reuse across shapes is bitwise identical to fresh plans" begin +@testset "ContractWorkspace: plan reuse refills every buffer before reading it" begin kernel = ScalarKernel(Val(4), Val(3), Float64) - m_block, k_block, n_block = 8, 6, 7 - alpha, beta = 1.75, -0.5 - - # Not monotone: sized by the first, grown by the second, then oversized. - shapes = ((13, 11, 10), (23, 19, 17), (4, 3, 2), (9, 10, 8), (1, 1, 1), (16, 5, 6)) - - ws = nothing - lengths_before = nothing - for (idx, (Ma, Ka, Na)) in enumerate(shapes) - Random.seed!(31_000 + idx) - Amat, Bmat = randn(Ma, Ka), randn(Ka, Na) - Cstart = randn(Ma, Na) - - Cfresh = copy(Cstart) - execute!(_mm_plan(Cfresh, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block), alpha, beta) - - Creuse = copy(Cstart) - plan = _mm_plan( - Creuse, Amat, Bmat; - kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, workspace = ws - ) - execute!(plan, alpha, beta) - - @test Creuse == Cfresh - - if ws !== nothing - @test plan.workspace === ws - @test all(_ws_lengths(ws) .>= lengths_before) - end - ws = plan.workspace - lengths_before = _ws_lengths(ws) - end - - Amat32, Bmat32, Cmat32 = randn(Float32, 4, 4), randn(Float32, 4, 4), zeros(Float32, 4, 4) - @test_throws ArgumentError _mm_plan( - Cmat32, Amat32, Bmat32; - kernel = ScalarKernel(Val(4), Val(3), Float32), workspace = ws - ) -end - -@testset "ContractWorkspace: an oversized reused buffer is not read beyond its live region" begin - kernel = ScalarKernel(Val(4), Val(3), Float64) - m_block, k_block, n_block = 8, 6, 7 alpha, beta = 2.5, -0.75 Random.seed!(606) - Abig, Bbig = randn(23, 19), randn(19, 17) - Cbig = zeros(23, 17) - big = _mm_plan(Cbig, Abig, Bbig; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block) - execute!(big, 1.0, 0.0) - ws = big.workspace - - Ma, Ka, Na = 5, 4, 3 - Amat, Bmat = randn(Ma, Ka), randn(Ka, Na) - Cstart = randn(Ma, Na) - + Amat, Bmat = randn(23, 19), randn(19, 17) + Cstart = randn(23, 17) Cref = copy(Cstart) - execute!(_mm_plan(Cref, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block), alpha, beta) + execute!(_mm_plan(Cref, Amat, Bmat; kernel = kernel, m_block = 8, k_block = 6, n_block = 7), alpha, beta) + + C = copy(Cstart) + plan = _mm_plan(C, Amat, Bmat; kernel = kernel, m_block = 8, k_block = 6, n_block = 7) + execute!(plan, alpha, beta) + copyto!(C, Cstart) - # Poison every reused buffer: a stale packed slot yields NaN, a stale - # offset or descriptor addresses far outside the operand. + # A stale packed slot yields NaN, a stale offset or descriptor addresses far + # outside the operand. + ws = plan.workspace fill!(ws.packed_a, NaN) fill!(ws.packed_b, NaN) - for buf in (ws.m_buf_A, ws.m_buf_C, ws.n_buf_B, ws.n_buf_C, ws.k_buf_A, ws.k_buf_B) + for buf in (ws.m.offsets..., ws.n.offsets..., ws.k...) fill!(buf, typemin(Int) ÷ 4) end poison = BlockDescriptor(typemin(Int) ÷ 4, 0, 1, true) - for desc in (ws.m_desc_A, ws.m_desc_C, ws.n_desc_B, ws.n_desc_C) - fill!(desc, poison) - end - - lengths_before = _ws_lengths(ws) - Cpoisoned = copy(Cstart) - plan = _mm_plan( - Cpoisoned, Amat, Bmat; - kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, workspace = ws - ) - @test _ws_lengths(ws) == lengths_before - @test length(ws.packed_a) > cld(Ma, 4) * 4 * min(k_block, Ka) - @test length(ws.packed_b) > cld(Na, 3) * 3 * min(k_block, Ka) + foreach(d -> fill!(d, poison), (ws.m.descriptors..., ws.n.descriptors...)) execute!(plan, alpha, beta) - @test all(isfinite, Cpoisoned) - @test Cpoisoned == Cref -end - -@testset "plan_contract: oracle = false skips execute_tilewise!'s buffers" begin - Random.seed!(909) - kernel = ScalarKernel(Val(4), Val(3), Float64) - Ma, Ka, Na = 9, 10, 8 - Amat, Bmat = randn(Ma, Ka), randn(Ka, Na) + @test C == Cref - Cmat = zeros(Ma, Na) - plan = _mm_plan(Cmat, Amat, Bmat; kernel = kernel, k_block = 4, oracle = false) - ws = plan.workspace - - @test isempty(ws.tw_packed_a) - @test isempty(ws.tw_packed_b) - @test isempty(ws.tw_k_buf_A) - @test isempty(ws.tw_k_buf_B) - # The tile buffers serve the beta-only pass too. - @test length(ws.tile_m_buf_A) == 4 - @test length(ws.tile_n_buf_C) == 3 - - execute!(plan, 1.0, 0.0) - @test Cmat ≈ Amat * Bmat - - Cstart = randn(Ma, Na) - Cbeta = copy(Cstart) - execute!(_mm_plan(Cbeta, Amat, Bmat; kernel = kernel, oracle = false), 0.0, 0.5) - @test Cbeta ≈ 0.5 .* Cstart - - @test_throws ArgumentError execute_tilewise!(plan, 1.0, 0.0) - - Ctw = zeros(Ma, Na) - plan_tw = _mm_plan(Ctw, Amat, Bmat; kernel = kernel, k_block = 4, workspace = ws, oracle = true) - @test plan_tw.workspace === ws - @test !isempty(ws.tw_packed_a) - execute_tilewise!(plan_tw, 1.0, 0.0) - @test Ctw ≈ Amat * Bmat + # The beta-only pass borrows the M/N offset buffers. + foreach(buf -> fill!(buf, typemin(Int) ÷ 4), (ws.m.offsets..., ws.n.offsets...)) + execute!(plan, 0.0, 0.5) + @test C == 0.5 .* Cref end -@testset "plan_contract: explicit allocators size the packed panels exactly once" begin +@testset "plan_contract: explicit allocators serve the packed panels" begin Random.seed!(1717) kernel = SIMDKernel(Val(4), Val(3), Float64) Ma, Ka, Na = 19, 23, 17 @@ -171,7 +77,7 @@ end plan = _mm_plan( Cmat, Amat, Bmat; kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, - allocator = allocator, oracle = false + allocator = allocator ) ws = plan.workspace @@ -179,9 +85,8 @@ end @test all(isconcretetype, fieldtypes(typeof(ws))) @test length(ws.packed_a) == cld(plan.blocking.m_block, 4) * 4 * plan.blocking.k_block @test length(ws.packed_b) == cld(plan.blocking.n_block, 3) * 3 * plan.blocking.k_block - @test isempty(ws.tw_packed_a) - @test ws.m_buf_A isa Vector{Int} - @test ws.k_buf_B isa Vector{Int} + @test ws.m.offsets[1] isa Vector{Int} + @test ws.k[2] isa Vector{Int} execute!(plan, 1.5, 0.0) @test Cmat == Cdefault @@ -189,24 +94,6 @@ end QuasiStrided.release!(ws, allocator) TO.allocator_reset!(allocator, checkpoint) end - - # Allocator-owned temporaries are never resized. - manual = TO.ManualAllocator() - Cmanual = zeros(Ma, Na) - manual_plan = _mm_plan( - Cmanual, Amat, Bmat; - kernel = kernel, m_block = m_block, k_block = k_block, n_block = n_block, allocator = manual, oracle = false - ) - @test_throws MethodError QuasiStrided.reserve!( - manual_plan.workspace, kernel, manual_plan.blocking, false - ) - QuasiStrided.release!(manual_plan.workspace, manual) - - @test_throws ArgumentError _mm_plan( - zeros(Ma, Na), Amat, Bmat; - kernel = kernel, allocator = TO.ManualAllocator(), - workspace = default_plan.workspace - ) end @testset "ContractWorkspace: storage type T, packed panels of real(T)" begin @@ -214,23 +101,17 @@ end k32 = SIMDKernel(Val(8), Val(6), Float32) b = Blocking(16, 8, 12) - ws = QuasiStrided.ContractWorkspace(Float64, k64, b; oracle = true) + ws = QuasiStrided.ContractWorkspace(Float64, k64, b) @test ws isa QuasiStrided.ContractWorkspace{Float64, Vector{Float64}} @test eltype(ws.packed_a) === Float64 - wsc = QuasiStrided.ContractWorkspace(ComplexF64, k64, b; oracle = true) + wsc = QuasiStrided.ContractWorkspace(ComplexF64, k64, b) @test wsc isa QuasiStrided.ContractWorkspace{ComplexF64, Vector{Float64}} - @test eltype(wsc.packed_a) === Float64 === eltype(wsc.tw_packed_b) + @test eltype(wsc.packed_a) === Float64 === eltype(wsc.packed_b) @test all(isconcretetype, fieldtypes(typeof(wsc))) @test !any(t -> t isa Union, fieldtypes(typeof(wsc))) @test_throws ArgumentError QuasiStrided.ContractWorkspace(ComplexF64, k32, b) @test_throws ArgumentError QuasiStrided.ContractWorkspace(Float64, k32, b) @test_throws ArgumentError QuasiStrided.ContractWorkspace(ComplexF32, k64, b) - - Amat, Bmat, Cmat = randn(6, 5), randn(5, 4), zeros(6, 4) - @test_throws ArgumentError plan_contract( - StridedView(Cmat), StridedView(Amat), (1, 2), StridedView(Bmat), (2, 3), (1, 3); - workspace = wsc - ) end diff --git a/test/helpers.jl b/test/helpers.jl index 23c7100..f4d4b69 100644 --- a/test/helpers.jl +++ b/test/helpers.jl @@ -7,7 +7,6 @@ using StridedViews: StridedView, offset const plan_contract = QuasiStrided.plan_contract const execute! = QuasiStrided.execute! const ContractPlan = QuasiStrided.ContractPlan -const execute_tilewise! = QuasiStrided.execute_tilewise! import TensorOperations as TO diff --git a/test/integrations/test_tensoroperations.jl b/test/integrations/test_tensoroperations.jl index 6076c8b..6d8f8ad 100644 --- a/test/integrations/test_tensoroperations.jl +++ b/test/integrations/test_tensoroperations.jl @@ -204,82 +204,28 @@ end @test c1 == c2 end -# Run `f` with an empty workspace pool on the current task (not a spawned one: -# `@test` finds its testset through task-local storage too), restoring it after. -function _qs_with_clean_pool(f) - key = QuasiStrided._QS_WORKSPACE_KEY - tls = task_local_storage() - prior = get(tls, key, nothing) - delete!(tls, key) - try - return f() - finally - prior === nothing ? delete!(tls, key) : (tls[key] = prior) - end -end - -_qs_haspool() = haskey(task_local_storage(), QuasiStrided._QS_WORKSPACE_KEY) - -@testset "workspace pooling: default path reuses a task-local workspace" begin - Random.seed!(90210) - A, B = randn(20, 30), randn(30, 25) - run_once() = (C = zeros(20, 25); @tensor backend = qsbackend C[i, j] = A[i, k] * B[k, j]; C) - - steady = _qs_with_clean_pool() do - @test !_qs_haspool() - @test run_once() ≈ A * B - pool = task_local_storage(QuasiStrided._QS_WORKSPACE_KEY) - ws = pool[Float64] - steady = @allocated run_once() - @test pool[Float64] === ws - steady - end - # Julia 1.10 does not keep SIMDKernel's accumulator in registers. - @test steady < 20_000 skip = (VERSION < v"1.11") - Cv, Av, Bv = StridedView(zeros(20, 25)), StridedView(A), StridedView(B) - cold = @allocated plan_contract(Cv, Av, (1, -1), Bv, (-1, 2), (1, 2); oracle = false) - @test cold - steady > 5_000 skip = (VERSION < v"1.11") -end - -@testset "workspace pooling: explicit allocators never touch the pool" begin +@testset "tensorcontract!: allocator-routed workspace, released per call" begin Random.seed!(13571113) A, B = randn(20, 30), randn(30, 25) - _qs_with_clean_pool() do - for _ in 1:2 - C1, C2, C3 = zeros(20, 25), zeros(20, 25), zeros(20, 25) - @tensor backend = qsbackend allocator = TensorOperations.ManualAllocator() C1[i, j] = A[i, k] * B[k, j] - @tensor backend = qsbackend allocator = TensorOperations.BufferAllocator() C2[i, j] = A[i, k] * B[k, j] - @no_escape begin - @tensor backend = qsbackend allocator = default_buffer() C3[i, j] = A[i, k] * B[k, j] - end - @test C1 ≈ C2 ≈ C3 ≈ A * B + buffer = TensorOperations.BufferAllocator() + for _ in 1:2 + C1, C2, C3 = zeros(20, 25), zeros(20, 25), zeros(20, 25) + @tensor backend = qsbackend allocator = TensorOperations.ManualAllocator() C1[i, j] = A[i, k] * B[k, j] + @tensor backend = qsbackend allocator = buffer C2[i, j] = A[i, k] * B[k, j] + @no_escape begin + @tensor backend = qsbackend allocator = default_buffer() C3[i, j] = A[i, k] * B[k, j] end - @test !_qs_haspool() - C = zeros(20, 25) - @tensor backend = qsbackend C[i, j] = A[i, k] * B[k, j] - @test _qs_haspool() + @test C1 ≈ C2 ≈ C3 ≈ A * B end + @test isempty(buffer) end -# One pooled workspace per eltype, grown and then reused oversized. -@testset "workspace pooling: correctness across shapes and eltypes on one task" begin - Random.seed!(2024) - _qs_with_clean_pool() do - for (m, k, n) in ((6, 8, 5), (37, 41, 29), (3, 3, 3), (17, 90, 2)) - A, B = randn(m, k), randn(k, n) - C = fill(NaN, m, n) - @tensor backend = qsbackend C[i, j] = A[i, k] * B[k, j] - @test C ≈ A * B - end - for T in (Float64, ComplexF64, Float32, ComplexF32, Float64, ComplexF64) - A, B = randn(T, (12, 7)), randn(T, (7, 9)) - C = fill(convert(T, NaN), (12, 9)) - @tensor backend = qsbackend C[i, j] = A[i, k] * B[k, j] - @test C ≈ A * B - end - pool = task_local_storage(QuasiStrided._QS_WORKSPACE_KEY) - for T in all_eltypes - @test pool[T] isa QuasiStrided.ContractWorkspace{T, Vector{real(T)}} - end - end +@testset "tensorcontract!: the default path allocates no more than planning" begin + A, B, C = randn(20, 30), randn(30, 25), zeros(20, 25) + run_once() = TO.tensorcontract!(C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1, 2), ()), 1.0, 0.0, qsbackend) + Cv, Av, Bv = StridedView(C), StridedView(A), StridedView(B) + plan_once() = plan_contract(Cv, Av, (1, -1), Bv, (-1, 2), (1, 2)) + run_once(); plan_once() + @test C ≈ A * B + @test (@allocated run_once()) <= (@allocated plan_once()) skip = (VERSION < v"1.11") end diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index a4a159b..d798003 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -171,7 +171,7 @@ using StridedViews: StridedView Bv = StridedView(randn(T, d, d, d, d)) plan = QuasiStrided.plan_contract( StridedView(zeros(T, d, d, d, d, d, d)), StridedView(randn(T, d, d, d, d)), - (i, j, m, b), Bv, (m, k, a, c), (a, b, c, i, j, k); oracle = false + (i, j, m, b), Bv, (m, k, a, c), (a, b, c, i, j, k) ) @test tile_size(plan.kernel, 1) == 16 @test plan.Astorage === parent(Bv) @@ -179,7 +179,7 @@ using StridedViews: StridedView B2 = StridedView(randn(T, d, d, d, d)) plan2 = QuasiStrided.plan_contract( StridedView(zeros(T, d, d, d, d)), StridedView(randn(T, d, d)), (q, b), - B2, (a, q, r, s), (a, b, r, s); oracle = false + B2, (a, q, r, s), (a, b, r, s) ) @test tile_size(plan2.kernel, 1) == 16 @test plan2.Astorage === parent(B2) diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index fe3f07f..5287545 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -6,7 +6,6 @@ using StridedViews: StridedView, offset const QS = QuasiStrided const _pcf_plan = QuasiStrided.plan_contract const _pcf_exec = QuasiStrided.execute! -const _pcf_exec_tw = QuasiStrided.execute_tilewise! const _pcf_Plan = QuasiStrided.ContractPlan @testset "per-call floor: classify_labels" begin @@ -77,10 +76,11 @@ end A = randn(64, 64); B = randn(64, 64); C = zeros(64, 64) Av, Bv, Cv = StridedView(A), StridedView(B), StridedView(C) p = _pcf_plan(Cv, Av, (1, 2), Bv, (2, 3), (1, 3)) - ws = p.workspace - f() = _pcf_plan(Cv, Av, (1, 2), Bv, (2, 3), (1, 3); workspace = ws) + ws() = QS.ContractWorkspace(Float64, p.kernel, p.blocking) + ws() + f() = _pcf_plan(Cv, Av, (1, 2), Bv, (2, 3), (1, 3)) f() - @test (@allocated f()) <= 3000 + @test (@allocated f()) <= 3000 + (@allocated ws()) end @testset "per-call floor: checked_span_bounds is equivalent to the per-tile check" begin @@ -142,7 +142,7 @@ end buf2 = [rand(-40:40) for _ in 1:blocklen] d1 = Vector{BlockDescriptor}(undef, nsliv) d2 = Vector{BlockDescriptor}(undef, nsliv) - (r1, r2) = QS._classify_slivers!(d1, d2, buf1, buf2, blocklen, reg, nsliv) + (r1, r2) = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) @test r1 == (minimum(buf1), maximum(buf1)) @test r2 == (minimum(buf2), maximum(buf2)) @test all(s -> d1[s].count == min(reg, blocklen - (s - 1) * reg), 1:nsliv) @@ -155,7 +155,7 @@ end buf2[7:12] .= [0, -7, -14, -21, -28, -35] # sliver 2, REGULAR, stride -7 d1 = Vector{BlockDescriptor}(undef, 3) d2 = Vector{BlockDescriptor}(undef, 3) - (r1, r2) = QS._classify_slivers!(d1, d2, buf1, buf2, blocklen, reg, 3) + (r1, r2) = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, 3) @test !d1[2].regular && d1[2].count == 6 # scan branch @test d2[2].regular && d2[2].stride == -7 # affine branch, negative stride @test r1 == (-500, 500) @@ -225,13 +225,6 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() @test cstore.n == 1 + forced * ntiles @test astore.n == 1 + forced * m_tiles @test bstore.n == 1 - - # The oracle checks per tile. - fill!(cstore.data, 0.0) - astore.n = 0; bstore.n = 0; cstore.n = 0 - _pcf_exec_tw(p, 1.0, 0.0) - @test reshape(cstore.data, Ma, Na) ≈ Amat * Bmat - @test cstore.n >= ntiles end @testset "per-call floor: a block that must be rejected is still rejected" begin @@ -259,7 +252,6 @@ end short_C = zeros(Ma * Na - 1) @test_throws BoundsError _pcf_exec(replan(Cstorage = short_C), 1.0, 0.0) @test all(iszero, short_C) - @test_throws BoundsError _pcf_exec_tw(replan(Cstorage = short_C), 1.0, 0.0) @test_throws BoundsError _pcf_exec(replan(Cbase = -1), 1.0, 0.0) @@ -332,12 +324,6 @@ end @test QS.checked_span_bounds(base.Cbase, mrange, firstonly, shortlen) === nothing @test QS.checked_span_bounds(base.Cbase, mrange, lastonly, shortlen) === nothing - # The oracle rejects too, but only at the offending tile, after writing - # the ones before it. - fill!(short_C, 0.0) - @test_throws BoundsError _pcf_exec_tw(pshort, 1.0, 0.0) - @test any(!iszero, short_C) - # An interior-sliver minimum: the base shifted down by one. plow = _pcf_Plan( base.kernel, base.mgroup, base.ngroup, base.kgroup, base.blocking, @@ -441,11 +427,11 @@ end d1 = Vector{BlockDescriptor}(undef, nsliv) d2 = Vector{BlockDescriptor}(undef, nsliv) fill_offsets!((buf1, buf2), g, first, blocklen) - want = QS._classify_slivers!(d1, d2, buf1, buf2, blocklen, reg, nsliv) + want = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) r1 = Vector{BlockDescriptor}(undef, nsliv) r2 = Vector{BlockDescriptor}(undef, nsliv) - got = QS._ramp_slivers!(r1, r2, steps[1], steps[2], first, blocklen, reg, nsliv) + got = QS._ramp_slivers!(QS.GroupBuffers((buf1, buf2), (r1, r2)), steps[1], steps[2], first, blocklen, reg, nsliv) for s in 1:nsliv @test r1[s].base == d1[s].base && r1[s].stride == d1[s].stride @@ -468,12 +454,7 @@ end StridedView(C), StridedView(A), indA, StridedView(B), indB, indC; kw... ) _pcf_exec(p, 2.0, -0.5) - ref = copy(C0) - p2 = _pcf_plan( - StridedView(ref), StridedView(A), indA, StridedView(B), indB, indC; kw... - ) - _pcf_exec_tw(p2, 2.0, -0.5) - @test C ≈ ref + @test C ≈ _lo_reference(C0, A, indA, B, indB, indC; alpha = 2.0, beta = -0.5) return p end diff --git a/test/planning/test_plan_contract.jl b/test/planning/test_plan_contract.jl index 489e8ca..b2b607d 100644 --- a/test/planning/test_plan_contract.jl +++ b/test/planning/test_plan_contract.jl @@ -30,9 +30,6 @@ end 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 @@ -111,9 +108,6 @@ end fill!(Cmat, 0.0) execute!(plan, 1.0, 0.0) @test Cmat ≈ Amat * Bmat - fill!(Cmat, 0.0) - execute_tilewise!(plan, 1.0, 0.0) - @test Cmat ≈ Amat * Bmat end # A real eltype collapses a view's `conj` op to identity. @@ -430,13 +424,6 @@ end end Cref = _lo_reference(C, Av, indA, Bv, indB, indC; alpha = 1.3, beta = -0.7) - Ctw = copy(C) - plan_tw = plan_contract( - StridedView(Ctw), Av, indA, Bv, indB, indC; kernel = plan.kernel - ) - execute_tilewise!(plan_tw, 1.3, -0.7) - @test Ctw ≈ Cref - Cex = copy(C) plan_ex = plan_contract(StridedView(Cex), Av, indA, Bv, indB, indC) execute!(plan_ex, 1.3, -0.7) @@ -581,12 +568,6 @@ end plan = plan_contract(Cv, Av, indA, Bv, indB, indC; conjA = conjA, conjB = conjB) execute!(plan, alpha, beta) @test isapprox(copy(Cv), Cref; rtol = rtol) - - copyto!(Csub, permutedims(Cstart, invperm(perm))) - plan_tw = plan_contract(Cv, Av, indA, Bv, indB, indC; conjA = conjA, conjB = conjB) - @test (plan_tw.Astorage === parent(Bv)) == (plan.Astorage === parent(Bv)) - execute_tilewise!(plan_tw, alpha, beta) - @test isapprox(copy(Cv), Cref; rtol = rtol) end end end @@ -742,7 +723,7 @@ end @test pt.kgroup.strides[2] == (14400, 360, 9) # then A's end -@testset "K order: correctness on scrambled K (all dtypes, alpha/beta, conj, both orientations, oracle)" begin +@testset "K order: correctness on scrambled K (all dtypes, alpha/beta, conj, both orientations)" begin Random.seed!(0x5C7A_0B1E) ext = _KO_BIG for T in (Float64, Float32, ComplexF64, ComplexF32) @@ -767,10 +748,6 @@ end plan = plan_contract(Cv, Av, indA, Bv, indB, indC; conjA = conjA, conjB = conjB) execute!(plan, alpha, beta) @test isapprox(C, Cref; rtol = rtol) - copyto!(C, Cstart) - plan_tw = plan_contract(Cv, Av, indA, Bv, indB, indC; conjA = conjA, conjB = conjB) - execute_tilewise!(plan_tw, alpha, beta) - @test isapprox(C, Cref; rtol = rtol) end end end From d0100e671bf4666b4432f739e93700ed9d25af68 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 2 Oct 2026 17:01:10 -0400 Subject: [PATCH 28/60] Benchmark: hybrid complex blocking variant; decision log: chunk 12 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 22 +++++++++++++++- benchmark/bench_complex_blocking.jl | 40 +++++++++++++++++++++-------- 2 files changed, 51 insertions(+), 11 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index bd521c7..59cedd8 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -24,7 +24,7 @@ TTFX when codegen is touched). | 9 | Labels | `planning/labels.jl` (`conjugation.jl` → plan.jl) | done | | 10 | Kernel selection | `planning/kernel_selection.jl` | done | | 11 | Blocking | `blocking.jl`, line packing (now `packing/line_packing.jl`) | done | -| 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | | +| 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | done | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | | 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | | 15 | Alternative paths | `oracle.jl`, `unpackedb.jl` | | @@ -203,6 +203,26 @@ size is used (`line_bytes`). The line-packing feature lives in `packing/line_packing.jl`; `pack_split` takes the `SliverSpec`; the packer is `pack_block_by_lines!`. +### D16. Workspace (chunk 12, applied) + +No workspace pooling or reuse API: a plan builds its workspace through the +TensorOperations allocator; reusing a plan is the way to reuse buffers +(allocation-free). Measured cost under `DefaultAllocator`: per-call floor +~3x, 128³–256³ +13–36%; none under Bumper. If this matters, design an +explicit reuse mechanism rather than a cache. The tile-wise oracle is gone +(tests compare against the brute-force reference). Offsets and descriptors +are bundled per group (`ws.m`, `ws.n`, `ws.k`), not per operand: a group's +two maps are filled together and share a lifetime. Plans are single-task. + +### Complex blocking benchmark (after chunk 11) + +Same-node A/B, `bench_complex_blocking.jl`, jobs 7160147–9: the kernel's own +model is faster on Icelake for large sizes (0.58–0.92; the old scaled row was +rounded *up* to MR past the L2 budget), but 1–10% slower on Genoa/Rome and at +small M, following the smaller `k_block`. Rerun with a third "hybrid" +variant (real row's `k_block`, own `m_block`/`n_block` rounded down) to pick +one. + ## Possible improvements - Piecewise-affine block descriptions instead of block-sized offset buffers: diff --git a/benchmark/bench_complex_blocking.jl b/benchmark/bench_complex_blocking.jl index 2cc7c25..e5d8498 100644 --- a/benchmark/bench_complex_blocking.jl +++ b/benchmark/bench_complex_blocking.jl @@ -1,7 +1,10 @@ -# Same-node A/B of complex cache blocking: the kernel's own analytical blocking -# (`default_blocking`) against the real row of `real(T)`'s default kernel -# divided by packed reals per element, interleaved per case, for every menu -# shape of the planar, 1m and fmaddsub kernels that fits the host's vectors. +# Same-node comparison of complex cache blocking, interleaved per case, for +# every menu shape of the planar, 1m and fmaddsub kernels that fits the host's +# vectors: +# old - the real row of `real(T)`'s default kernel divided by packed reals +# new - the kernel's own analytical blocking (`default_blocking`) +# hybrid - the real row's `k_block`, `m_block`/`n_block` from the kernel's +# byte budget at that `k_block`, rounded down to MR/NR # # julia --project=benchmark benchmark/bench_complex_blocking.jl # [--dtypes ComplexF64,ComplexF32] [--rounds 3] [--reps 3] [--outdir DIR] @@ -10,7 +13,7 @@ include(joinpath(@__DIR__, "harness.jl")) using QuasiStrided: PlanarKernel, OneMKernel, FMAddSubKernel, Blocking, default_blocking, kernel_shapes, kernel_from_shape, derived_shape, target_profile, pack_formats, - reals_per_element, scalartype + reals_per_element, scalartype, tile_size const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32")) const ROUNDS = argopt("rounds", 3) @@ -33,6 +36,20 @@ function scaled_real_row(kernel) return Blocking(max(1, row.m_block ÷ a_reals), row.k_block, max(1, row.n_block ÷ b_reals)) end +function hybrid_blocking(kernel) + T = scalartype(kernel) + R = real(T) + profile = target_profile() + row = default_blocking(kernel_from_shape(derived_shape(profile, R), R, SIMDKernel)) + profile.l1d.bytes > 0 && profile.l2.bytes > 0 || return scaled_real_row(kernel) + MR, NR = tile_size(kernel) + SA, SB = map(f -> reals_per_element(f) * sizeof(R), pack_formats(typeof(kernel))) + k = row.k_block + m = max(MR, ((profile.l2_share ÷ 2) ÷ (k * SA)) ÷ MR * MR) + n = max(NR, ((profile.l2_share + profile.l3_share) ÷ (k * SB)) ÷ NR * NR) + return Blocking(m, k, n) +end + function time_blocking(fx, kernel, b::Blocking, ::Type{T}) where {T} plan = plan_contract( fx.Cv, fx.Av, fx.indA, fx.Bv, fx.indB, fx.indC; @@ -43,7 +60,7 @@ end function run(io) print_env_header(stdout, "bench_complex_blocking.jl") - println(io, "dtype,kernel,MR,NR,W,case,old_m,old_k,old_n,new_m,new_k,new_n,round,t_old,t_new") + println(io, "dtype,kernel,MR,NR,W,case,old_m,old_k,old_n,new_m,new_k,new_n,hyb_m,hyb_k,hyb_n,round,t_old,t_new,t_hybrid") rng = MersenneTwister(1) lanes(T) = target_profile().vector_bytes ÷ sizeof(real(T)) for T in RUN_DTYPES, K in (PlanarKernel, OneMKernel, FMAddSubKernel) @@ -52,7 +69,8 @@ function run(io) kernel = K(Val(MR), Val(NR), T, Val(W)) new = default_blocking(kernel) old = scaled_real_row(kernel) - if new == old + hybrid = hybrid_blocking(kernel) + if new == old == hybrid println("skip $T $(nameof(K)) ($MR,$NR,$W): blocking unchanged $new") continue end @@ -61,18 +79,20 @@ function run(io) for r in 1:ROUNDS t_old = time_blocking(fx, kernel, old, T) t_new = time_blocking(fx, kernel, new, T) + t_hybrid = time_blocking(fx, kernel, hybrid, T) println( io, join( ( T, nameof(K), MR, NR, W, spec.name, old.m_block, old.k_block, old.n_block, - new.m_block, new.k_block, new.n_block, r, t_old, t_new, + new.m_block, new.k_block, new.n_block, + hybrid.m_block, hybrid.k_block, hybrid.n_block, r, t_old, t_new, t_hybrid, ), "," ) ) flush(io) @printf( - "%s %s (%d,%d,%d) %s round %d: old %.4g s, new %.4g s, new/old %.3f\n", - T, nameof(K), MR, NR, W, spec.name, r, t_old, t_new, t_new / t_old + "%s %s (%d,%d,%d) %s round %d: old %.4g s, new/old %.3f, hybrid/old %.3f\n", + T, nameof(K), MR, NR, W, spec.name, r, t_old, t_new / t_old, t_hybrid / t_old ) end end From cab517e4bbc428d232a7a7f17c6962c32918d7a2 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 3 Oct 2026 16:43:35 -0400 Subject: [PATCH 29/60] Hybrid complex blocking A complex kernel keeps the k_block of real(T)'s default real kernel (B sliver at half L1d) and derives m_block/n_block from its own packed bytes at that depth, rounded down to MR/NR. The per-kernel model from chunk 11 let a complex B sliver halve k_block, which measured slower. Same-node A/B over the planar, 1m and FMAddSub menu shapes (512^3, 2048^3, 12/24x2048x2048; Genoa, Rome, Icelake): hybrid/old 0.99-1.00 (worst 1.02), per-kernel model/old up to 1.10, worst at small M. Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 23 +++++++++++++++++++++-- src/planning/blocking.jl | 21 +++++++++++++-------- test/planning/test_kernel_selection.jl | 9 +++++---- 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 59cedd8..41c2c53 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -220,8 +220,27 @@ Same-node A/B, `bench_complex_blocking.jl`, jobs 7160147–9: the kernel's own model is faster on Icelake for large sizes (0.58–0.92; the old scaled row was rounded *up* to MR past the L2 budget), but 1–10% slower on Genoa/Rome and at small M, following the smaller `k_block`. Rerun with a third "hybrid" -variant (real row's `k_block`, own `m_block`/`n_block` rounded down) to pick -one. +variant (jobs 7160543–5, commit d0100e6): hybrid/old 0.99–1.00 on all three +nodes (worst single case 1.02), new/old up to 1.10 at small M on AMD; the +Icelake 0.58 did not reproduce (noise). Decision: the hybrid — `k_block` from +the real default kernel's B sliver (`l1_k_block`), `m_block`/`n_block` from +the kernel's byte budget at that `k_block`, rounded down. Implemented in +blocking.jl (uncommitted until the full suite passes). + +## Resume here + +1. Chunk 13: `planning/plan.jl` (+ `test_plan_contract.jl`, + `test_per_call_overhead.jl`), including D3 (execution path as a plan type + parameter, removing barrier #2 and the hint machinery; kernel types cross + the barrier as `Val`, `::Type{K}` everywhere after it) and the open items + for chunks 13/14 below. +2. Workflow: present each chunk (purpose, reading order, design decisions, + proposed fixes, questions), then hand the agreed changes to an Opus agent + with the standard checks (Runic, full suite, per-call floor/allocations, + TTFX, equivalence script where behaviour must not change); perf claims + need Slurm runs (user submits; `--reservation=rocky8`), never poll Slurm. + Only quote commit SHAs read from `git log`. + ## Possible improvements diff --git a/src/planning/blocking.jl b/src/planning/blocking.jl index 152ad99..b2835f0 100644 --- a/src/planning/blocking.jl +++ b/src/planning/blocking.jl @@ -31,12 +31,14 @@ default_blocking(kernel, profile::TargetProfile = target_profile()) = kernel_blocking(profile, scalartype(kernel), typeof(kernel), tile_size(kernel)...) # The default blocking of a kernel of type `K` for `T` with an `MR x NR` tile. -# The fallback rows are divided by the packed reals per element of each -# operand, so every kernel type gets the same packed byte budget. +# `k_block` is the real default kernel's for `real(T)`: derived from a complex +# B sliver it halves or worse, which is slower, most at small M. The fallback rows are divided by the packed reals per element of +# each operand, so every kernel type gets the same packed byte budget. @inline function kernel_blocking(profile::TargetProfile, ::Type{T}, ::Type{K}, MR::Int, NR::Int) where {T, K} R = real(T) a_reals, b_reals = map(reals_per_element, pack_formats(K)) - modelled = modelled_blocking(profile, MR, NR, a_reals * sizeof(R), b_reals * sizeof(R)) + k_block = l1_k_block(profile, derived_shape(profile, R)[2], sizeof(R)) + modelled = modelled_blocking(profile, MR, NR, a_reals * sizeof(R), b_reals * sizeof(R), k_block) modelled === nothing || return modelled base = fallback_blocking(R) return Blocking(max(1, base.m_block ÷ a_reals), base.k_block, max(1, base.n_block ÷ b_reals)) @@ -52,7 +54,7 @@ fallback_blocking(::Type{Float32}) = Blocking(96, 768, 1152) # `m_block*k_block`, and one `NR x k_block` B sliver is reused by every A # sliver of the block: # -# NR*k_block*SB <= L1/2 the reused B sliver +# NR*k_block*SB <= L1/2 the reused B sliver (real kernels) # m_block*k_block*SA <= L2core/2 the A block # n_block*k_block*SB <= L2core + L3core the B panel # @@ -61,15 +63,18 @@ fallback_blocking(::Type{Float32}) = Blocking(96, 768, 1152) # oversized B panel only re-streams B, an undersized one repacks A once per N # block. # `nothing` when L1d or L2 is undetected. -@inline function modelled_blocking(profile::TargetProfile, MR::Int, NR::Int, SA::Int, SB::Int) - l1 = profile.l1d - (l1.bytes > 0 && profile.l2.bytes > 0) || return nothing +@inline function modelled_blocking( + profile::TargetProfile, MR::Int, NR::Int, SA::Int, SB::Int, k_block::Int = l1_k_block(profile, NR, SB) + ) + (profile.l1d.bytes > 0 && profile.l2.bytes > 0) || return nothing l2core, l3core = profile.l2_share, profile.l3_share - k_block = max(1, (l1.bytes ÷ 2) ÷ (NR * SB)) m_block = max(MR, ((l2core ÷ 2) ÷ (k_block * SA)) ÷ MR * MR) n_block = max(NR, ((l2core + l3core) ÷ (k_block * SB)) ÷ NR * NR) return Blocking(m_block, k_block, n_block) end +# The K depth at which an `NR`-wide B sliver of `SB`-byte elements fills half of L1d. +@inline l1_k_block(profile::TargetProfile, NR::Int, SB::Int) = max(1, (profile.l1d.bytes ÷ 2) ÷ (NR * SB)) + # Smallest multiple of `n` that is >= `x` (`x >= 0`, `n >= 1`). @inline roundup(x::Int, n::Int) = cld(x, n) * n diff --git a/test/planning/test_kernel_selection.jl b/test/planning/test_kernel_selection.jl index d798003..b64eb13 100644 --- a/test/planning/test_kernel_selection.jl +++ b/test/planning/test_kernel_selection.jl @@ -305,16 +305,17 @@ end @test modelled_blocking(nosmt, 8, 6, 8, 8) === Blocking(96, 341, 1728) @test modelled_blocking(nosmt, 16, 6, 4, 4) === Blocking(96, 682, 1728) @test modelled_blocking(smt2, 16, 6, 8, 8) === Blocking(192, 341, 1572) - # A complex kernel blocks by its packed bytes: 1m packs 4 reals per A - # element, planar 2 per B element. - @test default_blocking(kernel_from_shape((12, 8, 8), ComplexF64, OneMKernel), nosmt) === Blocking(60, 128, 2304) + # A complex kernel keeps the real default kernel's `k_block` and blocks M + # and N by its packed bytes: 1m packs 4 reals per A element, 2 per B element. + @test default_blocking(kernel_from_shape((12, 8, 8), ComplexF64, OneMKernel), nosmt) === Blocking(24, 341, 864) for p in (nosmt, smt2), T in (Float64, Float32, ComplexF64, ComplexF32), K in (SIMDKernel, PlanarKernel, OneMKernel, FMAddSubKernel) for shape in kernel_shapes(T, K) k = kernel_from_shape(shape, T, K) b = default_blocking(k, p) MR, NR = tile_size(k) @test b.m_block % MR == 0 && b.n_block % NR == 0 - @test sliver_width(k, 2) * b.k_block * sizeof(real(T)) <= p.l1d.bytes ÷ 2 + real_default = kernel_from_shape(derived_shape(p, real(T)), real(T), SIMDKernel) + @test b.k_block == default_blocking(real_default, p).k_block end end # No L3: the B panel is budgeted from the L2 alone. From 308c07b96f3714652ffd0aad0ad514f2e332a6e7 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 3 Oct 2026 18:01:09 -0400 Subject: [PATCH 30/60] Chunk 13: execution path decided at plan time Blocking resolution (defaults, rounding, clamping, C panel, line-by-line packing) and path selection move before the kernel barrier, computed from the kernel type and tile (`resolve_blocking`, `select_path`). The barrier crosses once with the kernel (`Val{Kern}()`, or a named kernel itself), the path singleton, the transforms and one `RefValue` holding the request plus the resolved blocking. `ContractPlan` gains a `path` field/type parameter; `execute!` keeps its runtime short-circuits and then calls `execute_path!(plan, ..., plan.path)` statically, so barrier #2 is gone. The dot path's capacity check is evaluated at plan time (`dot_fits`) from `workspace_sizes`, which now works on the kernel type and tile. `pack_split` takes (side, lanes, format) instead of a `SliverSpec`. Deleted: `_path_hint`, `_default_k_block`, `_continue`, `_execute_hinted!`, `_hint_holds`, `_splits`, `_split_path`, `_execute_across_barrier!`, `_execute_resolved!`, `_strip_storage`, `_with_storage`, `SlotCache`, `barrier_slot!`/`barrier_slot_slow!`, the workspace's `slots` field, `_dot_capacity_ok`, and the global `_DOT_MODE`/`_OUTER_MODE`/ `_UNPACKED_B_MODE` Refs (now an internal `path_modes::PathModes` keyword of `plan_contract`). `Execute` is a callable struct (execute! + release!), and plan construction ends with `req.f(plan)`. One keyword copy method `ContractPlan(p; ...)` replaces the spelled-out copies in the panel path and tests. Underscores dropped on touched names; barrier.jl -> paths.jl. Measurements (workstation, cascadelake; noisy, medians of 4 interleaved runs of per-call medians, indicative only), before -> after: execute! 64^3 Float64 6.8 us -> 6.8 us, 0 B execute! 2x2 Float64 127 ns -> 85-190 ns (bimodal), 0 B execute! 32^3 ComplexF64 5.4 us -> 5.3 us, 0 B execute! dot 1x256x64 Float64 1.77 us -> 1.74 us, 0 B execute! panel 32x600x32 F32/F64 30.5 us -> 30.3 us, 0 B plan_contract 64^3 Float64 1.94 us -> 2.06 us; 72240 -> 71952 B plan_contract 2x2 Float64 417 ns -> 405 ns; 2272 -> 1984 B tensorcontract! 2x2 Float64 486 ns -> 484 ns; 2096 -> 1808 B tensorcontract! 2x2 ComplexF64 687 ns -> 627 ns; 2432 -> 2144 B TTFX (fresh process, median of 3): first Float64 plan+execute 3.44 -> 3.52 s, then ComplexF64 1.84 -> 1.93 s; first @tensor call 2.55 -> 2.53 s. Codegen on a fixed 172-case workload (all paths, eltypes, conj, swap, split, panel, named kernels, TensorOperations): path/nest/micro-tile/dot/outer specializations identical (58/40/18/22/14/17/4), --trace-compile 715 -> 714 lines, except that k_length == 0 plans now compile their (never run) nest path: +4 nests with four such cases included. Results of all 172 cases are bitwise identical to before, with identical paths, blocking and splits. Co-Authored-By: Claude Opus 5.5 --- benchmark/bench_degenerate.jl | 36 ++-- src/QuasiStrided.jl | 2 +- src/execution/barrier.jl | 108 ---------- src/execution/dot.jl | 15 +- src/execution/execute.jl | 194 +++--------------- src/execution/outer.jl | 4 - src/execution/paths.jl | 68 +++++++ src/execution/unpackedb.jl | 12 +- src/execution/workspace.jl | 26 +-- src/integrations/tensoroperations.jl | 25 ++- src/packing/line_packing.jl | 20 +- src/planning/plan.jl | 253 ++++++++++++++---------- test/execution/test_dot.jl | 6 +- test/execution/test_execute.jl | 40 ++-- test/execution/test_mixed.jl | 26 +-- test/execution/test_outer.jl | 28 ++- test/execution/test_unpackedb.jl | 6 +- test/microkernels/test_mixed_kernel.jl | 25 +-- test/planning/test_per_call_overhead.jl | 25 +-- test/runtests.jl | 2 +- 20 files changed, 371 insertions(+), 550 deletions(-) delete mode 100644 src/execution/barrier.jl create mode 100644 src/execution/paths.jl diff --git a/benchmark/bench_degenerate.jl b/benchmark/bench_degenerate.jl index 7ae088f..eeb8aca 100644 --- a/benchmark/bench_degenerate.jl +++ b/benchmark/bench_degenerate.jl @@ -1,9 +1,9 @@ # A/B of the three automatically selected small/degenerate-shape paths against -# the five-loop nest, via their mode overrides: +# the five-loop nest, via the internal `path_modes` keyword of `plan_contract`: # -# outer K = 1 outer product C[a,b] = A[a] B[b] _OUTER_MODE -# dot M = 1 gemv C[cde] = A[ab] B[abcde] _DOT_MODE -# unpackedb small square matmul C[a,c] = A[a,b] B[b,c] _UNPACKED_B_MODE +# outer K = 1 outer product C[a,b] = A[a] B[b] outer +# dot M = 1 gemv C[cde] = A[ab] B[abcde] dot +# unpackedb small square matmul C[a,c] = A[a,b] B[b,c] unpacked_b # # `execute!` only (planning excluded), plus OpenBLAS `mul!` on the equivalent # dense matrices for scale. Modes are interleaved rep by rep; each rep times a @@ -39,17 +39,11 @@ end fmt((m, lo, hi)) = @sprintf("%9.0f [%6.0f,%6.0f]", m, lo, hi) -function withmode(f, ref::Ref{Symbol}, mode::Symbol) - return () -> begin - old = ref[]; ref[] = mode - f() - ref[] = old - end -end - -function row(name, ref, plan, blasf) - ex = () -> execute!(plan, 1.0, 0.0) - r = interleaved(Any[withmode(ex, ref, :never), withmode(ex, ref, :auto), blasf]) +# One plan with the path ruled out (`mode = :never`), one with it automatic. +function row(name, mode::Symbol, args, blasf) + nest = plan_contract(args...; path_modes = QuasiStrided.PathModes(; mode => :never)) + path = plan_contract(args...) + r = interleaved(Any[() -> execute!(nest, 1.0, 0.0), () -> execute!(path, 1.0, 0.0), blasf]) @printf( "%-28s nest %s path %s blas %s path/nest %.2f path/blas %.2f\n", name, fmt(r[1]), fmt(r[2]), fmt(r[3]), r[2][1] / r[1][1], r[2][1] / r[3][1] @@ -62,19 +56,19 @@ for T in (Float64, ComplexF64) println("## $T") for d in (16, 63, 128) a = randn(T, d); b = randn(T, d); C = zeros(T, d, d) - plan = plan_contract(StridedView(C), StridedView(a), (1,), StridedView(b), (2,), (1, 2)) + args = (StridedView(C), StridedView(a), (1,), StridedView(b), (2,), (1, 2)) am = reshape(a, d, 1); bm = reshape(b, 1, d) - row("outer $(d)x$(d)", QuasiStrided._OUTER_MODE, plan, () -> mul!(C, am, bm)) + row("outer $(d)x$(d)", :outer, args, () -> mul!(C, am, bm)) end for d in (6, 8, 12, 16) A = randn(T, d, d); B = randn(T, d, d, d, d, d); C = zeros(T, d, d, d) - plan = plan_contract(StridedView(C), StridedView(A), (1, 2), StridedView(B), (1, 2, 3, 4, 5), (3, 4, 5)) + args = (StridedView(C), StridedView(A), (1, 2), StridedView(B), (1, 2, 3, 4, 5), (3, 4, 5)) Bm = reshape(B, d^2, d^3); av = vec(A); cv = vec(C) - row("dot 1x$(d^2)x$(d^3)", QuasiStrided._DOT_MODE, plan, () -> mul!(cv, transpose(Bm), av)) + row("dot 1x$(d^2)x$(d^3)", :dot, args, () -> mul!(cv, transpose(Bm), av)) end for d in (16, 32, 64, 128) A = randn(T, d, d); B = randn(T, d, d); C = zeros(T, d, d) - plan = plan_contract(StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3)) - row("unpackedB $(d)^3", QuasiStrided._UNPACKED_B_MODE, plan, () -> mul!(C, A, B)) + args = (StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3)) + row("unpackedB $(d)^3", :unpacked_b, args, () -> mul!(C, A, B)) end end diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 3449b52..5810af9 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -37,11 +37,11 @@ include("microkernels/stores.jl") include("planning/labels.jl") include("planning/kernel_selection.jl") include("planning/blocking.jl") -include("execution/barrier.jl") include("execution/workspace.jl") include("planning/plan.jl") # --- Execution: the five-loop nest and the specialised paths --- +include("execution/paths.jl") include("execution/macrokernel.jl") include("execution/execute.jl") include("execution/unpackedb.jl") diff --git a/src/execution/barrier.jl b/src/execution/barrier.jl deleted file mode 100644 index 036d93f..0000000 --- a/src/execution/barrier.jl +++ /dev/null @@ -1,108 +0,0 @@ -# Dispatch barriers. The kernel and the execution path are runtime choices -# that everything below them is specialised on. A static branch over every -# choice would make Julia infer and compile every arm (the whole kernel menu x -# every path x, for complex, four transform pairs) on the first call. A -# dynamic call instead compiles only the callee actually reached. -# -# `jl_apply_generic` boxes every argument that is not already a heap object -# or a singleton, so each barrier passes: -# * static choices (kernel shape, method, transforms, path) as singletons; -# * operand storages (already heap objects) as themselves; -# * everything else through a typed slot, a `Base.RefValue{P}` (owned by the -# workspace's `SlotCache` once there is a workspace): written just before -# the call, read back first thing in the callee. -# The payload never holds an operand array, so a slot retains no user data -# between calls. A workspace belongs to one task at a time, and the callee -# copies the payload out before it could run another contraction. - -# Typed payload slots: one `Base.RefValue{P}` per payload type, plus the last -# one used. They live as long as the workspace, which is what lets a barrier -# call allocate nothing. -mutable struct SlotCache - last::Any - const slots::IdDict{Any, Any} -end -SlotCache() = SlotCache(nothing, IdDict{Any, Any}()) - -# The typed slot for payload type `P`; the common case (same payload type as -# the previous crossing) is one type-tag test. -@inline function barrier_slot!(cache::SlotCache, ::Type{P}) where {P} - slot = cache.last - slot isa Base.RefValue{P} && return slot - return barrier_slot_slow!(cache, P) -end - -@noinline function barrier_slot_slow!(cache::SlotCache, ::Type{P}) where {P} - slot = get!(() -> Base.RefValue{P}(), cache.slots, P)::Base.RefValue{P} - cache.last = slot - return slot -end - -# ---------------------------------------------------------------------------- -# Execution paths: singletons that select one `_execute_path!` method. -# ---------------------------------------------------------------------------- - -# 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 offset-buffer `view` 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. A split group is -# 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, panel::Bool = false - ) - split_b &= !unpacked_b - return _nest_path_from_flags( - 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(table::Vector{Any}, flags::Vararg{Bool, 9}) - index = 1 - for i in 1:9 - index += flags[i] << (i - 1) - end - return @inbounds table[index] -end - -const _NEST_PATHS = let - 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 - -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 -# (`m_length == 1`). -struct _DotPath{MATB, W} end - -# The outer-product path at lane width `W`. -struct _OuterPath{W} end - -# A runtime lane width to its path singleton, one literal per arm. -@inline function _lane_path(::Type{P}, W::Int) where {P} - W == 16 && return P{16}() - W == 8 && return P{8}() - W == 4 && return P{4}() - return P{2}() -end diff --git a/src/execution/dot.jl b/src/execution/dot.jl index de07e78..4cd9920 100644 --- a/src/execution/dot.jl +++ b/src/execution/dot.jl @@ -6,9 +6,6 @@ # and unpacked, one horizontal reduction per output. Only the K-vector is # gathered, once per K block, with its conjugation folded in. -# Forces the path off (`:never`) for benchmarks and tests; `:auto` otherwise. -const _DOT_MODE = Ref{Symbol}(:auto) - # One full hardware vector register of `real(T)`. @inline function _dot_lanewidth(::Type{T}) where {T} R = real(T) @@ -24,7 +21,6 @@ _dot_group_width(::Type{T}) where {T} = T <: Complex ? 4 : 8 # free extent, a matrix operand with unit-ramp K in dense storage (raw-pointer # loads), and at least one vector of K. function _dot_applicable(::Type{T}, Astorage, Bstorage, kgroup::AxisGroup, m_length::Int, n_length::Int, k_length::Int) where {T} - _DOT_MODE[] === :never && return false (m_length == 1 || n_length == 1) || return false k_length >= _dot_lanewidth(T) || return false if m_length == 1 @@ -38,13 +34,10 @@ function _dot_applicable(::Type{T}, Astorage, Bstorage, kgroup::AxisGroup, m_len end # The gathered vector (and, complex, its pair-swapped copy) lives in the -# packed-A buffer. Always true for an automatically chosen kernel; `nothing` -# is the prediction's "assume so". -function _dot_capacity_ok(plan::ContractPlan{T}) where {T} - need = (T <: Complex ? 2 : 1) * plan.blocking.k_block - return (length(plan.workspace.packed_a) * sizeof(real(T))) ÷ sizeof(T) >= need -end -_dot_capacity_ok(::Nothing) = true +# packed-A buffer of `packed_a` reals. Always true for an automatically chosen +# kernel. +dot_fits(::Type{T}, packed_a::Int, k_block::Int) where {T} = + packed_a * sizeof(real(T)) ÷ sizeof(T) >= (T <: Complex ? 2 : 1) * k_block function _execute_dot!(plan::ContractPlan{T}, alphaT::T, betaT::T, matB::Bool, ::Val{W}) where {T, W} ws = plan.workspace diff --git a/src/execution/execute.jl b/src/execution/execute.jl index 6d8fa0d..c1ab5a8 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -36,13 +36,13 @@ block. Allocation-free. function execute!(plan::ContractPlan{T}, alpha::Number, beta::Number) where {T} alphaT = convert(T, alpha) betaT = convert(T, beta) - _execute_short_circuit!(plan, alphaT, betaT) && return plan.Cstorage - _execute_across_barrier!(plan, alphaT, betaT, _select_path(plan)) + execute_short_circuit!(plan, alphaT, betaT) && return plan.Cstorage + execute_path!(plan, alphaT, betaT, plan.path) return plan.Cstorage end # The short-circuits every path shares; `true` when the call is finished. -@inline function _execute_short_circuit!(plan::ContractPlan{T}, alphaT::T, betaT::T) where {T} +@inline function execute_short_circuit!(plan::ContractPlan{T}, alphaT::T, betaT::T) where {T} m_length = axis_length(plan.mgroup) n_length = axis_length(plan.ngroup) (m_length == 0 || n_length == 0) && return true @@ -53,155 +53,17 @@ end return false end -# The path `execute!` runs past its short-circuits: the dot path, else the -# 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.blocking.k_block, 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. -@inline function _select_path( - ::Type{T}, unpack_ok::Bool, Astorage, Bstorage, Cstorage, - mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup, capacity, k_block, - split_a::Bool = false, split_b::Bool = false - ) where {T} - m_length = axis_length(mgroup) - n_length = axis_length(ngroup) - k_length = axis_length(kgroup) - panel = _c_panel_needed(T, Cstorage, k_length, k_block) - if (m_length == 1 || n_length == 1) && !panel && _dot_applicable(T, Astorage, Bstorage, kgroup, m_length, n_length, k_length) && - _dot_capacity_ok(capacity) - W = _dot_lanewidth(T) - return m_length == 1 ? _lane_path(_DotPath{true}, W) : _lane_path(_DotPath{false}, W) - end - k_length == 1 && _outer_applicable(T, Astorage, Cstorage, mgroup, m_length) && - return _lane_path(_OuterPath, _dot_lanewidth(T)) - return _nest_path( - 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. -@inline function _c_panel_needed(::Type{T}, Cstorage, k_length::Int, k_block::Int) where {T} - sizeof(real(eltype(Cstorage))) < sizeof(real(T)) || return false - return k_length > k_block -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} - core = _strip_storage(plan) - slot = barrier_slot!(plan.workspace.slots, Tuple{typeof(core), T, T}) - slot[] = (core, alphaT, betaT) - Base.inferencebarrier(_execute_resolved!)( - path, slot, plan.Astorage, plan.Bstorage, plan.Cstorage - ) - return nothing -end - -function _execute_resolved!( - path::P, slot::Base.RefValue{Tuple{R, T, T}}, Astorage::SA, Bstorage::SB, Cstorage::SC - ) where {P, R <: ContractPlan, T, SA, SB, SC} - core, alphaT, betaT = slot[] - _execute_path!(_with_storage(core, Astorage, Bstorage, Cstorage), alphaT, betaT, path) - return nothing -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.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.mpack, p.npack -) - -# The continuation the TensorOperations adapter hands `_planned`: `execute!` -# and `release!` inside the planning barrier. A struct, not a closure, so the -# fields are concretely typed. -struct _Execute{T, AL} - alpha::T - beta::T - allocator::AL -end - -# The path predicted before the kernel barrier, so its callee is specialised -# on the path as well as the kernel. -@inline _path_hint(::_Execute, req::_PlanRequest{T}, shape, ::Val{K}) where {T, K} = _select_path( - T, _unpacked_b_kernel_eligible(K), - req.Astorage, req.Bstorage, req.Cstorage, req.mgroup, req.ngroup, req.kgroup, nothing, - req.k_block === nothing ? _default_k_block(T, req.Cstorage, shape, K) : req.k_block -) - -# The kernel's default `k_block`, which decides the path only under a C -# narrower than `T`; elsewhere a static placeholder. -@inline _default_k_block(::Type{T}, Cstorage, shape, ::Type{K}) where {T, K} = - sizeof(real(eltype(Cstorage))) < sizeof(real(T)) ? - kernel_blocking(target_profile(), T, K, shape[1], shape[2]).k_block : 0 - -@inline function _continue(e::_Execute{T}, plan::ContractPlan{T}, hint) where {T} - if hint === nothing - execute!(plan, e.alpha, e.beta) - else - _execute_hinted!(plan, e.alpha, e.beta, hint) - end - release!(plan, e.allocator) - return nothing -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. 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_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, ::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 -# dot path's workspace capacity (assumed) and the panel decision (made at the -# default `k_block`; 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.k_block) === (path isa _PanelPath) -@inline _hint_holds(plan::ContractPlan, ::_DotPath) = _dot_capacity_ok(plan) -@inline _hint_holds(plan::ContractPlan, ::_OuterPath) = true - -_execute_path!(plan::ContractPlan, alphaT, betaT, ::_DotPath{MATB, W}) where {MATB, W} = +execute_path!(plan::ContractPlan, alphaT, betaT, ::DotPath{MATB, W}) where {MATB, W} = (_execute_dot!(plan, alphaT, betaT, MATB, Val(W)); nothing) -_execute_path!(plan::ContractPlan, alphaT, betaT, ::_OuterPath{W}) where {W} = ( +execute_path!(plan::ContractPlan, alphaT, betaT, ::OuterPath{W}) where {W} = ( _execute_outer!( plan, alphaT, betaT, axis_length(plan.mgroup), axis_length(plan.ngroup), Val(W) ); nothing ) -function _execute_path!( - plan::ContractPlan{T}, alphaT::T, betaT::T, path::_NestPath +function execute_path!( + plan::ContractPlan{T}, alphaT::T, betaT::T, path::NestPath ) where {T} kernel = plan.kernel ws = plan.workspace @@ -217,18 +79,17 @@ function _execute_path!( end # The nest over `pplan`, a copy of `plan` whose C is the workspace panel with -# dense M/N maps, one N 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} +# dense M/N maps, one N 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 m_length = axis_length(plan.mgroup) pplan = ContractPlan( - kernel, _dense_second_map(plan.mgroup, 1), _dense_second_map(plan.ngroup, m_length), - 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 + plan; mgroup = dense_second_map(plan.mgroup, 1), ngroup = dense_second_map(plan.ngroup, m_length), + Cstorage = ws.c_panel, Cbase = 0 ) GC.@preserve ws begin _execute_nest!( @@ -241,31 +102,26 @@ function _execute_path!( 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} +@inline function dense_second_map(g::AxisGroup{D, 2}, step::Int) where {D} dense = ntuple(d -> step * prod(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, n_block_start, n_block_length, betaT) = plan -@inline _panel_exit!(::Nothing, n_block_start, n_block_length) = nothing - -function _panel_enter!(target::ContractPlan, plan::ContractPlan, n_block_start::Int, n_block_length::Int, betaT) - iszero(betaT) || _panel_copy!(target, n_block_start, n_block_length, true) - return ContractPlan( - plan.kernel, plan.mgroup, plan.ngroup, plan.kgroup, plan.blocking, - plan.Astorage, plan.Abase, plan.Bstorage, plan.Bbase, plan.Cstorage, - -n_block_start * axis_length(plan.mgroup), plan.atransform, plan.btransform, plan.workspace, - plan.mpack, plan.npack - ) +@inline panel_enter!(::Nothing, plan, n_block_start, n_block_length, betaT) = plan +@inline panel_exit!(::Nothing, n_block_start, n_block_length) = nothing + +function panel_enter!(target::ContractPlan, plan::ContractPlan, n_block_start::Int, n_block_length::Int, betaT) + iszero(betaT) || panel_copy!(target, n_block_start, n_block_length, true) + return ContractPlan(plan; Cbase = -n_block_start * axis_length(plan.mgroup)) end -_panel_exit!(target::ContractPlan, n_block_start::Int, n_block_length::Int) = _panel_copy!(target, n_block_start, n_block_length, false) +panel_exit!(target::ContractPlan, n_block_start::Int, n_block_length::Int) = panel_copy!(target, n_block_start, n_block_length, false) # Columns `n_block_start .+ (0:n_block_length-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, n_block_start::Int, n_block_length::Int, load::Bool) +function panel_copy!(target::ContractPlan, n_block_start::Int, n_block_length::Int, load::Bool) ws = target.workspace panel = ws.c_panel C = target.Cstorage @@ -295,7 +151,7 @@ end function _execute_nest!( plan::ContractPlan{T}, ws, kernel::K, m_tile::Int, n_tile::Int, m_length::Int, n_length::Int, k_length::Int, m_block::Int, k_block::Int, n_block::Int, - alphaT::T, betaT::T, ::_NestPath{UNPACKED_B, AFF, SPLIT}, target + 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; # `m_tile`/`n_tile` count register-tile rows. They differ for complex @@ -327,7 +183,7 @@ function _execute_nest!( n_block_start = 0 while n_block_start < n_length n_block_length = min(n_block, n_length - n_block_start) - cplan = _panel_enter!(target, plan, n_block_start, n_block_length, betaT) + cplan = panel_enter!(target, plan, n_block_start, n_block_length, betaT) n_tiles = cld(n_block_length, n_tile) (rng_nB, rng_nC) = if n_ramp _ramp_slivers!( @@ -445,7 +301,7 @@ function _execute_nest!( k_block_start += k_block_length end - _panel_exit!(target, n_block_start, n_block_length) + panel_exit!(target, n_block_start, n_block_length) n_block_start += n_block_length end diff --git a/src/execution/outer.jl b/src/execution/outer.jl index 1e802fe..ea106f2 100644 --- a/src/execution/outer.jl +++ b/src/execution/outer.jl @@ -6,12 +6,8 @@ # complex K = 1 tiles are not the bottleneck, and the path would need its own # interleaved arithmetic. -# `:never` forces the path off, for benchmarks and tests. -const _OUTER_MODE = Ref{Symbol}(:auto) - # N and B may have any layout: B is read one scalar per column. function _outer_applicable(::Type{T}, Astorage, Cstorage, mgroup::AxisGroup, m_length::Int) where {T} - _OUTER_MODE[] === :never && return false T <: Real || return false (Astorage isa DenseVector{T} && Cstorage isa DenseVector{T}) || return false m_length >= _dot_lanewidth(T) || return false diff --git a/src/execution/paths.jl b/src/execution/paths.jl new file mode 100644 index 0000000..4258650 --- /dev/null +++ b/src/execution/paths.jl @@ -0,0 +1,68 @@ +# Execution paths: singletons that select one `execute_path!` method. +# `plan_contract` decides the path (`select_path`) and stores it in the plan, +# so `execute!` reaches the path's code statically. + +# 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 offset-buffer `view` 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. A split group is +# 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, panel::Bool = false + ) + split_b &= !unpacked_b + return nest_path_from_flags( + 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(table::Vector{Any}, flags::Vararg{Bool, 9}) + index = 1 + for i in 1:9 + index += flags[i] << (i - 1) + end + return @inbounds table[index] +end + +const _NEST_PATHS = let + 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 + +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 +# (`m_length == 1`). +struct DotPath{MATB, W} end + +# The outer-product path at lane width `W`. +struct OuterPath{W} end + +# A runtime lane width to its path singleton, one literal per arm. +@inline function lane_path(::Type{P}, W::Int) where {P} + W == 16 && return P{16}() + W == 8 && return P{8}() + W == 4 && return P{4}() + return P{2}() +end diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index 70004d7..e9224ab 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -37,21 +37,17 @@ end # The kernels whose K step reads B through `b_scalar`/`b_complex` # (not 1m, whose inner kernel walks the planar panel, nor the scalar one); by -# type, to predict the path before the kernel exists. +# type, to decide the path before the kernel exists. @inline _unpacked_b_kernel_eligible(::Type{<:Union{SIMDKernel, PlanarKernel, FMAddSubKernel, ComplexRealKernel}}) = true @inline _unpacked_b_kernel_eligible(::Type) = false -@inline _unpacked_b_kernel_eligible(kernel::Microkernel) = _unpacked_b_kernel_eligible(typeof(kernel)) - -# `:always`/`:never` override the rule below, for benchmarks and tests. -const _UNPACKED_B_MODE = Ref{Symbol}(:auto) # Whether B is read in place (given an eligible kernel): small M, and every B # column contiguous along K. Packing pays off only through reuse across M # slivers, while `n_tile` contiguous columns read in place cost the kernel # nothing. A large K stride in B (each step its own cache line, and for a -# power of two only a few L1 sets) is what packing exists for. -@inline function _unpacked_b_rule(mgroup::AxisGroup, kgroup::AxisGroup) - mode = _UNPACKED_B_MODE[] +# power of two only a few L1 sets) is what packing exists for. `mode`: +# `:always`/`:never` override the rule (`PathModes`). +@inline function _unpacked_b_rule(mode::Symbol, mgroup::AxisGroup, kgroup::AxisGroup) mode === :always && return true mode === :never && return false axis_length(mgroup) <= _UNPACKED_B_MMAX || return false diff --git a/src/execution/workspace.jl b/src/execution/workspace.jl index 1051332..44120fa 100644 --- a/src/execution/workspace.jl +++ b/src/execution/workspace.jl @@ -19,8 +19,8 @@ be [`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 every GC-tracked field. +# A `mutable struct` with `const` fields so a plan holds one pointer to it +# instead of copying every GC-tracked field. mutable struct ContractWorkspace{T, VT <: AbstractVector, PT <: AbstractVector} const m::GroupBuffers const n::GroupBuffers @@ -31,11 +31,9 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector, PT <: AbstractVector} const packed_a::VT const packed_b::VT # The compute-type panel of C for an eltype of C narrower than `T` (see - # `_PanelPath`); empty otherwise. + # `PanelPath`); empty otherwise. const c_panel::PT - const slots::SlotCache - # GUARDRAIL: the packed panels hold `real(T)`, never `T`. function ContractWorkspace{T}( m::GroupBuffers, n::GroupBuffers, k::NTuple{2, Vector{Int}}, @@ -47,20 +45,22 @@ mutable struct ContractWorkspace{T, VT <: AbstractVector, PT <: AbstractVector} "(the real type of the compute type $T), got $(eltype(VT))" ) ) - return new{T, VT, PT}(m, n, k, packed_a, packed_b, c_panel, SlotCache()) + return new{T, VT, PT}(m, n, k, packed_a, packed_b, c_panel) end end -# GUARDRAIL: `packed_a_length` already counts reals, the sliver counts are in -# logical `m_tile`/`n_tile`; do not rescale either. -@inline function workspace_sizes(kernel, blocking::Blocking) - m_tile, n_tile = tile_size(kernel) +# For a kernel of type `K` with an `m_tile x n_tile` tile, from its type so +# that planning can size the workspace before the kernel exists. GUARDRAIL: +# the packed lengths count reals, the sliver counts logical `m_tile`/`n_tile`; +# do not rescale either. +@inline function workspace_sizes(::Type{K}, m_tile::Int, n_tile::Int, blocking::Blocking) where {K} + a_reals, b_reals = map(reals_per_element, pack_formats(K)) m_tiles = cld(blocking.m_block, m_tile) n_tiles = cld(blocking.n_block, n_tile) return ( m_tiles = m_tiles, n_tiles = n_tiles, - packed_a = m_tiles * packed_a_length(kernel, blocking.k_block), - packed_b = n_tiles * packed_b_length(kernel, blocking.k_block), + packed_a = m_tiles * m_tile * a_reals * blocking.k_block, + packed_b = n_tiles * n_tile * b_reals * blocking.k_block, ) end @@ -78,7 +78,7 @@ function ContractWorkspace( ::Type{T}, kernel, blocking::Blocking; allocator = TO.DefaultAllocator(), panel::Int = 0 ) where {T} - s = workspace_sizes(kernel, blocking) + s = workspace_sizes(typeof(kernel), tile_size(kernel)..., blocking) R = realtype(kernel) # `release!` frees in exactly the reverse order, for arena allocators. packed_a = TO.tensoralloc(Vector{R}, (s.packed_a,), Val(true), allocator) diff --git a/src/integrations/tensoroperations.jl b/src/integrations/tensoroperations.jl index 3a00477..e69d3d7 100644 --- a/src/integrations/tensoroperations.jl +++ b/src/integrations/tensoroperations.jl @@ -58,7 +58,7 @@ end end @inline function _qs_prepare(C, A, pA, B, pB, pAB, α, β, accumulator) - T = _compute_type(eltype(A), eltype(B), eltype(C), 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) @@ -80,7 +80,22 @@ end return Cv, Av, Bv, indA, indB, indC, convert(T, α), convert(T, β) end -# `_planned` builds, runs and releases the plan behind the kernel dispatch +# The continuation handed to `planned`: `execute!` and `release!` on the +# concretely typed plan. A struct, not a closure, so the fields are concretely +# typed. +struct Execute{T, AL} + alpha::T + beta::T + allocator::AL +end + +function (e::Execute)(plan::ContractPlan) + execute!(plan, e.alpha, e.beta) + release!(plan, e.allocator) + return nothing +end + +# `planned` builds, runs and releases the plan behind the kernel dispatch # barrier, so the plan is never boxed (as `execute!(plan_contract(...), ...)` # would be). Bracketed with checkpoint/reset like TO's own `blas_contract!`. function TO.tensorcontract!( @@ -93,9 +108,9 @@ function TO.tensorcontract!( ) where {AC} Cv, Av, Bv, indA, indB, indC, α′, β′ = _qs_prepare(C, A, pA, B, pB, pAB, α, β, AC) checkpoint = TO.allocator_checkpoint!(allocator) - _planned( - _Execute(α′, β′, allocator), Cv, Av, indA, Bv, indB, indC, - nothing, conjA, conjB, nothing, nothing, nothing, allocator, AC + planned( + Execute(α′, β′, allocator), Cv, Av, indA, Bv, indB, indC, + nothing, conjA, conjB, nothing, nothing, nothing, allocator, AC, PathModes() ) TO.allocator_reset!(allocator, checkpoint) return C diff --git a/src/packing/line_packing.jl b/src/packing/line_packing.jl index a2fd1b0..e94505b 100644 --- a/src/packing/line_packing.jl +++ b/src/packing/line_packing.jl @@ -29,17 +29,17 @@ const _NO_SPLIT = PackSplit(0, 0, 0, false, 0, 0) is_split(s::PackSplit) = s.L != 0 -# The split for the free group `g` of the operand packed as `spec` (A for M, B -# for N), whose storage eltype takes `S` bytes, and the block extent it needs, -# or `(eff, _NO_SPLIT)`. `eff` is the plan's extent, `rounded` the requested -# one rounded to the tile; a split block may take the budget the blocking -# reserved for `k_block_requested` when the K extent clamps `k_block`. -# `l2bytes = nothing`: `split_capacity`, doubled for complex, whose block walk -# costs more per element. +# The split for the free group `g` of operand `side` (1: A for M, 2: B for N), +# packed in slivers of `lanes` elements of `format`, whose storage eltype takes +# `S` bytes, and the block extent it needs, or `(eff, _NO_SPLIT)`. `eff` is the +# plan's extent, `rounded` the requested one rounded to the tile; a split block +# may take the budget the blocking reserved for `k_block_requested` when the K +# extent clamps `k_block`. `l2bytes = nothing`: `split_capacity`, doubled for +# complex, whose block walk costs more per element. @inline function pack_split( - g::AxisGroup{D}, kg::AxisGroup, spec::SliverSpec{I, R, F}, S::Int, k_block::Int, + g::AxisGroup{D}, kg::AxisGroup, side::Int, lanes::Int, format::PackFormat, S::Int, k_block::Int, eff::Int, rounded::Int, k_block_requested::Int; l2bytes::Union{Int, Nothing} = nothing - ) where {D, I, R, F} + ) where {D} D < 256 || return (eff, _NO_SPLIT) line = line_bytes(target_profile()) d1 = 0 @@ -52,7 +52,7 @@ is_split(s::PackSplit) = s.L != 0 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, I, d1, dj, R, S, F !== RealFormat, line, k_block, eff, rounded, k_block_requested, l2bytes) + return pack_split_window(g, kg, side, d1, dj, lanes, S, !(format isa RealFormat), line, k_block, eff, rounded, k_block_requested, l2bytes) end # Out of line: shared by every eltype, and most plans never get here. diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 2b332c4..2ae4d62 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -1,23 +1,25 @@ # The plan: everything `execute!` needs, resolved once so it can be reused. # `plan_contract` runs the planning stages (labels, conjugation, kernel -# selection, blocking) and sizes the `ContractWorkspace`. +# selection, blocking, execution path) and sizes the `ContractWorkspace`. """ ContractPlan Reusable, concretely typed plan from [`plan_contract`](@ref): the M/N/K -`AxisGroup`s, kernel, operand storages, effective [`Blocking`](@ref), packing -transforms and the [`ContractWorkspace`](@ref) holding every buffer -[`execute!`](@ref) needs. Reusing a plan reuses its buffers. A plan, and so -its workspace, must be used by one task at a time. Field layout is not part -of the public interface; after an M/N orientation swap the `A*` fields -describe the original `B`. +`AxisGroup`s, kernel, execution path, operand storages, effective +[`Blocking`](@ref), packing transforms and the [`ContractWorkspace`](@ref) +holding every buffer [`execute!`](@ref) needs. Reusing a plan reuses its +buffers. A plan, and so its workspace, must be used by one task at a time. +Field layout is not part of the public interface; 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, + T, Kern, P, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, TA, TB, VT <: AbstractVector, PT <: AbstractVector, } kernel::Kern + # The singleton selecting the `execute_path!` method (src/execution/paths.jl). + path::P mgroup::GM ngroup::GN kgroup::GK @@ -42,6 +44,25 @@ struct ContractPlan{ npack::PackSplit end +# `p` with some fields replaced. +@inline ContractPlan( + p::ContractPlan; path = p.path, mgroup = p.mgroup, ngroup = p.ngroup, + Astorage = p.Astorage, Abase = p.Abase, Bstorage = p.Bstorage, Bbase = p.Bbase, + Cstorage = p.Cstorage, Cbase = p.Cbase, atransform = p.atransform, btransform = p.btransform, + mpack = p.mpack, npack = p.npack +) = ContractPlan( + p.kernel, path, mgroup, ngroup, p.kgroup, p.blocking, Astorage, Abase, Bstorage, Bbase, + Cstorage, Cbase, atransform, btransform, p.workspace, mpack, npack +) + +# Internal, for tests and benchmarks: `:never` rules out the dot, outer-product +# or unpacked-B path, `:always` forces unpacked B where the kernel admits it. +Base.@kwdef struct PathModes + dot::Symbol = :auto + outer::Symbol = :auto + unpacked_b::Symbol = :auto +end + """ plan_contract(C::StridedView, A::StridedView, indA::NTuple{NA,Int}, B::StridedView, indB::NTuple{NB,Int}, @@ -102,18 +123,19 @@ function plan_contract( k_block::Union{Int, Nothing} = nothing, n_block::Union{Int, Nothing} = nothing, allocator = TO.DefaultAllocator(), - accumulator::Union{Nothing, Type{Float32}, Type{Float64}} = nothing + accumulator::Union{Nothing, Type{Float32}, Type{Float64}} = nothing, + path_modes::PathModes = PathModes() ) where {NA, NB, NC} - return _planned( + return planned( identity, C, A, indA, B, indB, indC, - kernel, conjA, conjB, m_block, k_block, n_block, allocator, accumulator + kernel, conjA, conjB, m_block, k_block, n_block, allocator, accumulator, path_modes ) 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 compute type. -struct _PlanRequest{ +# Everything planning needs that is concretely typed before the kernel is +# known (`run` is the chosen M composite's unit-stride run in C). `T` is the +# phantom compute type. +struct PlanRequest{ T, F, GM <: AxisGroup, GN <: AxisGroup, GK <: AxisGroup, SA, SB, SC, AL, } f::F @@ -131,16 +153,17 @@ struct _PlanRequest{ k_block::Union{Int, Nothing} n_block::Union{Int, Nothing} allocator::AL + modes::PathModes end -@inline function _plan_request( +@inline function plan_request( ::Type{T}, f::F, mgroup::GM, ngroup::GN, kgroup::GK, Astorage::SA, Abase::Int, Bstorage::SB, Bbase::Int, Cstorage::SC, Cbase::Int, - run::Int, m_block, k_block, n_block, allocator::AL + run::Int, m_block, k_block, n_block, allocator::AL, modes::PathModes ) where {T, F, GM, GN, GK, SA, SB, SC, AL} - return _PlanRequest{T, F, GM, GN, GK, SA, SB, SC, AL}( + return PlanRequest{T, F, GM, GN, GK, SA, SB, SC, AL}( f, mgroup, ngroup, kgroup, Astorage, Abase, Bstorage, Bbase, Cstorage, Cbase, - run, m_block, k_block, n_block, allocator + run, m_block, k_block, n_block, allocator, modes ) end @@ -168,18 +191,18 @@ isconj(v::StridedView{T}, flag::Bool) where {T} = (T <: Complex) && (flag ⊻ op_conjugates(v.op)) # `plan_contract`'s body, positional, with a continuation `f` applied to the -# plan inside the barrier, where its type is concrete (the TensorOperations -# adapter passes an executor). -function _planned( +# plan behind the kernel barrier, where its type is concrete (the +# TensorOperations adapter passes an executor). +function planned( f::F, C::StridedView, A::StridedView, indA::NTuple{NA, Int}, B::StridedView, indB::NTuple{NB, Int}, indC::NTuple{NC, Int}, kernel, conjA::Bool, conjB::Bool, m_block::Union{Int, Nothing}, k_block::Union{Int, Nothing}, n_block::Union{Int, Nothing}, - allocator, accumulator::AC + allocator, accumulator::AC, modes::PathModes ) where {F, NA, NB, NC, AC} - T = _compute_type(eltype(A), eltype(B), eltype(C), accumulator) + T = compute_type(eltype(A), eltype(B), eltype(C), accumulator) K = default_kernel_type(T, eltype(A), eltype(B)) - _check_kernel_domain(kernel, 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`. @@ -212,57 +235,57 @@ function _planned( k_length = axis_length(kgroup) - m_tile_asis, m_tile_swapped = _candidate_m_tiles(T, K, kernel, m_length, n_length, run_m, run_n) + m_tile_asis, m_tile_swapped = candidate_m_tiles(T, K, kernel, m_length, n_length, run_m, run_n) if T <: Real && prefer_swap(run_m, run_n, m_tile_asis, m_tile_swapped) # B takes the M role: groups, K maps, storages, run and transforms move # together; the sum is unchanged. kgroup_swapped = AxisGroup(korder, (indB, B), (indA, A)) # maps: (B, A) - req_swapped = _plan_request( + req_swapped = plan_request( T, f, ngroup, mgroup, kgroup_swapped, parent(B), offset(B), parent(A), offset(A), parent(C), offset(C), - run_n, m_block, k_block, n_block, allocator + run_n, m_block, k_block, n_block, allocator, modes ) - return _plan_with_kernel(kernel, K, btransform, atransform, req_swapped) + return plan_with_kernel(kernel, K, btransform, atransform, req_swapped) end - req = _plan_request( + req = plan_request( T, f, mgroup, ngroup, kgroup, parent(A), offset(A), parent(B), offset(B), parent(C), offset(C), - run_m, m_block, k_block, n_block, allocator + run_m, m_block, k_block, n_block, allocator, modes ) - return _plan_with_kernel(kernel, K, atransform, btransform, req) + return plan_with_kernel(kernel, K, 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) +@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( +@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} +@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) + 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 +@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( +@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( +@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") ) @@ -270,107 +293,137 @@ end # named kernel either way, else `select_shape`'s pick at that orientation's # extent and C run. The swap is judged before `fit_to_run`'s K-dependent # divisor step, which an unbounded `k_length` skips. -@inline _candidate_m_tiles(::Type{T}, ::Type, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = +@inline candidate_m_tiles(::Type{T}, ::Type, kernel, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T} = (tile_size(kernel, 1), tile_size(kernel, 1)) -@inline function _candidate_m_tiles(::Type{T}, ::Type{K}, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T, K} +@inline function candidate_m_tiles(::Type{T}, ::Type{K}, ::Nothing, m_length::Int, n_length::Int, run_m::Int, run_n::Int) where {T, K} m_tile_asis = select_shape(T, K, m_length, typemax(Int), run_m)[1][1] m_tile_swapped = T <: Real ? select_shape(T, K, n_length, typemax(Int), run_n)[1][1] : m_tile_asis return m_tile_asis, m_tile_swapped end # Kernel resolution. A named kernel goes straight through, never demoted. An -# automatic one is chosen as a `(shape, kernel type)` value and the plan is -# built across a dispatch barrier specialised on that one concrete kernel type, -# so only the chosen 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 (the request -# in a `Ref`): one method-cache hit. -@inline _plan_with_kernel(kernel, ::Type, atransform, btransform, req::_PlanRequest) = - _plan_contract(kernel, atransform, btransform, req, nothing) -@inline function _plan_with_kernel(::Nothing, ::Type{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} +# automatic one is chosen as a `(shape, kernel type)` value. The blocking and +# execution path follow from the kernel type and tile, so they are resolved +# here, before the kernel exists; the plan is then built across a dispatch +# barrier specialised on the one concrete kernel and path, so only their 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 (the request in a `Ref`): one method-cache hit. +@inline function plan_with_kernel(kernel, ::Type, atransform, btransform, req::PlanRequest{T}) where {T} + scalartype(kernel) === T || + throw(ArgumentError("kernel scalar type $(scalartype(kernel)) does not match the compute type $T")) + return plan_across_barrier(kernel, typeof(kernel), tile_size(kernel)..., atransform, btransform, req) +end +@inline function plan_with_kernel(::Nothing, ::Type{K}, atransform, btransform, req::PlanRequest{T}) where {K, T} shape, kernel_type = select_shape(T, K, axis_length(req.mgroup), axis_length(req.kgroup), req.run) - return _plan_at_shape(shape, kernel_type, atransform, btransform, req) + return plan_at_shape(shape, kernel_type, atransform, btransform, req) end -@inline function _plan_at_shape(shape, ::Val{K}, atransform, btransform, req::_PlanRequest{T}) where {K, T} - vkernel = menu_val(shape, T, K) - # The execution path, predicted so the callee is specialised on it. From - # the shape and kernel type, not the kernel: a call union-split on the - # kernel type is emitted out of line and boxes `req`. - hint = _path_hint(req.f, req, shape, Val(K)) - return Base.inferencebarrier(_plan_resolved)( - vkernel, atransform, btransform, hint, Base.RefValue{typeof(req)}(req) +# Nothing here is called on the kernel `Val`: a call union-split on the kernel +# type is emitted out of line and boxes `req`. +@inline plan_at_shape(shape, ::Val{K}, atransform, btransform, req::PlanRequest{T}) where {K, T} = + plan_across_barrier(menu_val(shape, T, K), K, shape[1], shape[2], atransform, btransform, req) + +@inline function plan_across_barrier( + kernel, ::Type{K}, m_tile::Int, n_tile::Int, atransform, btransform, req::PlanRequest + ) where {K} + resolved = resolve_blocking(K, m_tile, n_tile, req) + path = select_path(K, m_tile, n_tile, req, resolved) + return Base.inferencebarrier(build_plan)( + kernel, path, atransform, btransform, Base.RefValue((req, resolved)) ) end -function _plan_resolved( - ::Val{Kern}, atransform::TA, btransform::TB, hint::H, slot::Base.RefValue{R} - ) where {Kern, TA, TB, H, R <: _PlanRequest} - return _plan_contract(Kern(), atransform, btransform, slot[], hint) -end - -# `nothing` for a continuation that does not execute (see src/execution/execute.jl). -@inline _path_hint(f, req::_PlanRequest, shape, kernel_type) = nothing - -# Plan construction on a concrete kernel and transform pair (a named kernel's -# transform Unions die here). `req.f` runs in here, on the concrete plan type. -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 the compute type $T")) - +# The effective blocking for a kernel of type `K` with an `m_tile x n_tile` +# tile, its line-by-line packing (`mpack`/`npack`) and whether partial sums +# live in a panel of C (`panel`). +function resolve_blocking(::Type{K}, m_tile::Int, n_tile::Int, req::PlanRequest{T}) where {K, T} m_length = axis_length(req.mgroup) n_length = axis_length(req.ngroup) k_length = axis_length(req.kgroup) - defaults = default_blocking(kernel) + defaults = kernel_blocking(target_profile(), T, K, m_tile, n_tile) requested = Blocking( - req.m_block === nothing ? defaults.m_block : req.m_block, - req.k_block === nothing ? defaults.k_block : req.k_block, - req.n_block === nothing ? defaults.n_block : req.n_block + something(req.m_block, defaults.m_block), + something(req.k_block, defaults.k_block), + something(req.n_block, defaults.n_block) ) - m_tile, n_tile = tile_size(kernel) - # Empty extents: the drivers never read these; the floors keep them valid. m_block_rounded = roundup(requested.m_block, m_tile) n_block_rounded = roundup(requested.n_block, n_tile) m_block = m_length == 0 ? m_tile : min(m_block_rounded, roundup(m_length, m_tile)) n_block = n_length == 0 ? n_tile : min(n_block_rounded, roundup(n_length, n_tile)) k_block = k_length == 0 ? 1 : min(requested.k_block, k_length) - panel = _c_panel_needed(T, req.Cstorage, k_length, k_block) + panel = c_panel_needed(T, req.Cstorage, k_length, k_block) 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 N blocks in C's own N order, which a split N group does # not enumerate contiguously. if m_length > 0 && n_length > 0 && k_length > 0 + a_format, b_format = pack_formats(K) m_block, mpack = pack_split( - req.mgroup, req.kgroup, sliver_spec(kernel, 1), sizeof(eltype(req.Astorage)), + req.mgroup, req.kgroup, 1, m_tile, a_format, sizeof(eltype(req.Astorage)), k_block, m_block, m_block_rounded, requested.k_block ) if !panel n_block, npack = pack_split( - req.ngroup, req.kgroup, sliver_spec(kernel, 2), sizeof(eltype(req.Bstorage)), + req.ngroup, req.kgroup, 2, n_tile, b_format, sizeof(eltype(req.Bstorage)), k_block, n_block, n_block_rounded, requested.k_block ) end end - blocking = Blocking(m_block, k_block, n_block) - ws = ContractWorkspace( - T, kernel, blocking; allocator = req.allocator, panel = panel ? m_length * min(n_block, n_length) : 0 - ) + return (blocking = Blocking(m_block, k_block, n_block), mpack, npack, 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. +@inline function c_panel_needed(::Type{T}, Cstorage, k_length::Int, k_block::Int) where {T} + sizeof(real(eltype(Cstorage))) < sizeof(real(T)) || return false + return k_length > k_block +end + +# The path `execute!` runs past its short-circuits: the dot path, else the +# outer-product path, else the nest (B packed or read in place). +function select_path(::Type{K}, m_tile::Int, n_tile::Int, req::PlanRequest{T}, resolved) where {K, T} + (; mgroup, ngroup, kgroup, modes) = req + (; blocking, mpack, npack, panel) = resolved + m_length = axis_length(mgroup) + n_length = axis_length(ngroup) + k_length = axis_length(kgroup) + if (m_length == 1 || n_length == 1) && !panel && modes.dot !== :never && + _dot_applicable(T, req.Astorage, req.Bstorage, kgroup, m_length, n_length, k_length) && + dot_fits(T, workspace_sizes(K, m_tile, n_tile, blocking).packed_a, blocking.k_block) + W = _dot_lanewidth(T) + return m_length == 1 ? lane_path(DotPath{true}, W) : lane_path(DotPath{false}, W) + end + k_length == 1 && modes.outer !== :never && _outer_applicable(T, req.Astorage, req.Cstorage, mgroup, m_length) && + return lane_path(OuterPath, _dot_lanewidth(T)) + unpacked_b = _unpacked_b_kernel_eligible(K) && _unpacked_b_rule(modes.unpacked_b, mgroup, kgroup) + return nest_path(unpacked_b, mgroup, ngroup, kgroup, is_split(mpack), is_split(npack), panel) +end +# Type parameters on the transforms force specialisation on them: the +# compiler does not specialise on a `Function` argument it only passes on. +build_plan(::Val{Kern}, path, atransform::TA, btransform::TB, slot::Base.RefValue) where {Kern, TA, TB} = + build_plan(Kern(), path, atransform, btransform, slot) + +# The plan on a concrete kernel, path and transform pair (a named kernel's +# transform Unions die here). `req.f` runs in here, on the concrete plan type. +function build_plan( + kernel::Microkernel, path::P, atransform::TA, btransform::TB, slot::Base.RefValue{R} + ) where {P, TA, TB, R} + T = scalartype(kernel) + req, (; blocking, mpack, npack, panel) = slot[] + panel_length = panel ? axis_length(req.mgroup) * min(blocking.n_block, axis_length(req.ngroup)) : 0 + ws = ContractWorkspace(T, kernel, blocking; allocator = req.allocator, panel = panel_length) plan = ContractPlan( - kernel, req.mgroup, req.ngroup, req.kgroup, blocking, + kernel, path, req.mgroup, req.ngroup, req.kgroup, blocking, req.Astorage, req.Abase, req.Bstorage, req.Bbase, req.Cstorage, req.Cbase, atransform, btransform, ws, mpack, npack ) - return _continue(req.f, plan, hint) + return req.f(plan) end -# `hint` is used only by an executing continuation (src/execution/execute.jl). -@inline _continue(f::F, plan::ContractPlan, hint) where {F} = f(plan) - release!(plan::ContractPlan, allocator) = release!(plan.workspace, allocator) diff --git a/test/execution/test_dot.jl b/test/execution/test_dot.jl index c3a0227..4615cab 100644 --- a/test/execution/test_dot.jl +++ b/test/execution/test_dot.jl @@ -2,7 +2,7 @@ # matrix operand runs a K-vectorized gemv. Results match the nest and the # reference to a tolerance (the K summation order differs). -_dot_takes(plan) = _path_of(plan) isa QuasiStrided._DotPath +_dot_takes(plan) = _path_of(plan) isa QuasiStrided.DotPath const _DOT_TYPES = (Float64, Float32, ComplexF64, ComplexF32) @@ -37,7 +37,7 @@ end # The dot path against the nest and the reference. function _dot_check(mk, alpha, beta; conjA = false, conjB = false, plankw...) C_dot, plan = _run_fresh(execute!, mk, alpha, beta; conjA, conjB, plankw...) - C_nest, _ = _run_fresh(_run_nest!, mk, alpha, beta; conjA, conjB, plankw...) + C_nest, _ = _run_nest(mk, alpha, beta; conjA, conjB, plankw...) @test _dot_takes(plan) @test C_dot ≈ C_nest return @test C_dot ≈ _ref_of(mk, alpha, beta; conjA, conjB) @@ -117,7 +117,7 @@ end execute!(plan, 1.0, 0.0) @test (@allocated execute!(plan, 1.0, 0.0)) == 0 skip = (VERSION < v"1.11") end - # The backend runs the path predicted before planning (`_path_hint`). + # The backend plans and runs in one call. d = 6 A, B, C = randn(T, d, d), randn(T, d, d, d, d, d), zeros(T, d, d, d) pA, pB, pAB = TO.contract_indices((:a, :b), (:a, :b, :c, :d, :e), (:c, :d, :e)) diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index 57ab405..b8b177e 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -43,20 +43,12 @@ function _run_fresh(run!, mk, alpha, beta; plankw...) return Array(Cv), plan end -# `execute!` forced onto the nest with B packed, through the mode switches -# the benchmarks use: the baseline for the dedicated paths. -const _PATH_MODES = (QuasiStrided._DOT_MODE, QuasiStrided._OUTER_MODE, QuasiStrided._UNPACKED_B_MODE) -function _run_nest!(plan, alpha, beta) - old = map(getindex, _PATH_MODES) - foreach(m -> m[] = :never, _PATH_MODES) - try - return execute!(plan, alpha, beta) - finally - foreach(setindex!, _PATH_MODES, old) - end -end +# `_run_fresh` planned onto the nest with B packed, through the path modes the +# benchmarks use: the baseline for the dedicated paths. +const _NEST_ONLY = QuasiStrided.PathModes(dot = :never, outer = :never, unpacked_b = :never) +_run_nest(mk, alpha, beta; plankw...) = _run_fresh(execute!, mk, alpha, beta; path_modes = _NEST_ONLY, plankw...) -_path_of(plan) = QuasiStrided._select_path(plan) +_path_of(plan) = plan.path # A DenseMatrix that is not a DenseVector: a plan keeps it as storage, so paths # that need raw-pointer loads must decline it. @@ -182,11 +174,7 @@ end plans = ( p, _mm_plan(Cmat, Amat, Bmat; kernel = kernel, m_block = 16, k_block = 5, n_block = 8), - 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, p.mpack, p.npack, - ), + ContractPlan(p; atransform = conj, btransform = conj), ) @test parent(StridedView(Cmat)) isa DenseVector{Float64} for plan in plans @@ -253,28 +241,26 @@ function _sp_forced_plan(Cv, Av, iA, Bv, iB, iC; kw...) k, k_block, d, T = p.kernel, p.blocking.k_block, default_blocking(p.kernel), eltype(Cv) function split(g, i, eff, req) R = tile_size(k, i) - return pack_split(g, p.kgroup, sliver_spec(k, i), sizeof(T), k_block, eff, cld(req, R) * R, d.k_block; l2bytes = 0) + return pack_split(g, p.kgroup, i, R, pack_formats(typeof(k))[i], sizeof(T), k_block, eff, cld(req, R) * R, d.k_block; l2bytes = 0) end (m_block, ms) = split(p.mgroup, 1, p.blocking.m_block, something(get(kw, :m_block, nothing), d.m_block)) (n_block, ns) = split(p.ngroup, 2, p.blocking.n_block, d.n_block) q = plan_contract(Cv, Av, iA, Bv, iB, iC; kw..., kernel = k, m_block, n_block) @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 - ) + path = QuasiStrided.nest_path(false, q.mgroup, q.ngroup, q.kgroup, is_split(ms), is_split(ns)) + return ContractPlan(q; path, mpack = ms, npack = 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(l2bytes) = is_split(pack_split(p.mgroup, p.kgroup, sliver_spec(k, 1), 8, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block; l2bytes)[2]) + splits(l2bytes) = is_split(pack_split(p.mgroup, p.kgroup, 1, tile_size(k, 1), pack_formats(typeof(k))[1], 8, b.k_block, b.m_block, b.m_block, default_blocking(k).k_block; l2bytes)[2]) @test splits(2^20) && !splits(2^24) # a 4 MB reuse window host = splits(QuasiStrided.split_capacity(target_profile(), true)) @test 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)} + 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) @@ -288,7 +274,7 @@ end 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, m_block, kernel) - @test _path_of(plan) isa _NestPath{false, <:Any, (true, ind === _SP_BOTH)} + @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 diff --git a/test/execution/test_mixed.jl b/test/execution/test_mixed.jl index 62299ca..0679546 100644 --- a/test/execution/test_mixed.jl +++ b/test/execution/test_mixed.jl @@ -25,16 +25,16 @@ 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), + (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 @@ -84,10 +84,10 @@ end allocator = case === :manual_n_block_16 ? TO.ManualAllocator() : TO.DefaultAllocator() n_block = case === :manual_n_block_16 ? 16 : nothing plan = plan_contract(Cv, Av, (1, 2), Bv, (2, 3), (1, 3); accumulator = acc, k_block = 4, n_block, allocator) - @test _path_of(plan) isa QuasiStrided._PanelPath + @test _path_of(plan) isa QuasiStrided.PanelPath @test axis_length(plan.mgroup) == M case === :manual_n_block_16 && @test N > plan.blocking.n_block && !(plan.workspace.c_panel isa Vector) - case === :packed_b && @test _path_of(plan) isa QuasiStrided._PanelPath{<:QuasiStrided._NestPath{false}} + 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) @@ -101,6 +101,6 @@ end rng = MersenneTwister(3) Amat, Bmat, Cmat = randn(rng, TA, 64, K), randn(rng, TB, K, 40), zeros(TC, 64, 40) plan = _mm_plan(Cmat, Amat, Bmat; accumulator = acc, k_block = 256) - @test _path_of(plan) isa (TC === Float32 ? QuasiStrided._PanelPath : QuasiStrided._NestPath) + @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/execution/test_outer.jl b/test/execution/test_outer.jl index b1fb921..076d35a 100644 --- a/test/execution/test_outer.jl +++ b/test/execution/test_outer.jl @@ -2,7 +2,7 @@ # in A and C streams `alpha * A * B[n]` column by column. Results match the # nest and the reference (`≈`: the two differ on a signed zero). -_outer_takes(plan) = _path_of(plan) isa QuasiStrided._OuterPath +_outer_takes(plan) = _path_of(plan) isa QuasiStrided.OuterPath # `C[a, b] = A[a] * B[b]` at extents `(M, N)`: # :bstrided B every other element of a longer vector @@ -55,7 +55,7 @@ end (alpha, beta) in ((1.0, 0.0), (2.5, -0.75), (1.0, 1.0)) mk = _outer_maker(T, M, N, 10 + idx; variant) C_out, plan = _run_fresh(execute!, mk, alpha, beta) - C_nest, _ = _run_fresh(_run_nest!, mk, alpha, beta) + C_nest, _ = _run_nest(mk, alpha, beta) @test _outer_takes(plan) == (axis_length(plan.mgroup) >= W) @test C_out ≈ C_nest @test C_out ≈ _ref_of(mk, alpha, beta) @@ -96,19 +96,13 @@ end @test C ≈ a * transpose(b) end -@testset "the mode switches override path selection" begin - d, o, u = QuasiStrided._DOT_MODE, QuasiStrided._OUTER_MODE, QuasiStrided._UNPACKED_B_MODE - pd = plan_contract(_mm_maker(Float64, 1, 64, 9, 1)()...) - po = plan_contract(_outer_maker(Float64, 16, 9, 1)()...) - pu = plan_contract(_mm_maker(Float64, 20, 12, 9, 1)()...) - ps = plan_contract(_mm_maker(Float64, 20, 12, 9, 1; B = :transposed)()...) - @test _dot_takes(pd) && _outer_takes(po) && _ub_takes(pu) && !_ub_takes(ps) - try - d[] = o[] = u[] = :never - @test !_dot_takes(pd) && !_outer_takes(po) && !_ub_takes(pu) - u[] = :always - @test _ub_takes(ps) - finally - d[] = o[] = u[] = :auto - end +@testset "the path modes override path selection" begin + pd(; kw...) = plan_contract(_mm_maker(Float64, 1, 64, 9, 1)()...; kw...) + po(; kw...) = plan_contract(_outer_maker(Float64, 16, 9, 1)()...; kw...) + pu(; kw...) = plan_contract(_mm_maker(Float64, 20, 12, 9, 1)()...; kw...) + ps(; kw...) = plan_contract(_mm_maker(Float64, 20, 12, 9, 1; B = :transposed)()...; kw...) + @test _dot_takes(pd()) && _outer_takes(po()) && _ub_takes(pu()) && !_ub_takes(ps()) + @test !_dot_takes(pd(; path_modes = _NEST_ONLY)) && !_outer_takes(po(; path_modes = _NEST_ONLY)) && + !_ub_takes(pu(; path_modes = _NEST_ONLY)) + @test _ub_takes(ps(; path_modes = QuasiStrided.PathModes(unpacked_b = :always))) end diff --git a/test/execution/test_unpackedb.jl b/test/execution/test_unpackedb.jl index ec9c9df..3ac2f86 100644 --- a/test/execution/test_unpackedb.jl +++ b/test/execution/test_unpackedb.jl @@ -3,7 +3,7 @@ # (same arithmetic, same values, different address), and the packed-B buffer # must stay untouched. -_ub_takes(plan) = _path_of(plan) isa QuasiStrided._NestPath{true} +_ub_takes(plan) = _path_of(plan) isa QuasiStrided.NestPath{true} # Poison the packed-B buffer, then run. function _ub_run!(plan, alpha, beta) @@ -16,7 +16,7 @@ _ub_untouched(plan) = all(isnan, plan.workspace.packed_b) # and the reference. function _ub_check(mk, alpha, beta; conjA = false, conjB = false, plankw...) C_unp, plan = _run_fresh(_ub_run!, mk, alpha, beta; conjA, conjB, plankw...) - C_packed, _ = _run_fresh(_run_nest!, mk, alpha, beta; conjA, conjB, plankw...) + C_packed, _ = _run_nest(mk, alpha, beta; conjA, conjB, plankw...) @test _ub_takes(plan) @test _ub_untouched(plan) @test C_unp == C_packed @@ -98,7 +98,7 @@ end allocs = _steady_allocs!(execute!, plan, Cmat) @test allocs == 0 skip = (VERSION < v"1.11") @test Cmat ≈ Amat * Bmat - # The backend runs the path predicted before planning (`_path_hint`). + # The backend plans and runs in one call. fill!(Cmat, 0) TO.tensorcontract!(Cmat, Amat, ((1,), (2,)), false, Bmat, ((1,), (2,)), false, ((1, 2), ()), 1.0, 0.0, QuasiStridedBackend()) @test Cmat ≈ Amat * Bmat diff --git a/test/microkernels/test_mixed_kernel.jl b/test/microkernels/test_mixed_kernel.jl index ac78c9f..e985f3e 100644 --- a/test/microkernels/test_mixed_kernel.jl +++ b/test/microkernels/test_mixed_kernel.jl @@ -88,21 +88,16 @@ end 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, k_block = 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 + 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, k_block = 16, conjA = cA, conjB = cB, accumulator = acc, + path_modes = QuasiStrided.PathModes(unpacked_b = mode) + ) + 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 end end diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index 5287545..bb20778 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -210,11 +210,7 @@ Base.IndexStyle(::Type{<:CountingStorage}) = IndexLinear() cstore = CountingStorage(zeros(Ma * Na)) astore = CountingStorage(copy(vec(Amat))) bstore = CountingStorage(copy(vec(Bmat))) - 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.mpack, base.npack, - ) + p = _pcf_Plan(base; Astorage = astore, Abase = 0, Bstorage = bstore, Bbase = 0, Cstorage = cstore, Cbase = 0) astore.n = 0; bstore.n = 0; cstore.n = 0 _pcf_exec(p, 1.0, 0.0) @test reshape(cstore.data, Ma, Na) ≈ Amat * Bmat @@ -241,11 +237,7 @@ end Cstorage = vec(Cmat), Astorage = vec(Amat), Bstorage = vec(Bmat), Cbase = 0, Abase = 0, Bbase = 0 ) - 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.mpack, base.npack, - ) + return _pcf_Plan(base; Astorage, Abase, Bstorage, Bbase, Cstorage, Cbase) end # One element short: rejected before anything is written. @@ -302,11 +294,7 @@ end # One element short: only the address at `binding` overflows. short_C = zeros(M * N1 * N2 - 1) - 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.mpack, base.npack, - ) + pshort = _pcf_Plan(base; Cstorage = short_C) @test_throws BoundsError _pcf_exec(pshort, 1.0, 0.0) @test all(iszero, short_C) # nothing written before the throw @@ -325,12 +313,7 @@ end @test QS.checked_span_bounds(base.Cbase, mrange, lastonly, shortlen) === nothing # An interior-sliver minimum: the base shifted down by one. - plow = _pcf_Plan( - 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.mpack, base.npack, - ) + plow = _pcf_Plan(base; Cstorage = zeros(M * N1 * N2), Cbase = base.Cbase - 1) @test_throws BoundsError _pcf_exec(plow, 1.0, 0.0) end diff --git a/test/runtests.jl b/test/runtests.jl index 85c20ce..15f15ea 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,7 +14,7 @@ using QuasiStrided: AxisGroup, axis_length, fill_offsets!, BlockDescriptor, rule_applies, isa_nregisters, sliver_width, realtype, PlanarKernel, OneMKernel, FMAddSubKernel, accumulator_planes, pack_formats, reals_per_element, modelled_blocking, kernel_blocking, - planar_pressure, pack_split, is_split, _NestPath + planar_pressure, pack_split, is_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. From 8a0cfbaff5bdcab0cb133e146b4fd90f6618e532 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 3 Oct 2026 18:02:01 -0400 Subject: [PATCH 31/60] Decision log: chunk 13 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 53 +++++++++++++++++++++++++++++------------------------ 1 file changed, 29 insertions(+), 24 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 41c2c53..18112c6 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -25,7 +25,7 @@ TTFX when codegen is touched). | 10 | Kernel selection | `planning/kernel_selection.jl` | done | | 11 | Blocking | `blocking.jl`, line packing (now `packing/line_packing.jl`) | done | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | done | -| 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | | +| 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | done | | 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | | 15 | Alternative paths | `oracle.jl`, `unpackedb.jl` | | | 16 | Degenerate paths | `dot.jl`, `outer.jl` | | @@ -66,14 +66,25 @@ where the underscore meant "internal helper of `foo`" (known clashes: LinearAlgebra, SIMD or the package's own names. Explicit imports instead of `using LinearAlgebra`, so later shadowing is an error (applied up front). -### D3. Dispatch barriers (chunks 12–14) - -Keep barrier #1 (kernel choice). Remove barrier #2 (execution path) by -deciding the path at plan time and making it a `ContractPlan` type parameter; -this deletes `_path_hint`/`_continue`/`_execute_hinted!`/`_hint_holds`/ -`_splits`/`_split_path`. To verify first: the path is fully fixed at plan -time, including dot-path workspace capacity under a reused grow-only -workspace. Measured: one crossing ~13 ns; a 2×2 `tensorcontract!` ~320 ns. +### D3. Dispatch barriers (chunk 13, applied: 308c07b) + +One barrier, at plan time. Blocking (including line packing and the C panel) +and the execution path are resolved before the kernel barrier, from the kernel +type and shape; the barrier crosses once with `Val{Kern}()` and the path, and +the path is a `ContractPlan` field/type parameter, so `execute!` dispatches on +it statically. Deletes the hint family, `_strip_storage`/`_with_storage` and +the workspace's `SlotCache`. The TensorOperations route keeps its continuation +(one crossing per call); a reused plan has none. Accepted: a named-kernel +`plan_contract` is no longer inferable. The dot/outer/unpacked-B mode switches +become an internal `plan_contract` keyword instead of global `Ref`s. Also: a +keyword copy method for `ContractPlan`, blocking resolution as its own +function, underscores dropped, `barrier.jl` → `paths.jl`. Tests stay where they +are until chunk 18. As applied: named kernels cross as their (singleton) +instance; `pack_split` takes the pack format instead of a `SliverSpec` (whose +`MR` is runtime before the barrier); the dot-path fit is checked at plan time +from `workspace_sizes(K, …)`. Results and chosen paths bitwise identical to +before on 172 cases; per-call floor, TTFX and nest specialisations unchanged +within noise. ### D4. Precompile workload (later, separate) @@ -229,11 +240,8 @@ blocking.jl (uncommitted until the full suite passes). ## Resume here -1. Chunk 13: `planning/plan.jl` (+ `test_plan_contract.jl`, - `test_per_call_overhead.jl`), including D3 (execution path as a plan type - parameter, removing barrier #2 and the hint machinery; kernel types cross - the barrier as `Val`, `::Type{K}` everywhere after it) and the open items - for chunks 13/14 below. +1. Chunk 14: five-loop nest (`execution/macrokernel.jl`, `execute.jl`), + including the open items marked chunk 14 below. 2. Workflow: present each chunk (purpose, reading order, design decisions, proposed fixes, questions), then hand the agreed changes to an Opus agent with the standard checks (Runic, full suite, per-call floor/allocations, @@ -270,15 +278,12 @@ blocking.jl (uncommitted until the full suite passes). - Internal constants still use `_UPPER` names; drop the underscores in one pass at the end (D2). -- Detected `l1d.line` is unused: `_K_LINE_BYTES = 64` in `labels.jl` (K-order - cost model) and `pack_split.jl`; thread the detected line size through - planning (chunks 9/11). -- 16-argument `ContractPlan(...)` spelled out five times, `_PlanRequest` - twice more: one "copy with fields replaced" helper (chunk 13/14). -- `execute.jl` mixes public API, barrier/hint machinery, the C-panel path and - the nest; `_c_panel_needed` lives there but is used by `plan.jl` (chunk 14). -- `_select_path` has 12 positional arguments, two optional (chunk 14). -- `oracle = true` default makes every user plan allocate test-path buffers; - consider defaulting to `false` (chunk 13/15). +- `execute.jl` mixes public API, the C-panel path and the nest (chunk 14). +- Plans with empty extents or K = 0 now compile their path although the + short-circuit always fires first; `select_path` could return a trivial + path for them (chunk 14). +- `_execute_dot!` takes `matB` as a runtime `Bool` although `DotPath{MATB}` + makes it static (chunk 16). +- `plan_request` is spelled out twice in `planned` (swapped and not). - `workspace.jl`/`barrier.jl` live in `execution/` but are included in the planning section (chunk 12). From 5c08bdf838c0c22bba15f2ec4cda0e275fdb76a4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 00:06:48 -0400 Subject: [PATCH 32/60] Chunk 14: nest cleanup - select_path returns EmptyPath (M or N empty: nothing to do) and ScalePath (K = 0: C *= beta) before the dot/outer/nest choice; execute! keeps only the runtime alpha == 0 short-circuit. Such plans no longer compile a nest. - nest!(plan, alphaT, betaT, path, n_range) replaces the 15-argument _execute_nest!: kernel, workspace, tile sizes, lengths and blocking come from the plan. - One pack_block!(plan, path, Val(side), ...) per operand: line by line when the path splits it, else sliver by sliver; same @inbounds coverage under the hoisted span checks. - Ramp decisions come from the path's AFF flags (M: both M flags, K: both K flags, N: both N flags, except that N's B flag is false under unpacked B, so there the B map is still tested at runtime). The flags stay. - The C panel path loops over N blocks itself: load the block into the panel (beta != 0), run the nest over that N range on a plan whose C is the panel holding just that block, copy back. No target plan, no panel hooks in the nest. - execute.jl (execute!, contract!, execute_path! dispatch, scaling C), nest.jl (the nest and the former macrokernel.jl helpers), c_panel.jl. - Underscores dropped on the nest helpers; firstpanel -> first_k_block; the nest returns nothing. Results: bitwise identical to the parent commit on 408 cases (the 172-case script plus K = 0, empty M/N/K+N, alpha = 0 and panel cases spanning 3-5 N blocks, each with beta = 0 and beta != 0); paths identical except Empty/ScalePath for the 80 degenerate cases. Full suite passes (32155). Per-call floor, execute! on reused plans, medians of 3 interleaved runs (noisy workstation), before -> after, all 0 allocations: 64^3 Float64 6862 -> 6387 ns; 2x2 Float64 85 -> 80 ns; 32^3 ComplexF64 5239 -> 4937 ns; dot 1x256x64 1722 -> 1556 ns; panel 32x600x32 F32/acc F64 29810 -> 29860 ns; tensorcontract! 2x2 Float64 523 -> 461 ns, ComplexF64 642 -> 640 ns. TTFX (fresh process, medians of 3): f64 3.25 -> 3.31 s, c64 1.79 -> 1.84 s, @tensor after those 1.14 -> 1.15 s, @tensor first 2.39 -> 2.31 s. Specialisations on the equivalence workload: nest 46 -> 42 and packed-B micro-tile loops 23 -> 19 (the degenerate plans' nests are gone; 42/19 is also the count without the degenerate cases, before and after); execute_path! 64 -> 76, build_plan 136 -> 160, scale_all_of_C! 64 -> 76, since Empty/ScalePath plans are now their own plan types. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 3 +- src/execution/c_panel.jl | 66 +++++ src/execution/execute.jl | 331 +++--------------------- src/execution/macrokernel.jl | 116 --------- src/execution/nest.jl | 301 +++++++++++++++++++++ src/execution/paths.jl | 6 + src/execution/unpackedb.jl | 10 +- src/planning/plan.jl | 7 +- test/execution/test_execute.jl | 8 +- test/execution/test_manual_pipeline.jl | 2 +- test/planning/test_per_call_overhead.jl | 14 +- 11 files changed, 433 insertions(+), 431 deletions(-) create mode 100644 src/execution/c_panel.jl delete mode 100644 src/execution/macrokernel.jl create mode 100644 src/execution/nest.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 5810af9..277e989 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -42,8 +42,9 @@ include("planning/plan.jl") # --- Execution: the five-loop nest and the specialised paths --- include("execution/paths.jl") -include("execution/macrokernel.jl") +include("execution/nest.jl") include("execution/execute.jl") +include("execution/c_panel.jl") include("execution/unpackedb.jl") include("execution/dot.jl") include("execution/outer.jl") diff --git a/src/execution/c_panel.jl b/src/execution/c_panel.jl new file mode 100644 index 0000000..aba926f --- /dev/null +++ b/src/execution/c_panel.jl @@ -0,0 +1,66 @@ +# The panel of C: for an eltype of C narrower than the compute type, partial +# sums across K blocks live in a compute-type panel in the workspace, one N +# block at a time. + +# Per N block: load the block of C into the panel (unless `beta == 0`), run the +# nest on a plan whose C is the panel holding just that block, with dense M/N +# maps, and round the block back into C. +function execute_path!( + plan::ContractPlan{T}, alphaT::T, betaT::T, ::PanelPath{P} + ) where {T, P} + ws = plan.workspace + m_length = axis_length(plan.mgroup) + n_length = axis_length(plan.ngroup) + n_block = plan.blocking.n_block + panel_plan = ContractPlan( + plan; mgroup = dense_second_map(plan.mgroup, 1), ngroup = dense_second_map(plan.ngroup, m_length), + Cstorage = ws.c_panel + ) + n_block_start = 0 + while n_block_start < n_length + n_block_length = min(n_block, n_length - n_block_start) + n_range = n_block_start:(n_block_start + n_block_length - 1) + iszero(betaT) || panel_copy!(plan, n_range, true) + block_plan = ContractPlan(panel_plan; Cbase = -n_block_start * m_length) + GC.@preserve ws nest!(block_plan, alphaT, betaT, P(), n_range) + panel_copy!(plan, n_range, false) + n_block_start += n_block_length + 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 * prod(ntuple(i -> i < d ? g.lengths[i] : 1, Val(D))), Val(D)) + return AxisGroup(g.lengths, (g.strides[1], dense)) +end + +# Columns `n_range` of `plan`'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!(plan::ContractPlan, n_range::UnitRange{Int}, load::Bool) + ws = plan.workspace + panel = ws.c_panel + C = plan.Cstorage + m_length = axis_length(plan.mgroup) + m_block = plan.blocking.m_block + n_bufC = ws.n.offsets[2] + fill_offsets!(ws.n.offsets, plan.ngroup, first(n_range), length(n_range)) + m_block_start = 0 + while m_block_start < m_length + m_block_length = min(m_block, m_length - m_block_start) + m_bufC = ws.m.offsets[2] + fill_offsets!(ws.m.offsets, plan.mgroup, m_block_start, m_block_length) + for j in 1:length(n_range), i in 1:m_block_length + c = plan.Cbase + m_bufC[i] + n_bufC[j] + 1 + q = m_block_start + i + (j - 1) * m_length + if load + panel[q] = C[c] + else + C[c] = panel[q] + end + end + m_block_start += m_block_length + end + return nothing +end diff --git a/src/execution/execute.jl b/src/execution/execute.jl index c1ab5a8..20db1ec 100644 --- a/src/execution/execute.jl +++ b/src/execution/execute.jl @@ -1,28 +1,3 @@ -# `C *= beta` tile by tile, never reading A or B: the -# `k_length == 0 || alpha == 0` pass of `execute!`. Borrows the first tile of -# the M/N offset buffers. -function _scale_all_of_C!(plan, betaT::T, m_tile::Int, n_tile::Int, m_length::Int, n_length::Int) where {T} - ws = plan.workspace - m_bufs = ws.m.offsets - n_bufs = ws.n.offsets - m_tile_start = 0 - while m_tile_start < m_length - m_tile_length = min(m_tile, m_length - m_tile_start) - (_, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) - n_tile_start = 0 - while n_tile_start < n_length - n_tile_length = min(n_tile, n_length - n_tile_start) - (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) - rowsC = _axis_of(dM_C, m_bufs[2], 0) - colsC = _axis_of(dN_C, n_bufs[2], 0) - GC.@preserve ws _scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) - n_tile_start += n_tile_length - end - m_tile_start += m_tile_length - end - return nothing -end - """ execute!(plan::ContractPlan, alpha::Number, beta::Number) @@ -36,22 +11,17 @@ block. Allocation-free. function execute!(plan::ContractPlan{T}, alpha::Number, beta::Number) where {T} alphaT = convert(T, alpha) betaT = convert(T, beta) - execute_short_circuit!(plan, alphaT, betaT) && return plan.Cstorage - execute_path!(plan, alphaT, betaT, plan.path) + if iszero(alphaT) + scale_all_of_C!(plan, betaT) + else + execute_path!(plan, alphaT, betaT, plan.path) + end return plan.Cstorage end -# The short-circuits every path shares; `true` when the call is finished. -@inline function execute_short_circuit!(plan::ContractPlan{T}, alphaT::T, betaT::T) where {T} - m_length = axis_length(plan.mgroup) - n_length = axis_length(plan.ngroup) - (m_length == 0 || n_length == 0) && return true - if axis_length(plan.kgroup) == 0 || iszero(alphaT) - _scale_all_of_C!(plan, betaT, tile_size(plan.kernel)..., m_length, n_length) - return true - end - return false -end +execute_path!(plan::ContractPlan, alphaT, betaT, ::EmptyPath) = nothing + +execute_path!(plan::ContractPlan, alphaT, betaT, ::ScalePath) = scale_all_of_C!(plan, betaT) execute_path!(plan::ContractPlan, alphaT, betaT, ::DotPath{MATB, W}) where {MATB, W} = (_execute_dot!(plan, alphaT, betaT, MATB, Val(W)); nothing) @@ -62,275 +32,44 @@ execute_path!(plan::ContractPlan, alphaT, betaT, ::OuterPath{W}) where {W} = ( ); nothing ) -function execute_path!( - plan::ContractPlan{T}, alphaT::T, betaT::T, path::NestPath - ) where {T} - kernel = plan.kernel +function execute_path!(plan::ContractPlan{T}, alphaT::T, betaT::T, path::NestPath) where {T} ws = plan.workspace - # Packed panels and scatter axes borrow pointers into `ws`. - GC.@preserve ws begin - _execute_nest!( - plan, ws, kernel, tile_size(kernel)..., - axis_length(plan.mgroup), axis_length(plan.ngroup), axis_length(plan.kgroup), - plan.blocking.m_block, plan.blocking.k_block, plan.blocking.n_block, alphaT, betaT, path, nothing - ) - end + GC.@preserve ws nest!(plan, alphaT, betaT, path, 0:(axis_length(plan.ngroup) - 1)) return nothing end -# The nest over `pplan`, a copy of `plan` whose C is the workspace panel with -# dense M/N maps, one N 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 +# `C *= beta` tile by tile, never reading A or B. Borrows the first tile of the +# M/N offset buffers. +function scale_all_of_C!(plan::ContractPlan{T}, betaT::T) where {T} ws = plan.workspace + m_tile, n_tile = tile_size(plan.kernel) m_length = axis_length(plan.mgroup) - pplan = ContractPlan( - plan; mgroup = dense_second_map(plan.mgroup, 1), ngroup = dense_second_map(plan.ngroup, m_length), - Cstorage = ws.c_panel, Cbase = 0 - ) - GC.@preserve ws begin - _execute_nest!( - pplan, ws, kernel, tile_size(kernel)..., - m_length, axis_length(plan.ngroup), axis_length(plan.kgroup), - plan.blocking.m_block, plan.blocking.k_block, plan.blocking.n_block, 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 * prod(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, n_block_start, n_block_length, betaT) = plan -@inline panel_exit!(::Nothing, n_block_start, n_block_length) = nothing - -function panel_enter!(target::ContractPlan, plan::ContractPlan, n_block_start::Int, n_block_length::Int, betaT) - iszero(betaT) || panel_copy!(target, n_block_start, n_block_length, true) - return ContractPlan(plan; Cbase = -n_block_start * axis_length(plan.mgroup)) -end - -panel_exit!(target::ContractPlan, n_block_start::Int, n_block_length::Int) = panel_copy!(target, n_block_start, n_block_length, false) - -# Columns `n_block_start .+ (0:n_block_length-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, n_block_start::Int, n_block_length::Int, load::Bool) - ws = target.workspace - panel = ws.c_panel - C = target.Cstorage - m_length = axis_length(target.mgroup) - m_block = target.blocking.m_block - n_bufC = ws.n.offsets[2] - fill_offsets!(ws.n.offsets, target.ngroup, n_block_start, n_block_length) - m_block_start = 0 - while m_block_start < m_length - m_block_length = min(m_block, m_length - m_block_start) - m_bufC = ws.m.offsets[2] - fill_offsets!(ws.m.offsets, target.mgroup, m_block_start, m_block_length) - for j in 1:n_block_length, i in 1:m_block_length - c = target.Cbase + m_bufC[i] + n_bufC[j] + 1 - q = m_block_start + i + (j - 1) * m_length - if load - panel[q] = C[c] - else - C[c] = panel[q] - end + n_length = axis_length(plan.ngroup) + m_bufs = ws.m.offsets + n_bufs = ws.n.offsets + m_tile_start = 0 + while m_tile_start < m_length + m_tile_length = min(m_tile, m_length - m_tile_start) + (_, dM_C) = block_descriptors!(m_bufs, plan.mgroup, m_tile_start, m_tile_length) + n_tile_start = 0 + while n_tile_start < n_length + n_tile_length = min(n_tile, n_length - n_tile_start) + (_, dN_C) = block_descriptors!(n_bufs, plan.ngroup, n_tile_start, n_tile_length) + rowsC = axis_of(dM_C, m_bufs[2], 0) + colsC = axis_of(dN_C, n_bufs[2], 0) + GC.@preserve ws scale_micro_tile!(plan.Cstorage, plan.Cbase, rowsC, colsC, betaT) + n_tile_start += n_tile_length end - m_block_start += m_block_length + m_tile_start += m_tile_length end return nothing end -function _execute_nest!( - plan::ContractPlan{T}, ws, kernel::K, m_tile::Int, n_tile::Int, - m_length::Int, n_length::Int, k_length::Int, m_block::Int, k_block::Int, n_block::Int, - 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; - # `m_tile`/`n_tile` count register-tile rows. They differ for complex - # kernels. - a_sliver_width, b_sliver_width = sliver_width(kernel) - - atransform = plan.atransform - btransform = plan.btransform - - 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) - lenB = length(plan.Bstorage) - lenC = length(plan.Cstorage) - - aff_mA, aff_mC, aff_nB, aff_nC, aff_kA, aff_kB = map(Val, AFF) - - # --- loop over N blocks --- - n_block_start = 0 - while n_block_start < n_length - n_block_length = min(n_block, n_length - n_block_start) - cplan = panel_enter!(target, plan, n_block_start, n_block_length, betaT) - n_tiles = cld(n_block_length, n_tile) - (rng_nB, rng_nC) = if n_ramp - _ramp_slivers!( - ws.n, n_step[1], n_step[2], n_block_start, n_block_length, n_tile, n_tiles - ) - else - fill_offsets!(ws.n.offsets, ngroup, n_block_start, n_block_length) - _classify_slivers!(ws.n, n_block_length, n_tile, n_tiles) - end - - # --- loop over K blocks --- - k_block_start = 0 - firstpanel = true - while k_block_start < k_length - k_block_length = min(k_block, k_length - k_block_start) - dK_A, dK_B, rng_kA, rng_kB = if k_ramp - ( - _ramp_descriptor(k_step[1], k_block_start, k_block_length), - _ramp_descriptor(k_step[2], k_block_start, k_block_length), - _ramp_offset_range(k_step[1], k_block_start, k_block_length), - _ramp_offset_range(k_step[2], k_block_start, k_block_length), - ) - else - fill_offsets!(ws.k, plan.kgroup, k_block_start, k_block_length) - dA = describe_block(ws.k[1], 0, k_block_length) - dB = describe_block(ws.k[2], 0, k_block_length) - ( - dA, dB, - descriptor_offset_range(dA, ws.k[1], 0), - descriptor_offset_range(dB, ws.k[2], 0), - ) - end - - colsA_k = _axis_of(dK_A, ws.k[1], 0, aff_kA) - rowsB_k = _axis_of(dK_B, ws.k[2], 0, aff_kB) - - # Hoisted bounds checks (B here, A and C per M block): each - # rectangle is exactly the union of the per-sliver/per-tile - # regions, and the storage check compares only range extremes, so - # the `@inbounds` pack and tile calls below touch nothing unchecked. - checked_span_bounds(plan.Bbase, rng_kB, rng_nB, lenB) - - beta_eff = firstpanel ? betaT : one(T) - - if split_b - pack_block_by_lines!( - packed_panel(ws.packed_b, 1, b_sliver_width * k_block_length * n_tiles), sliver_spec(kernel, 2), - plan.Bstorage, plan.Bbase, ws.n.offsets[1], rowsB_k, btransform, n_block_length, k_block_length, - plan.npack - ) - elseif !UNPACKED_B - for n_tile_index in 0:(n_tiles - 1) - n_tile_start = n_tile_index * n_tile - bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - colsB = _axis_of(ws.n.descriptors[1][n_tile_index + 1], ws.n.offsets[1], n_tile_start, aff_nB) - @inbounds _pack_sliver!( - bpanel, plan.Bstorage, plan.Bbase, colsB, rowsB_k, - sliver_spec(kernel, 2), btransform - ) - end - end - - # --- loop over M blocks --- - m_block_start = 0 - while m_block_start < m_length - m_block_length = min(m_block, m_length - m_block_start) - m_tiles = cld(m_block_length, m_tile) - (rng_mA, rng_mC) = if m_ramp - _ramp_slivers!( - ws.m, m_step[1], m_step[2], m_block_start, m_block_length, m_tile, m_tiles - ) - else - fill_offsets!(ws.m.offsets, mgroup, m_block_start, m_block_length) - _classify_slivers!(ws.m, m_block_length, m_tile, m_tiles) - end - - checked_span_bounds(plan.Abase, rng_mA, rng_kA, lenA) - checked_span_bounds(cplan.Cbase, rng_mC, rng_nC, lenC) - - if split_a - pack_block_by_lines!( - packed_panel(ws.packed_a, 1, a_sliver_width * k_block_length * m_tiles), sliver_spec(kernel, 1), - plan.Astorage, plan.Abase, ws.m.offsets[1], colsA_k, atransform, m_block_length, k_block_length, - plan.mpack - ) - else - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsA = _axis_of(ws.m.descriptors[1][m_tile_index + 1], ws.m.offsets[1], m_tile_start, aff_mA) - @inbounds _pack_sliver!( - apanel, plan.Astorage, plan.Abase, rowsA, colsA_k, - sliver_spec(kernel, 1), atransform - ) - end - end - - # --- loops over N tiles and M tiles --- - if UNPACKED_B - _micro_tiles_unpacked_b!( - kernel, cplan, ws, rowsB_k, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC - ) - else - _micro_tiles_packed_b!( - kernel, cplan, ws, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, b_sliver_width, k_block_length, alphaT, beta_eff, aff_mC, aff_nC - ) - end - - m_block_start += m_block_length - end - - firstpanel = false - k_block_start += k_block_length - end - - panel_exit!(target, n_block_start, n_block_length) - n_block_start += n_block_length - end - - return plan.Cstorage -end - -# The tile loops over one (N, K, M) block. `@noinline`: one call per block, and -# the inlined microkernel is most of a nest's size; compile cost grows faster -# than linearly with function size. -@noinline function _micro_tiles_packed_b!( - kernel::K, plan::ContractPlan, ws, m_tiles::Int, n_tiles::Int, - m_tile::Int, n_tile::Int, a_sliver_width::Int, b_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, - aff_mC::Val{MC}, aff_nC::Val{NC} - ) where {K, MC, NC} - for n_tile_index in 0:(n_tiles - 1) - n_tile_start = n_tile_index * n_tile - bpanel = _sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) - colsC = _axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsC = _axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) - # Inside the caller's C check. - @inbounds _execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bpanel, k_block_length, alphaT, beta_eff - ) - end - end +@inline function scale_micro_tile!( + storage::S, base::Int, rows::R, cols::C, beta + ) where {S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} + destination = Tile(storage, base, rows, cols) + scale_tile!(destination, beta) return nothing end diff --git a/src/execution/macrokernel.jl b/src/execution/macrokernel.jl deleted file mode 100644 index dd0536a..0000000 --- a/src/execution/macrokernel.jl +++ /dev/null @@ -1,116 +0,0 @@ -# Macro-kernel helpers for the five-loop nest: function barriers over the -# tile axis types, packed-sliver addressing, and sliver classification. - -# The axis `d` describes. GUARDRAIL: a `Union{AffineAxis, ScatterAxis}` of -# `isbits` types, never boxed; each consumer below binds each axis type as its -# own parameter, so it builds a concretely typed `Tile`. -@inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int) = - d.regular ? AffineAxis(d.base, d.stride, d.count) : ScatterAxis(pointer(buffer, first + 1), d.count) - -# The same with the axis type fixed by the path: `Val(true)` for a ramp map, -# whose descriptors are always regular (checked). -@inline _axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{false}) = - _axis_of(d, buffer, first) -@inline function _axis_of(d::BlockDescriptor, ::Vector{Int}, ::Int, ::Val{true}) - d.regular || _throw_irregular_ramp_descriptor() - return AffineAxis(d.base, d.stride, d.count) -end -@noinline _throw_irregular_ramp_descriptor() = - throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) - -# B passes its axes swapped (its `transpose`). GUARDRAIL: `transform` needs its -# own bound type parameter, or it costs a dynamic dispatch per call. -Base.@propagate_inbounds function _pack_sliver!( - packed::PK, storage::S, base::Int, - lanes::R, steps::C, spec, transform::TF - ) where {PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} - pack!(packed, Tile(storage, base, lanes, steps), spec, transform) - return nothing -end - -Base.@propagate_inbounds function _execute_micro_tile!( - kernel, storage::S, base::Int, rows::R, cols::C, - packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta - ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} - destination = Tile(storage, base, rows, cols) - execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) - return nothing -end - -@inline function _scale_micro_tile!( - storage::S, base::Int, rows::R, cols::C, beta - ) where {S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} - destination = Tile(storage, base, rows, cols) - scale_tile!(destination, beta) - return nothing -end - -# Sliver `tile_index` of a packed panel at the current block's depth -# `k_block_length`, shared by the packing and the consuming step so the two -# cannot disagree. GUARDRAIL: `width` counts reals per K step -# (`sliver_width`), not the tile size; they differ for complex kernels. -@inline function _sliver_panel(buffer, width::Int, k_block_length::Int, tile_index::Int) - stride = width * k_block_length - return packed_panel(buffer, tile_index * stride + 1, stride) -end - -# Classify each register sliver of a just-filled macro block, and return the -# two maps' offset ranges over the whole block (the slivers partition it) for -# the hoisted bounds checks. -@inline function _classify_slivers!(g::GroupBuffers, block_length::Int, tile::Int, tile_count::Int) - desc1, desc2 = g.descriptors - buf1, buf2 = g.offsets - lo1 = typemax(Int); hi1 = typemin(Int) - lo2 = typemax(Int); hi2 = typemin(Int) - for tile_index in 0:(tile_count - 1) - tile_start = tile_index * tile - tile_length = min(tile, block_length - tile_start) - d1 = describe_block(buf1, tile_start, tile_length) - d2 = describe_block(buf2, tile_start, tile_length) - desc1[tile_index + 1] = d1 - desc2[tile_index + 1] = d2 - (l1, h1) = descriptor_offset_range(d1, buf1, tile_start) - if h1 >= l1 - lo1 = min(lo1, l1); hi1 = max(hi1, h1) - end - (l2, h2) = descriptor_offset_range(d2, buf2, tile_start) - if h2 >= l2 - lo2 = min(lo2, l2); hi2 = max(hi2, h2) - end - end - # `hi < lo` means empty to `checked_span_bounds`. - return ((lo1, hi1), (lo2, hi2)) -end - -# Closed-form block descriptors for an affine-ramp composite (offset of `q` -# is `q * step`), with no offset buffer or scan. `==` to what `describe_block` -# gives for the materialized interval, including its `stride == 0` for a -# count-1 block. -@inline _ramp_descriptor(step::Int, first::Int, count::Int) = - count == 0 ? BlockDescriptor(0, 0, 0, true) : - count == 1 ? BlockDescriptor(first * step, 0, 1, true) : - BlockDescriptor(first * step, step, count, true) - -# Offset range of `[first, first+count)`; `(0, -1)` if empty. Cannot overflow: -# `AxisGroup` validated every in-domain offset at construction. -@inline _ramp_offset_range(step::Int, first::Int, count::Int) = - count == 0 ? (0, -1) : minmax(first * step, (first + count - 1) * step) - -# `_classify_slivers!`'s closed-form twin. -@inline function _ramp_slivers!( - g::GroupBuffers, step1::Int, step2::Int, first::Int, - block_length::Int, tile::Int, tile_count::Int - ) - desc1, desc2 = g.descriptors - for tile_index in 0:(tile_count - 1) - tile_start = tile_index * tile - tile_length = min(tile, block_length - tile_start) - q0 = first + tile_start - desc1[tile_index + 1] = _ramp_descriptor(step1, q0, tile_length) - desc2[tile_index + 1] = _ramp_descriptor(step2, q0, tile_length) - end - return ( - _ramp_offset_range(step1, first, block_length), - _ramp_offset_range(step2, first, block_length), - ) -end diff --git a/src/execution/nest.jl b/src/execution/nest.jl new file mode 100644 index 0000000..7209391 --- /dev/null +++ b/src/execution/nest.jl @@ -0,0 +1,301 @@ +# The five-loop nest and its helpers: function barriers over the tile axis +# types, block packing, packed-sliver addressing and sliver classification. + +# The nest over the N coordinates `n_range`, in N blocks from its start. The +# caller preserves `plan.workspace`: packed panels and scatter axes borrow +# pointers into it. +function nest!( + plan::ContractPlan{T}, alphaT::T, betaT::T, path::NestPath{UNPACKED_B, AFF, SPLIT}, + n_range::UnitRange{Int} + ) where {T, UNPACKED_B, AFF, SPLIT} + kernel = plan.kernel + ws = plan.workspace + m_tile, n_tile = tile_size(kernel) + # GUARDRAIL: reals per sliver per K step address the packed panels; + # `m_tile`/`n_tile` count register-tile rows. They differ for complex + # kernels. + a_sliver_width, b_sliver_width = sliver_width(kernel) + (; m_block, k_block, n_block) = plan.blocking + m_length = axis_length(plan.mgroup) + k_length = axis_length(plan.kgroup) + + 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. A group is a ramp when both its maps are, + # which the path's flags record, except N's B map under unpacked B. + aff_mA, aff_mC, aff_nB, aff_nC, aff_kA, aff_kB = AFF + m_ramp = aff_mA && aff_mC + n_ramp = (UNPACKED_B ? is_ramp_map(ngroup, 1) : aff_nB) && aff_nC + k_ramp = aff_kA && aff_kB + m_step = m_ramp ? affine_ramp(mgroup)[2] : (0, 0) + n_step = n_ramp ? affine_ramp(ngroup)[2] : (0, 0) + k_step = k_ramp ? affine_ramp(plan.kgroup)[2] : (0, 0) + + lenA = length(plan.Astorage) + lenB = length(plan.Bstorage) + lenC = length(plan.Cstorage) + + # --- loop over N blocks --- + n_block_start = first(n_range) + while n_block_start <= last(n_range) + n_block_length = min(n_block, last(n_range) + 1 - n_block_start) + n_tiles = cld(n_block_length, n_tile) + (rng_nB, rng_nC) = if n_ramp + ramp_slivers!( + ws.n, n_step[1], n_step[2], n_block_start, n_block_length, n_tile, n_tiles + ) + else + fill_offsets!(ws.n.offsets, ngroup, n_block_start, n_block_length) + classify_slivers!(ws.n, n_block_length, n_tile, n_tiles) + end + + # --- loop over K blocks --- + k_block_start = 0 + first_k_block = true + while k_block_start < k_length + k_block_length = min(k_block, k_length - k_block_start) + dK_A, dK_B, rng_kA, rng_kB = if k_ramp + ( + ramp_descriptor(k_step[1], k_block_start, k_block_length), + ramp_descriptor(k_step[2], k_block_start, k_block_length), + ramp_offset_range(k_step[1], k_block_start, k_block_length), + ramp_offset_range(k_step[2], k_block_start, k_block_length), + ) + else + fill_offsets!(ws.k, plan.kgroup, k_block_start, k_block_length) + dA = describe_block(ws.k[1], 0, k_block_length) + dB = describe_block(ws.k[2], 0, k_block_length) + ( + dA, dB, + descriptor_offset_range(dA, ws.k[1], 0), + descriptor_offset_range(dB, ws.k[2], 0), + ) + end + + colsA_k = axis_of(dK_A, ws.k[1], 0, Val(aff_kA)) + rowsB_k = axis_of(dK_B, ws.k[2], 0, Val(aff_kB)) + + # Hoisted bounds checks (B here, A and C per M block): each + # rectangle is exactly the union of the per-sliver/per-tile + # regions, and the storage check compares only range extremes, so + # the `@inbounds` pack and tile calls below touch nothing unchecked. + checked_span_bounds(plan.Bbase, rng_kB, rng_nB, lenB) + + beta_eff = first_k_block ? betaT : one(T) + + UNPACKED_B || pack_block!(plan, path, Val(2), rowsB_k, n_block_length, k_block_length, n_tiles) + + # --- loop over M blocks --- + m_block_start = 0 + while m_block_start < m_length + m_block_length = min(m_block, m_length - m_block_start) + m_tiles = cld(m_block_length, m_tile) + (rng_mA, rng_mC) = if m_ramp + ramp_slivers!( + ws.m, m_step[1], m_step[2], m_block_start, m_block_length, m_tile, m_tiles + ) + else + fill_offsets!(ws.m.offsets, mgroup, m_block_start, m_block_length) + classify_slivers!(ws.m, m_block_length, m_tile, m_tiles) + end + + checked_span_bounds(plan.Abase, rng_mA, rng_kA, lenA) + checked_span_bounds(plan.Cbase, rng_mC, rng_nC, lenC) + + pack_block!(plan, path, Val(1), colsA_k, m_block_length, k_block_length, m_tiles) + + # --- loops over N tiles and M tiles --- + if UNPACKED_B + _micro_tiles_unpacked_b!( + kernel, plan, ws, rowsB_k, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, Val(aff_mC), Val(aff_nC) + ) + else + micro_tiles_packed_b!( + kernel, plan, ws, m_tiles, n_tiles, + m_tile, n_tile, a_sliver_width, b_sliver_width, k_block_length, alphaT, beta_eff, + Val(aff_mC), Val(aff_nC) + ) + end + + m_block_start += m_block_length + end + + first_k_block = false + k_block_start += k_block_length + end + + n_block_start += n_block_length + end + return nothing +end + +# Pack one macro block of operand `I` (1: A over M, 2: B over N): line by line +# when the path splits it, else sliver by sliver. Inside the caller's bounds +# checks. +@inline function pack_block!( + plan::ContractPlan, ::NestPath{UNPACKED_B, AFF, SPLIT}, side::Val{I}, kaxis::KA, + block_length::Int, k_block_length::Int, tiles::Int + ) where {UNPACKED_B, AFF, SPLIT, I, KA <: AbstractVector{Int}} + kernel = plan.kernel + spec = sliver_spec(kernel, I) + width = sliver_width(kernel, I) + tile = tile_size(kernel, I) + storage, base, transform, buffer, g, split = packed_operand(plan, side) + if SPLIT[I] + pack_block_by_lines!( + packed_panel(buffer, 1, width * k_block_length * tiles), spec, + storage, base, g.offsets[1], kaxis, transform, block_length, k_block_length, split + ) + else + for tile_index in 0:(tiles - 1) + tile_start = tile_index * tile + panel = sliver_panel(buffer, width, k_block_length, tile_index) + lanes = axis_of(g.descriptors[1][tile_index + 1], g.offsets[1], tile_start, Val(AFF[2I - 1])) + @inbounds pack_sliver!(panel, storage, base, lanes, kaxis, spec, transform) + end + end + return nothing +end + +# Operand `I`'s storage, base and transform, packed buffer, free-group buffers +# and split. +@inline packed_operand(plan::ContractPlan, ::Val{1}) = + (plan.Astorage, plan.Abase, plan.atransform, plan.workspace.packed_a, plan.workspace.m, plan.mpack) +@inline packed_operand(plan::ContractPlan, ::Val{2}) = + (plan.Bstorage, plan.Bbase, plan.btransform, plan.workspace.packed_b, plan.workspace.n, plan.npack) + +# The tile loops over one (N, K, M) block. `@noinline`: one call per block, and +# the inlined microkernel is most of a nest's size; compile cost grows faster +# than linearly with function size. +@noinline function micro_tiles_packed_b!( + kernel::K, plan::ContractPlan, ws, m_tiles::Int, n_tiles::Int, + m_tile::Int, n_tile::Int, a_sliver_width::Int, b_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, + aff_mC::Val{MC}, aff_nC::Val{NC} + ) where {K, MC, NC} + for n_tile_index in 0:(n_tiles - 1) + n_tile_start = n_tile_index * n_tile + bpanel = sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) + colsC = axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) + for m_tile_index in 0:(m_tiles - 1) + m_tile_start = m_tile_index * m_tile + apanel = sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + rowsC = axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) + # Inside the caller's C check. + @inbounds execute_micro_tile!( + kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, + apanel, bpanel, k_block_length, alphaT, beta_eff + ) + end + end + return nothing +end + +# The axis `d` describes. GUARDRAIL: a `Union{AffineAxis, ScatterAxis}` of +# `isbits` types, never boxed; each consumer below binds each axis type as its +# own parameter, so it builds a concretely typed `Tile`. +@inline axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int) = + d.regular ? AffineAxis(d.base, d.stride, d.count) : ScatterAxis(pointer(buffer, first + 1), d.count) + +# The same with the axis type fixed by the path: `Val(true)` for a ramp map, +# whose descriptors are always regular (checked). +@inline axis_of(d::BlockDescriptor, buffer::Vector{Int}, first::Int, ::Val{false}) = + axis_of(d, buffer, first) +@inline function axis_of(d::BlockDescriptor, ::Vector{Int}, ::Int, ::Val{true}) + d.regular || throw_irregular_ramp_descriptor() + return AffineAxis(d.base, d.stride, d.count) +end +@noinline throw_irregular_ramp_descriptor() = + throw(AssertionError("an affine-ramp map produced an irregular block descriptor")) + +# B passes its axes swapped (its `transpose`). GUARDRAIL: `transform` needs its +# own bound type parameter, or it costs a dynamic dispatch per call. +Base.@propagate_inbounds function pack_sliver!( + packed::PK, storage::S, base::Int, + lanes::R, steps::C, spec, transform::TF + ) where {PK, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}, TF} + pack!(packed, Tile(storage, base, lanes, steps), spec, transform) + return nothing +end + +Base.@propagate_inbounds function execute_micro_tile!( + kernel, storage::S, base::Int, rows::R, cols::C, + packed_a::PA, packed_b::PB, k_block_length::Int, alpha, beta + ) where {PA, PB, S, R <: AbstractVector{Int}, C <: AbstractVector{Int}} + destination = Tile(storage, base, rows, cols) + execute_tile!(kernel, destination, packed_a, packed_b, k_block_length, alpha, beta) + return nothing +end + +# Sliver `tile_index` of a packed panel at the current block's depth +# `k_block_length`, shared by the packing and the consuming step so the two +# cannot disagree. GUARDRAIL: `width` counts reals per K step +# (`sliver_width`), not the tile size; they differ for complex kernels. +@inline function sliver_panel(buffer, width::Int, k_block_length::Int, tile_index::Int) + stride = width * k_block_length + return packed_panel(buffer, tile_index * stride + 1, stride) +end + +# Classify each register sliver of a just-filled macro block, and return the +# two maps' offset ranges over the whole block (the slivers partition it) for +# the hoisted bounds checks. +@inline function classify_slivers!(g::GroupBuffers, block_length::Int, tile::Int, tile_count::Int) + desc1, desc2 = g.descriptors + buf1, buf2 = g.offsets + lo1 = typemax(Int); hi1 = typemin(Int) + lo2 = typemax(Int); hi2 = typemin(Int) + for tile_index in 0:(tile_count - 1) + tile_start = tile_index * tile + tile_length = min(tile, block_length - tile_start) + d1 = describe_block(buf1, tile_start, tile_length) + d2 = describe_block(buf2, tile_start, tile_length) + desc1[tile_index + 1] = d1 + desc2[tile_index + 1] = d2 + (l1, h1) = descriptor_offset_range(d1, buf1, tile_start) + if h1 >= l1 + lo1 = min(lo1, l1); hi1 = max(hi1, h1) + end + (l2, h2) = descriptor_offset_range(d2, buf2, tile_start) + if h2 >= l2 + lo2 = min(lo2, l2); hi2 = max(hi2, h2) + end + end + # `hi < lo` means empty to `checked_span_bounds`. + return ((lo1, hi1), (lo2, hi2)) +end + +# Closed-form block descriptors for an affine-ramp composite (offset of `q` +# is `q * step`), with no offset buffer or scan. `==` to what `describe_block` +# gives for the materialized interval, including its `stride == 0` for a +# count-1 block. +@inline ramp_descriptor(step::Int, first::Int, count::Int) = + count == 0 ? BlockDescriptor(0, 0, 0, true) : + count == 1 ? BlockDescriptor(first * step, 0, 1, true) : + BlockDescriptor(first * step, step, count, true) + +# Offset range of `[first, first+count)`; `(0, -1)` if empty. Cannot overflow: +# `AxisGroup` validated every in-domain offset at construction. +@inline ramp_offset_range(step::Int, first::Int, count::Int) = + count == 0 ? (0, -1) : minmax(first * step, (first + count - 1) * step) + +# `classify_slivers!`'s closed-form twin. +@inline function ramp_slivers!( + g::GroupBuffers, step1::Int, step2::Int, first::Int, + block_length::Int, tile::Int, tile_count::Int + ) + desc1, desc2 = g.descriptors + for tile_index in 0:(tile_count - 1) + tile_start = tile_index * tile + tile_length = min(tile, block_length - tile_start) + q0 = first + tile_start + desc1[tile_index + 1] = ramp_descriptor(step1, q0, tile_length) + desc2[tile_index + 1] = ramp_descriptor(step2, q0, tile_length) + end + return ( + ramp_offset_range(step1, first, block_length), + ramp_offset_range(step2, first, block_length), + ) +end diff --git a/src/execution/paths.jl b/src/execution/paths.jl index 4258650..81f9fb7 100644 --- a/src/execution/paths.jl +++ b/src/execution/paths.jl @@ -2,6 +2,12 @@ # `plan_contract` decides the path (`select_path`) and stores it in the plan, # so `execute!` reaches the path's code statically. +# Nothing to compute: M or N is empty. +struct EmptyPath end + +# `C *= beta`: K is empty. +struct ScalePath 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 diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl index e9224ab..b0fc938 100644 --- a/src/execution/unpackedb.jl +++ b/src/execution/unpackedb.jl @@ -73,7 +73,7 @@ const _UNPACKED_B_MMAX = 256 return UnpackedBView(storage, colbase, ksteps, sliver_width(kernel, 2), transform) end -# The tile loops of `_execute_nest!` with B read in place. A barrier over the K +# The tile loops of `nest!` with B read in place. A barrier over the K # axis type, so each view is concretely typed. `@noinline` so that the # packed-B nest does not grow by this body. The view's addresses are within # the caller's B check. @@ -90,13 +90,13 @@ end bview = _unpacked_b_view( kernel, Bstorage, Bbase, ws.n.descriptors[1][n_tile_index + 1], ws.n.offsets[1], n_tile_start, rowsB_k, btransform ) - colsC = _axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) + colsC = axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile - apanel = _sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsC = _axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) + apanel = sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) + rowsC = axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) # Inside the caller's C check. - @inbounds _execute_micro_tile!( + @inbounds execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, apanel, bview, k_block_length, alphaT, beta_eff ) diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 2ae4d62..1020b8f 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -384,14 +384,17 @@ end return k_length > k_block end -# The path `execute!` runs past its short-circuits: the dot path, else the -# outer-product path, else the nest (B packed or read in place). +# The path `execute!` runs for a nonzero `alpha`: none for an empty C, scaling C +# for an empty K, else the dot path, the outer-product path or the nest (B +# packed or read in place). function select_path(::Type{K}, m_tile::Int, n_tile::Int, req::PlanRequest{T}, resolved) where {K, T} (; mgroup, ngroup, kgroup, modes) = req (; blocking, mpack, npack, panel) = resolved m_length = axis_length(mgroup) n_length = axis_length(ngroup) k_length = axis_length(kgroup) + (m_length == 0 || n_length == 0) && return EmptyPath() + k_length == 0 && return ScalePath() if (m_length == 1 || n_length == 1) && !panel && modes.dot !== :never && _dot_applicable(T, req.Astorage, req.Bstorage, kgroup, m_length, n_length, k_length) && dot_fits(T, workspace_sizes(K, m_tile, n_tile, blocking).packed_a, blocking.k_block) diff --git a/test/execution/test_execute.jl b/test/execution/test_execute.jl index b8b177e..9db710a 100644 --- a/test/execution/test_execute.jl +++ b/test/execution/test_execute.jl @@ -122,16 +122,18 @@ end for (Ka, alpha) in ((0, 1.0), (Ka, 0.0)), beta in (0.5, 0.0) # A and B are poisoned: neither may be read. mk = _mm_maker(Float64, Ma, Ka, Na, 1; ABfill = NaN) - C, _ = _run_fresh(execute!, mk, alpha, beta) + C, plan = _run_fresh(execute!, mk, alpha, beta) @test C == beta .* Array(mk()[1]) + Ka == 0 && @test _path_of(plan) isa QuasiStrided.ScalePath end # beta = 0 never reads a NaN C, across several blocks. C, _ = _run_fresh(execute!, _mm_maker(Float64, 11, 10, 9, 2; Cfill = NaN), 2.5, 0.0; m_block = 4, k_block = 4, n_block = 4) @test C ≈ _ref_of(_mm_maker(Float64, 11, 10, 9, 2), 2.5, 0.0) # Empty M or N: a no-op. for (M, N) in ((3, 0), (0, 3)) - C, _ = _run_fresh(execute!, _mm_maker(Float64, M, 4, N, 3; Cfill = NaN), 1.0, 2.0) + C, plan = _run_fresh(execute!, _mm_maker(Float64, M, 4, N, 3; Cfill = NaN), 1.0, 2.0) @test size(C) == (M, N) + @test _path_of(plan) isa QuasiStrided.EmptyPath end # Singleton extents, and two K labels against a singleton free label. mk2K = () -> ( @@ -168,7 +170,7 @@ end kernel = SIMDKernel(Val(8), Val(6), Float64, Val(4)) # Default everything; several blocks with tail tiles in M and N (the # vectorized store's row tail); a conj transform, which the real path never - # builds but which must still cross `_pack_sliver!` concretely typed. + # builds but which must still cross `pack_sliver!` concretely typed. Cmat = zeros(Ma, Na) p = _mm_plan(Cmat, Amat, Bmat) plans = ( diff --git a/test/execution/test_manual_pipeline.jl b/test/execution/test_manual_pipeline.jl index c0cf5af..20c10bf 100644 --- a/test/execution/test_manual_pipeline.jl +++ b/test/execution/test_manual_pipeline.jl @@ -16,7 +16,7 @@ (dN_B, dN_C) = block_descriptors!(nb, N, 0, 4) # Scatter axes borrow pointers into the offset buffers. GC.@preserve mb nb kb begin - axis(d, buf) = QuasiStrided._axis_of(d, buf, 0) + axis(d, buf) = QuasiStrided.axis_of(d, buf, 0) row_A, row_C = axis(dM_A, mb[1]), axis(dM_C, mb[2]) col_B, col_C = axis(dN_B, nb[1]), axis(dN_C, nb[2]) diff --git a/test/planning/test_per_call_overhead.jl b/test/planning/test_per_call_overhead.jl index bb20778..b727f1f 100644 --- a/test/planning/test_per_call_overhead.jl +++ b/test/planning/test_per_call_overhead.jl @@ -130,7 +130,7 @@ end @test_throws BoundsError QS.checked_span_bounds(0, (0, 0), (0, 0), 0) end -@testset "per-call floor: _classify_slivers! accumulates the TRUE block range" begin +@testset "per-call floor: classify_slivers! accumulates the TRUE block range" begin # The range handed to `checked_span_bounds` is the true min/max over the # whole block; a too-small range would silently under-validate. Random.seed!(20260922) @@ -142,7 +142,7 @@ end buf2 = [rand(-40:40) for _ in 1:blocklen] d1 = Vector{BlockDescriptor}(undef, nsliv) d2 = Vector{BlockDescriptor}(undef, nsliv) - (r1, r2) = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) + (r1, r2) = QS.classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) @test r1 == (minimum(buf1), maximum(buf1)) @test r2 == (minimum(buf2), maximum(buf2)) @test all(s -> d1[s].count == min(reg, blocklen - (s - 1) * reg), 1:nsliv) @@ -155,7 +155,7 @@ end buf2[7:12] .= [0, -7, -14, -21, -28, -35] # sliver 2, REGULAR, stride -7 d1 = Vector{BlockDescriptor}(undef, 3) d2 = Vector{BlockDescriptor}(undef, 3) - (r1, r2) = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, 3) + (r1, r2) = QS.classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, 3) @test !d1[2].regular && d1[2].count == 6 # scan branch @test d2[2].regular && d2[2].stride == -7 # affine branch, negative stride @test r1 == (-500, 500) @@ -171,7 +171,7 @@ end buf = [rand(-30:30) for _ in 1:n] rand() < 0.4 && (buf = [3 + 5 * (t - 1) for t in 1:n]) # force a regular run d = describe_block(buf, 0, n) - GC.@preserve buf @test QS.descriptor_offset_range(d, buf, 0) == extrema(QS._axis_of(d, buf, 0)) + GC.@preserve buf @test QS.descriptor_offset_range(d, buf, 0) == extrema(QS.axis_of(d, buf, 0)) end @test QS.descriptor_offset_range(BlockDescriptor(0, 0, 0, true), Int[], 0) == (0, -1) end @@ -377,7 +377,7 @@ end end end -@testset "per-call floor: _ramp_slivers! reproduces _classify_slivers! exactly" begin +@testset "per-call floor: ramp_slivers! reproduces classify_slivers! exactly" begin Random.seed!(31337) for trial in 1:150 D = rand(1:3) @@ -410,11 +410,11 @@ end d1 = Vector{BlockDescriptor}(undef, nsliv) d2 = Vector{BlockDescriptor}(undef, nsliv) fill_offsets!((buf1, buf2), g, first, blocklen) - want = QS._classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) + want = QS.classify_slivers!(QS.GroupBuffers((buf1, buf2), (d1, d2)), blocklen, reg, nsliv) r1 = Vector{BlockDescriptor}(undef, nsliv) r2 = Vector{BlockDescriptor}(undef, nsliv) - got = QS._ramp_slivers!(QS.GroupBuffers((buf1, buf2), (r1, r2)), steps[1], steps[2], first, blocklen, reg, nsliv) + got = QS.ramp_slivers!(QS.GroupBuffers((buf1, buf2), (r1, r2)), steps[1], steps[2], first, blocklen, reg, nsliv) for s in 1:nsliv @test r1[s].base == d1[s].base && r1[s].stride == d1[s].stride From 5bd0a74a05cc331250b1cd456b7a916900bf6387 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 00:10:45 -0400 Subject: [PATCH 33/60] Benchmark: nest with and without ramp flags; decision log: chunk 14 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 40 +++++++++++++----- benchmark/bench_ramp_flags.jl | 79 +++++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 10 deletions(-) create mode 100644 benchmark/bench_ramp_flags.jl diff --git a/CLEANUP.md b/CLEANUP.md index 18112c6..65caed3 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -26,7 +26,7 @@ TTFX when codegen is touched). | 11 | Blocking | `blocking.jl`, line packing (now `packing/line_packing.jl`) | done | | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | done | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | done | -| 14 | Five-loop nest | `execution/macrokernel.jl`, `execute.jl` | | +| 14 | Five-loop nest | `execution/macrokernel.jl` (now `nest.jl`), `execute.jl`, `c_panel.jl` | done (ramp flags pending) | | 15 | Alternative paths | `oracle.jl`, `unpackedb.jl` | | | 16 | Degenerate paths | `dot.jl`, `outer.jl` | | | 17 | TensorOperations backend | `integrations/tensoroperations.jl` | | @@ -225,6 +225,27 @@ explicit reuse mechanism rather than a cache. The tile-wise oracle is gone are bundled per group (`ws.m`, `ws.n`, `ws.k`), not per operand: a group's two maps are filled together and share a lifetime. Plans are single-task. +### D17. Five-loop nest (chunk 14, applied: 5c08bdf; ramp-flag benchmark pending) + +Behaviour-preserving first: `EmptyPath`/`ScalePath` for empty M/N and K = 0 +(decided at plan time; `execute!` keeps only the `alpha == 0` check); the +nest reads its parameters from the plan; one `pack_block!` per operand +(per-sliver or by lines); ramp decisions from the path's static flags; the C +panel path loops over N blocks itself (copy in, nest on the panel, copy +out), so the nest has no `target`/hooks; `execute.jl` split into +`execute.jl`, `nest.jl` (absorbs `macrokernel.jl`) and `c_panel.jl`. +Then, measured separately (Slurm A/B): drop the six `NestPath` ramp flags in +favour of the runtime `axis_of` choice; if they don't pay off, `NestPath` +shrinks to `{UNPACKED_B, SPLIT}` and the 512-entry table goes, and callable +paths get reconsidered. +As applied: `nest!(plan, alphaT, betaT, path, n_range)`; `pack_block!` lives in +`nest.jl` (needs the plan and group buffers); N's ramp under unpacked B stays +a runtime test of B's map. 408 equivalence cases bitwise identical; the extra +nests of degenerate plans are gone. Ramp-flag A/B: `benchmark/bench_ramp_flags.jl` +(flags cleared also disables the closed-form ramp descriptors, so it bounds +the cost of dropping them). Workstation, indicative: large cases equal, 24³ +and 64³ 20–50% slower without flags. + ### Complex blocking benchmark (after chunk 11) Same-node A/B, `bench_complex_blocking.jl`, jobs 7160147–9: the kernel's own @@ -240,9 +261,10 @@ blocking.jl (uncommitted until the full suite passes). ## Resume here -1. Chunk 14: five-loop nest (`execution/macrokernel.jl`, `execute.jl`), - including the open items marked chunk 14 below. -2. Workflow: present each chunk (purpose, reading order, design decisions, +1. Ramp-flag benchmark: user submits `bench_ramp_flags.jl` on Icelake/Genoa/ + Rome; decide on the `NestPath` flags (D17) from it. +2. Chunk 15: alternative paths (`unpackedb.jl`; the oracle is already gone). +3. Workflow: present each chunk (purpose, reading order, design decisions, proposed fixes, questions), then hand the agreed changes to an Opus agent with the standard checks (Runic, full suite, per-call floor/allocations, TTFX, equivalence script where behaviour must not change); perf claims @@ -272,16 +294,14 @@ blocking.jl (uncommitted until the full suite passes). ## Open items (to revisit in their chunk) -- One `pack_block!(panel, block, spec, transform, order)` for per-sliver and - by-lines packing, so the nest makes one call per block (chunk 14). - Internal constants still use `_UPPER` names; drop the underscores in one pass at the end (D2). -- `execute.jl` mixes public API, the C-panel path and the nest (chunk 14). -- Plans with empty extents or K = 0 now compile their path although the - short-circuit always fires first; `select_path` could return a trivial - path for them (chunk 14). +- Path types carry compile-time parameters (lane width, nest flags), which is + why they are structs, not functions. Measure whether the six `NestPath` ramp + flags pay off (regular vs scattered layouts, specialisation count); if they + go, consider callable paths instead of `execute_path!` (chunk 14). - `_execute_dot!` takes `matB` as a runtime `Bool` although `DotPath{MATB}` makes it static (chunk 16). - `plan_request` is spelled out twice in `planned` (swapped and not). diff --git a/benchmark/bench_ramp_flags.jl b/benchmark/bench_ramp_flags.jl new file mode 100644 index 0000000..8a2a542 --- /dev/null +++ b/benchmark/bench_ramp_flags.jl @@ -0,0 +1,79 @@ +# Same-node comparison of the nest with and without its static ramp flags, +# interleaved per case: `on` is the plan as built; `off` is the same plan with +# every `NestPath` ramp flag cleared, so every sliver axis takes the runtime +# affine/scatter choice and ramp composites go through the offset buffers. `off` +# is an upper bound on what dropping the flags costs: without them the ramp +# descriptors could still be computed in closed form at runtime. +# +# julia --project=benchmark benchmark/bench_ramp_flags.jl +# [--dtypes Float64,ComplexF64] [--rounds 3] [--reps 5] [--outdir DIR] + +include(joinpath(@__DIR__, "harness.jl")) + +using QuasiStrided: NestPath, ContractPlan + +const RUN_DTYPES = parse_dtypes(argopt("dtypes", "Float64,ComplexF64")) +const ROUNDS = argopt("rounds", 3) +const REPS = argopt("reps", 5) +const OUTDIR = outdir() +const CSV_PATH = joinpath(OUTDIR, "bench_ramp_flags.csv") + +# (name, size of each label, indA, indB, indC) +const CASES = [ + ("matmul_24^3", (24, 24, 24), (1, 2), (2, 3), (1, 3)), + ("matmul_64^3", (64, 64, 64), (1, 2), (2, 3), (1, 3)), + ("matmul_512^3", (512, 512, 512), (1, 2), (2, 3), (1, 3)), + ("smallM_12x512x512", (12, 512, 512), (1, 2), (2, 3), (1, 3)), + # Composites that are ramps in some maps only. + ("partial_32", (32, 32, 32, 32), (1, 3, 4), (4, 2), (1, 2, 3)), + ("ccsd_t_16", (16, 16, 16, 16, 16, 16, 16), (7, 4, 5, 6), (1, 2, 3, 7), (1, 2, 3, 4, 5, 6)), +] + +# Median over `REPS` batches, each long enough (>= 1 ms) to time reliably. +function batch_time_s(f!) + f!() + n = max(1, ceil(Int, 1.0e-3 / max(@elapsed(f!()), 1.0e-9))) + return median_time_s( + () -> ( + for _ in 1:n + f!() + end + ); reps = REPS + ) / n +end + +operand(T, sizes, ind, rng) = randn(rng, T, map(l -> sizes[l], ind)...) + +function without_ramp_flags(plan) + plan.path isa NestPath || return nothing + U, _, S = typeof(plan.path).parameters + return ContractPlan(plan; path = NestPath{U, ntuple(_ -> false, 6), S}()) +end + +function run(io) + print_env_header(stdout, "bench_ramp_flags.jl") + println(io, "dtype,case,path,round,t_on,t_off") + rng = MersenneTwister(1) + for T in RUN_DTYPES, (name, sizes, indA, indB, indC) in CASES + A = operand(T, sizes, indA, rng) + B = operand(T, sizes, indB, rng) + C = operand(T, sizes, indC, rng) + on = plan_contract(StridedView(C), StridedView(A), indA, StridedView(B), indB, indC) + off = without_ramp_flags(on) + if off === nothing || off.path === on.path + println("skip $T $name: path $(typeof(on.path)) has no ramp flags set") + continue + end + for r in 1:ROUNDS + t_on = batch_time_s(() -> execute!(on, one(T), zero(T))) + t_off = batch_time_s(() -> execute!(off, one(T), zero(T))) + println(io, join((T, name, repr(string(typeof(on.path))), r, t_on, t_off), ',')) + flush(io) + @printf("%s %s round %d: on %.4g s, off/on %.3f\n", T, name, r, t_on, t_off / t_on) + end + end + return nothing +end + +open(run, CSV_PATH, "w") +println("wrote ", CSV_PATH) From 337dc4a13e87f9d963fd8a4deb42b09af922111e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 11:24:50 -0400 Subject: [PATCH 34/60] Chunk 15: unpacked B as a B source - One tile loop: `micro_tiles!(plan, path, rowsB_k, ...)` in nest.jl replaces `micro_tiles_packed_b!` and `_micro_tiles_unpacked_b!`. Kernel, workspace, tile sizes and sliver widths come from the plan; the B sliver from `b_sliver`, which dispatches on the path's `UNPACKED_B` (packed panel or `UnpackedBView`). B's K axis is passed only when B is read in place (`nothing` otherwise), so the packed loop does not specialise on it. - src/execution/unpackedb.jl is gone: `UnpackedBView`, its accessors and `unpacked_b_view` move to src/packing/unpacked_b.jl (a B source beside `PackedPanel`); `unpacked_b_rule` and `_UNPACKED_B_MMAX` to plan.jl next to `select_path`. - `reads_b_by_element(K)` in kernels.jl replaces the hard-coded eligible kernel list (same truth table). - Underscores dropped: `unpacked_b_element`, `unpacked_b_view`, `unpacked_b_rule`. Full suite: 32155 pass. Behaviour-preserving: 408 equivalence cases bitwise identical to 5bd0a74, paths identical. Tile-loop specialisations on that workload: 42 (old: 19 packed + 23 unpacked); nest! 42, execute_micro_tile! 35, unchanged. TTFX (s, using/first Float64 64^3/first ComplexF64/first @tensor), interleaved fresh processes, 3 each: before 0.57/3.45/1.83/1.20, 0.48/3.27/1.85/1.15, 0.52/3.51/2.04/1.17 after 0.60/3.31/1.81/1.18, 0.52/3.34/1.89/1.28, 0.51/3.61/1.93/1.18 execute! per call (ns, medians of 5 interleaved runs, noisy workstation, indicative), before -> after, all 0 allocations: 64^3 Float64 (unpacked B) 6709 -> 6460 32^3 ComplexF64 (unpacked B) 5211 -> 5144 2x2 Float64 80 -> 76 12x512x512 Float64 240117 -> 227535 12x512x512 ComplexF64 (fmaddsub) 428252 -> 419345 12x512x512 ComplexF64*Float64 283307 -> 276729 300x64x64 Float64 (packed B) 33769 -> 33062 300x32x32 ComplexF64 (packed B) 32332 -> 32289 64^3 Float64, B transposed (packed) 7874 -> 7752 Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 3 +- src/execution/nest.jl | 57 ++++++++++------- src/execution/unpackedb.jl | 106 ------------------------------- src/microkernels/kernels.jl | 12 ++++ src/microkernels/steps.jl | 2 +- src/packing/unpacked_b.jl | 52 +++++++++++++++ src/planning/plan.jl | 19 +++++- test/execution/test_unpackedb.jl | 2 +- 8 files changed, 119 insertions(+), 134 deletions(-) delete mode 100644 src/execution/unpackedb.jl create mode 100644 src/packing/unpacked_b.jl diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 277e989..8c5e23d 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -32,6 +32,8 @@ include("microkernels/kernels.jl") include("microkernels/vecops.jl") include("microkernels/steps.jl") include("microkernels/stores.jl") +# B read in place: a B source beside `PackedPanel`, after the kernel interface it extends. +include("packing/unpacked_b.jl") # --- Planning: labels, kernel and blocking choice, the plan --- include("planning/labels.jl") @@ -45,7 +47,6 @@ include("execution/paths.jl") include("execution/nest.jl") include("execution/execute.jl") include("execution/c_panel.jl") -include("execution/unpackedb.jl") include("execution/dot.jl") include("execution/outer.jl") diff --git a/src/execution/nest.jl b/src/execution/nest.jl index 7209391..928e9d8 100644 --- a/src/execution/nest.jl +++ b/src/execution/nest.jl @@ -11,10 +11,6 @@ function nest!( kernel = plan.kernel ws = plan.workspace m_tile, n_tile = tile_size(kernel) - # GUARDRAIL: reals per sliver per K step address the packed panels; - # `m_tile`/`n_tile` count register-tile rows. They differ for complex - # kernels. - a_sliver_width, b_sliver_width = sliver_width(kernel) (; m_block, k_block, n_block) = plan.blocking m_length = axis_length(plan.mgroup) k_length = axis_length(plan.kgroup) @@ -108,18 +104,10 @@ function nest!( pack_block!(plan, path, Val(1), colsA_k, m_block_length, k_block_length, m_tiles) # --- loops over N tiles and M tiles --- - if UNPACKED_B - _micro_tiles_unpacked_b!( - kernel, plan, ws, rowsB_k, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, k_block_length, alphaT, beta_eff, Val(aff_mC), Val(aff_nC) - ) - else - micro_tiles_packed_b!( - kernel, plan, ws, m_tiles, n_tiles, - m_tile, n_tile, a_sliver_width, b_sliver_width, k_block_length, alphaT, beta_eff, - Val(aff_mC), Val(aff_nC) - ) - end + micro_tiles!( + plan, path, UNPACKED_B ? rowsB_k : nothing, m_tiles, n_tiles, k_block_length, + alphaT, beta_eff, Val(aff_mC), Val(aff_nC) + ) m_block_start += m_block_length end @@ -170,15 +158,23 @@ end # The tile loops over one (N, K, M) block. `@noinline`: one call per block, and # the inlined microkernel is most of a nest's size; compile cost grows faster -# than linearly with function size. -@noinline function micro_tiles_packed_b!( - kernel::K, plan::ContractPlan, ws, m_tiles::Int, n_tiles::Int, - m_tile::Int, n_tile::Int, a_sliver_width::Int, b_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, - aff_mC::Val{MC}, aff_nC::Val{NC} - ) where {K, MC, NC} +# than linearly with function size. `rowsB_k`: B's K axis when B is read in +# place (a barrier over its type, so each view is concretely typed), else +# `nothing`, so that the packed-B loop does not specialise on it. +@noinline function micro_tiles!( + plan::ContractPlan, path::NestPath, rowsB_k::KB, m_tiles::Int, n_tiles::Int, k_block_length::Int, + alphaT, beta_eff, aff_mC::Val{MC}, aff_nC::Val{NC} + ) where {KB, MC, NC} + kernel = plan.kernel + ws = plan.workspace + m_tile, n_tile = tile_size(kernel) + # GUARDRAIL: reals per sliver per K step address the packed panels; + # `m_tile`/`n_tile` count register-tile rows. They differ for complex + # kernels. + a_sliver_width = sliver_width(kernel, 1) for n_tile_index in 0:(n_tiles - 1) n_tile_start = n_tile_index * n_tile - bpanel = sliver_panel(ws.packed_b, b_sliver_width, k_block_length, n_tile_index) + bsliver = b_sliver(plan, path, rowsB_k, k_block_length, n_tile_index, n_tile_start) colsC = axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) for m_tile_index in 0:(m_tiles - 1) m_tile_start = m_tile_index * m_tile @@ -187,13 +183,26 @@ end # Inside the caller's C check. @inbounds execute_micro_tile!( kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bpanel, k_block_length, alphaT, beta_eff + apanel, bsliver, k_block_length, alphaT, beta_eff ) end end return nothing end +# The B sliver of N tile `n_tile_index`: its packed panel, or B in place. +@inline b_sliver(plan::ContractPlan, ::NestPath{false}, ::Nothing, k_block_length::Int, n_tile_index::Int, ::Int) = + sliver_panel(plan.workspace.packed_b, sliver_width(plan.kernel, 2), k_block_length, n_tile_index) +@inline function b_sliver( + plan::ContractPlan, ::NestPath{true}, rowsB_k::KB, ::Int, n_tile_index::Int, n_tile_start::Int + ) where {KB <: AbstractVector{Int}} + n = plan.workspace.n + return unpacked_b_view( + plan.kernel, plan.Bstorage, plan.Bbase, n.descriptors[1][n_tile_index + 1], n.offsets[1], + n_tile_start, rowsB_k, plan.btransform + ) +end + # The axis `d` describes. GUARDRAIL: a `Union{AffineAxis, ScatterAxis}` of # `isbits` types, never boxed; each consumer below binds each axis type as its # own parameter, so it builds a concretely typed `Tile`. diff --git a/src/execution/unpackedb.jl b/src/execution/unpackedb.jl deleted file mode 100644 index b0fc938..0000000 --- a/src/execution/unpackedb.jl +++ /dev/null @@ -1,106 +0,0 @@ -# "Unpacked B": the microkernel reads B in place instead of from a packed -# panel. The kernels consume B one (complex) scalar at a time and broadcast -# it, so any K axis, column offsets and `AbstractVector` storage work. -# Packing B buys locality and padding columns at an O(N*K) gather per -# (N, K) block, which at small M costs as much as the kernel itself. - -# Stands in for a packed B micro-panel of `NR` columns: column `j`, K step -# `p` is storage address `colbase[j] + ksteps[p]` (base included). Padding -# columns of a tail sliver alias the last valid column, so every address read -# is inside the hoisted B check; their results are never stored. `length` is -# the packed-equivalent count that `execute_tile!` checks, not the -# storage length. -struct UnpackedBView{S, K <: AbstractVector{Int}, NR, F} - storage::S - colbase::NTuple{NR, Int} - ksteps::K - per_k::Int - transform::F -end - -Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) - -# `@inbounds`: the caller has checked the whole B panel region. -@inline function _unpacked_b_element(u::UnpackedBView{S, K, NR, F}, j::Int, p::Int) where {S, K, NR, F} - @inbounds z = u.storage[u.colbase[j] + u.ksteps[p] + 1] - return u.transform(z) -end - -@inline b_scalar(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_complex(u::UnpackedBView, kernel, j::Int, p::Int) - z = convert(scalartype(kernel), _unpacked_b_element(u, j, p)) - return (real(z), imag(z)) -end - -# The kernels whose K step reads B through `b_scalar`/`b_complex` -# (not 1m, whose inner kernel walks the planar panel, nor the scalar one); by -# type, to decide the path before the kernel exists. -@inline _unpacked_b_kernel_eligible(::Type{<:Union{SIMDKernel, PlanarKernel, FMAddSubKernel, ComplexRealKernel}}) = true -@inline _unpacked_b_kernel_eligible(::Type) = false - -# Whether B is read in place (given an eligible kernel): small M, and every B -# column contiguous along K. Packing pays off only through reuse across M -# slivers, while `n_tile` contiguous columns read in place cost the kernel -# nothing. A large K stride in B (each step its own cache line, and for a -# power of two only a few L1 sets) is what packing exists for. `mode`: -# `:always`/`:never` override the rule (`PathModes`). -@inline function _unpacked_b_rule(mode::Symbol, mgroup::AxisGroup, kgroup::AxisGroup) - mode === :always && return true - mode === :never && return false - axis_length(mgroup) <= _UNPACKED_B_MMAX || return false - (k_ramp, k_step) = affine_ramp(kgroup) - return k_ramp && abs(k_step[2]) == 1 -end - -# Beyond this M the packed B's reuse wins. -const _UNPACKED_B_MMAX = 256 - -# The view for one N sliver; `buf`/`n_tile_start` locate an irregular sliver's -# offsets in `ws.n.offsets[1]`. -@inline function _unpacked_b_view( - kernel::Microkernel{MR, NR, T}, storage::S, base::Int, - d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F - ) where {MR, NR, T, S, K <: AbstractVector{Int}, F} - last = d.count - 1 - colbase = ntuple(Val(NR)) do j1 - jj = min(j1 - 1, last) - off = d.regular ? d.base + jj * d.stride : (@inbounds buf[n_tile_start + jj + 1]) - base + off - end - return UnpackedBView(storage, colbase, ksteps, sliver_width(kernel, 2), transform) -end - -# The tile loops of `nest!` with B read in place. A barrier over the K -# axis type, so each view is concretely typed. `@noinline` so that the -# packed-B nest does not grow by this body. The view's addresses are within -# the caller's B check. -@noinline function _micro_tiles_unpacked_b!( - kernel::K, plan::ContractPlan, ws, rowsB_k::KA, m_tiles::Int, n_tiles::Int, - m_tile::Int, n_tile::Int, a_sliver_width::Int, k_block_length::Int, alphaT, beta_eff, - aff_mC::Val{MC}, aff_nC::Val{NC} - ) where {K, KA <: AbstractVector{Int}, MC, NC} - Bstorage = plan.Bstorage - Bbase = plan.Bbase - btransform = plan.btransform - for n_tile_index in 0:(n_tiles - 1) - n_tile_start = n_tile_index * n_tile - bview = _unpacked_b_view( - kernel, Bstorage, Bbase, ws.n.descriptors[1][n_tile_index + 1], ws.n.offsets[1], n_tile_start, rowsB_k, btransform - ) - colsC = axis_of(ws.n.descriptors[2][n_tile_index + 1], ws.n.offsets[2], n_tile_start, aff_nC) - for m_tile_index in 0:(m_tiles - 1) - m_tile_start = m_tile_index * m_tile - apanel = sliver_panel(ws.packed_a, a_sliver_width, k_block_length, m_tile_index) - rowsC = axis_of(ws.m.descriptors[2][m_tile_index + 1], ws.m.offsets[2], m_tile_start, aff_mC) - # Inside the caller's C check. - @inbounds execute_micro_tile!( - kernel, plan.Cstorage, plan.Cbase, rowsC, colsC, - apanel, bview, k_block_length, alphaT, beta_eff - ) - end - end - return nothing -end diff --git a/src/microkernels/kernels.jl b/src/microkernels/kernels.jl index 6d030a4..61086c8 100644 --- a/src/microkernels/kernels.jl +++ b/src/microkernels/kernels.jl @@ -184,6 +184,18 @@ pack_formats(::Type{<:FMAddSubKernel}) = (InterleavedFormat(), PlanarFormat()) pack_formats(::Type{<:ComplexRealKernel}) = (InterleavedFormat(), RealFormat()) pack_formats(::Type{<:RealComplexKernel}) = (RealFormat(), InterleavedFormat()) +""" + reads_b_by_element(K::Type{<:Microkernel}) -> Bool + +Whether the K step of every kernel of type `K` reads B only through +`b_scalar`/`b_complex`, so that any B source implementing those (such as +`UnpackedBView`, B read in place) can stand in for its packed B panel. +False otherwise: `ScalarKernel` loads from the panel directly, and 1m and +`RealComplexKernel` run a real kernel over the packed panel's reals. +""" +reads_b_by_element(::Type{<:Microkernel}) = false +reads_b_by_element(::Type{<:Union{SIMDKernel, PlanarKernel, FMAddSubKernel, ComplexRealKernel}}) = true + # Accumulator planes held live: planar keeps separate re/im planes. accumulator_planes(::Type{<:Microkernel}) = 1 accumulator_planes(::Type{<:PlanarKernel}) = 2 diff --git a/src/microkernels/steps.jl b/src/microkernels/steps.jl index 4e5334a..43eca00 100644 --- a/src/microkernels/steps.jl +++ b/src/microkernels/steps.jl @@ -54,7 +54,7 @@ end end # B column `j` at K step `p`, as a real or as `(re, im)`; `UnpackedBView` -# (src/execution/unpackedb.jl) overrides both to read B in place, resolved at +# (src/packing/unpacked_b.jl) overrides both to read B in place, resolved at # compile time. @inline b_scalar(packed_b::PB, kernel, j::Int, p::Int) where {PB} = panel_load(packed_b, packed_b_offset(kernel, j, p)) diff --git a/src/packing/unpacked_b.jl b/src/packing/unpacked_b.jl new file mode 100644 index 0000000..b34305e --- /dev/null +++ b/src/packing/unpacked_b.jl @@ -0,0 +1,52 @@ +# "Unpacked B": the microkernel reads B in place instead of from a packed +# panel. The kernels that read B by element (`reads_b_by_element`) broadcast +# one (complex) scalar at a time, so any K axis, column offsets and +# `AbstractVector` storage work. Packing B buys locality and padding columns +# at an O(N*K) gather per (N, K) block, which at small M costs as much as the +# kernel itself. + +# Stands in for a packed B micro-panel of `NR` columns: column `j`, K step +# `p` is storage address `colbase[j] + ksteps[p]` (base included). Padding +# columns of a tail sliver alias the last valid column, so every address read +# is inside the hoisted B check; their results are never stored. `length` is +# the packed-equivalent count that `execute_tile!` checks, not the +# storage length. +struct UnpackedBView{S, K <: AbstractVector{Int}, NR, F} + storage::S + colbase::NTuple{NR, Int} + ksteps::K + per_k::Int + transform::F +end + +Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) + +# `@inbounds`: the caller has checked the whole B panel region. +@inline function unpacked_b_element(u::UnpackedBView{S, K, NR, F}, j::Int, p::Int) where {S, K, NR, F} + @inbounds z = u.storage[u.colbase[j] + u.ksteps[p] + 1] + return u.transform(z) +end + +@inline b_scalar(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_complex(u::UnpackedBView, kernel, j::Int, p::Int) + z = convert(scalartype(kernel), unpacked_b_element(u, j, p)) + return (real(z), imag(z)) +end + +# The view for one N sliver; `buf`/`n_tile_start` locate an irregular sliver's +# offsets in `ws.n.offsets[1]`. +@inline function unpacked_b_view( + kernel::Microkernel{MR, NR, T}, storage::S, base::Int, + d::BlockDescriptor, buf::Vector{Int}, n_tile_start::Int, ksteps::K, transform::F + ) where {MR, NR, T, S, K <: AbstractVector{Int}, F} + last = d.count - 1 + colbase = ntuple(Val(NR)) do j1 + jj = min(j1 - 1, last) + off = d.regular ? d.base + jj * d.stride : (@inbounds buf[n_tile_start + jj + 1]) + base + off + end + return UnpackedBView(storage, colbase, ksteps, sliver_width(kernel, 2), transform) +end diff --git a/src/planning/plan.jl b/src/planning/plan.jl index 1020b8f..1d8cdb0 100644 --- a/src/planning/plan.jl +++ b/src/planning/plan.jl @@ -403,10 +403,27 @@ function select_path(::Type{K}, m_tile::Int, n_tile::Int, req::PlanRequest{T}, r end k_length == 1 && modes.outer !== :never && _outer_applicable(T, req.Astorage, req.Cstorage, mgroup, m_length) && return lane_path(OuterPath, _dot_lanewidth(T)) - unpacked_b = _unpacked_b_kernel_eligible(K) && _unpacked_b_rule(modes.unpacked_b, mgroup, kgroup) + unpacked_b = reads_b_by_element(K) && unpacked_b_rule(modes.unpacked_b, mgroup, kgroup) return nest_path(unpacked_b, mgroup, ngroup, kgroup, is_split(mpack), is_split(npack), panel) end +# Whether B is read in place (given a kernel that reads B by element): small +# M, and every B column contiguous along K. Packing pays off only through +# reuse across M slivers, while `n_tile` contiguous columns read in place cost +# the kernel nothing. A large K stride in B (each step its own cache line, and +# for a power of two only a few L1 sets) is what packing exists for. `mode`: +# `:always`/`:never` override the rule (`PathModes`). +@inline function unpacked_b_rule(mode::Symbol, mgroup::AxisGroup, kgroup::AxisGroup) + mode === :always && return true + mode === :never && return false + axis_length(mgroup) <= _UNPACKED_B_MMAX || return false + (k_ramp, k_step) = affine_ramp(kgroup) + return k_ramp && abs(k_step[2]) == 1 +end + +# Beyond this M the packed B's reuse wins. +const _UNPACKED_B_MMAX = 256 + # Type parameters on the transforms force specialisation on them: the # compiler does not specialise on a `Function` argument it only passes on. build_plan(::Val{Kern}, path, atransform::TA, btransform::TB, slot::Base.RefValue) where {Kern, TA, TB} = diff --git a/test/execution/test_unpackedb.jl b/test/execution/test_unpackedb.jl index 3ac2f86..7d14297 100644 --- a/test/execution/test_unpackedb.jl +++ b/test/execution/test_unpackedb.jl @@ -1,5 +1,5 @@ # The unpacked-B path: `execute!` reads B in place through an `UnpackedBView` -# where `_unpacked_b_rule` says so. The result must be bitwise the packed nest's +# where `unpacked_b_rule` says so. The result must be bitwise the packed nest's # (same arithmetic, same values, different address), and the packed-B buffer # must stay untouched. From 3ce5f9acc25b1a9488c145887f5bba606aaede53 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 11:25:23 -0400 Subject: [PATCH 35/60] Decision log: chunk 15, ramp-flag round 1 Co-Authored-By: Claude Opus 5.5 --- CLEANUP.md | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/CLEANUP.md b/CLEANUP.md index 65caed3..6cef01b 100644 --- a/CLEANUP.md +++ b/CLEANUP.md @@ -27,7 +27,7 @@ TTFX when codegen is touched). | 12 | Workspace | `execution/workspace.jl`, `barrier.jl` | done | | 13 | The plan | `planning/plan.jl` (+ `test_plan_contract.jl`, `test_per_call_overhead.jl`) | done | | 14 | Five-loop nest | `execution/macrokernel.jl` (now `nest.jl`), `execute.jl`, `c_panel.jl` | done (ramp flags pending) | -| 15 | Alternative paths | `oracle.jl`, `unpackedb.jl` | | +| 15 | Alternative paths | `unpackedb.jl` (oracle gone in chunk 12) | done | | 16 | Degenerate paths | `dot.jl`, `outer.jl` | | | 17 | TensorOperations backend | `integrations/tensoroperations.jl` | | | 18 | Test infrastructure | `runtests.jl`, `helpers.jl`, `forced_isa_runner.jl`, `quality/` | | @@ -245,6 +245,20 @@ nests of degenerate plans are gone. Ramp-flag A/B: `benchmark/bench_ramp_flags.j (flags cleared also disables the closed-form ramp descriptors, so it bounds the cost of dropping them). Workstation, indicative: large cases equal, 24³ and 64³ 20–50% slower without flags. +Slurm round 1 (jobs 7183225–7, Icelake/Genoa/Rome; Icelake and Rome ran on a +tree with chunk-15 edits in progress, ratios are within-process): off/on +1.07–1.35 at 24³, up to 1.10 at 64³, 1.00 at 512³, ccsd(t) up to 1.07. But +off is *faster* in three cases: C64 12×512×512 on Icelake (0.52), F64 +12×512×512 (0.75) and 64³ (0.91) on Genoa, all unpacked-B. Next: explain +that anomaly, then a round with a runtime-closed-form-ramp variant. + +### D18. Unpacked B (chunk 15, applied: 337dc4a) + +`UnpackedBView` is a B source next to `PackedPanel` (`packing/unpacked_b.jl`); +one tile loop in `nest.jl` takes its B sliver from either, by the path's +`UNPACKED_B`; the rule moves to `plan.jl`; eligibility is a kernel trait +(`reads_b_by_element`) defined beside each kernel. `unpackedb.jl` goes. The +`M <= 256` cutoff stays until measured (see possible improvements). ### Complex blocking benchmark (after chunk 11) @@ -261,9 +275,10 @@ blocking.jl (uncommitted until the full suite passes). ## Resume here -1. Ramp-flag benchmark: user submits `bench_ramp_flags.jl` on Icelake/Genoa/ - Rome; decide on the `NestPath` flags (D17) from it. -2. Chunk 15: alternative paths (`unpackedb.jl`; the oracle is already gone). +1. Ramp flags (D17): investigate why clearing them is faster on the + unpacked-B small-M cases, then round 2 with a variant keeping closed-form + ramps at runtime but no static axis-type flags. +2. Chunk 16: degenerate paths (`dot.jl`, `outer.jl`). 3. Workflow: present each chunk (purpose, reading order, design decisions, proposed fixes, questions), then hand the agreed changes to an Opus agent with the standard checks (Runic, full suite, per-call floor/allocations, @@ -292,6 +307,10 @@ blocking.jl (uncommitted until the full suite passes). while B's pack (and the unpacked-B path) would prefer B's strides. Measure on suite layouts where B's and C's N strides disagree. +- Unpacked-B cutoff `M <= 256` is a constant; a cutoff from the reuse a + packed B sliver would get (about `m_length / MR` M slivers) may be more + principled. Benchmark before changing. + ## Open items (to revisit in their chunk) From ea97e6d276e5f2765dc5613f129991ca2828080b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 12:56:14 -0400 Subject: [PATCH 36/60] Unpacked B: avoid vectorised B addresses and unrolling With B read in place, the K loop addresses B as storage[colbase[j] + ksteps[p] + 1] for each of the NR columns. Two codegen problems made it spill: - LLVM's SLP vectoriser can turn the NR per-column indices into vector arithmetic inside the K loop, which LSR cannot strength-reduce. On AVX-512 the FMAddSub 12x8 kernel then needs a 33rd zmm register and spills 12 accumulators per K step. Whether SLP fires depended on the surrounding code: it hit micro_tiles! specialised on both C maps being ramps, i.e. ordinary dense plans, but not the runtime-axis variants, which is why clearing the ramp flags made these plans faster. - Where the loop is unrolled (znver4 codegen unrolls it 8x), the runtime K stride gives each unrolled step its own column addresses, spilling general-purpose registers in every variant. For DenseVector storage, B is now loaded from a column pointer plus one K index shared by all columns, and the add_tile K loop disables unrolling when B is an UnpackedBView. The packed-B loop is unchanged. UnpackedBView is included before the K steps so the generator sees it. Hot K loop of micro_tiles! for C maps (ramp, ramp) / (ramp, runtime) / (runtime, runtime), (instructions, spill+reload lines), before -> after: - ComplexF64 12x512x512, native AVX-512: (111,29)/(74,0)/(74,0) -> (74,0) for all three. znver4: (85,2)/(74,0)/(74,0) -> (74,0). - Float64 12x512x512, znver4: (257,19)/(212,42)/(207,41), unrolled 8x -> (24,0) for all three. Float64 64^3 (znver4 kernel) likewise. - Packed B (300x64x64): identical code up to register names and stack slots on native, znver4 and znver2. AVX2 (znver2) unpacked loops unchanged except ComplexF64 12x512x512 with both C maps ramps, (40,4) -> (41,5), a PlanarKernel that spills in every variant. Cascade Lake, pinned, 3 alternating processes, medians of 7 batches: ComplexF64 12x512x512 392-408 us -> 315-325 us (with the ramp flags cleared: 321-325 us both before and after). Float64 12x512x512, 64^3 in both types and packed 300x64x64 in both types unchanged within noise. Results are bitwise identical on 408 cases; no allocations; TTFX unchanged. Co-Authored-By: Claude Opus 5.5 --- src/QuasiStrided.jl | 4 ++-- src/microkernels/steps.jl | 19 +++++++++++++------ src/packing/unpacked_b.jl | 8 ++++++++ 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/src/QuasiStrided.jl b/src/QuasiStrided.jl index 8c5e23d..b9c6587 100644 --- a/src/QuasiStrided.jl +++ b/src/QuasiStrided.jl @@ -30,10 +30,10 @@ include("packing/line_packing.jl") include("microkernels/interface.jl") include("microkernels/kernels.jl") include("microkernels/vecops.jl") +# B read in place: a B source beside `PackedPanel`, ahead of the K steps' generators. +include("packing/unpacked_b.jl") include("microkernels/steps.jl") include("microkernels/stores.jl") -# B read in place: a B source beside `PackedPanel`, after the kernel interface it extends. -include("packing/unpacked_b.jl") # --- Planning: labels, kernel and blocking choice, the plan --- include("planning/labels.jl") diff --git a/src/microkernels/steps.jl b/src/microkernels/steps.jl index 43eca00..d17d4fa 100644 --- a/src/microkernels/steps.jl +++ b/src/microkernels/steps.jl @@ -41,16 +41,23 @@ function zero_accumulator(kernel::VectorKernel{MR, NR, T, W}) where {MR, NR, T, return ntuple(_ -> z, Val(accumulator_length(layout, MR ÷ rows_per_vector(layout, W), NR))) end -@inline function add_tile( +# Not unrolled over an `UnpackedBView`: B's K stride is a runtime value, so +# every unrolled step needs its own column addresses, which spill. +@generated function add_tile( kernel::VectorKernel, acc::NTuple{NA, Vec{W, R}}, packed_a::PA, packed_b::PB, k_block_length::Int ) where {NA, W, R, PA <: PackedPanel, PB} - k_block_length == 0 && return acc - k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) - @inbounds for p in 1:k_steps(kernel, k_block_length) - acc = accumulate_step(kernel, acc, packed_a, packed_b, p) + no_unroll = PB <: UnpackedBView ? Expr(:loopinfo, (Symbol("llvm.loop.unroll.disable"),)) : nothing + return quote + Base.@_inline_meta + k_block_length == 0 && return acc + k_block_length > 0 || throw_negative_k_block_length(:add_tile, k_block_length) + @inbounds for p in 1:k_steps(kernel, k_block_length) + acc = accumulate_step(kernel, acc, packed_a, packed_b, p) + $no_unroll + end + return acc end - return acc end # B column `j` at K step `p`, as a real or as `(re, im)`; `UnpackedBView` diff --git a/src/packing/unpacked_b.jl b/src/packing/unpacked_b.jl index b34305e..c59b4c2 100644 --- a/src/packing/unpacked_b.jl +++ b/src/packing/unpacked_b.jl @@ -26,6 +26,14 @@ Base.length(u::UnpackedBView) = u.per_k * length(u.ksteps) @inbounds z = u.storage[u.colbase[j] + u.ksteps[p] + 1] return u.transform(z) end +# Column pointer plus one K index shared by all columns: indexed, LLVM's SLP +# vectoriser can turn the `NR` per-column indices into vector arithmetic inside +# the K loop, which costs a vector register and spills accumulators. +@inline function unpacked_b_element(u::UnpackedBView{S, K, NR, F}, j::Int, p::Int) where {S <: DenseVector, K, NR, F} + s = u.storage + z = GC.@preserve s unsafe_load(pointer(s) + u.colbase[j] * sizeof(eltype(s)), @inbounds(u.ksteps[p]) + 1) + return u.transform(z) +end @inline b_scalar(u::UnpackedBView, kernel, j::Int, p::Int) = convert(scalartype(kernel), unpacked_b_element(u, j, p)) From 4aa1dab01f6d3f06295d52fbf663bbabf0f22809 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 6 Oct 2026 12:56:14 -0400 Subject: [PATCH 37/60] Benchmark: same-node A/B of two revisions benchmark/submit_ab.sh runs one benchmark script for two git revisions on the same node, each from a temporary worktree with its own copy of the script, alternating a, b, a, b for QS_AB_ROUNDS rounds (default 2). Outputs go to benchmark/results/