From 7ae96197b82857c93504b33de1d6d8a842d70053 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 25 Oct 2026 04:48:49 +0100 Subject: [PATCH] Make plans a tree instead of pairing them with their nested algorithms A plan recorded the resolved algorithms of the operations it calls (the scan in radix sort and findall, the key range reduction) in a separate `nested` NamedTuple next to their sizes, and a workspace copied every field of the plan. A plan's sizes now hold the plans of the operations it calls, and a workspace holds its plan, device and buffers. A workspace serves a call whose plan is structurally equal to its own (`same_plan`): the same backend, algorithms and sizes down the tree, which is what was compared before. The nested reduction of `ViaReduce` predicates now also has its algorithm compared; it was the only nested operation recorded by its sizes alone. A mismatch says which algorithms or which buffers differ, as before. --- src/findall.jl | 10 ++-- src/predicates.jl | 2 +- src/sort/radix_sort.jl | 5 +- src/workspace.jl | 99 ++++++++++++++++++++++----------------- test/generic/workspace.jl | 16 ++++--- 5 files changed, 74 insertions(+), 58 deletions(-) diff --git a/src/findall.jl b/src/findall.jl index 10dd9274..a3c99dbf 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -280,14 +280,14 @@ function findall_setup(v, items, predicate::Bool, backend, alg) a = resolve_findall(alg, backend, eltype(v)) n = length(v) mask = predicate || !(v isa AbstractArray{Bool}) ? (; mask=scratch(Bool, size(v))) : (;) - counts, nested = if a isa ScanScatter + counts = if a isa ScanScatter num_blocks = cld(n, a.block_size * a.items_per_thread) - scan = scan_setup(+, Int, Int, (num_blocks,), backend; init=0).plan - (; counts=scratch(Int, num_blocks), scan=scan.sizes), (; scan=scan.alg) + (; counts=scratch(Int, num_blocks), + scan=scan_setup(+, Int, Int, (num_blocks,), backend; init=0).plan) else - (; counts=scratch(Int, task_count(n, a.max_tasks, a.min_elems))), (;) + (; counts=scratch(Int, task_count(n, a.max_tasks, a.min_elems))) end - return (; plan=Plan(backend, a, nested, merge(mask, counts))) + return (; plan=Plan(backend, a, merge(mask, counts))) end function findall_run(pred, v, items, s, bufs) diff --git a/src/predicates.jl b/src/predicates.jl index 3405ade8..06a96ee0 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -191,7 +191,7 @@ function predicate_setup(pred, op, init, v, backend, alg) (; flag=scratch(Int8, 1)) elseif a isa ViaReduce (; reduce=mapreduce_setup(BoolResult(pred), op, v, backend, init, nothing, nothing, :, - a.reduce).plan.sizes) + a.reduce).plan) else (;) end diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index 8aa223f9..32339115 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -446,9 +446,8 @@ function radix_plan(a, v::AbstractArray{T}, backend, dims) where T scan = scan_setup(+, UInt32, UInt32, (hist_len,), backend; init=UInt32(0), inclusive=false).plan key_range = first(radix_key_range_setup(v, backend)).plan - return Plan(backend, a, (; scan=scan.alg, key_range=key_range.alg), - (; temp=scratch(T, size(v)), hist=scratch(UInt32, hist_len), scan=scan.sizes, - key_range=key_range.sizes)) + return Plan(backend, a, (; temp=scratch(T, size(v)), hist=scratch(UInt32, hist_len), scan, + key_range)) end diff --git a/src/workspace.jl b/src/workspace.jl index e901d179..3ee8be71 100644 --- a/src/workspace.jl +++ b/src/workspace.jl @@ -11,19 +11,54 @@ end scratch(::Type{T}, dims::Integer...) where {T} = Scratch{T, length(dims)}(Int.(dims)) scratch(::Type{T}, dims::Tuple) where {T} = Scratch{T, length(dims)}(Int.(dims)) -# The public form of the sizes: `(eltype, dims)` pairs +# An operation's scratch: its backend, its resolved algorithm (which the operation runs with), +# and its buffers, as a `NamedTuple` of `Scratch` requirements and of the plans of the +# operations it calls +struct Plan{B <: Backend, A, S <: NamedTuple} + backend::B + alg::A + sizes::S +end + +# `plan(op, args...; kwargs...)` has one method per operation that takes a workspace, with the +# operation's arguments (without `workspace`) +plan(op, args...; kwargs...) = throw(ArgumentError( + "`$op` takes no workspace for arguments of types $(join(Base.map(Core.Typeof, args), ", "))")) + +# The public form of the sizes: `(eltype, dims)` pairs, nested for the operations a plan calls size_pairs(sizes::NamedTuple) = Base.map(size_pairs, sizes) +size_pairs(p::Plan) = size_pairs(p.sizes) size_pairs(b::Scratch{T}) where {T} = (T, b.dims) Base.show(io::IO, b::Scratch) = show(io, size_pairs(b)) +nbytes(sizes::NamedTuple) = Base.sum(nbytes, values(sizes); init=0) +nbytes(p::Plan) = nbytes(p.sizes) +nbytes(b::Scratch{T}) where {T} = Base.prod(b.dims; init=1) * Base.elsize(Array{T}) + +# Whether a workspace made from plan `a` serves a call with plan `b`: the same backend, +# algorithms and buffer sizes, also in the plans of the operations they call. A plan holds +# nothing of the call's arrays, so that a workspace serves any call that needs the same buffers. +same_plan(a::Plan, b::Plan) = + a.backend == b.backend && a.alg == b.alg && same_plan(a.sizes, b.sizes) +same_plan(a::NamedTuple{K}, b::NamedTuple{K}) where {K} = Base.all(Base.map(same_plan, a, b)) +same_plan(a::Scratch, b::Scratch) = a == b +same_plan(a, b) = false + +# The resolved algorithms of a plan, with those of the operations it calls +function algorithms(p::Plan) + nested = (; (k => algorithms(v) for (k, v) in pairs(p.sizes) if v isa Plan)...) + return isempty(nested) ? p.alg : (p.alg, nested) +end + """ Workspace Scratch memory for one operation, made by [`workspace`](@ref) and passed to the operation with its `workspace` keyword. A workspace records the backend and device it was made for, the -algorithm the operation resolved to, and its buffers; a call whose arguments need a different -algorithm or other buffers throws an `ArgumentError` instead of using it. +algorithms the operation and the operations it calls resolved to, and its buffers; a call whose +arguments need a different algorithm or other buffers throws an `ArgumentError` instead of using +it. A workspace holds no state between calls, so it can be reused by any number of calls, but not by calls that may run at the same time (on several tasks or streams). It covers the operations' @@ -32,40 +67,17 @@ per-task bookkeeping, and `Base.sort!` its own scratch for the per-task sorts of `CPUThreads.SampleSort` and for the slices of a sort along `dims`; and before Julia 1.12, a reduction of several arrays or of a `Broadcasted` object materializes it first. """ -struct Workspace{B <: Backend, A, N <: NamedTuple, S <: NamedTuple, T <: NamedTuple} - backend::B +struct Workspace{P <: Plan, T <: NamedTuple} + plan::P device::Int - alg::A - nested::N - sizes::S buffers::T end function Base.show(io::IO, ws::Workspace) - print(io, "Workspace(", backend_name(ws.backend), ", ", algname(ws.alg), ", ", - Base.format_bytes(nbytes(ws.sizes)), ")") + print(io, "Workspace(", backend_name(ws.plan.backend), ", ", algname(ws.plan.alg), ", ", + Base.format_bytes(nbytes(ws.plan)), ")") end -nbytes(sizes::NamedTuple) = Base.sum(nbytes, values(sizes); init=0) -nbytes(b::Scratch{T}) where {T} = Base.prod(b.dims; init=1) * Base.elsize(Array{T}) - - -# An operation's scratch: its backend, its resolved algorithm (which the operation runs with), -# the resolved algorithms of the operations it calls, and its buffers, as a `NamedTuple` of -# `(eltype, dims)` requirements, or of the `NamedTuple`s of nested operations. -struct Plan{B <: Backend, A, N <: NamedTuple, S <: NamedTuple} - backend::B - alg::A - nested::N - sizes::S -end -Plan(backend, alg, sizes::NamedTuple) = Plan(backend, alg, (;), sizes) - -# `plan(op, args...; kwargs...)` has one method per operation that takes a workspace, with the -# operation's arguments (without `workspace`) -plan(op, args...; kwargs...) = throw(ArgumentError( - "`$op` takes no workspace for arguments of types $(join(Base.map(Core.Typeof, args), ", "))")) - """ @@ -85,7 +97,7 @@ julia> AK.workspace_size(AK.sum, CuArray(rand(Float32, 10_000))) (partials = (Float32, (4,)),) ``` """ -workspace_size(op, args...; kwargs...) = size_pairs(plan(op, args...; kwargs...).sizes) +workspace_size(op, args...; kwargs...) = size_pairs(plan(op, args...; kwargs...)) """ workspace(op, args...; kwargs...) -> Workspace @@ -115,31 +127,34 @@ operation's arrays. """ function workspace(op, args...; kwargs...) p = plan(op, args...; kwargs...) - return Workspace(p.backend, KernelAbstractions.device(p.backend), p.alg, p.nested, p.sizes, - allocate_scratch(p.backend, p.sizes)) + return Workspace(p, KernelAbstractions.device(p.backend), allocate_scratch(p.backend, p)) end allocate_scratch(backend, sizes::NamedTuple) = Base.map(s -> allocate_scratch(backend, s), sizes) +allocate_scratch(backend, p::Plan) = allocate_scratch(backend, p.sizes) allocate_scratch(backend, b::Scratch{T}) where {T} = KernelAbstractions.allocate(backend, T, b.dims) # The buffers of the plan: the workspace's, checked against the plan and the operation's arrays, # or new ones -scratch_buffers(p::Plan, ::Nothing, arrays...) = allocate_scratch(p.backend, p.sizes) +scratch_buffers(p::Plan, ::Nothing, arrays...) = allocate_scratch(p.backend, p) function scratch_buffers(p::Plan, ws::Workspace, arrays...) - ws.backend == p.backend && ws.device == KernelAbstractions.device(p.backend) || - throw(ArgumentError("the workspace was made for another backend or device")) - ws.alg == p.alg && ws.nested == p.nested || throw(ArgumentError( - "the workspace was made for $(ws.alg) $(ws.nested), but this call resolves to " * - "$(p.alg) $(p.nested)")) - ws.sizes == p.sizes || throw(ArgumentError( - "the workspace holds the buffers $(size_pairs(ws.sizes)), but this call needs " * - "$(size_pairs(p.sizes))")) + ws.device == KernelAbstractions.device(p.backend) && same_plan(ws.plan, p) || + throw(plan_mismatch(ws, p)) for b in leaves(ws.buffers), a in arrays mayalias(b, a) && throw(ArgumentError("the workspace must not alias the operation's arrays")) end return ws.buffers end +function plan_mismatch(ws::Workspace, p::Plan) + ws.plan.backend == p.backend && ws.device == KernelAbstractions.device(p.backend) || + return ArgumentError("the workspace was made for another backend or device") + algorithms(ws.plan) == algorithms(p) || return ArgumentError( + "the workspace was made for $(algorithms(ws.plan)), but this call resolves to " * + "$(algorithms(p))") + return ArgumentError("the workspace holds the buffers $(size_pairs(ws.plan)), but this " * + "call needs $(size_pairs(p))") +end scratch_buffers(p::Plan, ws, arrays...) = throw(ArgumentError("`workspace` must be a `Workspace` from `workspace`, got $(typeof(ws))")) diff --git a/test/generic/workspace.jl b/test/generic/workspace.jl index 59d77930..afe6173d 100644 --- a/test/generic/workspace.jl +++ b/test/generic/workspace.jl @@ -5,9 +5,8 @@ buffer_sizes(buffers::NamedTuple) = map(buffer_sizes, buffers) buffer_sizes(buffer::AbstractArray) = (eltype(buffer), size(buffer)) -# `ws` with other buffers or nested algorithms, which the operations must reject -forge(ws; buffers=ws.buffers, nested=ws.nested) = - AK.Workspace(ws.backend, ws.device, ws.alg, nested, ws.sizes, buffers) +# `ws` with other buffers or another plan, which the operations must reject +forge(ws; buffers=ws.buffers, plan=ws.plan) = AK.Workspace(plan, ws.device, buffers) @testset "workspace: sizes" begin v = array_from_host(rand(Float32, 10_000)) @@ -197,12 +196,15 @@ end bools = array_from_host(rand(Bool, 10_000)) alg = AK.ScanScatter() ws = AK.workspace(AK.findall, bools; alg) - @test haskey(ws.nested, :scan) - other = forge(ws; nested=(; ws.nested..., scan=AK.ScanPrefixes(16, 1))) - @test_throws ArgumentError AK.findall(bools; alg, workspace=other) + p = ws.plan + @test p.sizes.scan isa AK.Plan + scan = AK.Plan(p.backend, AK.ScanPrefixes(16, 1), p.sizes.scan.sizes) + other = forge(ws; plan=AK.Plan(p.backend, p.alg, (; p.sizes..., scan))) + err = try AK.findall(bools; alg, workspace=other) catch e; e end + @test err isa ArgumentError && occursin("ScanPrefixes(16, 1)", err.msg) @test Array(AK.findall(bools; alg, workspace=ws)) == findall(Array(bools)) ws = AK.workspace(AK.sort!, v; alg=AK.RadixSort()) - @test Base.all(k -> haskey(ws.nested, k), (:scan, :key_range)) + @test Base.all(k -> ws.plan.sizes[k] isa AK.Plan, (:scan, :key_range)) else # A workspace of reference elements, on the host ws = AK.workspace(AK.sort!, fill("a", 100); alg=AK.CPUThreads.SampleSort(max_tasks=4))