From 5502b40cd494a953838fa2050e2e74e34ccab991 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 21 Oct 2025 14:48:46 -0700 Subject: [PATCH 01/43] DArray: Use Finch for sparse arrays --- Project.toml | 6 ++- ext/FinchExt.jl | 49 ++++++++++++++++++++++++ ext/SparseArraysExt.jl | 47 +++++++++++++++++++++++ src/Dagger.jl | 3 +- src/array/alloc.jl | 14 ------- src/array/mul.jl | 71 ++++++++++++++++++++++++++++++++++ src/array/sparse.jl | 72 +++++++++++++++++++++++++++++++++++ src/array/sparse_partition.jl | 37 ------------------ src/datadeps/remainders.jl | 14 ------- src/file-io.jl | 1 - 10 files changed, 245 insertions(+), 69 deletions(-) create mode 100644 ext/FinchExt.jl create mode 100644 ext/SparseArraysExt.jl create mode 100644 src/array/sparse.jl delete mode 100644 src/array/sparse_partition.jl diff --git a/Project.toml b/Project.toml index 21eb4a00d..62ed8252c 100644 --- a/Project.toml +++ b/Project.toml @@ -23,7 +23,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Requires = "ae029012-a4dd-5104-9daa-d747884805df" ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34" @@ -37,6 +36,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" @@ -45,12 +45,14 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] AbstractFFTsExt = "AbstractFFTs" CUDAExt = "CUDA" DistributionsExt = "Distributions" +FinchExt = "Finch" GraphVizExt = "GraphViz" GraphVizSimpleExt = "Colors" IntelExt = "oneAPI" @@ -62,6 +64,7 @@ OpenCLExt = "OpenCL" PlotsExt = ["DataFrames", "Plots"] PythonExt = "PythonCall" ROCExt = "AMDGPU" +SparseArraysExt = "SparseArrays" [compat] AMDGPU = "1, 2" @@ -74,6 +77,7 @@ DataStructures = "0.18, 0.19" DistributedNext = "1.0.0" Distributions = "0.25" FillArrays = "1.13.0" +Finch = "1.2.9" GPUArraysCore = "0.2.0" GraphViz = "0.2" Graphs = "1" diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl new file mode 100644 index 000000000..3d51937d2 --- /dev/null +++ b/ext/FinchExt.jl @@ -0,0 +1,49 @@ +module FinchExt + +import Finch +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix + +Dagger.sparse_mode(::Finch.Tensor) = :finch +Dagger._sparse_alloc(::Val{:finch}, T::Type, dims::Dims) = + Finch.fspzeros(T, dims...) +Dagger._sparse_collect(A::Finch.Tensor) = Array(A) + +function Finch.fsprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) + return Dagger._to_darray(a) +end +Finch.fsprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims, sparsity; assignment) +Finch.fsprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Finch.Tensor, + B::Finch.Tensor, + alpha, + beta +) + if !isa(C.mat, Finch.Tensor) + # Not supported here, forward to generic matmatmul! + return Dagger.matmatmul!(C.mat, transA, transB, A, B, alpha, beta) + end + + # Use optimized Finch operations + Cm = C.mat + Finch.@einsum Cm[i,j] += A[i,k] * B[k,j] + display(Cm) + C.mat = Cm + + return C +end + +end # module FinchExt \ No newline at end of file diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl new file mode 100644 index 000000000..42af139e8 --- /dev/null +++ b/ext/SparseArraysExt.jl @@ -0,0 +1,47 @@ +module SparseArraysExt + +import SparseArrays +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix + +Dagger.sparse_mode(::SparseArrays.SparseMatrixCSC) = :sparsearrays +Dagger._sparse_alloc(::Val{:sparsearrays}, T::Type, dims::Dims) = + SparseArrays.spzeros(T, dims...) + +function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(SparseArrays.sprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) + return Dagger._to_darray(a) +end +SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims, sparsity; assignment) +SparseArrays.sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) + +_apply_trans(X, t::Char) = + t == 'N' ? X : + t == 'T' ? transpose(X) : + t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::SparseArrays.SparseMatrixCSC, + B::SparseArrays.SparseMatrixCSC, + alpha, + beta +) + # Use fallback implementation + # TODO: Optimize this further + C.mat = alpha * A * B + beta * C.mat + + return C +end + +end # module SparseArraysExt diff --git a/src/Dagger.jl b/src/Dagger.jl index 81881716e..bc80217ac 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -2,7 +2,6 @@ module Dagger import Serialization import Serialization: AbstractSerializer, serialize, deserialize -import SparseArrays: sprand, SparseMatrixCSC import MemPool import MemPool: DRef, FileRef, poolget, poolset @@ -132,7 +131,7 @@ include("array/operators.jl") include("array/indexing.jl") include("array/setindex.jl") include("array/matrix.jl") -include("array/sparse_partition.jl") +include("array/sparse.jl") include("array/sort.jl") include("array/permute.jl") include("array/linalg.jl") diff --git a/src/array/alloc.jl b/src/array/alloc.jl index 93862d545..e90b375e2 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -86,20 +86,6 @@ Base.randn(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) Base.randn(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = randn(auto_blocks(dims), T, dims; assignment) -function sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) - d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, (T, _dims) -> sprand(T, _dims..., sparsity), false, d, partition(p, d), p, assignment) - return _to_darray(a) -end -sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims, sparsity; assignment) -sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(auto_blocks(dims), T, dims, sparsity; assignment) - function Base.ones(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = ArrayDomain(map(x->1:x, dims)) a = AllocateArray(T, ones, false, d, partition(p, d), p, assignment) diff --git a/src/array/mul.jl b/src/array/mul.jl index 14ed9f0ba..6f36d2747 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -1,3 +1,74 @@ +""" + matmatmul!(C, transA::Char, transB::Char, A, B, alpha, beta) + +A general-purpose matrix-matrix multiply, like `LinearAlgebra.generic_matmatmul!`, +but with extra functionality. May internally convert `A` and `B` to a type that +better matches `C` and provides optimal portability and, when possible, +better performance. The actual matrix multiply operation should happen in +`LinearAlgebra.generic_matmatmul!` or an equivalent call. + +The following automatic conversions are performed: +- If no `LinearAlgebra.generic_matmatmul!` method is available, convert `A` and `B` to dense Array-like +- If `C` is a `DSparseMatrix`, perform the operation out-of-place and then update `C` in-place +""" +function matmatmul!( + C, + transA::Char, + transB::Char, + A, + B, + alpha, + beta +) + EC = eltype(C) + EA = eltype(A) + EB = eltype(B) + + TC = typeof(C) + TA = typeof(A) + TB = typeof(B) + + mam = LinearAlgebra.MulAddMul(alpha, beta) + + # Check if C doesn't support in-place operations (e.g. DSparseMatrix) + # We'll get here if A and B don't have equivalent types + if isa(C, DSparseMatrix) + C.mat = alpha * A * B + beta * C.mat + return C + end + + # Check if the call will fail due to MethodError + sig = Tuple{TC, Char, Char, TA, TB, typeof(mam)} + if !hasmethod(LinearAlgebra.generic_matmatmul!, sig) + # Convert to Array-like + # FIXME: GPU support + C_new = C + A_new = collect(A) + B_new = collect(B) + alpha_new = alpha + beta_new = beta + # FIXME: Re-check hasmethod, and if no method, then convert and bounce C + @goto ready + end + + C_new = C + A_new = A + B_new = B + alpha_new = alpha + beta_new = beta + + @label ready + mam_new = LinearAlgebra.MulAddMul(alpha_new, beta_new) + return LinearAlgebra.generic_matmatmul!( + C_new, + transA, + transB, + A_new, + B_new, + mam_new + ) +end + function LinearAlgebra.generic_matmatmul!( C::DMatrix{T}, transA::Char, diff --git a/src/array/sparse.jl b/src/array/sparse.jl new file mode 100644 index 000000000..8f61a88b0 --- /dev/null +++ b/src/array/sparse.jl @@ -0,0 +1,72 @@ +""" + DSparseMatrix{T} <: AbstractMatrix{T} + +A sparse matrix container, for which the contained matrix may be replaced with +a new one to support in-place operations. Designed to work well with +Datadeps algorithms. +""" +mutable struct DSparseMatrix{T} <: AbstractMatrix{T} + mat +end +function DSparseMatrix{T}(undef, dims::NTuple{N, Int}) where {T,N} + M = sparse_mode() + return DSparseMatrix{T}(_sparse_alloc(Val(M), T, dims)) +end + +function _sparse_not_loaded(::Val{M}) where M + if M == :sparsearrays + throw(ArgumentError("SparseArrays must be loaded to use SparseMatrixCSC")) + elseif M == :finch + throw(ArgumentError("Finch must be loaded to use Finch.Tensor")) + elseif M == :none + throw(ArgumentError("Sparse mode not set\nSet it with `sparse_mode!(M)` where M is :sparsearrays or :finch, and load the corresponding package")) + else + throw(ArgumentError("Unknown sparse mode $M\nOptions are :sparsearrays and :finch")) + end +end +_sparse_alloc(::Val{M}, T::Type, dims::Dims) where M = _sparse_not_loaded(Val(M)) +_sparse_collect(M) = collect(M) +const SPARSE_MODE = TaskLocalValue{Symbol}(()->:none) +sparse_mode() = SPARSE_MODE[] +sparse_mode(::T) where T = error("Unknown sparse mode for type $T") +set_sparse_mode!(mode::Symbol) = SPARSE_MODE[] = mode + +Base.eltype(M::DSparseMatrix{T}) where T = T +Base.size(M::DSparseMatrix) = size(M.mat) +Base.ndims(M::DSparseMatrix) = 2 +Base.iterate(M::DSparseMatrix) = iterate(M.mat) +Base.iterate(M::DSparseMatrix, state) = iterate(M.mat, state) +Base.similar(M::DSparseMatrix, ::Type{T}, dims::Tuple{Int, Int}) where T = + DSparseMatrix{T}(similar(M.mat, T, dims)) +Base.collect(M::DSparseMatrix) = _sparse_collect(M.mat) + +# N.B. hash and aliasing shouldn't change even if M.mat changes +Base.hash(M::DSparseMatrix, h::UInt) = hash(objectid(M), hash(DSparseMatrix, h)) +aliasing(M::DSparseMatrix, _=identity) = + ObjectAliasing(M) + +Base.getindex(M::DSparseMatrix{T}, I::Int) where T = + getindex(M.mat, I) +Base.getindex(M::DSparseMatrix{T}, I::Vararg{Int,N}) where {T,N} = + getindex(M.mat, I...) +Base.setindex!(M::DSparseMatrix{T}, v, I::Int) where T = + setindex!(m.mat, v, I) +Base.setindex!(M::DSparseMatrix{T}, v, I::Vararg{Int,N}) where {T,N} = + setindex!(M.mat, v, I...) + +#Base.BroadcastStyle(::Type{<:DSparseMatrix}) = FinchStyle() + +LinearAlgebra.norm2(M::DSparseMatrix) = LinearAlgebra.norm2(M.mat) + +function matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::DSparseMatrix, + B::DSparseMatrix, + alpha, + beta +) + # Forward to sparse matmatmul!, but allow in-place update of C + return matmatmul!(C, transA, transB, A.mat, B.mat, alpha, beta) +end \ No newline at end of file diff --git a/src/array/sparse_partition.jl b/src/array/sparse_partition.jl deleted file mode 100644 index 26efe3e6a..000000000 --- a/src/array/sparse_partition.jl +++ /dev/null @@ -1,37 +0,0 @@ -export partition_sparse - -function partition_sparse(colptr, nnz, sz, nparts) - nnz_per_chunk = nnz/nparts - dom = ArrayDomain(map(x->1:x, sz)) - - nxt = nnz_per_chunk - colstart = 1 - splits = UnitRange[] - - for i=1:nparts-1 - idxs = searchsorted(colptr, nxt) - if length(idxs) == 0 - # the exact index was not found. latch on to the nearest column boundary - idx = first(idxs) - if abs(colptr[idx]-nxt) < abs(colptr[idx-1]-nxt) - nextidx = idx - else - nextidx = idx-1 - end - else - nextidx = last(idxs) - end - push!(splits, colstart:nextidx) - colstart=nextidx+1 - nxt=nxt + nnz_per_chunk - end - push!(splits, colstart:(length(colptr)-1)) - - - cumlength = cumsum(map(length, splits)) - subdomains = DomainBlocks((1,1), ([size(dom, 1)], cumlength)) - DomainSplit(dom, subdomains) -end - -partition_sparse(S::SparseMatrixCSC, nparts) = - partition_sparse(S.colptr, length(S.nzval), size(S), nparts) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 1000471b2..a5885a579 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -548,17 +548,3 @@ function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::Base write_remainder!(copies, copies_offset, to[], to_ptr, n) end end - -function find_object_holding_ptr(A::SparseMatrixCSC, ptr::UInt64) - span = LocalMemorySpan(pointer(A.nzval), length(A.nzval)*sizeof(eltype(A.nzval))) - if span_start(span) <= ptr <= span_end(span) - return A.nzval - end - span = LocalMemorySpan(pointer(A.colptr), length(A.colptr)*sizeof(eltype(A.colptr))) - if span_start(span) <= ptr <= span_end(span) - return A.colptr - end - span = LocalMemorySpan(pointer(A.rowval), length(A.rowval)*sizeof(eltype(A.rowval))) - @assert span_start(span) <= ptr <= span_end(span) "Pointer $ptr not found in SparseMatrixCSC" - return A.rowval -end diff --git a/src/file-io.jl b/src/file-io.jl index 1407fbced..15405d049 100644 --- a/src/file-io.jl +++ b/src/file-io.jl @@ -1,5 +1,4 @@ export save, load -using SparseArrays import MemPool: GenericFileDevice, FileRef struct File From 6dc624eb34966e8b411faaa999c477f92553405c Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 21 Apr 2024 21:23:13 -0400 Subject: [PATCH 02/43] DArray: Add matmul sparse support --- ext/FinchExt.jl | 42 ++++++++- ext/SparseArraysExt.jl | 11 ++- src/array/mul.jl | 168 +++++++++++++----------------------- src/array/sparse.jl | 10 +++ test/array/linalg/matmul.jl | 35 ++++++++ 5 files changed, 150 insertions(+), 116 deletions(-) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index 3d51937d2..9bc99d1ce 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -39,11 +39,47 @@ function Dagger.matmatmul!( # Use optimized Finch operations Cm = C.mat - Finch.@einsum Cm[i,j] += A[i,k] * B[k,j] - display(Cm) - C.mat = Cm + C_out = Finch.fspzeros(eltype(Cm), size(Cm)...) + # N.B. BEWARE: @einsum doesn't work correctly with Cm[i,j] = ... assignments. + #Finch.@einsum Cm[i,j] = alpha * (A[i,k] * B[k,j]) + beta * Cm[i,j] + Finch.@einsum C_out[i,j] += alpha * (A[i,k] * B[k,j]) + if beta != 0 + C_out += beta * Cm + end + # TODO: Remove this once we're sure @einsum works correctly + @assert Array(C_out) ≈ (alpha * Array(A) * Array(B)) .+ (beta * Array(Cm)) + C.mat = C_out return C end +function Dagger.transpose_tile(B::Finch.Tensor, uplo=nothing) + if uplo == 'U' + Finch.@finch begin + C .= 0 + for i in _, j in _ + if i <= j + C[i, j] = B[j, i] + end + end + return C + end + return C + elseif uplo == 'L' + Finch.@finch begin + C .= 0 + for i in _, j in _ + if i >= j + C[i, j] = B[j, i] + end + end + return C + end + return C + else + Finch.@einsum C[i,j] = B[j,i] + return C + end +end + end # module FinchExt \ No newline at end of file diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 42af139e8..b03b53f33 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -1,10 +1,11 @@ module SparseArraysExt import SparseArrays +import SparseArrays: SparseMatrixCSC import Dagger import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix -Dagger.sparse_mode(::SparseArrays.SparseMatrixCSC) = :sparsearrays +Dagger.sparse_mode(::SparseMatrixCSC) = :sparsearrays Dagger._sparse_alloc(::Val{:sparsearrays}, T::Type, dims::Dims) = SparseArrays.spzeros(T, dims...) @@ -32,8 +33,8 @@ function Dagger.matmatmul!( C::DSparseMatrix, transA::Char, transB::Char, - A::SparseArrays.SparseMatrixCSC, - B::SparseArrays.SparseMatrixCSC, + A::SparseMatrixCSC, + B::SparseMatrixCSC, alpha, beta ) @@ -44,4 +45,8 @@ function Dagger.matmatmul!( return C end +function Dagger.transpose_tile(B::SparseMatrixCSC) + return SparseArrays.sparse(B') +end + end # module SparseArraysExt diff --git a/src/array/mul.jl b/src/array/mul.jl index 6f36d2747..c74c19128 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -112,7 +112,8 @@ function LinearAlgebra.generic_matmatmul!( return gemm_dagger!(C, transA, transB, A, B, alpha, beta) end end -function _repartition_matmatmul(C, A, B, transA::Char, transB::Char)::Tuple{Blocks{2}, Blocks{2}, Blocks{2}} +# FIXME: Mixed-precision methods +function _repartition_matmatmul(C, A, B, transA::Char, transB::Char) partA = A.partitioning.blocksize partB = B.partitioning.blocksize istransA = transA == 'T' || transA == 'C' @@ -212,28 +213,28 @@ function gemm_dagger!( # A: NoTrans / B: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: NoTrans / B: [Conj]Trans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -242,28 +243,28 @@ function gemm_dagger!( # A: [Conj]Trans / B: NoTrans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: [Conj]Trans / B: [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -311,6 +312,7 @@ function syrk_dagger!( iscomplex = T <: Complex transs = iscomplex ? 'C' : 'T' + anti_transs = trans == 'N' ? transs : 'N' Dagger.spawn_datadeps() do for n in range(1, Cnt) @@ -318,116 +320,63 @@ function syrk_dagger!( # NoTrans for k in range(1, Ant) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + trans, + anti_transs, + In(Ac[n, k]), + In(Ac[n, k]), + _alpha, + mzone, + ) + end + # NoTrans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Ant) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), trans, - real(alpha), + transs, In(Ac[n, k]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[m, k]), alpha, - In(Ac[n, k]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # NoTrans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[m, k]), - In(Ac[n, k]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # NoTrans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[n, k]), - In(Ac[m, k]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end else # [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + _trans = iscomplex ? transs : trans + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + _trans, + anti_transs, + In(Ac[k, n]), + In(Ac[k, n]), + _alpha, + mzone, + ) + end + # [Conj]Trans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Amt) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), transs, - real(alpha), + 'N', In(Ac[k, n]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[k, m]), alpha, - In(Ac[k, n]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # [Conj]Trans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, m]), - In(Ac[k, n]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # [Conj]Trans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, n]), - In(Ac[k, m]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end end end end @@ -469,7 +418,7 @@ end return A end -@inline function copytile!(A::AbstractMatrix{T}, B::AbstractMatrix{T})::Nothing where {T} +function copytile!(A, B) m, n = size(A) C = B' @@ -479,15 +428,14 @@ end return nothing end -@inline function copydiagtile!(A::AbstractMatrix{T}, uplo::AbstractChar)::Nothing where {T} +function copydiagtile!(A, uplo) m, n = size(A) - Acpy = copy(A) if uplo == 'U' - C = UpperTriangular(Acpy)' + UpperTriangular(Acpy) + C = UpperTriangular(A)' + UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] elseif uplo == 'L' - C = LowerTriangular(Acpy)' + Acpy - UpperTriangular(Acpy) + C = LowerTriangular(A)' + A - UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] end diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 8f61a88b0..ad3ec79f3 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -69,4 +69,14 @@ function matmatmul!( ) # Forward to sparse matmatmul!, but allow in-place update of C return matmatmul!(C, transA, transB, A.mat, B.mat, alpha, beta) +end + +function transpose_tile end +function copytile!(A::DSparseMatrix, B::DSparseMatrix) + A.mat = transpose_tile(B.mat) + return +end +function copydiagtile!(A::DSparseMatrix, uplo) + A.mat = transpose_tile(A.mat, uplo) + return end \ No newline at end of file diff --git a/test/array/linalg/matmul.jl b/test/array/linalg/matmul.jl index 26fd0afe3..f0ec73287 100644 --- a/test/array/linalg/matmul.jl +++ b/test/array/linalg/matmul.jl @@ -29,36 +29,71 @@ function test_gemm!(T, szA, szB, partA, partB) DA = distribute(A, partA) DB = distribute(B, partB) + SA = sprand(T, szA..., 0.1) + SB = sprand(T, szA..., 0.1) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + ## Out-of-place gemm # No transA, No transB + # Dense DC = DA * DB C = A * B @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB + SC = SA * SB + @test collect(DSC) ≈ SC if szA == szB # No transA, transB + # Dense DC = DA * DB' C = A * B' @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB' + SC = SA * SB' + @test collect(DSC) ≈ SC # transA, No transB + # Dense DC = DA' * DB C = A' * B @test collect(DC) ≈ C + # Sparse + DSC = DSA' * DSB + SC = SA' * SB + @test collect(DSC) ≈ SC end # transA, transB + # Dense DC = DA' * DB' C = A' * B' @test collect(DC) ≈ C + #= Sparse + DSC = DSA' * DSB' + SC = SA' * SB' + @test collect(DSC) ≈ SC + =# ## In-place gemm # No transA, No transB + # Dense C = zeros(T, szC...) DC = distribute(C, partC) mul!(C, A, B) mul!(DC, DA, DB) @test collect(DC) ≈ C + #= Sparse + SC = zeros(T, szC...) + DSC = distribute(SC, partC) + mul!(SC, SA, SB) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC + =# if szA == szB # No transA, transB From 6472128400b7d85303d362955775279b3c96aafb Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 28 Oct 2025 12:00:28 -0700 Subject: [PATCH 03/43] DArray: Add spzeros/fspzeros --- ext/FinchExt.jl | 14 ++++++++++++++ ext/SparseArraysExt.jl | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index 9bc99d1ce..354c0b2c7 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -9,6 +9,20 @@ Dagger._sparse_alloc(::Val{:finch}, T::Type, dims::Dims) = Finch.fspzeros(T, dims...) Dagger._sparse_collect(A::Finch.Tensor) = Array(A) +function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment) + return Dagger._to_darray(a) +end +Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, T, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) +Finch.fspzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(Dagger.auto_blocks(dims), T, dims; assignment) + function Finch.fsprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index b03b53f33..ca291d63c 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -9,6 +9,20 @@ Dagger.sparse_mode(::SparseMatrixCSC) = :sparsearrays Dagger._sparse_alloc(::Val{:sparsearrays}, T::Type, dims::Dims) = SparseArrays.spzeros(T, dims...) +function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(SparseArrays.spzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment) + return Dagger._to_darray(a) +end +SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, T, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) +SparseArrays.spzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(Dagger.auto_blocks(dims), T, dims; assignment) + function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(SparseArrays.sprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) From 7a9aae42a0f25d94318e214e95ee06b73d2706f9 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 26 Nov 2025 11:32:49 -0700 Subject: [PATCH 04/43] sparse: Fix similar allocation --- src/array/alloc.jl | 33 +++++++++++++++++++++++++-------- src/array/darray.jl | 12 ++++++++++-- src/array/sparse.jl | 2 +- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/src/array/alloc.jl b/src/array/alloc.jl index e90b375e2..9710b78b9 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -10,13 +10,20 @@ mutable struct AllocateArray{T,N} <: ArrayOp{T,N} domainchunks partitioning::AbstractBlocks{N} procgrid::Union{AbstractProcGrid{N}, Nothing} + # Optional concrete tile type for staged `@spawn`s. When `nothing`, leave + # `return_type` unset so inference (and MPI `post_stage_array_chunks!`, + # which defaults dense tiles to `Array{T,N}`) decide. Sparse allocators + # pass e.g. `DSparseArray{T,N}` so MPI does not overwrite with `Array`. + return_type::Union{Type,Nothing} - function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, domainchunks, p::AbstractBlocks{N}, assignment::Union{AssignmentType{N},Nothing} = nothing) where {T,N} + function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, + domainchunks, p::AbstractBlocks{N}, + assignment::Union{AssignmentType{N},Nothing} = nothing; + return_type::Union{Type,Nothing} = nothing) where {T,N} sizeA = map(length, d.indexes) procgrid = build_procgrid(something(assignment, :arbitrary), Tuple(sizeA), p.blocksize, current_acceleration()) - return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid) + return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid, return_type) end - end size(a::AllocateArray) = size(a.domain) @@ -45,12 +52,22 @@ function stage(ctx, A::AllocateArray) tasks = emit_chunk_tasks!(A.domainchunks, A.procgrid, A.eltype, (scope, I, i) -> begin x = A.domainchunks[I] - N = ndims(A.domainchunks) - ret_type = Array{A.eltype, N} - if A.want_index - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, i, size(x)) + if A.f isa DArray + # `similar(::DArray)` carries the source tile type (e.g. sparse) forward. + chunk = A.f.chunks[I] + Dagger.@spawn compute_scope=scope similar(chunk, A.eltype, size(x)) + elseif A.want_index + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, i, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, i, size(x)) + end else - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, size(x)) + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, size(x)) + end end end) return DArray(A.eltype, A.domain, A.domainchunks, tasks, A.partitioning) diff --git a/src/array/darray.jl b/src/array/darray.jl index bd0e5efc5..72e2208a5 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -342,8 +342,16 @@ aliasing(x::DArray) = memory_space(x::DArray) = throw(ConcurrencyViolationError("DArray memory spaces may be mixed and unstable")) -Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} = - DArray{S,N}(undef, D.partitioning, dims) +function Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} + if dims == size(D) + domain = ArrayDomain(map(x->1:x, dims)) + subdomains = partition(D.partitioning, domain) + a = AllocateArray(S, D, false, domain, subdomains, D.partitioning, nothing) + return _to_darray(a) + else + return DArray{S,N}(undef, D.partitioning, dims) + end +end Base.similar(D::DArray{T,N1} where T, ::Type{S}, dims::Dims{N2}) where {S,N1,N2} = DArray{S,N2}(undef, auto_blocks(dims), dims) diff --git a/src/array/sparse.jl b/src/array/sparse.jl index ad3ec79f3..980b43e4f 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -37,7 +37,7 @@ Base.ndims(M::DSparseMatrix) = 2 Base.iterate(M::DSparseMatrix) = iterate(M.mat) Base.iterate(M::DSparseMatrix, state) = iterate(M.mat, state) Base.similar(M::DSparseMatrix, ::Type{T}, dims::Tuple{Int, Int}) where T = - DSparseMatrix{T}(similar(M.mat, T, dims)) + DSparseMatrix{T}(_sparse_alloc(Val(sparse_mode(M.mat)), T, dims)) Base.collect(M::DSparseMatrix) = _sparse_collect(M.mat) # N.B. hash and aliasing shouldn't change even if M.mat changes From beea74a1b6cd0ccb543bb791695fbfd2a84b5ec5 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 3 Jun 2026 08:12:45 -0700 Subject: [PATCH 05/43] test: Allow --test with directory --- test/runtests.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index bfc347ce0..799a05c2f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -109,7 +109,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ parsed_args = parse_args(s) to_test = String[] if isempty(parsed_args["test"]) - to_test = copy(all_test_names) + to_test = filter(!in(optin_test_names), copy(all_test_names)) else for _test in parsed_args["test"] test = only(_test) @@ -166,7 +166,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.offline(true) end else - to_test = all_test_names + to_test = filter(!in(optin_test_names), all_test_names) @info "Running all tests" end From 057f002a7e3fe980157e9ae9a7690ed81c4e1d3a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 3 Jun 2026 08:14:24 -0700 Subject: [PATCH 06/43] DArray: Add view(A, AutoBlocks()) and fill! --- src/array/alloc.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/array/alloc.jl b/src/array/alloc.jl index 9710b78b9..6ffbc8f50 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -170,3 +170,14 @@ function unsafe_free!(A::DArray) end end end + +# Initializers + +function Base.fill!(A::DArray, x) + spawn_datadeps() do + for chunk in A.chunks + Dagger.@spawn fill!(chunk, x) + end + end + return A +end From 286f903b17f99f32e91e368b0122c5bfd4022a30 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 4 Jun 2026 12:51:49 -0700 Subject: [PATCH 07/43] DArray/sparse: Improve copy/indexing and aliasing --- ext/FinchExt.jl | 12 ++++++ ext/SparseArraysExt.jl | 4 ++ src/array/copy.jl | 10 +++++ src/array/darray.jl | 11 ++++-- src/array/sparse.jl | 87 ++++++++++++++++++++++++++++++++++++------ 5 files changed, 110 insertions(+), 14 deletions(-) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index 354c0b2c7..d534b49d1 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -8,6 +8,7 @@ Dagger.sparse_mode(::Finch.Tensor) = :finch Dagger._sparse_alloc(::Val{:finch}, T::Type, dims::Dims) = Finch.fspzeros(T, dims...) Dagger._sparse_collect(A::Finch.Tensor) = Array(A) +Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseMatrix{eltype(x)}(x) function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) @@ -37,6 +38,17 @@ Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment:: Finch.fsprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = Finch.fsprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) +# Finch tensors do not define `Base.copy`; build a fresh equivalent tensor. +Dagger._sparse_copy(mat::Finch.Tensor) = _finch_materialize(mat, eltype(mat), size(mat)) + +# Finch tensors do not support `setindex!`, so copy-buffering writeback into a +# Finch tile densifies, applies the (partial-range) copy, then re-sparsifies. +function Dagger._sparse_copyto_view!(mat::Finch.Tensor, Brange, src) + dense = Array(mat) + copyto!(view(dense, Brange), src) + return _finch_materialize(dense, eltype(mat), size(mat)) +end + function Dagger.matmatmul!( C::DSparseMatrix, transA::Char, diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index ca291d63c..45f7d8524 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -8,6 +8,10 @@ import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix Dagger.sparse_mode(::SparseMatrixCSC) = :sparsearrays Dagger._sparse_alloc(::Val{:sparsearrays}, T::Type, dims::Dims) = SparseArrays.spzeros(T, dims...) +# Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. +Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) +# Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. +Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseMatrix{eltype(x)}(x) function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) diff --git a/src/array/copy.jl b/src/array/copy.jl index 7d92566e6..422bec7c9 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -154,3 +154,13 @@ StridedDArray{T,N} = Union{<:DArray{T,N}, SubArray{T,N,<:DArray{T,NP}} where NP} Base.copyto!(B::StridedDArray, A::StridedDArray) = darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +function Base.copyto!(B::Array, A::StridedDArray) + DB = view(B, AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(DB), parentindices(A)) + return B +end +function Base.copyto!(B::SubArray, A::StridedDArray) + DB = view(parent(B), AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(B), parentindices(A)) + return B +end diff --git a/src/array/darray.jl b/src/array/darray.jl index 72e2208a5..278be6098 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -180,6 +180,11 @@ stage(ctx, c::DArray) = c # closures get non-uniform ArgumentWrapper hashes). _collect_cat(concat, dims::Int, xs...) = concat(xs...; dims) +# Finalize a gathered DArray to a dense `Array{T,N}`. Sparse tile cats may yield +# a `DSparseArray` / `SparseMatrixCSC`, and `Base.collect` does not densify those. +_collect_dense(::Type{T}, ::Val{N}, x::Array{T,N}) where {T,N} = x +_collect_dense(::Type{T}, ::Val{N}, x) where {T,N} = Array{T,N}(x) + function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} a = fetch(d) if isempty(d.chunks) @@ -208,13 +213,13 @@ function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} result = Dagger.spawn_datadeps() do treereduce_nd(spawn_catfuncs, a.chunks) end - return collect(fetch(result)) + return _collect_dense(T, Val(N), fetch(result)) end # Distributed: fetch chunks directly and concat in-process. This avoids # routing chunk data through datadeps aliasing, which requires an # aliasing-resolvable (e.g. isbits) element type. dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] - return collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) + return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) end Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) @@ -511,7 +516,7 @@ function stage(ctx::Context, d::Distribute) cs = emit_chunk_tasks!(d.domainchunks, d.procgrid, T, (scope, I, i) -> begin c = d.domainchunks[I] - Dagger.@spawn compute_scope=scope identity(d.data[c]) + Dagger.@spawn compute_scope=scope maybe_wrap_tile(d.data[c]) end) end return DArray(eltype(d.data), diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 980b43e4f..4f1d9c21e 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -40,19 +40,84 @@ Base.similar(M::DSparseMatrix, ::Type{T}, dims::Tuple{Int, Int}) where T = DSparseMatrix{T}(_sparse_alloc(Val(sparse_mode(M.mat)), T, dims)) Base.collect(M::DSparseMatrix) = _sparse_collect(M.mat) -# N.B. hash and aliasing shouldn't change even if M.mat changes +# N.B. hash and aliasing shouldn't change even if M.mat changes. +# This is the key property that makes `DSparseMatrix` safe for Datadeps: +# the aliasing identity is tied to the (stable) wrapper object, not to the +# inner storage, which may be reallocated (and grow/shrink) on writes. Base.hash(M::DSparseMatrix, h::UInt) = hash(objectid(M), hash(DSparseMatrix, h)) -aliasing(M::DSparseMatrix, _=identity) = - ObjectAliasing(M) -Base.getindex(M::DSparseMatrix{T}, I::Int) where T = - getindex(M.mat, I) -Base.getindex(M::DSparseMatrix{T}, I::Vararg{Int,N}) where {T,N} = - getindex(M.mat, I...) -Base.setindex!(M::DSparseMatrix{T}, v, I::Int) where T = - setindex!(m.mat, v, I) -Base.setindex!(M::DSparseMatrix{T}, v, I::Vararg{Int,N}) where {T,N} = - setindex!(M.mat, v, I...) +# Sparse containers must alias as a single, indivisible unit. Their storage is +# reallocated (and may grow/shrink) on writes, so it is unsafe to alias any +# *part* of the container independently. We therefore force *all* access to a +# `DSparseMatrix` to resolve to the container's stable whole-object +# `ObjectAliasing` -- including access via views, transposes, adjoints, and +# reshapes -- rather than e.g. a `StridedAliasing` over storage that may have +# since moved. +aliasing(M::DSparseMatrix, _=identity) = ObjectAliasing(M) + +# Resolve the root `DSparseMatrix` beneath a stack of array wrappers (returns +# `nothing` if there is no sparse container at the root). +sparse_alias_root(@nospecialize(x)) = nothing +sparse_alias_root(M::DSparseMatrix) = M +sparse_alias_root(x::SubArray) = sparse_alias_root(parent(x)) +sparse_alias_root(x::Base.ReshapedArray) = sparse_alias_root(parent(x)) +sparse_alias_root(x::LinearAlgebra.Transpose) = sparse_alias_root(parent(x)) +sparse_alias_root(x::LinearAlgebra.Adjoint) = sparse_alias_root(parent(x)) +sparse_alias_root(x::Base.PermutedDimsArray) = sparse_alias_root(parent(x)) + +# Views/reshapes over a sparse container alias the *whole* container; any other +# (e.g. dense `Array`-backed) parent defers to the existing handling. Transpose +# and adjoint already forward aliasing to their parent, so the sparse-rooted case +# is handled there (and recursively via `sparse_alias_root`). +function aliasing(x::SubArray) + root = sparse_alias_root(x) + root === nothing && return invoke(aliasing, Tuple{Any}, x) + return aliasing(root) +end +function aliasing(x::Base.ReshapedArray) + root = sparse_alias_root(x) + root === nothing && return invoke(aliasing, Tuple{Any}, x) + return aliasing(root) +end + +# Forward indexing to the inner matrix. Ranges/colons are supported so that +# views (used by copy-buffering between mismatched partitionings) work. +Base.getindex(M::DSparseMatrix, I...) = getindex(M.mat, I...) +Base.setindex!(M::DSparseMatrix, v, I...) = setindex!(M.mat, v, I...) + +# Whole-tile copy. This is how Datadeps moves a tile between workers via +# `move!`: we *reallocate* the inner storage rather than mutating in place, +# since sparse storage cannot generally be updated in place. The wrapper's +# identity (and thus its aliasing) is preserved. +function Base.copyto!(dst::DSparseMatrix, src::DSparseMatrix) + dst.mat = _sparse_copy(src.mat) + return dst +end +# Backend-overridable deep copy of the inner storage (Finch tensors don't define +# `Base.copy`). +_sparse_copy(mat) = copy(mat) + +# Wrapping hook used when materializing tiles (e.g. in `distribute`). Backends +# overload this (e.g. in package extensions) to wrap freshly-created tiles in a +# container that Datadeps can track (such as `DSparseMatrix` for sparse tiles); +# everything else is left untouched. +maybe_wrap_tile(x) = x + +# Partial-range tile copy used by copy-buffering (e.g. when matmul operands need +# to be repartitioned). Some backends (Finch) forbid `setindex!`, so we route +# through a backend hook that returns the (possibly reallocated) inner storage. +# `DSparseMatrix` hides the reallocation from Datadeps. +function copyto_view!(Bpart::DSparseMatrix, Brange, Apart, Arange) + Bpart.mat = _sparse_copyto_view!(Bpart.mat, Brange, view(Apart, Arange)) + return +end +# Default: in-place element-wise copy (works for `setindex!`-capable backends +# such as `SparseArrays`). May reallocate and return new storage for backends +# that can't update in place, so callers must use the returned value. +function _sparse_copyto_view!(mat, Brange, src) + copyto!(view(mat, Brange), src) + return mat +end #Base.BroadcastStyle(::Type{<:DSparseMatrix}) = FinchStyle() From ecef8d27ab0faee6038b1d34cc56c8d1f226d6f3 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 4 Jun 2026 12:52:40 -0700 Subject: [PATCH 08/43] DArray/sparse: Fix sparse matmul implementations --- ext/FinchExt.jl | 111 +++++++++++++++++++++++++++++++---------- ext/SparseArraysExt.jl | 23 ++++++++- src/array/copy.jl | 14 +++--- src/array/mul.jl | 51 +++++++++++++++++++ 4 files changed, 164 insertions(+), 35 deletions(-) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index d534b49d1..a3ac5e7d4 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -1,6 +1,7 @@ module FinchExt import Finch +import LinearAlgebra import Dagger import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix @@ -38,6 +39,23 @@ Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment:: Finch.fsprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = Finch.fsprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) +# Materialize a (possibly lazy/`SwizzleArray`) result into a concrete sparse +# `Tensor`. `@einsum` returns a lazy `SwizzleArray` for transposed-output +# patterns, which is not itself a `Finch.Tensor` and corrupts subsequent +# accumulation steps; copying through `@finch` produces a clean `SparseCOO`. +function _finch_materialize(X, ::Type{T}, dims) where {T} + out = Finch.fspzeros(T, dims...) + Finch.@finch begin + out .= 0 + for j = _, i = _ + if X[i, j] != 0 + out[i, j] = X[i, j] + end + end + end + return out +end + # Finch tensors do not define `Base.copy`; build a fresh equivalent tensor. Dagger._sparse_copy(mat::Finch.Tensor) = _finch_materialize(mat, eltype(mat), size(mat)) @@ -49,6 +67,40 @@ function Dagger._sparse_copyto_view!(mat::Finch.Tensor, Brange, src) return _finch_materialize(dense, eltype(mat), size(mat)) end +# Compute `C_out += alpha * op(A) * op(B)` via Finch's `@einsum`, choosing the +# right index pattern for each transpose option. We deliberately use *explicit* +# index patterns (`A[k,i]`) rather than `Finch.swizzle`, because feeding a +# swizzled tensor into `@einsum` makes it return a lazy `SwizzleArray` (rather +# than a materialized `Tensor`), which breaks subsequent accumulation steps. +# Adjoint ('C') additionally applies `conj` inside the einsum. +function _finch_einsum_matmul!(C_out, A, B, transA::Char, transB::Char, alpha) + aT = transA == 'T' || transA == 'C' + aC = transA == 'C' + bT = transB == 'T' || transB == 'C' + bC = transB == 'C' + if !aT && !bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[k,j]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[k,j])) + elseif !aT && bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[j,k]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[j,k])) + elseif aT && !bT + aC ? Finch.@einsum(C_out[i,j] += alpha * (conj(A[k,i]) * B[k,j])) : + Finch.@einsum(C_out[i,j] += alpha * (A[k,i] * B[k,j])) + else # aT && bT + if aC && bC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * conj(B[j,k])) + elseif aC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * B[j,k]) + elseif bC + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * conj(B[j,k])) + else + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * B[j,k]) + end + end + return C_out +end + function Dagger.matmatmul!( C::DSparseMatrix, transA::Char, @@ -62,50 +114,55 @@ function Dagger.matmatmul!( # Not supported here, forward to generic matmatmul! return Dagger.matmatmul!(C.mat, transA, transB, A, B, alpha, beta) end + transA in ('N', 'T', 'C') || throw(ArgumentError("Invalid transA: $transA")) + transB in ('N', 'T', 'C') || throw(ArgumentError("Invalid transB: $transB")) - # Use optimized Finch operations + # `@einsum` cannot write into an existing sparse tensor in place, so we build + # a fresh output tensor and (re)assign it to the wrapper. `DSparseMatrix` + # hides this reallocation from Datadeps. + # + # N.B. BEWARE: `@einsum Cm[i,j] = ...` (assignment into an existing tensor) + # does not work; use a fresh `C_out` with `+=`. Cm = C.mat C_out = Finch.fspzeros(eltype(Cm), size(Cm)...) - # N.B. BEWARE: @einsum doesn't work correctly with Cm[i,j] = ... assignments. - #Finch.@einsum Cm[i,j] = alpha * (A[i,k] * B[k,j]) + beta * Cm[i,j] - Finch.@einsum C_out[i,j] += alpha * (A[i,k] * B[k,j]) + C_out = _finch_einsum_matmul!(C_out, A, B, transA, transB, alpha) + # `@einsum` may return a lazy `SwizzleArray`; materialize to a concrete + # sparse `Tensor` so accumulation and later iterations stay well-typed. + C_out = _finch_materialize(C_out, eltype(Cm), size(Cm)) if beta != 0 - C_out += beta * Cm + C_out = C_out + beta * Cm end - # TODO: Remove this once we're sure @einsum works correctly - @assert Array(C_out) ≈ (alpha * Array(A) * Array(B)) .+ (beta * Array(Cm)) C.mat = C_out return C end +# Tile transpose/symmetrization used by `copytri!`. Finch tensors don't support +# in-place `setindex!`, so we densify, build the result, and re-sparsify via +# `_finch_materialize`. +# - `uplo === nothing`: off-diagonal tile, return the conjugate transpose. +# - `uplo == 'U'`/`'L'`: diagonal tile, build the full Hermitian tile from the +# given triangle (matching the dense `copydiagtile!` semantics). function Dagger.transpose_tile(B::Finch.Tensor, uplo=nothing) + _B = Array(B) + if uplo === nothing + C = adjoint(_B) + return _finch_materialize(C, eltype(_B), size(C)) + end + n = LinearAlgebra.checksquare(_B) + C = zeros(eltype(_B), n, n) if uplo == 'U' - Finch.@finch begin - C .= 0 - for i in _, j in _ - if i <= j - C[i, j] = B[j, i] - end - end - return C + for j in 1:n, i in 1:n + C[i, j] = i <= j ? _B[i, j] : conj(_B[j, i]) end - return C elseif uplo == 'L' - Finch.@finch begin - C .= 0 - for i in _, j in _ - if i >= j - C[i, j] = B[j, i] - end - end - return C + for j in 1:n, i in 1:n + C[i, j] = i >= j ? _B[i, j] : conj(_B[j, i]) end - return C else - Finch.@einsum C[i,j] = B[j,i] - return C + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) end + return _finch_materialize(C, eltype(_B), size(C)) end end # module FinchExt \ No newline at end of file diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 45f7d8524..078b89b3c 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -2,6 +2,7 @@ module SparseArraysExt import SparseArrays import SparseArrays: SparseMatrixCSC +import LinearAlgebra import Dagger import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix @@ -58,13 +59,33 @@ function Dagger.matmatmul!( ) # Use fallback implementation # TODO: Optimize this further - C.mat = alpha * A * B + beta * C.mat + opA = _apply_trans(A, transA) + opB = _apply_trans(B, transB) + C.mat = alpha * (opA * opB) + beta * C.mat return C end +# Off-diagonal tile copy in `copytri!`: produce the (conjugate) transpose tile. function Dagger.transpose_tile(B::SparseMatrixCSC) return SparseArrays.sparse(B') end +# Diagonal tile symmetrization in `copytri!`: build the full Hermitian tile from +# its `uplo` triangle (matching the dense `copydiagtile!` semantics). +function Dagger.transpose_tile(B::SparseMatrixCSC, uplo::Char) + if uplo == 'U' + Bt = SparseArrays.triu(B) + elseif uplo == 'L' + Bt = SparseArrays.tril(B) + else + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + end + C = Bt + Bt' + # The shared diagonal was added twice; restore the original tile's diagonal. + for i in 1:LinearAlgebra.checksquare(B) + C[i, i] = B[i, i] + end + return C +end end # module SparseArraysExt diff --git a/src/array/copy.jl b/src/array/copy.jl index 422bec7c9..d703f49d1 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -36,6 +36,13 @@ function allocate_copy_buffer(part::Blocks{N}, A::DArray{T,N}) where {T,N} return DArray{T}(undef, part, size(A)) end +to_range(x::UnitRange) = x +to_range(x::Integer) = x:x +to_range(x::Base.OneTo{Int}) = UnitRange(x) +to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) +to_range(::StepRange) = throw(ArgumentError("Cannot convert StepRange to UnitRange")) +to_range(x) = throw(ArgumentError("Cannot convert $(typeof(x)) to UnitRange")) + function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(B), Ainds=parentindices(A)) where {TB,NB,TA,NA} Nmax = max(NA, NB) @@ -45,13 +52,6 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices( padNmax(x) = ntuple(i->pad1range(x, i), Nmax) padNmax(x::ArrayDomain) = padNmax(x.indexes) - to_range(x::UnitRange) = x - to_range(x::Integer) = x:x - to_range(x::Base.OneTo{Int}) = UnitRange(x) - to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) - to_range(::StepRange) = throw(ArgumentError("Non-continuous ranges are not yet supported for DArray copy")) - to_range(x) = throw(ArgumentError("Unsupported range type for DArray copy: $(typeof(x))")) - if any(x->x isa Vector, Binds) || any(x->x isa Vector, Ainds) # Split the copy into multiple copies dims_with_vector = findall(x->x[1] isa Vector || x[2] isa Vector, collect(zip(Binds, Ainds))) diff --git a/src/array/mul.jl b/src/array/mul.jl index c74c19128..94a8a235c 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -444,6 +444,7 @@ function copydiagtile!(A, uplo) end return nothing end + function LinearAlgebra.generic_matvecmul!( C::DVector{T}, transA::Char, @@ -560,3 +561,53 @@ function gemv_dagger!( return C end + +wrap_as_darray(A::DArray) = A +wrap_as_darray(A::Array) = view(A, AutoBlocks()) +function wrap_as_darray(A::SubArray{T,Nv,Array{T,Na}}) where {T,Nv,Na} + Ap = parent(A) + part = auto_blocks(map(last, parentindices(Ap))) + partsize = part.blocksize + inds = parentindices(A) + inds_ranges_parent = ntuple(i->to_range(inds[i]), Val(Na)) + inds_ranges_view = ntuple(i->to_range(inds[i]), Val(Nv)) + subdomains = partition(part, ArrayDomain(inds_ranges_parent)) + nparts = size(subdomains) + chunks = Array{Any,Na}(undef, nparts...) + for idx in CartesianIndices(nparts) + subdomain_view = subdomains[idx] + subdomain_parent = ArrayDomain(ntuple(i->Nv >= i ? subdomain_view.indexes[i] : inds_ranges_parent[i], Val(Na))) + subinds = ntuple(i->subdomain_parent.indexes[i], Val(Na)) + subA = view(Ap, subinds...) + chunks[idx] = tochunk(subA) + end + return DArray(T, ArrayDomain(inds_ranges_parent), subdomains, chunks, part) +end + +# Generate generic_matvecmul! methods for all combinations of DArray, Array, and SubArray +for CT in (DVector, Vector, SubArray{<:Any,1,<:Array}), + AT in (DMatrix, Matrix, SubArray{<:Any,2,<:Array}), + BT in (DVector, Vector, SubArray{<:Any,1,<:Array}) + + # Don't commit type piracy + CT isa DArray || AT isa DArray || BT isa DArray || continue + + @eval function LinearAlgebra.generic_matvecmul!( + C::$(CT), + transA::Char, + A::$(AT), + B::$(BT), + _add::LinearAlgebra.MulAddMul, + ) + new_C = wrap_as_darray(C) + new_A = wrap_as_darray(A) + new_B = wrap_as_darray(B) + return LinearAlgebra.generic_matvecmul!( + new_C, + transA, + new_A, + new_B, + _add, + ) + end +end \ No newline at end of file From 0177c5863af9461fa9e2cb0e090162326cc118e7 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 4 Jun 2026 12:53:35 -0700 Subject: [PATCH 09/43] test: Add explicit Finch/SparseArrays matmul tests --- test/Project.toml | 1 + test/array/linalg/matmul_finch.jl | 73 +++++++++++++++++++++++++++ test/array/linalg/matmul_sparse.jl | 81 ++++++++++++++++++++++++++++++ test/runtests.jl | 2 + 4 files changed, 157 insertions(+) create mode 100644 test/array/linalg/matmul_finch.jl create mode 100644 test/array/linalg/matmul_sparse.jl diff --git a/test/Project.toml b/test/Project.toml index edc738961..ffbd286a7 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" diff --git a/test/array/linalg/matmul_finch.jl b/test/array/linalg/matmul_finch.jl new file mode 100644 index 000000000..59ecc87a5 --- /dev/null +++ b/test/array/linalg/matmul_finch.jl @@ -0,0 +1,73 @@ +# Fast, focused sparse matrix-multiply tests for the Finch backend. +# +# Mirrors `matmul_sparse.jl` (SparseArrays) but exercises Finch tensor tiles via +# `DSparseMatrix`. Finch is an optional (weak) dependency, so this test loads it +# on all workers and skips gracefully if it cannot be loaded. +# +# Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_finch +# +# N.B. Finch precompilation is heavy (several minutes); this test is intended for +# local validation rather than fast CI. + +# Make the test-project environment (which provides Finch) resolvable on the +# workers, which are launched with Dagger's project. +const _FINCH_TEST_ENV = abspath(joinpath(@__DIR__, "..", "..")) +@everywhere pushfirst!(LOAD_PATH, $_FINCH_TEST_ENV) + +const FINCH_AVAILABLE = try + @eval using Finch + @everywhere using Finch + true +catch err + @warn "Finch unavailable; skipping Finch matmul tests" exception=err + false +end + +if !FINCH_AVAILABLE + @testset "Finch GEMM (quick)" begin + @test_skip false + end +else + const FINCH_DENSITY = 0.3 + + # Self-consistent: build distributed Finch operands, take the dense + # `collect` as the reference, and compare against the distributed result. + function test_finch_gemm!(T, sz, partA, partB) + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + DSA = Finch.fsprand(partA, T, sz, FINCH_DENSITY) + DSB = Finch.fsprand(partB, T, sz, FINCH_DENSITY) + A = collect(DSA) + B = collect(DSB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ A * B # N / N + @test collect(DSA * DSB') ≈ A * B' # N / C + @test collect(DSA' * DSB) ≈ A' * B # C / N + @test collect(DSA' * DSB') ≈ A' * B' # C / C + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ A' * A + @test collect(DSA * DSA') ≈ A * A' + + ## In-place + DSC = Finch.fspzeros(partC, T, sz) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ A * B + end + + const FINCH_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ] + + @testset "Finch GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in FINCH_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_gemm!(T, sz, partA, partB) + end + end + end +end diff --git a/test/array/linalg/matmul_sparse.jl b/test/array/linalg/matmul_sparse.jl new file mode 100644 index 000000000..afdc822f9 --- /dev/null +++ b/test/array/linalg/matmul_sparse.jl @@ -0,0 +1,81 @@ +# Fast, focused sparse matrix-multiply tests. +# +# This is intentionally a *small* subset of the full GEMM suite in `matmul.jl`, +# meant for quick iteration while developing sparse-array support (the full +# suite takes ~1hr). Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_sparse +# +# Keep this fast: only a couple of sizes/partitionings, just enough to exercise +# the out-of-place and in-place paths (incl. transposes) for both real and +# complex element types. + +const SPARSE_DENSITY = 0.3 + +function test_sparse_gemm!(T, sz, partA, partB) + rows, cols = sz + @assert rows == cols "sparse quick-test uses square matrices so transposes line up" + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + SA = sprand(T, sz..., SPARSE_DENSITY) + SB = sprand(T, sz..., SPARSE_DENSITY) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ SA * SB # N / N + @test collect(DSA * DSB') ≈ SA * SB' # N / T + @test collect(DSA' * DSB) ≈ SA' * SB # T / N + @test collect(DSA' * DSB') ≈ SA' * SB' # T / T + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ Array(SA)' * Array(SA) + @test collect(DSA * DSA') ≈ Array(SA) * Array(SA)' + + ## In-place + SC = SA * SB + DSC = distribute(sparse(zeros(T, sz...)), partC) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC +end + +const SPARSE_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ((8, 8), Blocks(4, 2), Blocks(2, 4)), +] + +@testset "Sparse GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in SPARSE_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_gemm!(T, sz, partA, partB) + end + end +end + +# Any *partial* or *reinterpreted* access to a sparse container (views, +# transposes, adjoints, reshapes, and combinations thereof) must alias the +# *entire* container, so that Datadeps never tracks stale sub-spans of storage +# that may have been reallocated on write. +@testset "Sparse whole-container aliasing" begin + M = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + a_full = Dagger.aliasing(M) + + @test Dagger.will_alias(a_full, Dagger.aliasing(view(M, 1:3, :))) + @test Dagger.will_alias(a_full, Dagger.aliasing(view(M, 4:6, 2:4))) + @test Dagger.will_alias(a_full, Dagger.aliasing(transpose(M))) + @test Dagger.will_alias(a_full, Dagger.aliasing(M')) + @test Dagger.will_alias(a_full, Dagger.aliasing(reshape(M, 36))) + @test Dagger.will_alias(a_full, Dagger.aliasing(view(transpose(M), 1:2, :))) + # Non-overlapping views of the *same* container still alias (whole-container). + @test Dagger.will_alias(Dagger.aliasing(view(M, 1:3, :)), Dagger.aliasing(view(M, 4:6, :))) + + # Distinct containers must not alias. + M2 = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + @test !Dagger.will_alias(a_full, Dagger.aliasing(M2)) + + # Dense-array views are unaffected (still strided, not whole-array). + A = rand(Float64, 6, 6) + @test Dagger.aliasing(view(A, 1:3, :)) isa Dagger.StridedAliasing +end diff --git a/test/runtests.jl b/test/runtests.jl index 799a05c2f..e0bc5f8b2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -32,6 +32,8 @@ tests = [ ("Array - LinearAlgebra - Core", "array/linalg/core.jl"), ("Array - LinearAlgebra - Arithmetic", "array/linalg/arithmetic.jl"), ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"), + ("Array - LinearAlgebra - Matmul (Sparse, quick)", "array/linalg/matmul_sparse.jl"), + ("Array - LinearAlgebra - Matmul (Finch, quick)", "array/linalg/matmul_finch.jl"), ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"), ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"), ("Array - LinearAlgebra - Solve", "array/linalg/solve.jl"), From 93195f8d15dc1a07ce512009f4a8aeff3d807948 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 6 Jun 2026 11:08:00 -0700 Subject: [PATCH 10/43] DArray/sparse: Remove global sparse mode switch --- ext/FinchExt.jl | 8 +++++--- ext/SparseArraysExt.jl | 3 --- src/array/sparse.jl | 29 ++++++++--------------------- 3 files changed, 13 insertions(+), 27 deletions(-) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index a3ac5e7d4..54767fa05 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -5,10 +5,12 @@ import LinearAlgebra import Dagger import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix -Dagger.sparse_mode(::Finch.Tensor) = :finch -Dagger._sparse_alloc(::Val{:finch}, T::Type, dims::Dims) = - Finch.fspzeros(T, dims...) Dagger._sparse_collect(A::Finch.Tensor) = Array(A) +# Finch's generic `similar` produces tensor formats that destabilize later +# `@finch`/`@einsum` kernels (observed as segfaults during tile moves), so +# allocate an empty COO-backed tile explicitly instead. +Dagger._sparse_similar(::Finch.Tensor, ::Type{T}, dims::Dims) where {T} = + Finch.fspzeros(T, dims...) Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseMatrix{eltype(x)}(x) function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 078b89b3c..31c570784 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -6,9 +6,6 @@ import LinearAlgebra import Dagger import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix -Dagger.sparse_mode(::SparseMatrixCSC) = :sparsearrays -Dagger._sparse_alloc(::Val{:sparsearrays}, T::Type, dims::Dims) = - SparseArrays.spzeros(T, dims...) # Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) # Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 4f1d9c21e..6060dd33e 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -8,28 +8,15 @@ Datadeps algorithms. mutable struct DSparseMatrix{T} <: AbstractMatrix{T} mat end -function DSparseMatrix{T}(undef, dims::NTuple{N, Int}) where {T,N} - M = sparse_mode() - return DSparseMatrix{T}(_sparse_alloc(Val(M), T, dims)) -end -function _sparse_not_loaded(::Val{M}) where M - if M == :sparsearrays - throw(ArgumentError("SparseArrays must be loaded to use SparseMatrixCSC")) - elseif M == :finch - throw(ArgumentError("Finch must be loaded to use Finch.Tensor")) - elseif M == :none - throw(ArgumentError("Sparse mode not set\nSet it with `sparse_mode!(M)` where M is :sparsearrays or :finch, and load the corresponding package")) - else - throw(ArgumentError("Unknown sparse mode $M\nOptions are :sparsearrays and :finch")) - end -end -_sparse_alloc(::Val{M}, T::Type, dims::Dims) where M = _sparse_not_loaded(Val(M)) _sparse_collect(M) = collect(M) -const SPARSE_MODE = TaskLocalValue{Symbol}(()->:none) -sparse_mode() = SPARSE_MODE[] -sparse_mode(::T) where T = error("Unknown sparse mode for type $T") -set_sparse_mode!(mode::Symbol) = SPARSE_MODE[] = mode +# Allocate a destination tile matching the inner storage's sparse backend. The +# default delegates to the storage type's own `similar`, which carries the +# concrete backend (e.g. `SparseMatrixCSC`) forward to the destination -- this is +# how `similar(::DArray)` propagates a tile's sparse type to a newly-allocated +# result DArray. Backends whose `similar` is unsuitable (e.g. Finch, whose generic +# `similar` yields tensor formats that destabilize later kernels) override this. +_sparse_similar(mat, ::Type{T}, dims::Dims) where {T} = similar(mat, T, dims) Base.eltype(M::DSparseMatrix{T}) where T = T Base.size(M::DSparseMatrix) = size(M.mat) @@ -37,7 +24,7 @@ Base.ndims(M::DSparseMatrix) = 2 Base.iterate(M::DSparseMatrix) = iterate(M.mat) Base.iterate(M::DSparseMatrix, state) = iterate(M.mat, state) Base.similar(M::DSparseMatrix, ::Type{T}, dims::Tuple{Int, Int}) where T = - DSparseMatrix{T}(_sparse_alloc(Val(sparse_mode(M.mat)), T, dims)) + DSparseMatrix{T}(_sparse_similar(M.mat, T, dims)) Base.collect(M::DSparseMatrix) = _sparse_collect(M.mat) # N.B. hash and aliasing shouldn't change even if M.mat changes. From b081f2ef90de453795177d2e4663c6e0e7f4eb74 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 6 Jun 2026 11:09:27 -0700 Subject: [PATCH 11/43] SparseArrays: Lightly optimize matmul kernel --- ext/SparseArraysExt.jl | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 31c570784..a764b04e9 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -54,11 +54,24 @@ function Dagger.matmatmul!( alpha, beta ) - # Use fallback implementation - # TODO: Optimize this further opA = _apply_trans(A, transA) opB = _apply_trans(B, transB) - C.mat = alpha * (opA * opB) + beta * C.mat + # Sparse*sparse yields a freshly-allocated sparse matrix, which we reassign + # into the wrapper (`DSparseMatrix` hides this reallocation from Datadeps). + # `SparseArrays` provides no efficient 5-arg `mul!` into a sparse `C` -- the + # output sparsity pattern is determined by the product -- so we form the + # product out-of-place and apply only the alpha/beta scaling that is actually + # needed. The transposed-operand products dispatch to specialized SparseArrays + # methods, so `opA`/`opB` are not materialized. + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = prod + C.mat + else + C.mat = prod + beta * C.mat + end return C end From bccf7fc30c9bacbc1cf9e982b8d13f3113139ba0 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 6 Jun 2026 11:12:02 -0700 Subject: [PATCH 12/43] Datadeps/sparse: Ignore array wrappers for objects like DSparseMatrix --- ext/MPIExt.jl | 7 ++- src/array/sparse.jl | 40 ++++++-------- src/datadeps/aliasing.jl | 7 ++- src/datadeps/chunkview.jl | 5 +- src/datadeps/remainders.jl | 31 ++++++++++- src/memory-spaces.jl | 87 ++++++++++++++++++++++++++++-- test/array/linalg/matmul_sparse.jl | 46 +++++++++++----- 7 files changed, 177 insertions(+), 46 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 9274496b4..71725742b 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -8,7 +8,8 @@ import Dagger: @dagdebug, @opcounter # extends with new methods (for MPI-specific types) or calls directly. import Dagger: AbstractAliasing, accelerate!, accel_matches_proc, aliased_object!, - AliasedObjectCache, AliasedObjectCacheStore, aliasing, bind_moved_argument, + AliasedObjectCache, AliasedObjectCacheStore, aliasing, aliasing_unwrapped, + bind_moved_argument, chunktype, ChunkView, check_uniform, check_uniformity!, CHECK_UNIFORMITY, cleanup_tasks_accel!, constrain, CPURAMMemorySpace, current_acceleration, CyclicProcGrid, datasize, default_enabled, @@ -1757,7 +1758,9 @@ function aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) if handle.rank == rank ainfo = _with_default_acceleration() do v = view(unwrap(x.chunk), x.slices...) - aliasing(v, dep_mod) + # Resolve whole-object containers (e.g. `DSparseArray`) where `v` + # lives; see `aliasing_unwrapped`. + aliasing_unwrapped(v, dep_mod) end ainfo = mpi_remap_ainfo(ainfo, handle.rank) @opcounter :aliasing_bcast_send_yield diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 6060dd33e..04687a67e 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -39,32 +39,24 @@ Base.hash(M::DSparseMatrix, h::UInt) = hash(objectid(M), hash(DSparseMatrix, h)) # `DSparseMatrix` to resolve to the container's stable whole-object # `ObjectAliasing` -- including access via views, transposes, adjoints, and # reshapes -- rather than e.g. a `StridedAliasing` over storage that may have -# since moved. +# since moved. `aliases_as_whole` opts the type into this behavior, and Datadeps' +# `aliasing_root` resolves any wrapper of a `DSparseMatrix` back to the container +# before computing aliasing. Note: this must return a *bare* `ObjectAliasing`, as +# `aliases_as_whole(::ObjectAliasing)` relies on it to drive whole-object copies. aliasing(M::DSparseMatrix, _=identity) = ObjectAliasing(M) +aliases_as_whole(::DSparseMatrix) = true -# Resolve the root `DSparseMatrix` beneath a stack of array wrappers (returns -# `nothing` if there is no sparse container at the root). -sparse_alias_root(@nospecialize(x)) = nothing -sparse_alias_root(M::DSparseMatrix) = M -sparse_alias_root(x::SubArray) = sparse_alias_root(parent(x)) -sparse_alias_root(x::Base.ReshapedArray) = sparse_alias_root(parent(x)) -sparse_alias_root(x::LinearAlgebra.Transpose) = sparse_alias_root(parent(x)) -sparse_alias_root(x::LinearAlgebra.Adjoint) = sparse_alias_root(parent(x)) -sparse_alias_root(x::Base.PermutedDimsArray) = sparse_alias_root(parent(x)) - -# Views/reshapes over a sparse container alias the *whole* container; any other -# (e.g. dense `Array`-backed) parent defers to the existing handling. Transpose -# and adjoint already forward aliasing to their parent, so the sparse-rooted case -# is handled there (and recursively via `sparse_alias_root`). -function aliasing(x::SubArray) - root = sparse_alias_root(x) - root === nothing && return invoke(aliasing, Tuple{Any}, x) - return aliasing(root) -end -function aliasing(x::Base.ReshapedArray) - root = sparse_alias_root(x) - root === nothing && return invoke(aliasing, Tuple{Any}, x) - return aliasing(root) +# A `DSparseMatrix` has no meaningful raw data pointer (its storage may be +# reallocated/resized). This trap ensures that if aliasing ever tries to treat a +# wrapper of a `DSparseMatrix` as strided memory -- e.g. a new array-wrapper type +# that `aliasing_root` failed to resolve via `Base.parent` -- it fails loudly +# instead of silently corrupting Datadeps' aliasing tracking. +function Base.pointer(::DSparseMatrix) + throw(ArgumentError("`pointer(::DSparseMatrix)` is intentionally unsupported: \ + a DSparseMatrix may reallocate its storage, so it must be aliased as a \ + whole object via `Dagger.aliasing`. If you reached here through Datadeps \ + aliasing of an array wrapper, ensure that wrapper implements `Base.parent` \ + so that `Dagger.aliasing_root` can resolve it to the DSparseMatrix.")) end # Forward indexing to the inner matrix. Ranges/colons are supported so that diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index dc6668486..a412182ea 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -610,7 +610,12 @@ function aliasing!(state::DataDepsState, target_space::MemorySpace, arg_w::Argum return state.ainfo_cache[remote_arg_w] end - # Calculate the ainfo + # Calculate the ainfo via the current acceleration so SPMD backends (MPI) + # can broadcast owner-computed aliasing and keep ranks uniform. The generic + # `aliasing(::Acceleration, ...)` fallback uses `aliasing_unwrapped`, which + # resolves wrappers of whole-object containers (e.g. a view of a + # `DSparseArray`) to the container itself. For `Chunk`s, unwrap+aliasing + # happen where the data lives (remotely / on the owning rank). ainfo = AliasingWrapper(aliasing(current_acceleration(), remote_arg, arg_w.dep_mod)) # Cache the result diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 903785f93..6bf37cb14 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -69,7 +69,10 @@ function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices x = unwrap(x) v = view(x, slices...) - return aliasing(accel, v, dep_mod) + # A view of a whole-object container (e.g. `DSparseMatrix`) must alias the + # entire container; `aliasing_unwrapped` resolves that (otherwise it just + # aliases the view), and crucially does so here where `v` lives. + return aliasing_unwrapped(v) end end aliasing(x::ChunkView) = aliasing(current_acceleration(), x, identity) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index a5885a579..c530e0766 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -43,6 +43,7 @@ memory_spans(mra::MultiRemainderAliasing) = vcat(memory_spans.(mra.remainders).. Base.hash(mra::MultiRemainderAliasing, h::UInt) = hash(mra.remainders, hash(MultiRemainderAliasing, h)) Base.:(==)(mra1::MultiRemainderAliasing, mra2::MultiRemainderAliasing) = mra1.remainders == mra2.remainders +# A whole-object copy from the argument's recorded owner space. struct FullCopy end """ @@ -87,9 +88,34 @@ function compute_remainder_for_arg!(state::DataDepsState, target_space::MemorySpace, arg_w::ArgumentWrapper, write_num::Int; compute_syncdeps::Bool=true) + owner_space = state.arg_owner[arg_w] + + # Whole-object containers (e.g. `DSparseMatrix`) reallocate their entire + # storage on every write, so they can never be *partially* current in a + # memory space -- the whole object is current in exactly one space: the most + # recent writer (or the origin owner if never written). There is thus never a + # meaningful span "remainder" to compute, so we short-circuit to a whole-object + # `FullCopy` from the owner space (or `NoAliasing` if the target already holds + # it). This avoids the span-diff computation and the construction + transfer of + # a `RemainderAliasing` entirely. We detect this from the aliasing info, which + # is a bare `ObjectAliasing` for such containers (see `aliases_as_whole`). + # + # `arg_owner` is the most-recent *direct* writer of `arg_w`, while the merged + # `arg_history` records the latest write through *any* overlapping argument. + # For every path Datadeps currently exercises a whole-object container is only + # ever reached through a single argument, so these agree -- the assertion below + # guards that invariant. They could only diverge under cross-argument aliasing + # of a whole-object container (the same container written through one argument + # and read through a different, overlapping one); supporting that would require + # sourcing the copy from `history[end].space` instead of `arg_owner`. + if aliases_as_whole(aliasing!(state, target_space, arg_w)) + history = state.arg_history[arg_w] + @assert isempty(history) || history[end].space == owner_space "whole-object container reached through a different overlapping argument than its owner; cross-argument aliasing of whole-object containers is unsupported" + return owner_space == target_space ? (NoAliasing(), 0) : (FullCopy(), 0) + end + spaces_set = Set{MemorySpace}() push!(spaces_set, target_space) - owner_space = state.arg_owner[arg_w] push!(spaces_set, owner_space) @label restart @@ -437,6 +463,9 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: multi_span_copy!(to_s, from_s, dep_mod.spans) return end + # N.B. Whole-object containers (e.g. `DSparseMatrix`) never reach here: their + # storage reallocates on write, so they can't be span-copied. They are routed + # to a whole-object `FullCopy` in `compute_remainder_for_arg!` instead. # Gather source spans into a host buffer (device arrays pack via KA first). # Pin when the source is device memory so DtoH hits page-locked host RAM. diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 617f7ea8b..a1eefe52e 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -375,7 +375,11 @@ function memory_spans(oa::ObjectAliasing{S}) where S return [span] end -aliasing(accel::Acceleration, x, T) = aliasing(x, T) +# Acceleration entry point used by Datadeps. Unwrap whole-object containers +# (e.g. views of `DSparseArray`) before computing aliasing. SPMD backends +# (MPI) overload this for `Chunk`/`ChunkView` to broadcast owner-computed +# aliasing; those methods call `aliasing_unwrapped` on the owning rank. +aliasing(accel::Acceleration, x, T) = aliasing_unwrapped(x, T) function aliasing(x, dep_mod) if dep_mod isa Symbol return aliasing(getfield(x, dep_mod)) @@ -411,22 +415,95 @@ end aliasing(::String) = NoAliasing() # FIXME: Not necessarily true aliasing(::Symbol) = NoAliasing() aliasing(::Type) = NoAliasing() + +""" + aliases_as_whole(x) -> Bool + +Whether `x` -- either an object, or an already-computed aliasing -- must be +treated as a single, indivisible unit, i.e. it is never safe to alias (or to +consider current) only *part* of it. + +For an object: containers whose backing storage may be reallocated or resized on +write (such as `DSparseMatrix`) return `true`. This causes [`aliasing_root`](@ref) +to resolve any view/wrapper of them to the whole container, so all access funnels +through the container's own `aliasing`. By contract, such a container's `aliasing` +must be a bare [`ObjectAliasing`](@ref). + +For an aliasing: a bare `ObjectAliasing` (the aliasing of a whole-object +container) reports `true`. Datadeps uses this to copy such arguments as a whole +(a `FullCopy`) rather than computing per-span remainders, since they can never be +*partially* current in a memory space. +""" +aliases_as_whole(@nospecialize x) = false +aliases_as_whole(ainfo::AliasingWrapper) = aliases_as_whole(ainfo.inner) +aliases_as_whole(::ObjectAliasing) = true + +""" + aliasing_root(x) + +Resolve `x` down to the whole-object container it wraps: peel array wrappers +(views, transposes, adjoints, reshapes, permutations, ...) off of `x` until +reaching a container that must alias as a whole (see [`aliases_as_whole`](@ref)), +and return that container. If no such container is wrapped, return `x` unchanged. + +This relies solely on the standard `Base.parent` interface, so it transparently +handles *any* array wrapper without per-wrapper `aliasing` methods: a new wrapper +type needs no special-casing here, and a new whole-object container only needs to +define [`aliases_as_whole`](@ref) (plus its own `aliasing` method, which must +return a bare `ObjectAliasing`). A trapping `Base.pointer` on such containers +guards against a wrapper slipping through and being misinterpreted as strided +memory. + +!!! warning + The resolved object's *identity* defines its aliasing, so this is only + meaningful in the memory space where `x` physically resides. Never return its + result across a worker boundary and *then* compute `aliasing` -- the transfer + copies the object and changes its aliasing. Use [`aliasing_unwrapped`](@ref), + which fuses the two operations so they always run together. +""" +aliasing_root(@nospecialize x) = x +function aliasing_root(x::AbstractArray) + aliases_as_whole(x) && return x + p = parent(x) + # `Base.parent` returns the argument itself for non-wrapper arrays, which + # terminates the recursion. + p === x && return x + root = aliasing_root(p) + return aliases_as_whole(root) ? root : x +end + +""" + aliasing_unwrapped(x[, dep_mod]) + +Compute the `aliasing` of `x`, first resolving (via [`aliasing_root`](@ref)) any +whole-object container that `x` wraps. This is the entry point Datadeps uses to +compute the aliasing of a (possibly wrapped) argument. + +Unwrapping and `aliasing` are intentionally fused into a single call so they +always execute in the same place. It MUST be evaluated where `x` physically lives +-- e.g. *inside* a `Chunk`'s `remotecall_fetch` block -- because the unwrapped +object's identity determines its aliasing; returning the unwrapped object across +a worker boundary first would copy it and silently change the result. +""" +aliasing_unwrapped(x) = aliasing(aliasing_root(x)) +aliasing_unwrapped(x, dep_mod) = aliasing(aliasing_root(x), dep_mod) + function aliasing(x::Chunk, T) if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x), T) + return aliasing_unwrapped(unwrap(x), T) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T - aliasing(unwrap(x), T) + aliasing_unwrapped(unwrap(x), T) end end function aliasing(x::Chunk) if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x)) + return aliasing_unwrapped(unwrap(x)) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x) do x - aliasing(unwrap(x)) + aliasing_unwrapped(unwrap(x)) end end aliasing(x::DTask, T) = aliasing(fetch(x; move_value=false, unwrap=false), T) diff --git a/test/array/linalg/matmul_sparse.jl b/test/array/linalg/matmul_sparse.jl index afdc822f9..0bbdb3c4a 100644 --- a/test/array/linalg/matmul_sparse.jl +++ b/test/array/linalg/matmul_sparse.jl @@ -57,25 +57,47 @@ end # Any *partial* or *reinterpreted* access to a sparse container (views, # transposes, adjoints, reshapes, and combinations thereof) must alias the # *entire* container, so that Datadeps never tracks stale sub-spans of storage -# that may have been reallocated on write. +# that may have been reallocated on write. This is enforced by `aliasing_root`, +# which resolves any such wrapper back to the container itself (via `Base.parent`) +# -- exposed to Datadeps through the fused `aliasing_unwrapped` -- plus a trapping +# `Base.pointer` that fires if a wrapper ever slips through to the strided path. @testset "Sparse whole-container aliasing" begin M = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) - a_full = Dagger.aliasing(M) - @test Dagger.will_alias(a_full, Dagger.aliasing(view(M, 1:3, :))) - @test Dagger.will_alias(a_full, Dagger.aliasing(view(M, 4:6, 2:4))) - @test Dagger.will_alias(a_full, Dagger.aliasing(transpose(M))) - @test Dagger.will_alias(a_full, Dagger.aliasing(M')) - @test Dagger.will_alias(a_full, Dagger.aliasing(reshape(M, 36))) - @test Dagger.will_alias(a_full, Dagger.aliasing(view(transpose(M), 1:2, :))) - # Non-overlapping views of the *same* container still alias (whole-container). - @test Dagger.will_alias(Dagger.aliasing(view(M, 1:3, :)), Dagger.aliasing(view(M, 4:6, :))) + @test Dagger.aliases_as_whole(M) + @test !Dagger.aliases_as_whole(rand(Float64, 6, 6)) + # The container's aliasing is a bare `ObjectAliasing`, which also reports + # whole-object (drives the Datadeps whole-object copy short-circuit). + @test Dagger.aliasing(M) isa Dagger.ObjectAliasing + @test Dagger.aliases_as_whole(Dagger.aliasing(M)) + @test !Dagger.aliases_as_whole(Dagger.aliasing(rand(Float64, 6, 6))) + + # `aliasing_root` peels any array wrapper back to the sparse container. + @test Dagger.aliasing_root(M) === M + @test Dagger.aliasing_root(view(M, 1:3, :)) === M + @test Dagger.aliasing_root(view(M, 4:6, 2:4)) === M + @test Dagger.aliasing_root(transpose(M)) === M + @test Dagger.aliasing_root(M') === M + @test Dagger.aliasing_root(reshape(M, 36)) === M + @test Dagger.aliasing_root(view(transpose(M), 1:2, :)) === M + + # After unwrapping, all wrappers alias the same whole container. + a_full = Dagger.aliasing(M) + @test Dagger.will_alias(a_full, Dagger.aliasing_unwrapped(view(M, 1:3, :))) + @test Dagger.will_alias(Dagger.aliasing_unwrapped(view(M, 1:3, :)), + Dagger.aliasing_unwrapped(view(M, 4:6, :))) # Distinct containers must not alias. M2 = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) @test !Dagger.will_alias(a_full, Dagger.aliasing(M2)) - # Dense-array views are unaffected (still strided, not whole-array). + # Dense arrays are untouched: not whole-object, views stay strided, and + # `aliasing_root` returns them unchanged. A = rand(Float64, 6, 6) - @test Dagger.aliasing(view(A, 1:3, :)) isa Dagger.StridedAliasing + Av = view(A, 1:3, :) + @test Dagger.aliasing_root(Av) === Av + @test Dagger.aliasing_unwrapped(Av) isa Dagger.StridedAliasing + + # The `pointer` trap fires if a sparse container is treated as raw memory. + @test_throws ArgumentError pointer(M) end From 35b383e4bd2f7d9bef78f910fb230174cf3ddf4d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 6 Jun 2026 17:46:52 -0700 Subject: [PATCH 13/43] test/sparse: Only run Finch tests on demand --- .buildkite/pipeline.yml | 12 ++++++++++++ test/Project.toml | 1 - test/runtests.jl | 23 +++++++++++++++++++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b3f72a3c9..b61fbc9b4 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -111,5 +111,17 @@ steps: julia --project=test/metalenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' julia --project=test/metalenv test/run_mpi.jl 2 2 test/mpi_metal.jl + - label: Julia 1.11 (Finch) + timeout_in_minutes: 90 + if: build.message !~ /\[skip tests\]/ + plugins: + - JuliaCI/julia#v1: + version: "1.11" + - JuliaCI/julia-test#v1: ~ + - JuliaCI/julia-coverage#v1: + codecov: true + env: + CI_TEST_FINCH: "1" + # env: # SECRET_CODECOV_TOKEN: "" diff --git a/test/Project.toml b/test/Project.toml index ffbd286a7..edc738961 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,7 +6,6 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" -Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" diff --git a/test/runtests.jl b/test/runtests.jl index e0bc5f8b2..0d9d7b945 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,6 +7,11 @@ USE_METAL = parse(Bool, get(ENV, "CI_USE_METAL", "0")) USE_OPENCL = parse(Bool, get(ENV, "CI_USE_OPENCL", "0")) USE_GPU = USE_CUDA || USE_ROCM || USE_ONEAPI || USE_METAL || USE_OPENCL +# Finch is a heavy optional dependency (precompilation takes minutes), so it is +# not a default test dependency. The dedicated Finch CI job signals via this +# variable to `Pkg.add` Finch and run only the Finch-backed tests. +USE_FINCH = parse(Bool, get(ENV, "CI_TEST_FINCH", "0")) + tests = [ ("Thunk", "thunk.jl"), ("Scheduler", "scheduler.jl"), @@ -59,8 +64,22 @@ if USE_GPU ("Array - Stencils", "array/stencil.jl"), ] end +if USE_FINCH + # Only run the Finch-backed tests in the dedicated Finch CI job. + tests = [ + ("Array - LinearAlgebra - Matmul (Finch)", "array/linalg/matmul_finch.jl"), + ] +end all_test_names = map(test -> replace(last(test), ".jl"=>""), tests) +# Tests excluded from default runs; they only run when explicitly requested via +# `--test` or enabled by an environment signal. Finch's precompilation is very +# heavy, so its tests are opt-in (except in the dedicated Finch job, where we +# explicitly want them to run). +optin_test_names = USE_FINCH ? String[] : String[ + "array/linalg/matmul_finch", +] + additional_workers::Int = 3 worker_threads::Int = 1 @@ -73,6 +92,10 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.instantiate() catch end + if USE_FINCH + # Finch is not a default test dependency; add it on demand when signaled. + Pkg.add("Finch") + end using ArgParse s = ArgParseSettings(description = "Dagger Testsuite") From 885dec211d9edc36b615e069c79cca8283f50471 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 6 Jun 2026 18:05:12 -0700 Subject: [PATCH 14/43] DArray/sparse: DSparseMatrix->DSparseArray, and add matrix-vector product support --- ext/FinchExt.jl | 26 +++++++-- ext/SparseArraysExt.jl | 24 +++++++-- src/array/linalg.jl | 65 ++++++++++++++++++++++ src/array/mul.jl | 27 +++++++--- src/array/sparse.jl | 86 ++++++++++++++++++------------ src/datadeps/chunkview.jl | 2 +- src/datadeps/remainders.jl | 4 +- src/memory-spaces.jl | 2 +- test/array/allocation.jl | 3 +- test/array/linalg/core.jl | 54 +++++++++++++++++++ test/array/linalg/matmul_finch.jl | 32 +++++++++++ test/array/linalg/matmul_sparse.jl | 36 +++++++++++++ 12 files changed, 307 insertions(+), 54 deletions(-) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index 54767fa05..be2f0e356 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -3,7 +3,7 @@ module FinchExt import Finch import LinearAlgebra import Dagger -import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix Dagger._sparse_collect(A::Finch.Tensor) = Array(A) # Finch's generic `similar` produces tensor formats that destabilize later @@ -11,11 +11,13 @@ Dagger._sparse_collect(A::Finch.Tensor) = Array(A) # allocate an empty COO-backed tile explicitly instead. Dagger._sparse_similar(::Finch.Tensor, ::Type{T}, dims::Dims) where {T} = Finch.fspzeros(T, dims...) -Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseMatrix{eltype(x)}(x) +Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseArray(x) function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = @@ -29,7 +31,9 @@ Finch.fspzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = : function Finch.fsprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end Finch.fsprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = @@ -139,6 +143,20 @@ function Dagger.matmatmul!( return C end +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# Finch tensor `A` and dense vectors `B`/`C`. Finch's `@einsum` does not reliably +# accumulate into a dense output, so we densify the single sparse tile and use a +# dense BLAS-backed `mul!`; this is a per-tile densification, not the whole matrix. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Finch.Tensor, B::AbstractVector, alpha, beta) + Ad = Array(A) + opA = transA == 'N' ? Ad : + transA == 'C' ? adjoint(Ad) : + transA == 'T' ? transpose(Ad) : + throw(ArgumentError("Invalid transA: $transA")) + LinearAlgebra.mul!(C, opA, B, alpha, beta) + return C +end + # Tile transpose/symmetrization used by `copytri!`. Finch tensors don't support # in-place `setindex!`, so we densify, build the result, and re-sparsify via # `_finch_materialize`. diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index a764b04e9..a6c6755fa 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -1,19 +1,22 @@ module SparseArraysExt import SparseArrays -import SparseArrays: SparseMatrixCSC +import SparseArrays: SparseMatrixCSC, SparseVector import LinearAlgebra import Dagger -import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseMatrix +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix # Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) # Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. -Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseMatrix{eltype(x)}(x) +Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::SparseVector) = DSparseArray(x) function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(SparseArrays.spzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(SparseArrays.spzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = @@ -27,7 +30,9 @@ SparseArrays.spzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentTy function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseMatrix{T}(SparseArrays.sprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(SparseArrays.sprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = @@ -76,6 +81,15 @@ function Dagger.matmatmul!( return C end +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# `SparseMatrixCSC` `A` and dense vectors `B`/`C`. SparseArrays provides an +# efficient 5-arg `mul!` (SpMV) into a dense output, including for transposed and +# adjoint operands, so this updates `C` in place with no allocation. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::SparseMatrixCSC, B::AbstractVector, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + # Off-diagonal tile copy in `copytri!`: produce the (conjugate) transpose tile. function Dagger.transpose_tile(B::SparseMatrixCSC) return SparseArrays.sparse(B') diff --git a/src/array/linalg.jl b/src/array/linalg.jl index 2fcaa42fa..89d0575e1 100644 --- a/src/array/linalg.jl +++ b/src/array/linalg.jl @@ -4,6 +4,71 @@ function LinearAlgebra.norm2(A::DArray{T,N}) where {T,N} zeroRT = zero(real(T)) return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT)) end + +# --- BLAS-1 vector operations --------------------------------------------- +# Efficient, distributed level-1 operations needed by iterative (Krylov) +# solvers. These run one BLAS-1 kernel per chunk (avoiding the generic +# scalar-indexing fallbacks, which are slow and trigger scalar access on GPUs). +# Operands may be partitioned differently: we align them with +# `maybe_copy_buffered` (a no-op when they already share a layout, and which +# copies results back to the originals for in-place operations), so users keep +# full control over their own (possibly sub-optimal) partitioning. + +function LinearAlgebra.dot(x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("dot: x has size $(size(x)), y has size $(size(y))")) + R = typeof(LinearAlgebra.dot(zero(Tx), zero(Ty))) + return maybe_copy_buffered(x => x.partitioning, y => x.partitioning) do x, y + xc, yc = x.chunks, y.chunks + parts = [Dagger.@spawn LinearAlgebra.dot(xc[i], yc[i]) for i in eachindex(xc)] + sum(fetch, parts; init=zero(R)) + end +end + +function LinearAlgebra.axpy!(a::Number, x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpy!: x has size $(size(x)), y has size $(size(y))")) + # Align the read-only `x` to `y`'s layout so only `x` may be buffered; `y` is + # updated in place (and copied back by `maybe_copy_buffered` if it was buffered). + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpy!(a, In(xc[i]), InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.axpby!(a::Number, x::DArray{Tx,N}, b::Number, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpby!: x has size $(size(x)), y has size $(size(y))")) + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpby!(a, In(xc[i]), b, InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.rmul!(x::DArray, a::Number) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.rmul!(InOut(c), a) + end + end + return x +end + +function LinearAlgebra.lmul!(a::Number, x::DArray) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.lmul!(a, InOut(c)) + end + end + return x +end function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T Ac = parent(A).chunks Ac_upper = [] diff --git a/src/array/mul.jl b/src/array/mul.jl index 94a8a235c..322b27760 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -491,6 +491,21 @@ function _repartition_matvecmul(C, A, B, transA::Char)::Tuple{Blocks{1}, Blocks{ partC = (dimA,) return Blocks(partC...), Blocks(partA...), Blocks(partB...) end +""" + matvecmul!(C, transA::Char, A, B, alpha, beta) + +Tile-level matrix-vector multiply computing `C = alpha*op(A)*B + beta*C` in +place on the (dense) output vector `C`, where `op` is determined by `transA` +(`'N'`, `'T'`, `'C'`). Dispatches on the tile types: dense tiles use BLAS, while +sparse tiles (e.g. `DSparseArray`) provide their own method (in a package +extension) using a sparse matrix-vector product. This is the matvec analogue of +[`matmatmul!`](@ref). +""" +function matvecmul!(C, transA::Char, A, B, alpha, beta) + BLAS.gemv!(transA, alpha, A, B, beta, C) + return C +end + function gemv_dagger!( C::DVector{T}, transA::Char, @@ -533,26 +548,26 @@ function gemv_dagger!( # A: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[m, k]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end else # A: [Conj]Trans — C's blocks index A's column-blocks for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[k, m]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end end diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 04687a67e..65fdf317c 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -1,13 +1,24 @@ """ - DSparseMatrix{T} <: AbstractMatrix{T} + DSparseArray{T,N} <: AbstractArray{T,N} -A sparse matrix container, for which the contained matrix may be replaced with -a new one to support in-place operations. Designed to work well with -Datadeps algorithms. +A sparse array container, for which the contained array may be replaced with a +new one to support in-place operations. Designed to work well with Datadeps +algorithms: writes that reallocate (and grow/shrink) the inner sparse storage +are hidden behind the wrapper's stable identity, so Datadeps aliasing tracking +remains valid (see `aliases_as_whole`). + +`DSparseVector{T}` and `DSparseMatrix{T}` are aliases for the 1- and 2-dimensional +cases. The wrapper is general over `N` so it can hold sparse vectors, matrices, +and (eventually) higher-order sparse tensors (e.g. Finch tensors). """ -mutable struct DSparseMatrix{T} <: AbstractMatrix{T} +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} mat end +const DSparseVector{T} = DSparseArray{T,1} +const DSparseMatrix{T} = DSparseArray{T,2} + +# Convenience constructor inferring element type and dimensionality from `x`. +DSparseArray(x) = DSparseArray{eltype(x),ndims(x)}(x) _sparse_collect(M) = collect(M) # Allocate a destination tile matching the inner storage's sparse backend. The @@ -18,57 +29,56 @@ _sparse_collect(M) = collect(M) # `similar` yields tensor formats that destabilize later kernels) override this. _sparse_similar(mat, ::Type{T}, dims::Dims) where {T} = similar(mat, T, dims) -Base.eltype(M::DSparseMatrix{T}) where T = T -Base.size(M::DSparseMatrix) = size(M.mat) -Base.ndims(M::DSparseMatrix) = 2 -Base.iterate(M::DSparseMatrix) = iterate(M.mat) -Base.iterate(M::DSparseMatrix, state) = iterate(M.mat, state) -Base.similar(M::DSparseMatrix, ::Type{T}, dims::Tuple{Int, Int}) where T = - DSparseMatrix{T}(_sparse_similar(M.mat, T, dims)) -Base.collect(M::DSparseMatrix) = _sparse_collect(M.mat) +Base.eltype(M::DSparseArray{T}) where T = T +Base.size(M::DSparseArray) = size(M.mat) +Base.iterate(M::DSparseArray) = iterate(M.mat) +Base.iterate(M::DSparseArray, state) = iterate(M.mat, state) +Base.similar(M::DSparseArray, ::Type{T}, dims::Dims) where T = + DSparseArray{T,length(dims)}(_sparse_similar(M.mat, T, dims)) +Base.collect(M::DSparseArray) = _sparse_collect(M.mat) # N.B. hash and aliasing shouldn't change even if M.mat changes. -# This is the key property that makes `DSparseMatrix` safe for Datadeps: +# This is the key property that makes `DSparseArray` safe for Datadeps: # the aliasing identity is tied to the (stable) wrapper object, not to the # inner storage, which may be reallocated (and grow/shrink) on writes. -Base.hash(M::DSparseMatrix, h::UInt) = hash(objectid(M), hash(DSparseMatrix, h)) +Base.hash(M::DSparseArray, h::UInt) = hash(objectid(M), hash(DSparseArray, h)) # Sparse containers must alias as a single, indivisible unit. Their storage is # reallocated (and may grow/shrink) on writes, so it is unsafe to alias any # *part* of the container independently. We therefore force *all* access to a -# `DSparseMatrix` to resolve to the container's stable whole-object +# `DSparseArray` to resolve to the container's stable whole-object # `ObjectAliasing` -- including access via views, transposes, adjoints, and # reshapes -- rather than e.g. a `StridedAliasing` over storage that may have # since moved. `aliases_as_whole` opts the type into this behavior, and Datadeps' -# `aliasing_root` resolves any wrapper of a `DSparseMatrix` back to the container +# `aliasing_root` resolves any wrapper of a `DSparseArray` back to the container # before computing aliasing. Note: this must return a *bare* `ObjectAliasing`, as # `aliases_as_whole(::ObjectAliasing)` relies on it to drive whole-object copies. -aliasing(M::DSparseMatrix, _=identity) = ObjectAliasing(M) -aliases_as_whole(::DSparseMatrix) = true +aliasing(M::DSparseArray, _=identity) = ObjectAliasing(M) +aliases_as_whole(::DSparseArray) = true -# A `DSparseMatrix` has no meaningful raw data pointer (its storage may be +# A `DSparseArray` has no meaningful raw data pointer (its storage may be # reallocated/resized). This trap ensures that if aliasing ever tries to treat a -# wrapper of a `DSparseMatrix` as strided memory -- e.g. a new array-wrapper type +# wrapper of a `DSparseArray` as strided memory -- e.g. a new array-wrapper type # that `aliasing_root` failed to resolve via `Base.parent` -- it fails loudly # instead of silently corrupting Datadeps' aliasing tracking. -function Base.pointer(::DSparseMatrix) - throw(ArgumentError("`pointer(::DSparseMatrix)` is intentionally unsupported: \ - a DSparseMatrix may reallocate its storage, so it must be aliased as a \ +function Base.pointer(::DSparseArray) + throw(ArgumentError("`pointer(::DSparseArray)` is intentionally unsupported: \ + a DSparseArray may reallocate its storage, so it must be aliased as a \ whole object via `Dagger.aliasing`. If you reached here through Datadeps \ aliasing of an array wrapper, ensure that wrapper implements `Base.parent` \ - so that `Dagger.aliasing_root` can resolve it to the DSparseMatrix.")) + so that `Dagger.aliasing_root` can resolve it to the DSparseArray.")) end -# Forward indexing to the inner matrix. Ranges/colons are supported so that +# Forward indexing to the inner array. Ranges/colons are supported so that # views (used by copy-buffering between mismatched partitionings) work. -Base.getindex(M::DSparseMatrix, I...) = getindex(M.mat, I...) -Base.setindex!(M::DSparseMatrix, v, I...) = setindex!(M.mat, v, I...) +Base.getindex(M::DSparseArray, I...) = getindex(M.mat, I...) +Base.setindex!(M::DSparseArray, v, I...) = setindex!(M.mat, v, I...) # Whole-tile copy. This is how Datadeps moves a tile between workers via # `move!`: we *reallocate* the inner storage rather than mutating in place, # since sparse storage cannot generally be updated in place. The wrapper's # identity (and thus its aliasing) is preserved. -function Base.copyto!(dst::DSparseMatrix, src::DSparseMatrix) +function Base.copyto!(dst::DSparseArray, src::DSparseArray) dst.mat = _sparse_copy(src.mat) return dst end @@ -78,15 +88,15 @@ _sparse_copy(mat) = copy(mat) # Wrapping hook used when materializing tiles (e.g. in `distribute`). Backends # overload this (e.g. in package extensions) to wrap freshly-created tiles in a -# container that Datadeps can track (such as `DSparseMatrix` for sparse tiles); +# container that Datadeps can track (such as `DSparseArray` for sparse tiles); # everything else is left untouched. maybe_wrap_tile(x) = x # Partial-range tile copy used by copy-buffering (e.g. when matmul operands need # to be repartitioned). Some backends (Finch) forbid `setindex!`, so we route # through a backend hook that returns the (possibly reallocated) inner storage. -# `DSparseMatrix` hides the reallocation from Datadeps. -function copyto_view!(Bpart::DSparseMatrix, Brange, Apart, Arange) +# `DSparseArray` hides the reallocation from Datadeps. +function copyto_view!(Bpart::DSparseArray, Brange, Apart, Arange) Bpart.mat = _sparse_copyto_view!(Bpart.mat, Brange, view(Apart, Arange)) return end @@ -98,10 +108,11 @@ function _sparse_copyto_view!(mat, Brange, src) return mat end -#Base.BroadcastStyle(::Type{<:DSparseMatrix}) = FinchStyle() +#Base.BroadcastStyle(::Type{<:DSparseArray}) = FinchStyle() -LinearAlgebra.norm2(M::DSparseMatrix) = LinearAlgebra.norm2(M.mat) +LinearAlgebra.norm2(M::DSparseArray) = LinearAlgebra.norm2(M.mat) +# Matrix-specific operations dispatch on the 2-D alias `DSparseMatrix`. function matmatmul!( C::DSparseMatrix, transA::Char, @@ -115,6 +126,13 @@ function matmatmul!( return matmatmul!(C, transA, transB, A.mat, B.mat, alpha, beta) end +# Sparse matrix-vector multiply: `C = alpha*op(A)*B + beta*C` with a sparse +# matrix tile `A` and dense vectors `B`/`C`. Unwrap to the inner storage so each +# backend can provide an efficient (in-place on dense `C`) method. +function matvecmul!(C, transA::Char, A::DSparseMatrix, B, alpha, beta) + return matvecmul!(C, transA, A.mat, B, alpha, beta) +end + function transpose_tile end function copytile!(A::DSparseMatrix, B::DSparseMatrix) A.mat = transpose_tile(B.mat) diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 6bf37cb14..86964e6a9 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -69,7 +69,7 @@ function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices x = unwrap(x) v = view(x, slices...) - # A view of a whole-object container (e.g. `DSparseMatrix`) must alias the + # A view of a whole-object container (e.g. `DSparseArray`) must alias the # entire container; `aliasing_unwrapped` resolves that (otherwise it just # aliases the view), and crucially does so here where `v` lives. return aliasing_unwrapped(v) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index c530e0766..3ee709e6d 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -90,7 +90,7 @@ function compute_remainder_for_arg!(state::DataDepsState, write_num::Int; compute_syncdeps::Bool=true) owner_space = state.arg_owner[arg_w] - # Whole-object containers (e.g. `DSparseMatrix`) reallocate their entire + # Whole-object containers (e.g. `DSparseArray`) reallocate their entire # storage on every write, so they can never be *partially* current in a # memory space -- the whole object is current in exactly one space: the most # recent writer (or the origin owner if never written). There is thus never a @@ -463,7 +463,7 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: multi_span_copy!(to_s, from_s, dep_mod.spans) return end - # N.B. Whole-object containers (e.g. `DSparseMatrix`) never reach here: their + # N.B. Whole-object containers (e.g. `DSparseArray`) never reach here: their # storage reallocates on write, so they can't be span-copied. They are routed # to a whole-object `FullCopy` in `compute_remainder_for_arg!` instead. diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index a1eefe52e..015ca502c 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -424,7 +424,7 @@ treated as a single, indivisible unit, i.e. it is never safe to alias (or to consider current) only *part* of it. For an object: containers whose backing storage may be reallocated or resized on -write (such as `DSparseMatrix`) return `true`. This causes [`aliasing_root`](@ref) +write (such as `DSparseArray`) return `true`. This causes [`aliasing_root`](@ref) to resolve any view/wrapper of them to the whole container, so all access funnels through the container's own `aliasing`. By contract, such a container's `aliasing` must be a bare [`ObjectAliasing`](@ref). diff --git a/test/array/allocation.jl b/test/array/allocation.jl index 83dba36d2..4548f3ed1 100644 --- a/test/array/allocation.jl +++ b/test/array/allocation.jl @@ -99,7 +99,8 @@ end @test size(Xsp) == dims AT = length(dims) == 2 ? SparseMatrixCSC : SparseVector Ach = fetch(Xsp.chunks[1]) - @test Ach isa AT{T} + @test Ach isa Dagger.DSparseArray{T,length(dims)} + @test Ach.mat isa AT{T} AXsp = collect(Xsp) @test AXsp isa Array{T,length(dims)} @test AXsp == collect(Xsp) diff --git a/test/array/linalg/core.jl b/test/array/linalg/core.jl index 4c66ac34a..5de6fefd1 100644 --- a/test/array/linalg/core.jl +++ b/test/array/linalg/core.jl @@ -23,3 +23,57 @@ end DA = DArray(A) @test isapprox(norm(A), norm(DA)) end + +@testset "BLAS-1 vector ops" begin + @testset "T=$T part=$(part.blocksize)" for T in (Float64, ComplexF64), part in (Blocks(16), Blocks(4)) + n = 16 + x = rand(T, n) + y = rand(T, n) + a = T(2) + b = T(3) + Dx = distribute(x, part) + + # dot (conjugates first arg for complex) + @test dot(Dx, distribute(y, part)) ≈ dot(x, y) + + # axpy!: y += a*x + Dy = distribute(copy(y), part) + axpy!(a, Dx, Dy) + @test collect(Dy) ≈ a .* x .+ y + + # axpby!: y = a*x + b*y + Dy = distribute(copy(y), part) + axpby!(a, Dx, b, Dy) + @test collect(Dy) ≈ a .* x .+ b .* y + + # rmul!/lmul!: scale in place + Dz = distribute(copy(x), part) + rmul!(Dz, a) + @test collect(Dz) ≈ a .* x + Dz = distribute(copy(x), part) + lmul!(a, Dz) + @test collect(Dz) ≈ a .* x + end + + # Operands with *different* partitionings must still work (aligned internally + # via maybe_copy_buffered), and the in-place result keeps its own layout. + @testset "mismatched partitionings T=$T" for T in (Float64, ComplexF64) + n = 16 + x = rand(T, n) + y = rand(T, n) + a, b = T(2), T(3) + Dx = distribute(x, Blocks(4)) + Dy = distribute(y, Blocks(8)) + + @test dot(Dx, Dy) ≈ dot(x, y) + + Dy2 = distribute(copy(y), Blocks(8)) + axpy!(a, Dx, Dy2) + @test collect(Dy2) ≈ a .* x .+ y + @test Dy2.partitioning == Blocks(8) # output layout preserved + + Dy3 = distribute(copy(y), Blocks(8)) + axpby!(a, Dx, b, Dy3) + @test collect(Dy3) ≈ a .* x .+ b .* y + end +end diff --git a/test/array/linalg/matmul_finch.jl b/test/array/linalg/matmul_finch.jl index 59ecc87a5..061874f1e 100644 --- a/test/array/linalg/matmul_finch.jl +++ b/test/array/linalg/matmul_finch.jl @@ -70,4 +70,36 @@ else end end end + + # Sparse matrix-vector multiply (SpMV): Finch sparse matrix tiles, dense vectors. + function test_finch_spmv!(T, n, part) + bs = part.blocksize[1] + DSA = Finch.fsprand(part, T, (n, n), FINCH_DENSITY) + A = collect(DSA) + x = rand(T, n) + Dx = distribute(x, Blocks(bs)) + + @test collect(DSA * Dx) ≈ A * x + @test collect(transpose(DSA) * Dx) ≈ transpose(A) * x + @test collect(DSA' * Dx) ≈ A' * x + + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (A * x) + beta * y + end + + const FINCH_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), + ] + + @testset "Finch SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in FINCH_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_spmv!(T, n, part) + end + end + end end diff --git a/test/array/linalg/matmul_sparse.jl b/test/array/linalg/matmul_sparse.jl index 0bbdb3c4a..cc2e8274e 100644 --- a/test/array/linalg/matmul_sparse.jl +++ b/test/array/linalg/matmul_sparse.jl @@ -54,6 +54,42 @@ const SPARSE_QUICK_CASES = [ end end +# Sparse matrix-vector multiply (SpMV): sparse matrix tiles, dense vectors. This +# is the core kernel for distributed iterative (Krylov) solvers. +function test_sparse_spmv!(T, n, part) + bs = part.blocksize[1] + SA = sprand(T, n, n, SPARSE_DENSITY) + x = rand(T, n) + + DSA = distribute(SA, part) + Dx = distribute(x, Blocks(bs)) + + ## Out-of-place (N / T / C) + @test collect(DSA * Dx) ≈ SA * x + @test collect(transpose(DSA) * Dx) ≈ transpose(SA) * x + @test collect(DSA' * Dx) ≈ SA' * x + + ## In-place 5-arg mul!: y = alpha*A*x + beta*y + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (SA * x) + beta * y +end + +const SPARSE_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), +] + +@testset "Sparse SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in SPARSE_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_spmv!(T, n, part) + end + end +end + # Any *partial* or *reinterpreted* access to a sparse container (views, # transposes, adjoints, reshapes, and combinations thereof) must alias the # *entire* container, so that Datadeps never tracks stale sub-spans of storage From 3b1c811407e18b9d28757f7327e2d61dbf18cb8c Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 7 Jun 2026 13:34:38 -0700 Subject: [PATCH 15/43] DArray/sparse: Add iterative solvers via Krylov.jl --- Project.toml | 3 + ext/KrylovExt.jl | 36 ++++ src/Dagger.jl | 1 + src/array/iterativesolvers.jl | 237 ++++++++++++++++++++++++++ src/array/sparse.jl | 4 + test/Project.toml | 1 + test/array/linalg/iterativesolvers.jl | 156 +++++++++++++++++ test/runtests.jl | 1 + 8 files changed, 439 insertions(+) create mode 100644 ext/KrylovExt.jl create mode 100644 src/array/iterativesolvers.jl create mode 100644 test/array/linalg/iterativesolvers.jl diff --git a/Project.toml b/Project.toml index 62ed8252c..2c953d185 100644 --- a/Project.toml +++ b/Project.toml @@ -39,6 +39,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" @@ -57,6 +58,7 @@ GraphVizExt = "GraphViz" GraphVizSimpleExt = "Colors" IntelExt = "oneAPI" JSON3Ext = "JSON3" +KrylovExt = "Krylov" LinuxPerfExt = "LinuxPerf" MPIExt = "MPI" MetalExt = "Metal" @@ -83,6 +85,7 @@ GraphViz = "0.2" Graphs = "1" JSON3 = "1" KernelAbstractions = "0.9" +Krylov = "0.10" LinuxPerf = "0.4.2" MPI = "0.20.22" MacroTools = "0.5" diff --git a/ext/KrylovExt.jl b/ext/KrylovExt.jl new file mode 100644 index 000000000..5d36ba2fa --- /dev/null +++ b/ext/KrylovExt.jl @@ -0,0 +1,36 @@ +module KrylovExt + +import Krylov +import Dagger +import Dagger: DVector, DMatrix +import LinearAlgebra + +# Distributed iterative solvers, implemented on top of Krylov.jl. +# +# We never form `A⁻¹`: Krylov only needs `mul!(y, A, x)` (and `mul!(y, A', x)` +# for two-sided methods), which Dagger provides over `DVector`s via its +# distributed SpMV/`gemv_dagger!` path. The remaining requirements -- `dot`, +# `norm`, `axpy!`, `axpby!`, `rmul!`, `copyto!`, `fill!`, broadcasting -- are the +# distributed BLAS-1 ops in `Dagger`. +# +# Workspace vectors are allocated through a `KrylovConstructor` built from +# `similar(b)`, so every internal vector inherits `b`'s element type *and* +# partitioning (`similar` preserves both). This keeps all `mul!`/`dot`/`axpy!` +# operands on a compatible chunk layout, which the distributed kernels require. + +function _dagger_krylov(method::Symbol, A, b::DVector; kwargs...) + kc = Krylov.KrylovConstructor(similar(b)) + workspace = Krylov.krylov_workspace(Val(method), kc) + Krylov.krylov_solve!(workspace, A, b; kwargs...) + return Krylov.solution(workspace), Krylov.statistics(workspace) +end + +Dagger.krylov_solve(method::Symbol, A, b::DVector; kwargs...) = + _dagger_krylov(method, A, b; kwargs...) + +Dagger.cg(A, b::DVector; kwargs...) = _dagger_krylov(:cg, A, b; kwargs...) +Dagger.minres(A, b::DVector; kwargs...) = _dagger_krylov(:minres, A, b; kwargs...) +Dagger.gmres(A, b::DVector; kwargs...) = _dagger_krylov(:gmres, A, b; kwargs...) +Dagger.bicgstab(A, b::DVector; kwargs...) = _dagger_krylov(:bicgstab, A, b; kwargs...) + +end # module KrylovExt diff --git a/src/Dagger.jl b/src/Dagger.jl index bc80217ac..15e5961a1 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -141,6 +141,7 @@ include("array/trsm.jl") include("array/lu.jl") include("array/qr.jl") include("array/svd.jl") +include("array/iterativesolvers.jl") # GPU include("gpu.jl") diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl new file mode 100644 index 000000000..f6b12752b --- /dev/null +++ b/src/array/iterativesolvers.jl @@ -0,0 +1,237 @@ +# Distributed iterative (Krylov) linear solvers. +# +# The user-facing entry points (`Dagger.cg`, `Dagger.minres`, `Dagger.gmres`, +# `Dagger.bicgstab`, and the generic `Dagger.krylov_solve`) live here so the +# solver backend can evolve without changing user code. The actual solves are +# implemented in `ext/KrylovExt.jl`, which is loaded when `Krylov` is available +# (`using Krylov`). The matrix-free building blocks they rely on -- distributed +# `mul!`/SpMV and the BLAS-1 vector ops -- live in `mul.jl`/`linalg.jl`. + +""" + cg(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric positive-definite system `A x = b` with the conjugate +gradient method, distributed over `A`'s and `b`'s chunks. `A` may be a +`DMatrix` (dense or sparse-backed) or any object supporting `mul!(y, A, x)` over +`DVector`s (matrix-free). Requires `Krylov.jl` to be loaded. + +See also [`minres`](@ref), [`gmres`](@ref), [`bicgstab`](@ref), and +[`krylov_solve`](@ref). Keyword arguments are forwarded to `Krylov.cg!`. +""" +function cg end + +""" + minres(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric (possibly indefinite) system `A x = b` with MINRES. +Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function minres end + +""" + gmres(A, b::DVector; M=I, N=I, restart=false, memory=20, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with restarted GMRES. +`M`/`N` are left/right preconditioners. Requires `Krylov.jl` to be loaded. +See [`cg`](@ref). +""" +function gmres end + +""" + bicgstab(A, b::DVector; M=I, N=I, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with BiCGStab (short +recurrence, low memory). Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function bicgstab end + +""" + krylov_solve(method::Symbol, A, b::DVector; kwargs...) -> (x::DVector, stats) + +Generic entry point dispatching to the iterative `method` (`:cg`, `:minres`, +`:gmres`, `:bicgstab`). Requires `Krylov.jl` to be loaded. +""" +function krylov_solve end + +# Friendly fallbacks: these generic methods are shadowed by the more specific +# `(A, b::DVector)` methods added in `ext/KrylovExt.jl` once Krylov is loaded. +_krylov_required(name) = throw(ArgumentError( + "Dagger.$name requires Krylov.jl. Run `using Krylov` (or `import Krylov`) \ + to enable distributed iterative solvers.")) +cg(A, b; kwargs...) = _krylov_required(:cg) +minres(A, b; kwargs...) = _krylov_required(:minres) +gmres(A, b; kwargs...) = _krylov_required(:gmres) +bicgstab(A, b; kwargs...) = _krylov_required(:bicgstab) +krylov_solve(method::Symbol, A, b; kwargs...) = _krylov_required(:krylov_solve) + +# --- Preconditioners ------------------------------------------------------ +# A preconditioner object `P` represents the (approximate) *inverse* operator +# `M⁻¹`: applying it (`mul!(y, P, x)`) computes `y = M⁻¹ x`, which is exactly an +# ordinary matrix-vector product by the operator `P` stands for (no inversion +# happens at apply time -- any reciprocals/factorizations are precomputed once). +# This matches Krylov's `ldiv=false` convention, where the object passed as `M` +# is applied via `mul!(y, M, x)` to compute `y ← M⁻¹ x`, so these are passed +# straight through as `M=P`. +# +# These preconditioners are backend-agnostic: they operate on `DVector`s (and a +# `DMatrix`'s diagonal tiles) and need no solver backend, so they live in core. +# +# A note on the square-tile requirement below: the iterative solvers allocate +# *all* workspace vectors as `similar(b)` (one partitioning), and each must +# serve as both the length-`n` input and length-`n` output of `mul!(y, A, x)`. +# `gemv_dagger!` requires the input to match `A`'s column blocks and the output +# to match `A`'s row blocks, which is only simultaneously possible when those +# block sizes are equal. So the operator must use square tiles regardless of +# preconditioning; the preconditioners simply share that constraint. (Any +# uniform square tile size is fine, including a ragged final block.) + +""" + AbstractDaggerPreconditioner + +Supertype for Dagger's distributed preconditioners. A preconditioner `P` +represents an (approximate) inverse operator `M⁻¹` and applies it via +`mul!(y, P, x)` (`y = M⁻¹ x`) over `DVector`s, so it can be passed to the +solvers as `M=P` (with `ldiv=false`, the default). +""" +abstract type AbstractDaggerPreconditioner end + +""" + JacobiPreconditioner(A::DMatrix) + +Diagonal (Jacobi) preconditioner. The object represents the inverse-diagonal +operator `M⁻¹ = inv(Diagonal(A))`; it precomputes and stores the reciprocal +diagonal `dinv = 1 ./ diag(A)`, so applying it (`mul!(y, P, x)`) is a single +elementwise multiply `y = dinv .* x` per chunk -- an ordinary product by the +stored operator, not an inversion. Cheap to build (one diagonal extraction per +diagonal tile) and to apply. + +Requires square diagonal tiles (equal row/column block sizes); see the note +above on why the solver itself imposes this. Repartition `A` with equal block +sizes if needed. +""" +struct JacobiPreconditioner{V<:DVector} <: AbstractDaggerPreconditioner + dinv::V +end + +JacobiPreconditioner(A::DMatrix) = JacobiPreconditioner(_jacobi_dinv(A)) + +# Per-tile kernel: write the reciprocal diagonal of `tile` into `out`. +_set_inv_diag!(out, tile) = (out .= inv.(LinearAlgebra.diag(tile)); return nothing) + +# Validate that `A` has square diagonal tiles and return (n, Ac, mt, blocksize). +function _check_square_tiles(A::DMatrix, who) + n = LinearAlgebra.checksquare(A) + Ac = A.chunks + mt, nt = size(Ac) + mb, nb = A.partitioning.blocksize + # For a square matrix with square tiles, mt == nt follows automatically; we + # only need to reject non-square tiles (mb != nb). + mb == nb || throw(ArgumentError( + "$who requires square diagonal tiles (got block size $(mb)x$(nb)); \ + repartition A with equal block sizes (this matches the iterative \ + solver's own requirement on the operator)")) + return n, Ac, mt, mb +end + +function _jacobi_dinv(A::DMatrix{T}) where T + n, Ac, mt, mb = _check_square_tiles(A, "JacobiPreconditioner") + dinv = DVector{T}(undef, Blocks(mb), n) + dc = dinv.chunks + Dagger.spawn_datadeps() do + for i in 1:mt + Dagger.@spawn _set_inv_diag!(Out(dc[i]), In(Ac[i, i])) + end + end + return dinv +end + +function LinearAlgebra.mul!(y::DVector, P::JacobiPreconditioner, x::DVector) + part = P.dinv.partitioning + maybe_copy_buffered(P.dinv => part, x => part, y => part) do dinv, x, y + dc, xc, yc = dinv.chunks, x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(yc) + Dagger.@spawn _jacobi_apply!(Out(yc[i]), In(dc[i]), In(xc[i])) + end + end + end + return y +end +_jacobi_apply!(y, dinv, x) = (y .= dinv .* x; return nothing) + +""" + BlockJacobiPreconditioner(A::DMatrix) + +Block-Jacobi preconditioner. The object represents the block-diagonal inverse +operator `M⁻¹ = blockdiag(A₁₁, …, A_kk)⁻¹`, where `Aᵢᵢ` are `A`'s diagonal tiles. +Applying it (`mul!(y, P, x)`) solves `Aᵢᵢ yᵢ = xᵢ` independently per block -- +embarrassingly parallel and a natural fit for the tiled layout. Stronger than +`JacobiPreconditioner` (it captures intra-block coupling); a single tile recovers +an exact solve. + +Requires square diagonal tiles (see the note above and `JacobiPreconditioner`). + + +Each diagonal tile is factorized *once* at construction. A factorization object +cannot be moved between workers (dense `LU` has no `move!`; sparse `UmfpackLU` +holds external/UMFPACK resources tied to its process), so each factor is pinned +(via `tochunk(..., ProcessScope)`) to the worker that owns its tile, and every +apply for that block is scheduled there (`compute_scope`). Datadeps then moves +only the (movable) vector chunks to the factor, never the factor itself. +""" +struct BlockJacobiPreconditioner{F,S} <: AbstractDaggerPreconditioner + factors::F # cached per-tile factorizations (pinned to their workers) + scopes::S # the `ProcessScope` each factor/apply is pinned to + part::Blocks{1} # partitioning of the vectors it applies to + n::Int +end + +# Factorize a diagonal tile. Backends override for their storage (e.g. sparse +# tiles factorize the inner `SparseMatrixCSC`); the default is a dense LU factor. +_factorize_tile(A) = LinearAlgebra.lu(A) + +# Factorize on the current worker and pin the result there: the returned `Chunk` +# is process-scoped, so the factor can never be moved off this worker. +function _factorize_tile_pinned(Aii) + F = _factorize_tile(Aii) + proc = Dagger.task_processor() + return tochunk(F, proc, ProcessScope(root_worker_id(proc))) +end + +_tile_scope(c::Chunk) = ProcessScope(root_worker_id(c)) +_tile_scope(t::DTask) = ProcessScope(root_worker_id(fetch(t; raw=true))) + +function BlockJacobiPreconditioner(A::DMatrix) + n, Ac, mt, mb = _check_square_tiles(A, "BlockJacobiPreconditioner") + factors = Vector{Any}(undef, mt) + scopes = Vector{ProcessScope}(undef, mt) + for i in 1:mt + tile = Ac[i, i] + scope = _tile_scope(tile) + scopes[i] = scope + # Factor where the tile lives; the result is pinned to that worker. + factors[i] = Dagger.@spawn compute_scope=scope _factorize_tile_pinned(tile) + end + return BlockJacobiPreconditioner(factors, scopes, Blocks(mb), n) +end + +function LinearAlgebra.mul!(y::DVector, P::BlockJacobiPreconditioner, x::DVector) + part = P.part + maybe_copy_buffered(x => part, y => part) do x, y + xc, yc = x.chunks, y.chunks + length(yc) == length(P.factors) || throw(DimensionMismatch( + "BlockJacobiPreconditioner has $(length(P.factors)) blocks but the \ + vector has $(length(yc)) chunks")) + Dagger.spawn_datadeps() do + for i in eachindex(yc) + # Pin the solve to the factor's worker; the factor is passed as an + # untracked arg (read-only, already-pinned) so datadeps never tries + # to move it -- only the vector chunks are moved to this worker. + Dagger.@spawn compute_scope=P.scopes[i] _blockjacobi_solve!(Out(yc[i]), P.factors[i], In(xc[i])) + end + end + end + return y +end +# Apply the cached block factorization: solve `Aᵢᵢ yᵢ = xᵢ`. +_blockjacobi_solve!(y, F, x) = (y .= F \ x; return nothing) diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 65fdf317c..608845fcd 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -133,6 +133,10 @@ function matvecmul!(C, transA::Char, A::DSparseMatrix, B, alpha, beta) return matvecmul!(C, transA, A.mat, B, alpha, beta) end +# Factorize the inner sparse storage (e.g. sparse LU/UMFPACK) for block-Jacobi, +# rather than the `DSparseArray` wrapper (whose generic `lu` would densify). +_factorize_tile(M::DSparseArray) = LinearAlgebra.lu(M.mat) + function transpose_tile end function copytile!(A::DSparseMatrix, B::DSparseMatrix) A.mat = transpose_tile(B.mat) diff --git a/test/Project.toml b/test/Project.toml index edc738961..1e04b44ef 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -11,6 +11,7 @@ GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94" diff --git a/test/array/linalg/iterativesolvers.jl b/test/array/linalg/iterativesolvers.jl new file mode 100644 index 000000000..fd7927104 --- /dev/null +++ b/test/array/linalg/iterativesolvers.jl @@ -0,0 +1,156 @@ +# Distributed iterative (Krylov) linear-solver tests. +# +# Exercises the matrix-free Krylov integration (`Dagger.cg`/`minres`/`gmres`/ +# `bicgstab` + the generic `krylov_solve`) over both dense and sparse-backed +# `DMatrix` operators, plus the Jacobi preconditioner. Reference solutions come +# from a dense direct solve. +# +# julia test/runtests.jl --test array/linalg/iterativesolvers + +using Krylov + +# Strongly diagonally-dominant tridiagonal SPD matrix. The large diagonal keeps +# the condition number small so the Krylov methods converge in a handful of +# iterations (keeping the distributed test fast). `inv(diag) == 1/4`. +const SPD_DIAG = 4.0 + +laplacian_1d(T, n) = SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(SPD_DIAG), n), + 1 => fill(-one(T), n - 1), +) + +# Add a first-order advection term -> nonsymmetric, still well-conditioned. +function advection_diffusion_1d(T, n) + return laplacian_1d(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end + +@testset "Iterative solvers (Krylov)" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + @testset "SPD operator ($(backend))" for backend in (:dense, :sparse) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Generic entry point. + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = 1e-12, rtol = 1e-10) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "nonsymmetric operator ($(backend))" for backend in (:dense, :sparse) + Asp = advection_diffusion_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "complex SPD (Hermitian) operator" begin + # Real SPD tridiagonal is Hermitian as a complex matrix. + Asp = SparseArrays.spdiagm( + -1 => fill(ComplexF64(-1), n - 1), + 0 => fill(ComplexF64(SPD_DIAG), n), + 1 => fill(ComplexF64(-1), n - 1), + ) + b = rand(ComplexF64, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, Db_part) + x, stats = Dagger.cg(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(1 / SPD_DIAG, n) # 1/diag + + # Apply: y = M⁻¹ x = dinv .* x. + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (1 / SPD_DIAG) .* b + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Non-square block grid must be rejected with a helpful error. + DA_ragged = distribute(Matrix(Asp), Blocks(k, k ÷ 2)) + @test_throws ArgumentError Dagger.JacobiPreconditioner(DA_ragged) + end + + @testset "block-Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockJacobiPreconditioner(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes block-Jacobi an *exact* solve, so PCG converges + # essentially immediately. + DA1 = distribute(Adense, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = Dagger.BlockJacobiPreconditioner(DA1) + x1, s1 = Dagger.cg(DA1, Db1; M = P1, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test s1.solved + @test s1.niter <= 2 + @test collect(x1) ≈ xref rtol = 1e-8 + + @test_throws ArgumentError Dagger.BlockJacobiPreconditioner(distribute(Adense, Blocks(k, k ÷ 2))) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 0d9d7b945..b02db950c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -44,6 +44,7 @@ tests = [ ("Array - LinearAlgebra - Solve", "array/linalg/solve.jl"), ("Array - LinearAlgebra - QR", "array/linalg/qr.jl"), ("Array - LinearAlgebra - SVD", "array/linalg/svd.jl"), + ("Array - LinearAlgebra - Iterative Solvers", "array/linalg/iterativesolvers.jl"), ("Array - Permute", "array/permute.jl"), ("Array - Random", "array/random.jl"), ("Array - Stencils", "array/stencil.jl"), From 66001b9c1fc651d1dc150a563ea273aac5c245dc Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 8 Jun 2026 14:27:17 -0700 Subject: [PATCH 16/43] DArray/sparse: Add AMG and ILU preconditioner integrations --- Project.toml | 6 + ext/AlgebraicMultigridExt.jl | 45 +++++++ ext/IncompleteLUExt.jl | 36 ++++++ src/array/iterativesolvers.jl | 171 ++++++++++++++++++-------- src/array/sparse.jl | 4 + test/Project.toml | 2 + test/array/linalg/iterativesolvers.jl | 50 ++++++++ 7 files changed, 264 insertions(+), 50 deletions(-) create mode 100644 ext/AlgebraicMultigridExt.jl create mode 100644 ext/IncompleteLUExt.jl diff --git a/Project.toml b/Project.toml index 2c953d185..54ccd7e50 100644 --- a/Project.toml +++ b/Project.toml @@ -32,11 +32,13 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" +AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce" +IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" @@ -51,11 +53,13 @@ oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] AbstractFFTsExt = "AbstractFFTs" +AlgebraicMultigridExt = ["AlgebraicMultigrid", "SparseArrays"] CUDAExt = "CUDA" DistributionsExt = "Distributions" FinchExt = "Finch" GraphVizExt = "GraphViz" GraphVizSimpleExt = "Colors" +IncompleteLUExt = ["IncompleteLU", "SparseArrays"] IntelExt = "oneAPI" JSON3Ext = "JSON3" KrylovExt = "Krylov" @@ -72,6 +76,7 @@ SparseArraysExt = "SparseArrays" AMDGPU = "1, 2" AbstractFFTs = "1.5.0" Adapt = "4" +AlgebraicMultigrid = "1" CUDA = "3, 4, 5" Colors = "0.12, 0.13" DataFrames = "1" @@ -83,6 +88,7 @@ Finch = "1.2.9" GPUArraysCore = "0.2.0" GraphViz = "0.2" Graphs = "1" +IncompleteLU = "0.2" JSON3 = "1" KernelAbstractions = "0.9" Krylov = "0.10" diff --git a/ext/AlgebraicMultigridExt.jl b/ext/AlgebraicMultigridExt.jl new file mode 100644 index 000000000..dd68bbcfa --- /dev/null +++ b/ext/AlgebraicMultigridExt.jl @@ -0,0 +1,45 @@ +module AlgebraicMultigridExt + +import AlgebraicMultigrid +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Algebraic-multigrid preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). The AMG +# hierarchy's expensive setup runs once per tile at construction; each apply runs +# one V-cycle (a sequence of SpMVs, smoother sweeps, and a coarse solve). + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _amg_operator(tile, method::Symbol; kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + ml = if method === :ruge_stuben + AlgebraicMultigrid.ruge_stuben(S; kwargs...) + elseif method === :smoothed_aggregation + AlgebraicMultigrid.smoothed_aggregation(S; kwargs...) + else + throw(ArgumentError("AMGPreconditioner: unknown method $(method); use \ + :ruge_stuben or :smoothed_aggregation")) + end + return AlgebraicMultigrid.aspreconditioner(ml) +end + +function Dagger.AMGPreconditioner(A::DMatrix; method::Symbol=:ruge_stuben, kwargs...) + build = tile -> _amg_operator(tile, method; kwargs...) + return Dagger._build_block_preconditioner(Dagger.AMGPreconditioner, A, + "AMGPreconditioner", build) +end + +# An AMG preconditioner applies a V-cycle via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, p::AlgebraicMultigrid.Preconditioner, x) = + (LinearAlgebra.ldiv!(y, p, x); return nothing) + +# The hierarchy is read-only and pinned to its worker; it is never written and +# its internals (a `Vector{Level}`) aren't introspectable by datadeps, so treat +# it as non-aliasing rather than recursing into it. +Dagger.aliasing(::AlgebraicMultigrid.Preconditioner) = Dagger.NoAliasing() + +end # module AlgebraicMultigridExt diff --git a/ext/IncompleteLUExt.jl b/ext/IncompleteLUExt.jl new file mode 100644 index 000000000..06ecab522 --- /dev/null +++ b/ext/IncompleteLUExt.jl @@ -0,0 +1,36 @@ +module IncompleteLUExt + +import IncompleteLU +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Block incomplete-LU preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). Each tile +# gets an ILU factorization (drop tolerance `τ`) once at construction; the apply +# is a forward/backward substitution per block. + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _ilu_operator(tile; τ=0.001, kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + return IncompleteLU.ilu(S; τ=τ, kwargs...) +end + +function Dagger.BlockILUPreconditioner(A::DMatrix; kwargs...) + build = tile -> _ilu_operator(tile; kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockILUPreconditioner, A, + "BlockILUPreconditioner", build) +end + +# An ILU factorization applies via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, F::IncompleteLU.ILUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::IncompleteLU.ILUFactorization) = Dagger.NoAliasing() + +end # module IncompleteLUExt diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl index f6b12752b..e499d2a1c 100644 --- a/src/array/iterativesolvers.jl +++ b/src/array/iterativesolvers.jl @@ -159,79 +159,150 @@ function LinearAlgebra.mul!(y::DVector, P::JacobiPreconditioner, x::DVector) end _jacobi_apply!(y, dinv, x) = (y .= dinv .* x; return nothing) -""" - BlockJacobiPreconditioner(A::DMatrix) - -Block-Jacobi preconditioner. The object represents the block-diagonal inverse -operator `M⁻¹ = blockdiag(A₁₁, …, A_kk)⁻¹`, where `Aᵢᵢ` are `A`'s diagonal tiles. -Applying it (`mul!(y, P, x)`) solves `Aᵢᵢ yᵢ = xᵢ` independently per block -- -embarrassingly parallel and a natural fit for the tiled layout. Stronger than -`JacobiPreconditioner` (it captures intra-block coupling); a single tile recovers -an exact solve. - -Requires square diagonal tiles (see the note above and `JacobiPreconditioner`). +# --- Block-diagonal preconditioners --------------------------------------- +# A family of preconditioners of the form `M⁻¹ = blockdiag(op₁, …, op_k)`, where +# each `opⱼ` approximates the inverse of `A`'s `j`-th diagonal tile `Aⱼⱼ`. They +# share all machinery and differ only in how a per-tile operator is *built*: +# +# - BlockJacobiPreconditioner: exact tile solve via `lu` (dense or sparse). +# - BlockILUPreconditioner: incomplete LU per tile (needs IncompleteLU.jl). +# - AMGPreconditioner: an AMG hierarchy per tile (needs AlgebraicMultigrid.jl). +# +# Building per tile is embarrassingly parallel and a natural fit for the tiled +# layout. A single tile (`Blocks(n, n)`) makes any of these a *global* +# preconditioner over the whole matrix; many tiles make it a scalable +# block-Jacobi / additive-Schwarz variant that trades some convergence for +# parallelism. All require square diagonal tiles (see `JacobiPreconditioner`). +# +# A per-tile operator (`lu`/ILU factor, AMG hierarchy) generally cannot be moved +# between workers (dense `LU` has no `move!`; `UmfpackLU`/AMG hold process-bound +# resources). So each operator is built *once* and pinned (via +# `tochunk(..., ProcessScope)`) to the worker that owns its tile, and every apply +# for that block is scheduled there (`compute_scope`); datadeps then moves only +# the (movable) vector chunks to the operator, never the operator itself. +""" + AbstractBlockPreconditioner <: AbstractDaggerPreconditioner -Each diagonal tile is factorized *once* at construction. A factorization object -cannot be moved between workers (dense `LU` has no `move!`; sparse `UmfpackLU` -holds external/UMFPACK resources tied to its process), so each factor is pinned -(via `tochunk(..., ProcessScope)`) to the worker that owns its tile, and every -apply for that block is scheduled there (`compute_scope`). Datadeps then moves -only the (movable) vector chunks to the factor, never the factor itself. +Block-diagonal preconditioner: holds one per-diagonal-tile operator `opⱼ` +(`y ← opⱼ⁻¹ x`), pinned to its tile's worker, and applies them independently per +block. Concrete subtypes (`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, +`AMGPreconditioner`) share these fields: `ops`, `scopes`, `part`, `n`. """ -struct BlockJacobiPreconditioner{F,S} <: AbstractDaggerPreconditioner - factors::F # cached per-tile factorizations (pinned to their workers) - scopes::S # the `ProcessScope` each factor/apply is pinned to - part::Blocks{1} # partitioning of the vectors it applies to - n::Int -end +abstract type AbstractBlockPreconditioner <: AbstractDaggerPreconditioner end -# Factorize a diagonal tile. Backends override for their storage (e.g. sparse -# tiles factorize the inner `SparseMatrixCSC`); the default is a dense LU factor. -_factorize_tile(A) = LinearAlgebra.lu(A) +_tile_scope(c::Chunk) = ProcessScope(root_worker_id(c)) +_tile_scope(t::DTask) = ProcessScope(root_worker_id(fetch(t; raw=true))) -# Factorize on the current worker and pin the result there: the returned `Chunk` -# is process-scoped, so the factor can never be moved off this worker. -function _factorize_tile_pinned(Aii) - F = _factorize_tile(Aii) +# Build the per-tile operator on the current worker and pin the result there: the +# returned `Chunk` is process-scoped, so it can never be moved off this worker. +function _build_tile_pinned(build, tile) + op = build(tile) proc = Dagger.task_processor() - return tochunk(F, proc, ProcessScope(root_worker_id(proc))) + return tochunk(op, proc, ProcessScope(root_worker_id(proc))) end -_tile_scope(c::Chunk) = ProcessScope(root_worker_id(c)) -_tile_scope(t::DTask) = ProcessScope(root_worker_id(fetch(t; raw=true))) - -function BlockJacobiPreconditioner(A::DMatrix) - n, Ac, mt, mb = _check_square_tiles(A, "BlockJacobiPreconditioner") - factors = Vector{Any}(undef, mt) +# Build a block preconditioner of type `Ctor` by applying `build` to each +# diagonal tile (pinned to the tile's worker). `build(tile) -> op` returns a +# per-tile operator supporting the block apply below. +function _build_block_preconditioner(Ctor, A::DMatrix, who, build) + n, Ac, mt, mb = _check_square_tiles(A, who) + ops = Vector{Any}(undef, mt) scopes = Vector{ProcessScope}(undef, mt) for i in 1:mt tile = Ac[i, i] scope = _tile_scope(tile) scopes[i] = scope - # Factor where the tile lives; the result is pinned to that worker. - factors[i] = Dagger.@spawn compute_scope=scope _factorize_tile_pinned(tile) + ops[i] = Dagger.@spawn compute_scope=scope _build_tile_pinned(build, tile) end - return BlockJacobiPreconditioner(factors, scopes, Blocks(mb), n) + return Ctor(ops, scopes, Blocks(mb), n) end -function LinearAlgebra.mul!(y::DVector, P::BlockJacobiPreconditioner, x::DVector) +function LinearAlgebra.mul!(y::DVector, P::AbstractBlockPreconditioner, x::DVector) part = P.part maybe_copy_buffered(x => part, y => part) do x, y xc, yc = x.chunks, y.chunks - length(yc) == length(P.factors) || throw(DimensionMismatch( - "BlockJacobiPreconditioner has $(length(P.factors)) blocks but the \ - vector has $(length(yc)) chunks")) + length(yc) == length(P.ops) || throw(DimensionMismatch( + "$(nameof(typeof(P))) has $(length(P.ops)) blocks but the vector has \ + $(length(yc)) chunks")) Dagger.spawn_datadeps() do for i in eachindex(yc) - # Pin the solve to the factor's worker; the factor is passed as an - # untracked arg (read-only, already-pinned) so datadeps never tries - # to move it -- only the vector chunks are moved to this worker. - Dagger.@spawn compute_scope=P.scopes[i] _blockjacobi_solve!(Out(yc[i]), P.factors[i], In(xc[i])) + # Pin the apply to the operator's worker; the operator is passed as + # an untracked arg (read-only, already-pinned) so datadeps never + # moves it -- only the vector chunks are moved to this worker. + Dagger.@spawn compute_scope=P.scopes[i] _block_apply!(Out(yc[i]), P.ops[i], In(xc[i])) end end end return y end -# Apply the cached block factorization: solve `Aᵢᵢ yᵢ = xᵢ`. -_blockjacobi_solve!(y, F, x) = (y .= F \ x; return nothing) + +# Apply one block operator: `y = op⁻¹ x`. Default uses `\`; backends override for +# operator types where `\` is unavailable (e.g. ILU/AMG use `ldiv!`). +_block_apply!(y, op, x) = (y .= op \ x; return nothing) + +# Underlying matrix of a tile (overridden for `DSparseArray` in `sparse.jl`). +_tile_matrix(A) = A + +""" + BlockJacobiPreconditioner(A::DMatrix) + +Block-Jacobi preconditioner: `M⁻¹ = blockdiag(A₁₁, …, A_kk)⁻¹`. Each diagonal +tile is factorized *once* with `lu` (sparse or dense) and the apply solves +`Aᵢᵢ yᵢ = xᵢ`. Stronger than `JacobiPreconditioner` (captures intra-block +coupling); a single tile recovers an exact solve. See +[`AbstractBlockPreconditioner`](@ref). Requires square diagonal tiles. +""" +struct BlockJacobiPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F # cached per-tile factorizations (pinned to their workers) + scopes::S # the `ProcessScope` each operator/apply is pinned to + part::Blocks{1} # partitioning of the vectors it applies to + n::Int +end + +# Factorize a diagonal tile. Backends override for their storage (e.g. sparse +# tiles factorize the inner `SparseMatrixCSC`); the default is a dense LU factor. +_factorize_tile(A) = LinearAlgebra.lu(A) + +BlockJacobiPreconditioner(A::DMatrix) = + _build_block_preconditioner(BlockJacobiPreconditioner, A, "BlockJacobiPreconditioner", _factorize_tile) + +""" + BlockILUPreconditioner(A::DMatrix; τ=0.001, kwargs...) + +Block incomplete-LU preconditioner: an ILU factorization (with drop tolerance +`τ`) of each diagonal tile, applied per block. Cheaper setup than a full block +solve, good as a general-purpose preconditioner. Requires `IncompleteLU.jl` to +be loaded and sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct BlockILUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/IncompleteLUExt.jl`). +BlockILUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockILUPreconditioner requires IncompleteLU.jl. Run `using IncompleteLU` \ + to enable block incomplete-LU preconditioning.")) + +""" + AMGPreconditioner(A::DMatrix; method=:ruge_stuben, kwargs...) + +Algebraic-multigrid preconditioner: builds an AMG hierarchy (`method` is +`:ruge_stuben` or `:smoothed_aggregation`) for each diagonal tile and applies a +V-cycle per block. Near mesh-independent convergence for elliptic (Poisson-like) +operators. With one tile this is global AMG; with many tiles it is a scalable +block/additive-Schwarz AMG. Requires `AlgebraicMultigrid.jl` to be loaded and +sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct AMGPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/AlgebraicMultigridExt.jl`). +AMGPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.AMGPreconditioner requires AlgebraicMultigrid.jl. Run \ + `using AlgebraicMultigrid` to enable algebraic-multigrid preconditioning.")) diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 608845fcd..3e56c37b0 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -137,6 +137,10 @@ end # rather than the `DSparseArray` wrapper (whose generic `lu` would densify). _factorize_tile(M::DSparseArray) = LinearAlgebra.lu(M.mat) +# Expose the inner sparse storage so preconditioner backends (ILU, AMG) build on +# the actual `SparseMatrixCSC` rather than the `DSparseArray` wrapper. +_tile_matrix(M::DSparseArray) = M.mat + function transpose_tile end function copytile!(A::DSparseMatrix, B::DSparseMatrix) A.mat = transpose_tile(B.mat) diff --git a/test/Project.toml b/test/Project.toml index 1e04b44ef..4cb663d5b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" @@ -9,6 +10,7 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" diff --git a/test/array/linalg/iterativesolvers.jl b/test/array/linalg/iterativesolvers.jl index fd7927104..ca7865f29 100644 --- a/test/array/linalg/iterativesolvers.jl +++ b/test/array/linalg/iterativesolvers.jl @@ -8,6 +8,8 @@ # julia test/runtests.jl --test array/linalg/iterativesolvers using Krylov +using AlgebraicMultigrid +using IncompleteLU # Strongly diagonally-dominant tridiagonal SPD matrix. The large diagonal keeps # the condition number small so the Krylov methods converge in a handful of @@ -153,4 +155,52 @@ end @test_throws ArgumentError Dagger.BlockJacobiPreconditioner(distribute(Adense, Blocks(k, k ÷ 2))) end + + @testset "AMG preconditioner ($(method))" for method in (:ruge_stuben, :smoothed_aggregation) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "$(backend)" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.AMGPreconditioner(DA; method = method) + # Apply is a V-cycle approximating M⁻¹ x; just check it runs + is finite + # (and repeatable, exercising the cached, pinned hierarchy). + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "block-ILU preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockILUPreconditioner(DA; τ = 0.01) + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @test_throws ArgumentError Dagger.BlockILUPreconditioner(distribute(Adense, Blocks(k, k ÷ 2)); τ = 0.01) + end end From 33e40cd0bb9be46107baf05c719fcf634fc40e46 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 23 Jul 2026 17:25:15 -0700 Subject: [PATCH 17/43] DArray/sparse: Add sparse direct solvers (KLU/UMFPACK) Add `Dagger.klu` (PureKLU) and `Dagger.splu` (PureUMFPACK) whole-matrix direct solves for sparse `DMatrix`, plus block direct preconditioners (`BlockKLUPreconditioner`/`BlockUMFPACKPreconditioner`). The pure-Julia factorizations are movable, so the factor is gathered/factored on the worker owning the most tiles and pinned there; solves move only O(n) vectors. `Dagger.splu` also exposes two opt-in parallel variants (PureUMFPACK): - `distributed=true, method=:trsv`: re-tile the L/U factors as sparse DMatrices and run a blocked datadeps forward/backward substitution. - `distributed=true, method=:schur`: single-level METIS vertex-separator domain decomposition (`ext/MetisExt.jl`) that factors interior blocks in parallel across workers and reduces/factors the Schur complement. Supporting changes: `_gather_sparse`/`_sparse_copy_of` hooks in SparseArraysExt, Metis/PureKLU/PureUMFPACK weak deps + extensions, and a multi-worker `array/linalg/sparsedirect` test suite (314 tests). Co-authored-by: Cursor --- Project.toml | 9 + ext/MetisExt.jl | 75 +++++ ext/PureKLUExt.jl | 38 +++ ext/PureUMFPACKExt.jl | 63 ++++ ext/SparseArraysExt.jl | 21 ++ src/Dagger.jl | 1 + src/array/iterativesolvers.jl | 41 +++ src/array/sparsedirect.jl | 518 ++++++++++++++++++++++++++++++ test/Project.toml | 3 + test/array/linalg/sparsedirect.jl | 249 ++++++++++++++ test/runtests.jl | 1 + 11 files changed, 1019 insertions(+) create mode 100644 ext/MetisExt.jl create mode 100644 ext/PureKLUExt.jl create mode 100644 ext/PureUMFPACKExt.jl create mode 100644 src/array/sparsedirect.jl create mode 100644 test/array/linalg/sparsedirect.jl diff --git a/Project.toml b/Project.toml index 54ccd7e50..d25a38dc2 100644 --- a/Project.toml +++ b/Project.toml @@ -45,8 +45,11 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" +Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01" +PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" @@ -66,8 +69,11 @@ KrylovExt = "Krylov" LinuxPerfExt = "LinuxPerf" MPIExt = "MPI" MetalExt = "Metal" +MetisExt = ["Metis", "SparseArrays"] OpenCLExt = "OpenCL" PlotsExt = ["DataFrames", "Plots"] +PureKLUExt = ["PureKLU", "SparseArrays"] +PureUMFPACKExt = ["PureUMFPACK", "SparseArrays"] PythonExt = "PythonCall" ROCExt = "AMDGPU" SparseArraysExt = "SparseArrays" @@ -97,12 +103,15 @@ MPI = "0.20.22" MacroTools = "0.5" MemPool = "0.4.17" Metal = "1.1" +Metis = "1" NextLA = "0.2.2" OnlineStats = "1" OpenCL = "0.10" Plots = "1" PrecompileTools = "1.2" Preferences = "1.4.3" +PureKLU = "1" +PureUMFPACK = "0.1" PythonCall = "0.9" Requires = "1" ScopedValues = "1.1" diff --git a/ext/MetisExt.jl b/ext/MetisExt.jl new file mode 100644 index 000000000..9f4a2623b --- /dev/null +++ b/ext/MetisExt.jl @@ -0,0 +1,75 @@ +module MetisExt + +import Metis +import SparseArrays +import SparseArrays: SparseMatrixCSC, nnz, rowvals, nzrange +import LinearAlgebra +import Dagger + +# Single-level (non-recursive) METIS k-way partition + vertex-separator +# construction for Stage 4c Schur-complement domain decomposition. +# Recursive / nested dissection is future work. + +# Symmetric adjacency pattern of `A + Aᵀ` (structure only, unit weights). +function _symmetrize_pattern(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + n = size(A, 1) + S = A + transpose(A) + return SparseMatrixCSC{Tv,Ti}(n, n, S.colptr, S.rowval, ones(Tv, nnz(S))) +end + +""" + Dagger._nested_dissection_partition(A, nparts) -> (; perm, interiors, separator) + +K-way METIS partition of the undirected adjacency of `A + Aᵀ`, then a vertex +separator `Γ` = every vertex with a neighbor in a different part. Interiors are +the remaining vertices grouped by part (empty parts dropped). The permutation is +`[I₁; …; Iₖ; Γ]`. + +For tiny graphs (or `nparts == 1`), returns a trivial partition: one interior +covering `1:n` and an empty separator (caller should fall back to a serial +factorization). +""" +function Dagger._nested_dissection_partition(A::SparseMatrixCSC, nparts::Integer) + n = LinearAlgebra.checksquare(A) + nparts = Int(nparts) + nparts >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $nparts")) + + # Too small to partition usefully → trivial (serial) partition. + if nparts == 1 || n < 4 || n <= nparts + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + pattern = _symmetrize_pattern(A) + g = Metis.graph(pattern) + part = Int.(Metis.partition(g, nparts)) + + is_sep = falses(n) + rows = rowvals(pattern) + @inbounds for v in 1:n + for k in nzrange(pattern, v) + u = rows[k] + u == v && continue + if part[u] != part[v] + is_sep[v] = true + break + end + end + end + + separator = findall(is_sep) + interiors = Vector{Vector{Int}}() + for p in 1:nparts + Ii = findall(i -> !is_sep[i] && part[i] == p, 1:n) + isempty(Ii) || push!(interiors, Ii) + end + + # Everything landed in the separator (or no interiors) → serial fallback. + if isempty(interiors) + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + perm = vcat(interiors..., separator) + return (; perm, interiors, separator) +end + +end # module MetisExt diff --git a/ext/PureKLUExt.jl b/ext/PureKLUExt.jl new file mode 100644 index 000000000..c2bbf22f5 --- /dev/null +++ b/ext/PureKLUExt.jl @@ -0,0 +1,38 @@ +module PureKLUExt + +import PureKLU +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (KLU) integration. KLU factorizations are pure-Julia data +# (serializable/movable), so they back both a whole-matrix direct solve +# (`Dagger.klu`) and per-tile block direct solves (`BlockKLUPreconditioner`), +# the latter via Dagger's block-preconditioner machinery. + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +function Dagger.klu(A::DMatrix; kwargs...) + return Dagger._spawn_direct_factorization(A, S -> PureKLU.klu(S; kwargs...)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockKLUPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureKLU.klu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockKLUPreconditioner, A, + "BlockKLUPreconditioner", build) +end + +# KLU supports in-place `ldiv!`, so apply each block solve without allocating. +Dagger._block_apply!(y, F::PureKLU.KLUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureKLU.KLUFactorization) = Dagger.NoAliasing() + +end # module PureKLUExt diff --git a/ext/PureUMFPACKExt.jl b/ext/PureUMFPACKExt.jl new file mode 100644 index 000000000..1817a5010 --- /dev/null +++ b/ext/PureUMFPACKExt.jl @@ -0,0 +1,63 @@ +module PureUMFPACKExt + +import PureUMFPACK +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (UMFPACK-style, pure-Julia) integration. Like KLU, these +# factorizations are plain Julia data (movable), backing both a whole-matrix +# direct solve (`Dagger.splu`) and per-tile block direct solves +# (`BlockUMFPACKPreconditioner`). With `distributed=true`: +# - `method=:trsv` (default) — Stage 4b tiled triangular solves +# - `method=:schur` — Stage 4c Schur-complement domain decomposition + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _umfpack_factorize(; kwargs...) + return S -> PureUMFPACK.splu(_as_sparse(S); kwargs...) +end + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +# Opt-in distributed paths select Stage 4b (`:trsv`) or Stage 4c (`:schur`). +function Dagger.splu(A::DMatrix; distributed::Bool=false, method::Symbol=:trsv, + blocksize=nothing, nparts=nothing, kwargs...) + factorize = _umfpack_factorize(; kwargs...) + if !distributed + return Dagger._spawn_direct_factorization(A, factorize) + end + if method === :trsv + F = Dagger._spawn_direct_factorization(A, factorize) + return Dagger._make_distributed_sparse_lu(F, A; blocksize) + elseif method === :schur + return Dagger._make_schur_sparse_lu(A, factorize; nparts) + else + throw(ArgumentError( + "Dagger.splu distributed method must be :trsv or :schur, got $(repr(method))")) + end +end + +# `(Rs .* A)[p, q] == L * U` with unit-lower `L` (explicit stored ones on the +# diagonal; solve treats it as unit diagonal via `UnitLowerTriangular`). +function Dagger._extract_lu_factors(F::PureUMFPACK.PureLU) + return (F.L, F.U, copy(F.p), copy(F.q), copy(F.Rs)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureUMFPACK.splu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockUMFPACKPreconditioner, A, + "BlockUMFPACKPreconditioner", build) +end + +# `PureLU` provides `\` (no `ldiv!`); the default `_block_apply!` (`y .= op \ x`) +# already uses it, so no apply override is needed. + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureUMFPACK.PureLU) = Dagger.NoAliasing() + +end # module PureUMFPACKExt diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index a6c6755fa..7ef7c28c9 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -8,6 +8,27 @@ import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, D # Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) + +# Assemble already-local tiles into one global `SparseMatrixCSC` without +# densifying: unwrap each tile, offset its (i,j) by the precomputed subdomain +# offsets, and build from triplets. Intended to run inside a worker-scoped task +# (the scheduler moves the tile chunks there); used by `Dagger.klu` / `Dagger.splu`. +function Dagger._gather_sparse(::Type{T}, tiles, row_offsets, col_offsets, m, n) where T + Is = Int[]; Js = Int[]; Vs = T[] + for k in 1:length(tiles) + tile = SparseMatrixCSC(Dagger._tile_matrix(tiles[k])) + ti, tj, tv = SparseArrays.findnz(tile) + append!(Is, ti .+ row_offsets[k]) + append!(Js, tj .+ col_offsets[k]) + append!(Vs, tv) + end + return SparseArrays.sparse(Is, Js, Vs, m, n) +end + +# Dense → sparse for Stage-4c Schur complements (fill-in is expected). +Dagger._sparse_copy_of(S::AbstractMatrix) = SparseArrays.sparse(S) +Dagger._sparse_copy_of(S::SparseMatrixCSC) = S + # Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseArray(x) Dagger.maybe_wrap_tile(x::SparseVector) = DSparseArray(x) diff --git a/src/Dagger.jl b/src/Dagger.jl index 15e5961a1..3fe2a54ba 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -142,6 +142,7 @@ include("array/lu.jl") include("array/qr.jl") include("array/svd.jl") include("array/iterativesolvers.jl") +include("array/sparsedirect.jl") # GPU include("gpu.jl") diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl index e499d2a1c..95e9100ac 100644 --- a/src/array/iterativesolvers.jl +++ b/src/array/iterativesolvers.jl @@ -306,3 +306,44 @@ end AMGPreconditioner(A; kwargs...) = throw(ArgumentError( "Dagger.AMGPreconditioner requires AlgebraicMultigrid.jl. Run \ `using AlgebraicMultigrid` to enable algebraic-multigrid preconditioning.")) + +""" + BlockKLUPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using KLU: an exact (pure-Julia) KLU factorization of +each diagonal tile, applied per block. With a single tile this is an exact direct +solve; with many tiles it is an exact-block-Jacobi preconditioner. Requires +`PureKLU.jl` to be loaded and sparse-backed tiles. See also +[`AbstractBlockPreconditioner`](@ref) and the whole-matrix solver +[`Dagger.klu`](@ref). +""" +struct BlockKLUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureKLUExt.jl`). +BlockKLUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockKLUPreconditioner requires PureKLU.jl. Run `using PureKLU` to \ + enable block KLU preconditioning.")) + +""" + BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using a pure-Julia UMFPACK-style LU of each diagonal +tile, applied per block. With a single tile this is an exact direct solve; with +many tiles it is an exact-block-Jacobi preconditioner. Requires `PureUMFPACK.jl` +to be loaded and sparse-backed tiles. See also [`AbstractBlockPreconditioner`](@ref) +and the whole-matrix solver [`Dagger.splu`](@ref). +""" +struct BlockUMFPACKPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureUMFPACKExt.jl`). +BlockUMFPACKPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockUMFPACKPreconditioner requires PureUMFPACK.jl. Run \ + `using PureUMFPACK` to enable block UMFPACK preconditioning.")) diff --git a/src/array/sparsedirect.jl b/src/array/sparsedirect.jl new file mode 100644 index 000000000..508b24ac4 --- /dev/null +++ b/src/array/sparsedirect.jl @@ -0,0 +1,518 @@ +# Sparse direct solvers. +# +# Pure-Julia sparse LU backends (KLU via `PureKLU.jl`, UMFPACK-style via +# `PureUMFPACK.jl`) provide *direct* solves for sparse systems. Unlike the +# C-bound `SparseArrays.UmfpackLU` (whose factorization holds process-local +# UMFPACK resources and can't be moved between workers), these factorizations are +# plain Julia data: serializable and movable, so Dagger can schedule them freely. +# +# Two entry points are exposed (implemented in `ext/PureKLUExt.jl` / +# `ext/PureUMFPACKExt.jl`): +# +# - A whole-matrix direct solve: `Dagger.klu(A)` / `Dagger.splu(A)` gather the +# (sparse) `DMatrix` onto a single worker (the one owning the most tiles), +# assemble one `SparseMatrixCSC`, factor it once, and return a +# `DaggerSparseLU` whose factor is pinned there. Solves move only the RHS / +# solution vectors. This is a correctness-first, single-worker factorization +# for systems that fit on one worker (e.g. coarse-grid solves). +# - An opt-in distributed solve path: `Dagger.splu(A; distributed=true)` +# (PureUMFPACK only) reuses that factorization, extracts the triangular +# factors into sparse `DMatrix`es, and solves with a tiled datadeps +# forward/backward substitution (`DistributedSparseLU`, `method=:trsv`). +# - A Schur-complement domain-decomposition path: +# `Dagger.splu(A; distributed=true, method=:schur)` factors interior blocks +# in parallel across workers and solves via a separator Schur system +# (`DistributedSchurLU`; requires Metis.jl). +# - Per-tile block direct solves (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`, in `iterativesolvers.jl`), which reuse the +# block-preconditioner machinery. + +""" + klu(A::DMatrix; kwargs...) -> DaggerSparseLU + +Direct sparse LU factorization of a sparse `DMatrix` using KLU. Tiles are +gathered onto one worker, assembled into a single `SparseMatrixCSC`, and factored +once; the factor stays pinned there. Solve with `F \\ b` or `ldiv!(x, F, b)` over +`DVector`s. Requires `PureKLU.jl` to be loaded. + +KLU is well suited to unsymmetric systems with near-triangular structure (e.g. +circuit simulation). See also [`splu`](@ref). +""" +function klu end + +""" + splu(A::DMatrix; distributed=false, method=:trsv, blocksize=nothing, nparts=nothing, kwargs...) + +Direct sparse LU factorization of a sparse `DMatrix` using a pure-Julia +UMFPACK-style multifrontal solver. + +- `distributed=false` → [`DaggerSparseLU`](@ref): gather+factor on one worker. +- `distributed=true, method=:trsv` (default) → [`DistributedSparseLU`](@ref): + Stage-4a factor plus tiled datadeps triangular solves (`blocksize` controls + the solve tiling). +- `distributed=true, method=:schur` → [`DistributedSchurLU`](@ref): single-level + Schur-complement domain decomposition (METIS separator; `nparts` defaults to + `max(2, nworkers())`). Requires `Metis.jl`. + +Requires `PureUMFPACK.jl`. See also [`klu`](@ref). +""" +function splu end + +_klu_required() = throw(ArgumentError( + "Dagger.klu requires PureKLU.jl. Run `using PureKLU` to enable sparse direct \ + (KLU) solves.")) +_splu_required() = throw(ArgumentError( + "Dagger.splu requires PureUMFPACK.jl. Run `using PureUMFPACK` to enable sparse \ + direct (UMFPACK) solves.")) +klu(A; kwargs...) = _klu_required() +splu(A; kwargs...) = _splu_required() + +""" + DaggerSparseLU + +A direct sparse factorization of a sparse `DMatrix`, produced by +[`Dagger.klu`](@ref) or [`Dagger.splu`](@ref). The underlying factorization is +pure-Julia and pinned to a single worker (`fact` is a `Chunk`/`DTask` with +`scope`). Solve `A x = b` with `F \\ b` (returns a `DVector` partitioned like +`b`) or `ldiv!(x, F, b)`. +""" +struct DaggerSparseLU{F,S,P} + fact::F # pinned factorization (`Chunk` / `DTask`) + scope::S # `ProcessScope` of the factor's worker + n::Int # system size (square) + part::P # default vector partitioning (column blocks of `A`) +end + +Base.size(F::DaggerSparseLU) = (F.n, F.n) +Base.size(F::DaggerSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +""" + DistributedSparseLU + +A sparse LU factorization with triangular factors stored as sparse `DMatrix`es, +produced by `Dagger.splu(A; distributed=true)`. Solves use a tiled datadeps +forward/backward substitution over `L` and `U` (PureUMFPACK only). +""" +struct DistributedSparseLU{L,U,P,Q,R,Part} + L::L # unit-lower triangular factor (`DMatrix` of sparse tiles) + U::U # upper triangular factor (`DMatrix` of sparse tiles) + p::P # row permutation + q::Q # column permutation + Rs::R # row scaling + n::Int # system size (square) + part::Part # vector partitioning matching `L`/`U` tile rows +end + +Base.size(F::DistributedSparseLU) = (F.n, F.n) +Base.size(F::DistributedSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +# Assemble local tiles into one `SparseMatrixCSC` (no densification). The real +# method is added by `SparseArraysExt`; without `SparseArrays` there are no +# sparse tiles to gather. +function _gather_sparse end +_gather_sparse(args...) = throw(ArgumentError( + "gathering a sparse DMatrix requires SparseArrays to be loaded")) + +# Backend-specific extraction of `(L, U, p, q, Rs)` from a factored object. +# Implemented for `PureUMFPACK.PureLU` in `PureUMFPACKExt`. +function _extract_lu_factors end +_extract_lu_factors(F) = throw(ArgumentError( + "distributed sparse LU requires a PureUMFPACK.PureLU factor (got $(typeof(F))). \ + Use `Dagger.splu(A; distributed=true)` with PureUMFPACK.jl loaded.")) + +# Choose the worker that already owns the most tiles of `A` (minimizes gather +# traffic). Ties break toward the lowest worker id. +function _select_factor_scope(A::DMatrix) + counts = Dict{Int,Int}() + for c in A.chunks + wid = _tile_scope(c).wid + counts[wid] = get(counts, wid, 0) + 1 + end + best_wid, best_count = 0, -1 + for (wid, cnt) in counts + if cnt > best_count || (cnt == best_count && wid < best_wid) + best_wid, best_count = wid, cnt + end + end + return ProcessScope(best_wid) +end + +# Runs on the chosen worker: assemble from tiles (already moved here by the +# scheduler), factorize, and pin the factor to this process. +function _assemble_and_factor_pinned(factorize, ::Type{T}, row_offsets, col_offsets, + m, n, tiles...) where T + S = _gather_sparse(T, tiles, row_offsets, col_offsets, m, n) + F = factorize(S) + proc = task_processor() + return tochunk(F, proc, ProcessScope(root_worker_id(proc))) +end + +# Shared path for `klu` / `splu`: pick a worker, spawn assemble+factor there. +function _spawn_direct_factorization(A::DMatrix{T}, factorize) where T + n = LinearAlgebra.checksquare(A) + scope = _select_factor_scope(A) + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = Ac[i, j] + idx += 1 + end + mA, nA = size(A) + part = Blocks(A.partitioning.blocksize[2]) + fact = spawn(_assemble_and_factor_pinned, Options(; compute_scope=scope), + factorize, T, row_offsets, col_offsets, mA, nA, tiles...) + return DaggerSparseLU(fact, scope, n, part) +end + +# Gather RHS chunks, solve against the pinned factor, return the dense solution. +function _direct_solve(fact, bparts...) + b = length(bparts) == 1 ? bparts[1] : reduce(vcat, bparts) + return fact \ b +end + +function Base.:\(F::DaggerSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Solve on the factor's worker (factor stays pinned); only the O(n) solution + # returns here to be redistributed like `b`. + x = fetch(spawn(_direct_solve, Options(; compute_scope=F.scope), + F.fact, b.chunks...)) + return distribute(x, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DaggerSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4b: distributed triangular solves -------------------------------- + +_fetch_lu_factors(fact) = _extract_lu_factors(fact) + +function _resolve_solve_blocksize(A::DMatrix, blocksize) + if blocksize === nothing + return Int(min(A.partitioning.blocksize...)) + end + bs = Int(blocksize) + bs > 0 || throw(ArgumentError("blocksize must be positive, got $bs")) + return bs +end + +# Re-tile PureLU factors into sparse `DMatrix`es after a Stage-4a factorization. +function _make_distributed_sparse_lu(F::DaggerSparseLU, A::DMatrix; blocksize=nothing) + bs = _resolve_solve_blocksize(A, blocksize) + L, U, p, q, Rs = fetch(spawn(_fetch_lu_factors, Options(; compute_scope=F.scope), F.fact)) + part2 = Blocks(bs, bs) + part1 = Blocks(bs) + return DistributedSparseLU(distribute(L, part2), distribute(U, part2), + p, q, Rs, F.n, part1) +end + +# Tile kernels: unwrap `DSparseArray` → `SparseMatrixCSC`, never densify. +function _sparse_unit_ltrsv!(Ltile, x::AbstractVector) + LinearAlgebra.ldiv!(UnitLowerTriangular(_tile_matrix(Ltile)), x) + return +end + +function _sparse_utrsv!(Utile, x::AbstractVector) + LinearAlgebra.ldiv!(UpperTriangular(_tile_matrix(Utile)), x) + return +end + +function _sparse_gemv_sub!(Atile, x::AbstractVector, y::AbstractVector) + matvecmul!(y, 'N', Atile, x, -one(eltype(y)), one(eltype(y))) + return +end + +# Blocked forward substitution for unit-lower `L` (mirrors `trsv!` lower/'N'). +function _sparse_trsv_forward!(L::DMatrix, c::DVector) + Lc, Cc = L.chunks, c.chunks + nt = length(Cc) + size(Lc, 1) == nt && size(Lc, 2) == nt || throw(DimensionMismatch( + "L tile grid $(size(Lc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in 1:nt + Dagger.@spawn _sparse_unit_ltrsv!(In(Lc[k, k]), InOut(Cc[k])) + for i in (k + 1):nt + Dagger.@spawn _sparse_gemv_sub!(In(Lc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +# Blocked backward substitution for upper `U` (mirrors `trsv!` upper/'N'). +function _sparse_trsv_backward!(U::DMatrix, c::DVector) + Uc, Cc = U.chunks, c.chunks + nt = length(Cc) + size(Uc, 1) == nt && size(Uc, 2) == nt || throw(DimensionMismatch( + "U tile grid $(size(Uc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in reverse(1:nt) + Dagger.@spawn _sparse_utrsv!(In(Uc[k, k]), InOut(Cc[k])) + for i in 1:(k - 1) + Dagger.@spawn _sparse_gemv_sub!(In(Uc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +function Base.:\(F::DistributedSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Permute + row-scale (driver O(n)); matches PureUMFPACK `_solve_factor`. + b_local = collect(b) + T = promote_type(eltype(b_local), eltype(F.Rs), eltype(F.U)) + c_local = Vector{T}(undef, F.n) + p, q, Rs = F.p, F.q, F.Rs + @inbounds for k in eachindex(p) + c_local[k] = Rs[p[k]] * b_local[p[k]] + end + c = distribute(c_local, F.part) + _sparse_trsv_forward!(F.L, c) + _sparse_trsv_backward!(F.U, c) + c_local = collect(c) + x_local = similar(c_local) + @inbounds for k in eachindex(q) + x_local[q[k]] = c_local[k] + end + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4c: Schur-complement domain decomposition ------------------------- +# +# Single-level (non-recursive) k-way Schur factorization. Recursive / nested +# dissection of the separator is future work. + +""" + DistributedSchurLU + +Sparse LU via single-level Schur-complement domain decomposition, produced by +`Dagger.splu(A; distributed=true, method=:schur)`. Interior blocks are factored +in parallel (pinned round-robin across workers); the separator Schur system is +factored on one worker. Solve with `F \\ b` / `ldiv!(x, F, b)`. +""" +struct DistributedSchurLU{F,W,G,SF,S,P} + interior_facts::Vector{F} # pinned `Aᵢᵢ` factors + W::Vector{W} # pinned `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (dense columns) + AΓi::Vector{G} # pinned `A_Γi` blocks + interior_scopes::Vector{S} + Sfact::SF # pinned Schur factor (`nothing` if `|Γ|==0`) + Sscope::S + interiors::Vector{Vector{Int}} + separator::Vector{Int} + perm::Vector{Int} + n::Int + part::P +end + +Base.size(F::DistributedSchurLU) = (F.n, F.n) +Base.size(F::DistributedSchurLU, i::Integer) = i <= 2 ? F.n : 1 + +# METIS partition + separator (implemented in `MetisExt`). +function _nested_dissection_partition end +_nested_dissection_partition(A, nparts) = throw(ArgumentError( + "Dagger.splu(...; method=:schur) requires Metis.jl. Run `using Metis` to \ + enable Schur-complement distributed sparse LU.")) + +# Gather a sparse `DMatrix` into one `SparseMatrixCSC` on the caller (pattern +# assembly for the partitioner; no densification). +function _collect_sparse_dmatrix(A::DMatrix{T}) where T + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = fetch(Ac[i, j]) + idx += 1 + end + mA, nA = size(A) + return _gather_sparse(T, tiles, row_offsets, col_offsets, mA, nA) +end + +_schur_worker_ids() = (w = workers(); isempty(w) ? Int[myid()] : Int[w...]) + +# Factor `Aᵢᵢ`, form dense `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (column solves; densifying the thin +# RHS is simpler than sparse multi-RHS), and local Schur contribution `A_Γi Wᵢ`. +# Pins `(fact, W, A_Γi)` to the current worker. +function _schur_factor_interior_pinned(factorize, Aii, AiΓ, AΓi) + F = factorize(Aii) + ni, ng = size(AiΓ) + T = eltype(AiΓ) + W = zeros(T, ni, ng) + for j in 1:ng + W[:, j] = F \ Vector{T}(AiΓ[:, j]) + end + Si = AΓi * W + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return (tochunk(F, proc, scope), tochunk(W, proc, scope), + tochunk(AΓi, proc, scope), Si) +end + +function _schur_factor_interface_pinned(factorize, S) + F = factorize(S) + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(F, proc, scope) +end + +# Forward: `y = Aᵢᵢ⁻¹ bᵢ`, `contrib = A_Γi y`. Pin `y` for the back-substitution. +function _schur_forward_pinned(fact, AΓi, bI) + y = fact \ bI + contrib = AΓi * y + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(y, proc, scope), contrib +end + +_schur_back_pinned(y, W, xΓ) = y - W * xΓ + +_schur_interface_solve(Sfact, g) = Sfact \ g + +function _make_schur_sparse_lu(A::DMatrix, factorize; nparts=nothing) + n = LinearAlgebra.checksquare(A) + np = nparts === nothing ? max(2, nworkers()) : Int(nparts) + np >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $np")) + + Asp = _collect_sparse_dmatrix(A) + nd = _nested_dissection_partition(Asp, np) + interiors = nd.interiors + separator = nd.separator + perm = nd.perm + + # Trivial partition → Stage-4a serial factorization. + if length(interiors) <= 1 && isempty(separator) + return _spawn_direct_factorization(A, factorize) + end + + wids = _schur_worker_ids() + k = length(interiors) + interior_facts = Vector{Any}(undef, k) + Ws = Vector{Any}(undef, k) + AΓis = Vector{Any}(undef, k) + scopes = Vector{ProcessScope}(undef, k) + Si_tasks = Vector{Any}(undef, k) + + Γ = separator + ng = length(Γ) + for i in 1:k + Ii = interiors[i] + scope = ProcessScope(wids[mod1(i, length(wids))]) + scopes[i] = scope + Aii = Asp[Ii, Ii] + AiΓ = Asp[Ii, Γ] + AΓi = Asp[Γ, Ii] + task = spawn(_schur_factor_interior_pinned, Options(; compute_scope=scope), + factorize, Aii, AiΓ, AΓi) + Si_tasks[i] = task + end + + # Fetch pinned chunks + local Schur contributions; reduce S on the driver. + T = eltype(Asp) + S = isempty(Γ) ? zeros(T, 0, 0) : Matrix{T}(Asp[Γ, Γ]) + for i in 1:k + Fi, Wi, AΓi_c, Si = fetch(Si_tasks[i]) + interior_facts[i] = Fi + Ws[i] = Wi + AΓis[i] = AΓi_c + isempty(Γ) || (S .-= Si) + end + + if isempty(Γ) + Sfact = nothing + Sscope = scopes[1] + else + Sscope = ProcessScope(wids[1]) + # Schur fill-in is expected; hand a sparse view to the backend factorizer. + Ssp = _sparse_copy_of(S) + Sfact = fetch(spawn(_schur_factor_interface_pinned, + Options(; compute_scope=Sscope), factorize, Ssp)) + end + + part = Blocks(A.partitioning.blocksize[2]) + return DistributedSchurLU(interior_facts, Ws, AΓis, scopes, Sfact, Sscope, + interiors, separator, perm, n, part) +end + +# Convert a dense Schur matrix to the sparse type the backend expects. Defined +# as a hook so the core package need not depend on SparseArrays; the real method +# lives in `SparseArraysExt`. +function _sparse_copy_of end +_sparse_copy_of(S) = throw(ArgumentError( + "Schur-complement sparse LU requires SparseArrays to be loaded")) + +function Base.:\(F::DistributedSchurLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + b_local = collect(b) + bp = b_local[F.perm] + k = length(F.interiors) + ng = length(F.separator) + + # Split permuted RHS into interior segments + separator. + ys = Vector{Any}(undef, k) + contrib_tasks = Vector{Any}(undef, k) + off = 0 + for i in 1:k + ni = length(F.interiors[i]) + bI = bp[off+1:off+ni] + off += ni + contrib_tasks[i] = spawn(_schur_forward_pinned, + Options(; compute_scope=F.interior_scopes[i]), + F.interior_facts[i], F.AΓi[i], bI) + end + + g = ng == 0 ? similar(bp, 0) : copy(bp[off+1:end]) + for i in 1:k + yi, contrib = fetch(contrib_tasks[i]) + ys[i] = yi + ng == 0 || (g .-= contrib) + end + + xΓ = if ng == 0 + similar(g) + else + fetch(spawn(_schur_interface_solve, Options(; compute_scope=F.Sscope), + F.Sfact, g)) + end + + x_parts = Vector{Any}(undef, k) + back_tasks = Vector{Any}(undef, k) + for i in 1:k + back_tasks[i] = spawn(_schur_back_pinned, + Options(; compute_scope=F.interior_scopes[i]), + ys[i], F.W[i], xΓ) + end + for i in 1:k + x_parts[i] = fetch(back_tasks[i]) + end + + xp = vcat(x_parts..., xΓ) + x_local = similar(b_local) + x_local[F.perm] = xp + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSchurLU, b::DVector) + return copyto!(x, F \ b) +end + diff --git a/test/Project.toml b/test/Project.toml index 4cb663d5b..47894c0bf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -17,10 +17,13 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94" +Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Preferences = "21216c6a-2e73-6563-6e65-726566657250" +PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01" +PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" diff --git a/test/array/linalg/sparsedirect.jl b/test/array/linalg/sparsedirect.jl new file mode 100644 index 000000000..e2f94ad50 --- /dev/null +++ b/test/array/linalg/sparsedirect.jl @@ -0,0 +1,249 @@ +# Sparse direct-solver tests (KLU via PureKLU, UMFPACK-style via PureUMFPACK). +# +# Exercises the whole-matrix direct solves (`Dagger.klu` / `Dagger.splu`) and the +# per-tile block direct preconditioners (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`). Reference solutions come from a dense direct +# solve. +# +# julia test/runtests.jl --test array/linalg/sparsedirect + +using PureKLU +using PureUMFPACK +using Metis +using Krylov + +# A nonsymmetric, well-conditioned sparse system (advection-diffusion-like). +function nonsym_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(3), n), + 1 => fill(T(-7) / 10, n - 1), + ) +end + +# SPD, diagonally dominant tridiagonal (stable under pivoting). +function spd_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end + +# 2D 5-point Laplacian on an N×N grid (SPD; clean METIS separators). +function lap2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(4)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, -one(T)); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, -one(T)); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, -one(T)); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, -one(T)); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +# Nonsymmetric, diagonally dominant 2D advection-diffusion stencil. +function advdiff2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(5)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, T(-12) / 10); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, T(-8) / 10); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, T(-11) / 10); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, T(-9) / 10); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +@testset "Sparse direct solvers" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + Asp = nonsym_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "whole-matrix direct ($(solver)) — $(backend)" for solver in (:klu, :splu), + backend in (:dense, :sparse, :singletile) + DA = if backend === :dense + distribute(Adense, A_part) + elseif backend === :sparse + distribute(Asp, A_part) + else + distribute(Asp, Blocks(n, n)) # single tile + end + Db = distribute(b, Db_part) + + F = solver === :klu ? Dagger.klu(DA) : Dagger.splu(DA) + @test size(F) == (n, n) + + # `F \ b` returns a DVector partitioned like `b`. + x = F \ Db + @test x isa Dagger.DVector + @test collect(x) ≈ xref + + # In-place `ldiv!` writes into a preallocated DVector. + y = similar(Db) + LinearAlgebra.ldiv!(y, F, Db) + @test collect(y) ≈ xref + + # Dimension mismatch is caught. + @test_throws DimensionMismatch F \ distribute(rand(n ÷ 2), Blocks(k)) + end + + @testset "block direct preconditioner ($(name))" for (name, build) in ( + ("KLU", Dagger.BlockKLUPreconditioner), + ("UMFPACK", Dagger.BlockUMFPACKPreconditioner), + ) + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + P = build(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + # Repeatable (exercises the cached, pinned factorization). + y2 = similar(Db) + mul!(y2, P, Db) + @test collect(y2) ≈ yref + + x, stats = Dagger.bicgstab(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes the block solve an *exact* whole-matrix solve. + DA1 = distribute(Asp, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = build(DA1) + y1 = similar(Db1) + mul!(y1, P1, Db1) + @test collect(y1) ≈ xref + + # Non-square block grid must be rejected with a helpful error. + @test_throws ArgumentError build(distribute(Adense, Blocks(k, k ÷ 2))) + end + + @testset "distributed triangular solve (splu)" begin + systems = ( + ("nonsym", nonsym_1d(Float64, n)), + ("spd", spd_1d(Float64, n)), + ) + # Include a case where A's tiling differs from the solve blocksize. + configs = ( + (Blocks(k, k), nothing), # default: match A's blocks + (Blocks(k, k), 8), # finer solve tiles than A + (Blocks(32, 16), 16), # non-square A tiling → square solve + (Blocks(n, n), 16), # single-tile A, multi-tile solve + ) + for (name, Asp_sys) in systems, (A_part_cfg, solve_bs) in configs + @testset "$name part=$(A_part_cfg.blocksize) bs=$(solve_bs)" begin + Adense_sys = Matrix(Asp_sys) + DA = distribute(Asp_sys, A_part_cfg) + Db = distribute(b, Db_part) + + F4a = Dagger.splu(DA) + F4b = if solve_bs === nothing + Dagger.splu(DA; distributed=true) + else + Dagger.splu(DA; distributed=true, blocksize=solve_bs) + end + @test F4b isa Dagger.DistributedSparseLU + @test size(F4b) == (n, n) + @test F4b.L isa Dagger.DMatrix + @test F4b.U isa Dagger.DMatrix + # Factors stay sparse (no densification). + @test fetch(F4b.L.chunks[1, 1]) isa Dagger.DSparseArray + @test fetch(F4b.U.chunks[1, 1]) isa Dagger.DSparseArray + + x4a = collect(F4a \ Db) + x = F4b \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + @test xc ≈ x4a rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xc - b) / norm(b) < 1e-10 + + # Multi-RHS / repeated solves against one factor. + for _ in 1:3 + bi = rand(n) + Dbi = distribute(bi, Db_part) + xi = collect(F4b \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xi - bi) / norm(bi) < 1e-10 + end + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4b, Db) + @test collect(y) ≈ x4a rtol=1e-12 atol=1e-14 + + @test_throws DimensionMismatch F4b \ distribute(rand(n ÷ 2), Blocks(k)) + end + end + end + + @testset "Schur distributed LU (4c)" begin + systems = ( + ("lap2d", N -> lap2d(Float64, N)), + ("advdiff2d", N -> advdiff2d(Float64, N)), + ) + for (name, buildA) in systems, N in (10, 20), nparts in (2, 4) + @testset "$name N=$N nparts=$nparts" begin + Asp_sys = buildA(N) + nsys = size(Asp_sys, 1) + bs = max(8, nsys ÷ 4) + DA = distribute(Asp_sys, Blocks(bs, bs)) + bsys = rand(nsys) + Db = distribute(bsys, Blocks(bs)) + + F4a = Dagger.splu(DA) + F4c = Dagger.splu(DA; distributed=true, method=:schur, nparts=nparts) + @test F4c isa Dagger.DistributedSchurLU + @test size(F4c) == (nsys, nsys) + + x = F4c \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + x4a = collect(F4a \ Db) + @test xc ≈ x4a rtol=1e-8 + @test norm(Asp_sys * xc - bsys) / norm(bsys) < 1e-9 + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4c, Db) + @test collect(y) ≈ x4a rtol=1e-8 + @test norm(Asp_sys * collect(y) - bsys) / norm(bsys) < 1e-9 + + # Factor-once / solve-many. + for _ in 1:3 + bi = rand(nsys) + Dbi = distribute(bi, Blocks(bs)) + xi = collect(F4c \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-8 + @test norm(Asp_sys * xi - bi) / norm(bi) < 1e-9 + end + + @test_throws DimensionMismatch F4c \ distribute(rand(nsys ÷ 2), Blocks(bs)) + end + end + end +end + diff --git a/test/runtests.jl b/test/runtests.jl index b02db950c..d327ee50a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -45,6 +45,7 @@ tests = [ ("Array - LinearAlgebra - QR", "array/linalg/qr.jl"), ("Array - LinearAlgebra - SVD", "array/linalg/svd.jl"), ("Array - LinearAlgebra - Iterative Solvers", "array/linalg/iterativesolvers.jl"), + ("Array - LinearAlgebra - Sparse Direct", "array/linalg/sparsedirect.jl"), ("Array - Permute", "array/permute.jl"), ("Array - Random", "array/random.jl"), ("Array - Stencils", "array/stencil.jl"), From 918e5b2e10e636b1969a5e522f2ccbd5d5d7f475 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 23 Jul 2026 17:26:46 -0700 Subject: [PATCH 18/43] docs: Document sparse DArrays --- docs/make.jl | 2 + docs/src/index.md | 70 +++++++++++ docs/src/iterative-solving.md | 226 ++++++++++++++++++++++++++++++++++ docs/src/sparse-arrays.md | 172 ++++++++++++++++++++++++++ 4 files changed, 470 insertions(+) create mode 100644 docs/src/iterative-solving.md create mode 100644 docs/src/sparse-arrays.md diff --git a/docs/make.jl b/docs/make.jl index 611f7301d..d7f6f0168 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -23,6 +23,8 @@ makedocs(; "Task Affinity" => "task-affinity.md", "Data Management" => "data-management.md", "Distributed Arrays" => "darray.md", + "Sparse Arrays" => "sparse-arrays.md", + "Iterative Solvers" => "iterative-solving.md", "Streaming Tasks" => "streaming.md", "Scopes" => "scopes.md", "Processors" => "processors.md", diff --git a/docs/src/index.md b/docs/src/index.md index 4b79aaa72..3e4b8ca59 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -363,6 +363,76 @@ collect(DA) # returns a `Matrix{Float64}` ----- +## Quickstart: Sparse Arrays + +A `DArray` can hold sparse tiles, giving a distributed, tiled sparse matrix. +Load a sparse backend (`SparseArrays`) to enable it. + +For more details: [Sparse Distributed Arrays](@ref) + +### Distribute or allocate a sparse `DArray` + +```julia +using SparseArrays + +# Partition an existing sparse matrix into sparse tiles +DA = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250)) + +# Or allocate sparse DArrays directly +Z = spzeros(Blocks(250, 250), Float64, 1000, 1000) +R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01) +``` + +### Multiply sparse arrays + +```julia +using LinearAlgebra +x = distribute(rand(1000), Blocks(250)) +y = DA * x # distributed sparse matrix-vector multiply (dense result) +C = DA * DA # distributed sparse-sparse matmul (sparse result) +``` + +`collect(DA)` returns a dense `Array`; operate on the `DArray` to stay sparse. + +----- + +## Quickstart: Iterative Solvers + +Dagger provides distributed, matrix-free Krylov solvers (`cg`, `minres`, +`gmres`, `bicgstab`) for `A x = b`. Load `Krylov` to enable them. + +For more details: [Iterative Solvers](@ref) + +### Solve a sparse system + +```julia +using SparseArrays, Krylov, LinearAlgebra + +n = 1000 +A = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.0, n-1)) +DA = distribute(A, Blocks(250, 250)) # square tiles +b = distribute(rand(n), Blocks(250)) + +x, stats = Dagger.cg(DA, b) +@show stats.solved, stats.niter +``` + +### Add a preconditioner + +```julia +using AlgebraicMultigrid # enables Dagger.AMGPreconditioner + +P = Dagger.AMGPreconditioner(DA) # build once +x, stats = Dagger.cg(DA, b; M = P) # pass as `M` +``` + +Other preconditioners: `Dagger.JacobiPreconditioner`, +`Dagger.BlockJacobiPreconditioner` (core), and `Dagger.BlockILUPreconditioner` +(load `IncompleteLU`). Any object implementing `mul!(y, A, x)` over `DVector`s +can be used as a matrix-free operator `A`. + +----- + ## Quickstart: Stencil Operations Dagger's `@stencil` macro allows for easy specification of stencil operations on `DArray`s, often used in simulations and image processing. These operations typically involve updating an element based on the values of its neighbors. diff --git a/docs/src/iterative-solving.md b/docs/src/iterative-solving.md new file mode 100644 index 000000000..31b62d755 --- /dev/null +++ b/docs/src/iterative-solving.md @@ -0,0 +1,226 @@ +# Iterative Solvers + +Dagger provides distributed, **matrix-free** Krylov solvers for linear systems +`A x = b`, built on top of [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl). +They run entirely over `DArray`s: the operator `A` may be a dense or +[sparse](@ref "Sparse Distributed Arrays") `DMatrix`, or any object that +implements distributed `mul!(y, A, x)` over `DVector`s. This makes them a natural +fit for large, sparse systems arising from ODE/PDE discretizations. + +Solvers are provided through a package extension; load `Krylov` to enable them: + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays, Krylov +``` + +If you call a solver without `Krylov` loaded, you'll get a clear error telling +you to `using Krylov`. + +## Solving a system + +```julia +using SparseArrays, Krylov + +# A symmetric positive-definite operator (1-D Laplacian) and a right-hand side +n = 1000 +A = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.0, n-1)) +DA = distribute(A, Blocks(250, 250)) # square tiles (see below) +b = distribute(rand(n), Blocks(250)) + +x, stats = Dagger.cg(DA, b) +@show stats.solved, stats.niter +``` + +Each solver returns a tuple `(x, stats)`, where `x` is a `DVector` and `stats` +is Krylov's statistics object (`stats.solved`, `stats.niter`, residual history, +etc.). Keyword arguments such as `atol`, `rtol`, and `itmax` are forwarded +to Krylov. + +### Available methods + +| Function | Operator class | Notes | +|---------------------|-------------------------------|------------------------------------| +| [`Dagger.cg`](@ref) | symmetric positive-definite | cheapest; the PDE workhorse | +| [`Dagger.minres`](@ref) | symmetric (indefinite ok) | saddle-point / indefinite systems | +| [`Dagger.gmres`](@ref) | general nonsymmetric | robust; `restart`/`memory` to bound memory | +| [`Dagger.bicgstab`](@ref) | general nonsymmetric | short recurrence, low memory | + +[`Dagger.krylov_solve`](@ref) is a generic entry point taking the method as a +symbol (`:cg`, `:minres`, `:gmres`, `:bicgstab`): + +```julia +x, stats = Dagger.krylov_solve(:gmres, DA, b; memory=50) +``` + +## Matrix-free operators + +The solvers never form `A⁻¹`, and they never require `A` to be a materialized +matrix. They only need `mul!(y, A, x)` (and `mul!(y, A', x)` for two-sided +methods like GMRES/BiCGStab) to work over `DVector`s. So you can pass any custom +operator type that implements distributed `mul!`: + +```julia +struct MyStencilOperator + # ...your distributed state... +end +function LinearAlgebra.mul!(y::Dagger.DVector, A::MyStencilOperator, x::Dagger.DVector) + # fill y with A*x using Dagger tasks / datadeps + return y +end + +x, stats = Dagger.cg(MyStencilOperator(...), b) +``` + +Workspace vectors are allocated via `similar(b)`, so every internal vector +inherits `b`'s element type **and** partitioning. The distributed BLAS-1 building +blocks the solvers rely on — `dot`, `norm`, `axpy!`, `axpby!`, `rmul!`, +`copyto!`, `fill!`, broadcasting — are all implemented for `DVector` and align +mismatched partitionings automatically. + +!!! note "Square tiles for `DMatrix` operators" + When the operator is a `DMatrix`, use **square tiles** (`Blocks(k, k)`). The + solver's workspace vectors are all allocated as `similar(b)` (one + partitioning), and each must serve as both the length-`n` input and the + length-`n` output of `mul!(y, A, x)`. The distributed SpMV requires the input + to match `A`'s column blocks and the output to match `A`'s row blocks, which + is only simultaneously possible when those block sizes are equal. Any uniform + square tile size works (a ragged final block is fine). This constraint is on + how *you* partition the operator — generic code calling `mul!`/`dot` never + needs to know about it. + +## Preconditioners + +A preconditioner accelerates convergence by approximating `A⁻¹`. Dagger's +preconditioners follow Krylov's `ldiv=false` convention: the object **represents +the inverse operator `M⁻¹`** and is applied via `mul!(y, P, x)` (computing +`y = M⁻¹ x`). Pass one as the `M` keyword: + +```julia +P = Dagger.JacobiPreconditioner(DA) +x, stats = Dagger.cg(DA, b; M = P) +``` + +The built-in preconditioners, from cheapest to strongest: + +| Preconditioner | Needs | Idea | +|--------------------------------------|------------------------|--------------------------------------------| +| [`Dagger.JacobiPreconditioner`](@ref) | (core) | scale by `1 ./ diag(A)` | +| [`Dagger.BlockJacobiPreconditioner`](@ref) | (core) | exact `lu` solve per diagonal tile | +| [`Dagger.BlockILUPreconditioner`](@ref) | `IncompleteLU` | incomplete-LU (drop tol `τ`) per tile | +| [`Dagger.AMGPreconditioner`](@ref) | `AlgebraicMultigrid` | AMG V-cycle per tile | + +```julia +using AlgebraicMultigrid, IncompleteLU + +x, _ = Dagger.cg(DA, b; M = Dagger.BlockJacobiPreconditioner(DA)) +x, _ = Dagger.cg(DA, b; M = Dagger.BlockILUPreconditioner(DA; τ = 0.01)) +x, _ = Dagger.gmres(DA, b; M = Dagger.AMGPreconditioner(DA; method = :ruge_stuben)) +``` + +### How block preconditioners are distributed + +`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, and `AMGPreconditioner` +are all **block-diagonal** preconditioners: they build one operator per diagonal +tile of `A` and apply them independently per block. They share a common +mechanism ([`Dagger.AbstractBlockPreconditioner`](@ref)): + +- The per-tile operator (an `lu`/ILU factorization, or an AMG hierarchy) is built + **once**, at construction. +- A factorization/hierarchy generally cannot be moved between workers (sparse + `lu` factors and AMG hierarchies hold process-bound resources). So each + operator is **pinned** to the worker owning its tile, and every apply for that + block is scheduled there — only the (small, movable) vector chunks are + transferred. + +A useful consequence of the per-tile design: with a **single tile** +(`Blocks(n, n)`), any of these becomes a *global* preconditioner over the whole +matrix (e.g. a global AMG, or an exact direct solve for block-Jacobi). With many +tiles, it becomes a scalable block-Jacobi / additive-Schwarz variant that trades +some convergence for parallelism. As with the operator, these require square +diagonal tiles. + +### Choosing a preconditioner + +- **SPD elliptic (Poisson-like) problems:** `AMGPreconditioner` gives near + mesh-independent convergence and is usually the best choice; `cg` as the + solver. +- **General sparse systems:** `BlockILUPreconditioner` is a solid, cheap-setup + general-purpose option; pair with `gmres` or `bicgstab`. +- **Quick baseline / very well-conditioned systems:** `JacobiPreconditioner` (or + none) may suffice. +- **Strong per-subdomain coupling:** `BlockJacobiPreconditioner` (exact tile + solves) is stronger than diagonal Jacobi. + +## Sparse direct solvers + +For systems that fit on a single worker, Dagger also offers **direct** sparse +solves via pure-Julia factorization backends. Unlike the C-bound `UmfpackLU`, +these factorizations are plain Julia data, so Dagger can move and schedule them +freely. + +Load `PureKLU` (KLU; good for unsymmetric/circuit systems) or `PureUMFPACK` +(UMFPACK-style multifrontal LU): + +```julia +using SparseArrays, PureKLU, PureUMFPACK + +A = distribute(sprand(2000, 2000, 0.005) + 10I, Blocks(500, 500)) +b = distribute(rand(2000), Blocks(500)) + +F = Dagger.klu(A) # or Dagger.splu(A) +x = F \ b # returns a DVector partitioned like b +``` + +`Dagger.klu`/`Dagger.splu` gather the sparse `DMatrix` into one +`SparseMatrixCSC` (without densifying), factor it once, and return a +[`Dagger.DaggerSparseLU`](@ref) supporting `F \ b` and `ldiv!(x, F, b)`. Factor +once, solve many right-hand sides cheaply. + +There are also **block direct preconditioners** that factor each diagonal tile +exactly (`Dagger.BlockKLUPreconditioner`, `Dagger.BlockUMFPACKPreconditioner`), +usable like the other block preconditioners. With a single tile they are exact +whole-matrix solves; with many tiles they are exact-block-Jacobi preconditioners +for the iterative solvers. + +## A worked example: implicit time stepping + +Implicit ODE/PDE integrators repeatedly solve systems with the same operator +`(I - Δt·L)` and changing right-hand sides. Build the preconditioner once and +reuse it across steps: + +```julia +using SparseArrays, Krylov, AlgebraicMultigrid + +L = distribute(laplacian, Blocks(k, k)) # discretized operator (square tiles) +A = I - Δt * L # or a custom matrix-free operator +P = Dagger.AMGPreconditioner(L) # build hierarchy once + +u = distribute(u0, Blocks(k)) +for step in 1:nsteps + rhs = ... # depends on current state + u, stats = Dagger.cg(A, rhs; M = P, rtol = 1e-8) +end +``` + +## API + +```@docs +Dagger.cg +Dagger.minres +Dagger.gmres +Dagger.bicgstab +Dagger.krylov_solve +Dagger.AbstractDaggerPreconditioner +Dagger.JacobiPreconditioner +Dagger.AbstractBlockPreconditioner +Dagger.BlockJacobiPreconditioner +Dagger.BlockILUPreconditioner +Dagger.AMGPreconditioner +Dagger.BlockKLUPreconditioner +Dagger.BlockUMFPACKPreconditioner +Dagger.klu +Dagger.splu +Dagger.DaggerSparseLU +``` diff --git a/docs/src/sparse-arrays.md b/docs/src/sparse-arrays.md new file mode 100644 index 000000000..8af6f73f0 --- /dev/null +++ b/docs/src/sparse-arrays.md @@ -0,0 +1,172 @@ +# Sparse Distributed Arrays + +Dagger's [`DArray`](@ref) can hold **sparse** tiles, giving you a distributed, +tiled sparse matrix (or vector) that participates in the same scheduling, +[Datadeps](@ref), and linear-algebra machinery as dense `DArray`s. This is the +foundation for distributed sparse matrix multiplication and the +[matrix-free iterative solvers](@ref "Iterative Solvers"). + +Sparse support is provided through package extensions, so you opt in by loading a +sparse backend: + +- **`SparseArrays`** (the standard library) — tiles are `SparseMatrixCSC` / + `SparseVector`. This is the default, well-supported backend. +- **`Finch`** — tiles are `Finch.Tensor`s, enabling a wider range of sparse + formats. This backend is more experimental. + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays +``` + +!!! note "Load order with workers" + As with all Dagger usage, add your workers *before* `using Dagger` and the + backend package, so the packages load on every worker. See the note at the + top of the [home page](@ref "Dagger: A framework for out-of-core and parallel execution"). + +## Creating a sparse `DArray` + +### From an existing sparse array + +`distribute` accepts a sparse matrix or vector and partitions it into sparse +tiles according to a `Blocks` specification: + +```julia +using SparseArrays +A = sprand(1000, 1000, 0.01) # a SparseMatrixCSC +DA = distribute(A, Blocks(250, 250)) # a 4×4 grid of sparse tiles +``` + +Each tile is a sparse matrix in its own right, stored on one of the workers. + +### Allocating directly + +You can also allocate a sparse `DArray` without first building a local sparse +array, using the `Blocks`-aware methods of `spzeros` and `sprand`: + +```julia +using SparseArrays + +# All-zeros sparse DArray, Float64, 1000×1000 in 250×250 tiles +Z = spzeros(Blocks(250, 250), Float64, 1000, 1000) + +# Random sparse DArray with ~1% nonzeros per tile +R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01) +``` + +These run the per-tile allocation on the owning worker, so no large sparse array +is ever materialized on a single process. + +### Converting back to a dense array + +`collect` gathers the tiles and returns a **dense** `Array`: + +```julia +M = collect(DA) # dense Matrix{Float64} +``` + +To keep data sparse and distributed, operate on the `DArray` directly rather +than collecting. + +## How it works + +### The `DSparseArray` wrapper + +Internally, each sparse tile is wrapped in a [`Dagger.DSparseArray`](@ref) — a +small mutable container holding the actual sparse storage (`mat`): + +```julia +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} + mat # e.g. a SparseMatrixCSC, SparseVector, or Finch.Tensor +end +``` + +`DSparseVector{T}` and `DSparseMatrix{T}` are the 1- and 2-dimensional aliases. + +The wrapper exists because **sparse storage is reallocated on writes**. Many +sparse operations (e.g. `A*B`, or anything that changes the sparsity pattern) +cannot update their result in place — they produce a brand-new sparse array of a +different size. Datadeps, however, tracks data dependencies by the *identity* of +the objects it manages, and it does not support objects that grow or shrink. The +`DSparseArray` wrapper solves this: its identity is stable, and a write simply +swaps the inner `mat` for the new storage: + +```julia +# Conceptually, how an in-place sparse update is modeled: +tile.mat = tile.mat * other # identity of `tile` is unchanged +``` + +### Aliasing as a whole + +Because the inner storage may move, it is never safe to alias *part* of a sparse +tile (e.g. via a `view` or a strided sub-region). Dagger therefore treats a +`DSparseArray` as an **indivisible aliasing unit**: any access — including +through `view`, `transpose`, `adjoint`, or `reshape` — resolves to the +container's stable whole-object aliasing. This is what keeps Datadeps correct +when sparse writes reallocate storage. (For the curious: the type opts in via +`Dagger.aliases_as_whole`, and Datadeps' `aliasing_root` unwraps any wrapper of a +`DSparseArray` before computing aliasing. Calling `pointer` on a `DSparseArray` +intentionally errors, to catch any code path that tries to treat it as raw +strided memory.) + +The practical upshot: you can pass sparse tiles, or views of sparse `DArray`s, +into `Dagger.spawn_datadeps` regions and trust that read/write ordering is +tracked correctly. + +## Operations + +Sparse `DArray`s support the array operations that have distributed +implementations, including: + +- **Matrix–matrix multiply** (`A * B`, `mul!`), sparse × sparse, producing a + sparse result. +- **Sparse matrix–vector multiply** (SpMV: `A * x`, `mul!(y, A, x)`) with a + sparse matrix and dense vectors — the workhorse of iterative solvers. +- **Transpose/adjoint**, **`collect`**, and elementwise/`norm` operations. + +```julia +using SparseArrays, LinearAlgebra +A = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250)) +x = distribute(rand(1000), Blocks(250)) + +y = A * x # distributed SpMV -> dense DVector +C = A * A # distributed sparse-sparse matmul -> sparse DArray +``` + +### Partitioning guidance + +- Choose tile sizes so each tile comfortably fits on a worker, and so the number + of tiles is at least the number of workers (for parallelism). +- For **square operators used with the iterative solvers**, use **square tiles** + (`Blocks(k, k)`); see [Iterative Solvers](@ref) for why. +- Operands with mismatched partitionings are aligned automatically (by buffered + copy) where needed, but matching partitionings avoid that overhead. + +## Backends + +### `SparseArrays` (recommended) + +Tiles are `SparseMatrixCSC` (matrices) or `SparseVector` (vectors). This backend +provides efficient SpMV (including transposed/adjoint operands) and uses +`SparseArrays`' own `*` for sparse–sparse products. + +### `Finch` (experimental) + +Loading `Finch` makes tiles `Finch.Tensor`s, supporting a broader set of sparse +and structured formats. Finch support is newer and exercised by a dedicated test +suite; prefer `SparseArrays` unless you specifically need a Finch format. + +## Limitations + +- `collect` densifies; there is no sparse-preserving global gather. +- A sparse tile is aliased as a whole — Datadeps cannot track independent writes + to disjoint sub-regions of a single sparse tile (use finer tiling instead). +- Not every dense `DArray` operation has a sparse counterpart yet; sparse support + focuses on multiplication and the building blocks needed for iterative solving. + +## API + +```@docs +Dagger.DSparseArray +``` From 644e6ca3043b99684d21c94dba95dc3b42b479dd Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 10:44:05 -0700 Subject: [PATCH 19/43] Datadeps: Drop unsound arg_current write-back elision Always compute region-end write-back from history. The MPI SPMD arg_current shortcut could skip needed copies and leave uninitialized Krylov workspace tiles under multi-worker Distributed. Co-authored-by: Cursor --- src/datadeps/queue.jl | 17 +++++++---------- test/mpi.jl | 7 +++++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 9441f3bbe..0ef6c6655 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -153,16 +153,13 @@ function distribute_tasks!(queue::DataDepsTaskQueue) check_uniform(arg_w) arg = arg_w.arg origin_space = state.arg_origin[arg] - # When the origin still holds a fully-current replica (the argument was - # only read, or copies merely propagated it), the write-back is elided. - # This is only safe here at region end: mid-region, the copy tasks also - # serialize readers against later writers, so they must not be skipped. - current = get(state.arg_current, arg_w, nothing) - if current !== nothing && origin_space in current - remainder = NoAliasing() - else - remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) - end + # Always compute write-back from history. Eliding via `arg_current` + # (`origin in arg_current` ⇒ skip) is unsound under multi-worker + # Distributed: Krylov `similar` workspaces can leave origin marked + # current while a remote true write holds the only initialized bytes. + # Read-only copy-ins may therefore schedule a redundant write-back + # (copy history advances `arg_owner`); that is intentional and harmless. + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) if remainder isa MultiRemainderAliasing origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) diff --git a/test/mpi.jl b/test/mpi.jl index f5e7ffd6f..d8d25921f 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -543,14 +543,17 @@ read_only(X) = (sum(X); nothing) r1 = min(1, nranks-1) cross = r1 != 0 # single-rank runs degenerate to zero copies - # Read-only cross-rank arg: one copy-in, no writeback + # Read-only cross-rank arg: copy-in plus region-end write-back. The + # write-back is redundant (origin still holds the bytes) but is retained + # so write-back planning can stay history-based; `arg_current` elision is + # unsound for Krylov `similar` workspaces under Distributed. A = ones(4, 4) logs = with_logs() do Dagger.spawn_datadeps() do Dagger.@spawn scope=rank_scope(r1) read_only(In(A)) end end - @test count_move_tasks(logs) == (cross ? 1 : 0) + @test count_move_tasks(logs) == (cross ? 2 : 0) @test Dagger.check_uniform(count_move_tasks(logs)) # Same-rank InOut: the origin slot aliases the original, zero copies From 7fb631859a6506219d3e1d693d9454951ee20d07 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 10:44:30 -0700 Subject: [PATCH 20/43] Datadeps: Record producer task on HistoryEntry for remainder syncdeps Remainder copies previously re-resolved syncdeps via live ainfos_owner / get_read_deps!, which can miss the true producer after ownership moves. Store the writer DTask on each HistoryEntry and wait on it directly. Co-authored-by: Cursor --- src/datadeps/aliasing.jl | 16 ++++++++++------ src/datadeps/remainders.jl | 13 ++++++++++++- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index a412182ea..22c7ac071 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -240,6 +240,11 @@ struct HistoryEntry ainfo::AliasingWrapper space::MemorySpace write_num::Int + # Producer of this write/copy. Remainder syncdeps wait on this task + # directly instead of re-resolving through live `ainfos_owner`, which may + # later name a different task for the same ainfo or miss the producer when + # only an overlapping ainfo is consulted. + task::DTask end struct AliasedObjectCacheStore @@ -666,10 +671,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ history = state.arg_history[arg_w] @opcounter :merge_history @opcounter :merge_history_complexity length(history) - origin_space = state.arg_origin[other_arg_w.arg] for other_entry in state.arg_history[other_arg_w] - write_num_tuple = HistoryEntry(AliasingWrapper(NoAliasing()), origin_space, other_entry.write_num) - range = searchsorted(history, write_num_tuple; by=x->x.write_num) + range = searchsorted(history, other_entry; by=x->x.write_num) if !isempty(range) # Find and skip duplicates match = false @@ -677,7 +680,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ source_entry = history[source_idx] if source_entry.ainfo == other_entry.ainfo && source_entry.space == other_entry.space && - source_entry.write_num == other_entry.write_num + source_entry.write_num == other_entry.write_num && + source_entry.task == other_entry.task match = true break end @@ -844,12 +848,12 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M empty!(state.arg_history[arg_w]) # Add our own history - push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) # Find overlapping arguments and update their history for other_arg_w in state.arg_overlaps[arg_w] other_arg_w == arg_w && continue - push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) end # Track which spaces hold a fully-current replica of this region diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 3ee709e6d..cbd494fef 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -178,6 +178,7 @@ function compute_remainder_for_arg!(state::DataDepsState, break end + other_entry = nothing if idx > 0 other_entry = state.arg_history[arg_w][idx] other_ainfo = other_entry.ainfo @@ -225,7 +226,17 @@ function compute_remainder_for_arg!(state::DataDepsState, has_overlap = schedule_remainder!(tracker_other_space[1], other_space_idx, target_space_idx, remainder, other_many_spans) if compute_syncdeps && has_overlap @assert haskey(state.ainfos_owner, other_ainfo) "[idx $idx] ainfo $(typeof(other_ainfo)) has no owner" - get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + if other_entry !== nothing + # Wait on the producer recorded in history. Re-resolving via + # live `ainfos_owner`/`get_read_deps!` is not equivalent: + # ownership of `other_ainfo` may have moved to a later copy, or + # the producer may only be registered on an overlapping ainfo + # missing from `ainfos_overlaps[other_ainfo]`. + push!(tracker_other_space[3], ThunkSyncdep(other_entry.task)) + else + # idx==0 owner fallback: no HistoryEntry, use live writer set. + get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + end push!(tracker_other_space[2], other_ainfo) end end From 77bdc4c124f48d0596c21fda97e74689ae294b4d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 10:48:14 -0700 Subject: [PATCH 21/43] MPIExt: Declare SparseArrays as an extension dependency MPIExt imports SparseMatrixCSC; Julia 1.12 requires SparseArrays in the extension trigger list so the extension can precompile. Co-authored-by: Cursor --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index d25a38dc2..4ee392650 100644 --- a/Project.toml +++ b/Project.toml @@ -67,7 +67,7 @@ IntelExt = "oneAPI" JSON3Ext = "JSON3" KrylovExt = "Krylov" LinuxPerfExt = "LinuxPerf" -MPIExt = "MPI" +MPIExt = ["MPI", "SparseArrays"] MetalExt = "Metal" MetisExt = ["Metis", "SparseArrays"] OpenCLExt = "OpenCL" From f7c0bcae05d90b5ceba7803a75c5d6258c639d55 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 11:18:08 -0700 Subject: [PATCH 22/43] test/mpienv: Add SparseArrays so MPIExt can load MPIExt now declares SparseArrays as an extension trigger; the MPI test environment must provide it. Co-authored-by: Cursor --- test/mpienv/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/mpienv/Project.toml b/test/mpienv/Project.toml index 71e617954..581efcb0e 100644 --- a/test/mpienv/Project.toml +++ b/test/mpienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] From 25bab3f09442faea197e73ba507cec34f04f952f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 13 Jul 2026 12:50:46 -0700 Subject: [PATCH 23/43] DArray: Base.fetch(::DArray) just converts DTask to Chunk --- src/array/darray.jl | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/src/array/darray.jl b/src/array/darray.jl index 278be6098..2df44afbf 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -426,31 +426,21 @@ end """ Base.fetch(c::DArray) -If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk. +If a `DArray` tree has a `DTask` in it, `fetch` each one (as a `Chunk`, +without moving its data) and return a new `DArray` with all partitions +resolved to `Chunk`s. """ function Base.fetch(c::DArray{T}) where T - if any(istask, chunks(c)) - if uniform_execution() - # SPMD (MPI): every rank holds every task locally, and chunk - # records are rank-local views (placeholders for non-owned data). - # Resolve them locally instead of assembling on one rank and - # broadcasting (which would clobber the local records with the - # assembling rank's placeholders). - thunks = chunks(c) - new_chunks = Any[istask(t) ? fetch(t; raw=true) : t for t in thunks] - return DArray(T, domain(c), domainchunks(c), - reshape(new_chunks, size(thunks)), - c.partitioning, c.concat) - end - thunks = chunks(c) - sz = size(thunks) + thunks = chunks(c) + if any(istask, thunks) dmn = domain(c) dmnchunks = domainchunks(c) - return fetch(Dagger.spawn(Options(meta=true), thunks...) do results... - t = eltype(fetch(results[1])) - DArray(t, dmn, dmnchunks, reshape(Any[results...], sz), - c.partitioning, c.concat) - end) + new_chunks = similar(thunks, Any) + for idx in eachindex(thunks) + part = thunks[idx] + new_chunks[idx] = istask(part) ? fetch(part; raw=true) : part + end + return DArray(T, dmn, dmnchunks, new_chunks, c.partitioning, c.concat) else return c end From 4b73abaeaa87d8756cbc87d81ab0f03eac9a9602 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 15:44:10 -0700 Subject: [PATCH 24/43] MPI: Respect specified return_type option Co-authored-by: Cursor --- ext/MPIExt.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 71725742b..f7521bd6e 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -2107,7 +2107,11 @@ function mpi_propagate_chunk_types!(tasks, accel::MPIAcceleration, expected_type for t in tasks if t isa Thunk if t.options !== nothing - t.options.return_type = expected_type + # Respect a more specific tile type already set by the spawner + # (e.g. `DSparseArray` from sparse allocators). + if t.options.return_type === nothing + t.options.return_type = expected_type + end else t.options = Options(return_type=expected_type) end From d1de06bcc7898348db5d16ba0b4889ab673c155d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 22:23:16 -0700 Subject: [PATCH 25/43] DArray/sparse: Add GPU sparse tile support (CUDA/ROC/OpenCL/Metal/oneAPI) Preserve sparsity across GPU moves and Datadeps for SpGEMM/SpMV, using vendor sparse libraries where available and DeviceSparseMatrixCSC otherwise. Co-authored-by: Cursor --- Project.toml | 5 + docs/src/sparse-arrays.md | 9 +- ext/CUDASparseArraysExt.jl | 177 +++++++++++++++++++++++++++++++++ ext/IntelSparseArraysExt.jl | 68 +++++++++++++ ext/MetalSparseArraysExt.jl | 68 +++++++++++++ ext/OpenCLSparseArraysExt.jl | 82 ++++++++++++++++ ext/ROCSparseArraysExt.jl | 185 +++++++++++++++++++++++++++++++++++ ext/SparseArraysExt.jl | 114 +++++++++++++++++++-- src/array/darray.jl | 10 +- src/array/sparse.jl | 102 ++++++++++++++++++- test/cudaenv/Project.toml | 1 + test/gpu.jl | 89 +++++++++++++++++ test/metalenv/Project.toml | 1 + test/oneapienv/Project.toml | 1 + test/openclenv/Project.toml | 2 + test/rocmenv/Project.toml | 1 + 16 files changed, 904 insertions(+), 11 deletions(-) create mode 100644 ext/CUDASparseArraysExt.jl create mode 100644 ext/IntelSparseArraysExt.jl create mode 100644 ext/MetalSparseArraysExt.jl create mode 100644 ext/OpenCLSparseArraysExt.jl create mode 100644 ext/ROCSparseArraysExt.jl diff --git a/Project.toml b/Project.toml index 4ee392650..3f64360ae 100644 --- a/Project.toml +++ b/Project.toml @@ -58,24 +58,29 @@ oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" AbstractFFTsExt = "AbstractFFTs" AlgebraicMultigridExt = ["AlgebraicMultigrid", "SparseArrays"] CUDAExt = "CUDA" +CUDASparseArraysExt = ["CUDA", "SparseArrays"] DistributionsExt = "Distributions" FinchExt = "Finch" GraphVizExt = "GraphViz" GraphVizSimpleExt = "Colors" IncompleteLUExt = ["IncompleteLU", "SparseArrays"] IntelExt = "oneAPI" +IntelSparseArraysExt = ["oneAPI", "SparseArrays"] JSON3Ext = "JSON3" KrylovExt = "Krylov" LinuxPerfExt = "LinuxPerf" MPIExt = ["MPI", "SparseArrays"] MetalExt = "Metal" +MetalSparseArraysExt = ["Metal", "SparseArrays"] MetisExt = ["Metis", "SparseArrays"] OpenCLExt = "OpenCL" +OpenCLSparseArraysExt = ["OpenCL", "SparseArrays"] PlotsExt = ["DataFrames", "Plots"] PureKLUExt = ["PureKLU", "SparseArrays"] PureUMFPACKExt = ["PureUMFPACK", "SparseArrays"] PythonExt = "PythonCall" ROCExt = "AMDGPU" +ROCSparseArraysExt = ["AMDGPU", "SparseArrays"] SparseArraysExt = "SparseArrays" [compat] diff --git a/docs/src/sparse-arrays.md b/docs/src/sparse-arrays.md index 8af6f73f0..9935b033c 100644 --- a/docs/src/sparse-arrays.md +++ b/docs/src/sparse-arrays.md @@ -10,9 +10,14 @@ Sparse support is provided through package extensions, so you opt in by loading sparse backend: - **`SparseArrays`** (the standard library) — tiles are `SparseMatrixCSC` / - `SparseVector`. This is the default, well-supported backend. + `SparseVector` on the CPU. This is the default, well-supported backend. +- **GPU + `SparseArrays`** — under a GPU compute scope (`cuda_gpu`, `rocm_gpu`, + `cl_device`, `metal_gpu`, `intel_gpu`), tiles use the vendor sparse type when + available (CUDA cuSPARSE / AMDGPU rocSPARSE) or else + `Dagger.DeviceSparseMatrixCSC` (OpenCL / Metal / oneAPI) with host SpGEMM/SpMV + fallbacks. Load the GPU package together with `SparseArrays`. - **`Finch`** — tiles are `Finch.Tensor`s, enabling a wider range of sparse - formats. This backend is more experimental. + formats. This backend is more experimental (CPU only for now). ```julia using Distributed diff --git a/ext/CUDASparseArraysExt.jl b/ext/CUDASparseArraysExt.jl new file mode 100644 index 000000000..072577ff3 --- /dev/null +++ b/ext/CUDASparseArraysExt.jl @@ -0,0 +1,177 @@ +module CUDASparseArraysExt + +# Sparse DArray support for CUDA / cuSPARSE. Loaded when both CUDA and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import CUDA +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import CUDA: CuArray +import CUDA.CUSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const CUDAExt = Base.get_extension(Dagger, :CUDAExt)::Module +using .CUDAExt: CuArrayDeviceProc +_with_context(f, proc) = CUDAExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::CuSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseVector) = Dagger.memory_space(x.nzVal) +Dagger.memory_space(x::CuSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + CuSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + CuSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::CuSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::CuSparseVector, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::CuSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(A)) +Dagger._sparse_collect(A::CuSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_cu_sparse(x::SparseMatrixCSC) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::SparseVector) = CuSparseVector(x) +_to_cu_sparse(x::CuSparseMatrixCSC) = x +_to_cu_sparse(x::CuSparseMatrixCSR) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::CuSparseVector) = x +_to_cu_sparse(x::Dagger.DeviceSparseMatrixCSC) = CuSparseMatrixCSC(SparseMatrixCSC(x)) +_to_cu_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_cu_sparse(x.mat)) + +_to_host_sparse(x::CuSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(x)) +_to_host_sparse(x::CuSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::CuSparseMatrixCSC) = A +_to_csc(A::CuSparseMatrixCSR) = CuSparseMatrixCSC(A) +_to_csc(A) = CuSparseMatrixCSC(A) + +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return CuSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return CuSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +#----- Move -------------------------------------------------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_cu_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + CUDA.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, + x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) + _with_context(from_proc) do + CUDA.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + _with_context(CUDA.synchronize, from_proc) + return x + end + _with_context(to_proc) do + return _to_cu_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +function Dagger.matvecmul!(C::CuArray, transA::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::CuArray, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::CuSparseMatrixCSC) + return CuSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::CuSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return CuSparseMatrixCSC(Ct) +end + +end # module CUDASparseArraysExt diff --git a/ext/IntelSparseArraysExt.jl b/ext/IntelSparseArraysExt.jl new file mode 100644 index 000000000..9cf2f2c01 --- /dev/null +++ b/ext/IntelSparseArraysExt.jl @@ -0,0 +1,68 @@ +module IntelSparseArraysExt + +# Sparse DArray support for oneAPI via DeviceSparseMatrixCSC (no vendor sparse +# library assumed). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import oneAPI +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import oneAPI: oneArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const IntelExt = Base.get_extension(Dagger, :IntelExt)::Module +using .IntelExt: oneArrayDeviceProc +_with_context(f, proc) = IntelExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:oneArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(oneArray, S) + +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + oneAPI.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end + +end # module IntelSparseArraysExt diff --git a/ext/MetalSparseArraysExt.jl b/ext/MetalSparseArraysExt.jl new file mode 100644 index 000000000..0d47f97b2 --- /dev/null +++ b/ext/MetalSparseArraysExt.jl @@ -0,0 +1,68 @@ +module MetalSparseArraysExt + +# Sparse DArray support for Metal via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import Metal +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import Metal: MtlArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const MetalExt = Base.get_extension(Dagger, :MetalExt)::Module +using .MetalExt: MtlArrayDeviceProc +_with_context(f, proc) = MetalExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:MtlArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(MtlArray, S) + +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + Metal.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end + +end # module MetalSparseArraysExt diff --git a/ext/OpenCLSparseArraysExt.jl b/ext/OpenCLSparseArraysExt.jl new file mode 100644 index 000000000..2a2f52772 --- /dev/null +++ b/ext/OpenCLSparseArraysExt.jl @@ -0,0 +1,82 @@ +module OpenCLSparseArraysExt + +# Sparse DArray support for OpenCL via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import OpenCL +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import OpenCL: CLArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const OpenCLExt = Base.get_extension(Dagger, :OpenCLExt)::Module +using .OpenCLExt: CLArrayDeviceProc +_with_context(f, proc) = OpenCLExt.with_context(f, proc) + +# Do not densify via Adapt (generic `adapt(CLArray, SparseMatrixCSC)` would). +Adapt.adapt_structure(::Type{<:CLArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(CLArray, S) + +_empty_device_sparse(::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.spzeros(T, dims...)) +_empty_device_sparse(::Type{T}, dims::Dims{1}) where T = + # Store 1-D sparse as a 1-column DeviceSparseMatrixCSC + Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(SparseArrays.spzeros(T, dims...))) + +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + _empty_device_sparse(T, dims) +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) # vectors stay host CSC; wrap still applies +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseVector) + # Keep SparseVector on host (wrapped); OpenCL SpMV host-falls-back anyway. + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + if mat isa SparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, mat)) + elseif mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + else + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + OpenCL.cl.finish(OpenCL.cl.queue()) + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + mat = x.mat + S = mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(mat) : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, S)) + end +end + +end # module OpenCLSparseArraysExt diff --git a/ext/ROCSparseArraysExt.jl b/ext/ROCSparseArraysExt.jl new file mode 100644 index 000000000..7340109d5 --- /dev/null +++ b/ext/ROCSparseArraysExt.jl @@ -0,0 +1,185 @@ +module ROCSparseArraysExt + +# Sparse DArray support for AMDGPU / rocSPARSE. Loaded when both AMDGPU and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import AMDGPU +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import AMDGPU: ROCArray +import AMDGPU.rocSPARSE: ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +# ROCExt is triggered by AMDGPU and provides the processor + context helpers. +const ROCExt = Base.get_extension(Dagger, :ROCExt)::Module +using .ROCExt: ROCArrayDeviceProc +_with_context(f, proc) = ROCExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::ROCSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseVector) = Dagger.memory_space(x.nzVal) +# Datadeps `aliased_object!` keys slots by `memory_space(x)` (see DSparseArray). +Dagger.memory_space(x::ROCSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + ROCSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + ROCSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::ROCSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::ROCSparseVector, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::ROCSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(A)) +Dagger._sparse_collect(A::ROCSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_roc_sparse(x::SparseMatrixCSC) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::SparseVector) = ROCSparseVector(x) +_to_roc_sparse(x::ROCSparseMatrixCSC) = x +_to_roc_sparse(x::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::ROCSparseVector) = x +_to_roc_sparse(x::Dagger.DeviceSparseMatrixCSC) = ROCSparseMatrixCSC(SparseMatrixCSC(x)) +_to_roc_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_roc_sparse(x.mat)) + +_to_host_sparse(x::ROCSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(x)) +_to_host_sparse(x::ROCSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::ROCSparseMatrixCSC) = A +_to_csc(A::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(A) +_to_csc(A) = ROCSparseMatrixCSC(A) + +# Materialize transpose/adjoint as CSC so SpGEMM stays on rocSPARSE / host CSC +# paths; `A * transpose(B)` would otherwise fall into scalar-indexing generic mul. +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return ROCSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return ROCSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +#----- Move (preserve sparsity; wrap in DSparseArray) -------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_roc_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + AMDGPU.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, + x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) + _with_context(from_proc) do + AMDGPU.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + # Same device: identity (like dense ROCArray). A copy here would discard + # in-place SpGEMM writes under Datadeps/Sch argument moves. + if from_proc == to_proc + _with_context(AMDGPU.synchronize, from_proc) + return x + end + # Distinct devices: stage through host (no peer-copy helper yet). + _with_context(to_proc) do + return _to_roc_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function Dagger.matvecmul!(C::ROCArray, transA::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::ROCArray, alpha, beta) + # rocSPARSE SpMV supports transpose/adjoint wrappers. + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::ROCSparseMatrixCSC) + return ROCSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::ROCSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return ROCSparseMatrixCSC(Ct) +end + +end # module ROCSparseArraysExt diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 7ef7c28c9..8820bd45a 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -32,11 +32,64 @@ Dagger._sparse_copy_of(S::SparseMatrixCSC) = S # Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseArray(x) Dagger.maybe_wrap_tile(x::SparseVector) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::Dagger.DeviceSparseMatrixCSC) = DSparseArray(x) + +# Host defaults for scoped sparse allocation (GPU Exts override per-processor). +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +# DeviceSparseMatrixCSC ↔ SparseMatrixCSC +function SparseArrays.SparseMatrixCSC(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + return SparseMatrixCSC{Tv,Ti}(A.m, A.n, Array(A.colptr), Array(A.rowval), Array(A.nzval)) +end +""" + device_sparse_from_host(Arr, S::SparseMatrixCSC) -> DeviceSparseMatrixCSC + +Upload a host CSC onto device vectors of type `Arr` (e.g. `CLArray`, `MtlArray`, +`oneArray`). Indices are converted to `Int32` for device friendliness. +""" +function Dagger.device_sparse_from_host(::Type{Arr}, S::SparseMatrixCSC{Tv}) where {Tv,Arr} + colptr = Arr(Int32.(S.colptr)) + rowval = Arr(Int32.(S.rowval)) + nzval = Arr(Array(S.nzval)) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end +Base.copy(A::Dagger.DeviceSparseMatrixCSC) = + Dagger.DeviceSparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), copy(A.nzval)) +Dagger._sparse_copy(A::Dagger.DeviceSparseMatrixCSC) = copy(A) +Dagger._sparse_collect(A::Dagger.DeviceSparseMatrixCSC) = SparseMatrixCSC(A) +function Dagger._sparse_similar(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}, ::Type{T}, dims::Dims{2}) where {Tv,Ti,T} + n = dims[2] + colptr = similar(A.colptr, Ti, n + 1) + fill!(colptr, one(Ti)) + rowval = similar(A.rowval, Ti, 0) + nzval = similar(A.nzval, T, 0) + return Dagger.DeviceSparseMatrixCSC(dims[1], n, colptr, rowval, nzval) +end + +# Rebuild a DeviceSparseMatrixCSC on the same device array type as `like`. +function _to_device_sparse(like::Dagger.DeviceSparseMatrixCSC, S::SparseMatrixCSC) + colptr = similar(like.colptr, eltype(like.colptr), length(S.colptr)) + rowval = similar(like.rowval, eltype(like.rowval), length(S.rowval)) + nzval = similar(like.nzval, eltype(S), length(S.nzval)) + copyto!(colptr, eltype(colptr).(S.colptr)) + copyto!(rowval, eltype(rowval).(S.rowval)) + copyto!(nzval, S.nzval) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(SparseArrays.spzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment; + # Route through `allocate_sparse_zeros` so a GPU compute scope yields + # device-resident sparse tiles (vendor sparse or DeviceSparseMatrixCSC). + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_zeros(Dagger.task_processor(), T, _dims)), false, d, Dagger.partition(p, d), p, assignment; return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end @@ -52,7 +105,7 @@ SparseArrays.spzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentTy function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(SparseArrays.sprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment; + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_rand(Dagger.task_processor(), T, _dims, sparsity)), false, d, Dagger.partition(p, d), p, assignment; return_type=DSparseArray{T,N}) return Dagger._to_darray(a) end @@ -71,6 +124,17 @@ _apply_trans(X, t::Char) = t == 'C' ? adjoint(X) : throw(ArgumentError("Invalid trans char: $t")) +function _sparse_gemm_assign!(C::DSparseMatrix, prod, beta) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = prod + C.mat + else + C.mat = prod + beta * C.mat + end + return C +end + function Dagger.matmatmul!( C::DSparseMatrix, transA::Char, @@ -91,14 +155,35 @@ function Dagger.matmatmul!( # methods, so `opA`/`opB` are not materialized. AB = opA * opB prod = isone(alpha) ? AB : alpha * AB + return _sparse_gemm_assign!(C, prod, beta) +end + +# DeviceSparseMatrixCSC SpGEMM: gather to host, multiply, scatter back. +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Dagger.DeviceSparseMatrixCSC, + B::Dagger.DeviceSparseMatrixCSC, + alpha, + beta +) + Ah = SparseMatrixCSC(A) + Bh = SparseMatrixCSC(B) + Ch = C.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(C.mat) : + C.mat isa SparseMatrixCSC ? C.mat : SparseMatrixCSC(C.mat) + opA = _apply_trans(Ah, transA) + opB = _apply_trans(Bh, transB) + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB if iszero(beta) - C.mat = prod + result = prod elseif isone(beta) - C.mat = prod + C.mat + result = prod + Ch else - C.mat = prod + beta * C.mat + result = prod + beta * Ch end - + C.mat = _to_device_sparse(A, SparseMatrixCSC(result)) return C end @@ -111,10 +196,24 @@ function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::SparseMatrixCSC, return C end +# DeviceSparseMatrixCSC SpMV: host fallback (works for any dense vector type +# that supports Array(::)/copyto!). +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Dagger.DeviceSparseMatrixCSC, B::AbstractVector, alpha, beta) + Ah = SparseMatrixCSC(A) + Bh = Array(B) + Ch = Array(C) + LinearAlgebra.mul!(Ch, _apply_trans(Ah, transA), Bh, alpha, beta) + copyto!(C, Ch) + return C +end + # Off-diagonal tile copy in `copytri!`: produce the (conjugate) transpose tile. function Dagger.transpose_tile(B::SparseMatrixCSC) return SparseArrays.sparse(B') end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC) + return _to_device_sparse(B, SparseArrays.sparse(SparseMatrixCSC(B)')) +end # Diagonal tile symmetrization in `copytri!`: build the full Hermitian tile from # its `uplo` triangle (matching the dense `copydiagtile!` semantics). function Dagger.transpose_tile(B::SparseMatrixCSC, uplo::Char) @@ -132,5 +231,8 @@ function Dagger.transpose_tile(B::SparseMatrixCSC, uplo::Char) end return C end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC, uplo::Char) + return _to_device_sparse(B, Dagger.transpose_tile(SparseMatrixCSC(B), uplo)) +end end # module SparseArraysExt diff --git a/src/array/darray.jl b/src/array/darray.jl index 2df44afbf..b774d913f 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -218,8 +218,16 @@ function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} # Distributed: fetch chunks directly and concat in-process. This avoids # routing chunk data through datadeps aliasing, which requires an # aliasing-resolvable (e.g. isbits) element type. + # + # Sparse tiles (`DSparseArray`) are unwrapped to host storage before `cat`: + # GPU sparse types (CuSparse/ROCSparse/DeviceSparseMatrixCSC) disallow the + # scalar indexing that generic `cat` would use. dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] - return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) + tiles = asyncmap(a.chunks) do c + x = fetch(c) + x isa DSparseArray ? _sparse_collect(x.mat) : x + end + return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, tiles)) end Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 3e56c37b0..147e73e84 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -53,8 +53,25 @@ Base.hash(M::DSparseArray, h::UInt) = hash(objectid(M), hash(DSparseArray, h)) # `aliasing_root` resolves any wrapper of a `DSparseArray` back to the container # before computing aliasing. Note: this must return a *bare* `ObjectAliasing`, as # `aliases_as_whole(::ObjectAliasing)` relies on it to drive whole-object copies. -aliasing(M::DSparseArray, _=identity) = ObjectAliasing(M) +# +# The ObjectAliasing is stamped with the inner storage's memory space so GPU- +# resident sparse tiles (e.g. CuSparseMatrixCSC / ROCSparseMatrixCSC, or +# DeviceSparseMatrixCSC) are tracked in VRAM rather than host RAM. The pointer +# itself is still the host wrapper's `pointer_from_objref` (identity only); +# Datadeps never span-copies these objects (see `aliases_as_whole` / `FullCopy`). +function aliasing(M::DSparseArray, _=identity) + space = value_memory_space(M.mat) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(M)), space) + return ObjectAliasing(ptr, sizeof(typeof(M))) +end aliases_as_whole(::DSparseArray) = true +# Chunk space follows the inner sparse storage (host CSC → CPURAM, GPU sparse → VRAM). +# `memory_space` (not only `value_memory_space`) must resolve correctly: Datadeps' +# `aliased_object!` uses `memory_space(x)` when deciding whether a same-space +# slot can share the original object. The generic fallback labels unknown values +# as CPURAM, which would force a discarded GPU copy and drop SpGEMM writes. +value_memory_space(M::DSparseArray) = value_memory_space(M.mat) +memory_space(M::DSparseArray) = value_memory_space(M) # A `DSparseArray` has no meaningful raw data pointer (its storage may be # reallocated/resized). This trap ensures that if aliasing ever tries to treat a @@ -86,6 +103,26 @@ end # `Base.copy`). _sparse_copy(mat) = copy(mat) +# Adapt preserves the `DSparseArray` wrapper while adapting the inner storage +# (e.g. SparseMatrixCSC → CuSparseMatrixCSC under `adapt(CuArray, …)`). GPU +# package extensions override `move` for sparse types so OpenCL/Metal/oneAPI do +# not densify via the generic `adapt(DeviceArray, SparseMatrixCSC)` path. +Adapt.adapt_structure(to, M::DSparseArray) = DSparseArray(Adapt.adapt(to, M.mat)) + +# Whole-object in-place move for Datadeps `FullCopy`. Reassigns `to.mat` from a +# space-converted copy of `from.mat` rather than span-copying storage. +function move!(to_space::MemorySpace, from_space::MemorySpace, to::DSparseArray, from::DSparseArray) + if to_space == from_space + to.mat = _sparse_copy(from.mat) + else + to_proc = first(processors(to_space)) + from_proc = first(processors(from_space)) + moved = move(from_proc, to_proc, from) + to.mat = moved.mat + end + return +end + # Wrapping hook used when materializing tiles (e.g. in `distribute`). Backends # overload this (e.g. in package extensions) to wrap freshly-created tiles in a # container that Datadeps can track (such as `DSparseArray` for sparse tiles); @@ -149,4 +186,65 @@ end function copydiagtile!(A::DSparseMatrix, uplo) A.mat = transpose_tile(A.mat, uplo) return -end \ No newline at end of file +end + +#============================================================================== + Device-resident CSC for GPU backends without a vendor sparse library + (OpenCL / Metal / oneAPI). Storage vectors live in device memory; SpGEMM/SpMV + kernels fall back to host SparseArrays (see SparseArraysExt). CUDA/ROCm use + their native CuSparse/ROCSparse types instead. +==============================================================================# + +""" + DeviceSparseMatrixCSC{Tv,Ti,V,I} + +A CSC sparse matrix whose `colptr`/`rowval`/`nzval` live in device arrays +(`CLArray`, `MtlArray`, `oneArray`, …). Used as the inner storage of a +[`DSparseArray`](@ref) on GPU backends that lack a vendor sparse library. +""" +struct DeviceSparseMatrixCSC{Tv,Ti,V<:AbstractVector{Tv},I<:AbstractVector{Ti}} <: AbstractMatrix{Tv} + m::Int + n::Int + colptr::I + rowval::I + nzval::V +end +Base.size(A::DeviceSparseMatrixCSC) = (A.m, A.n) +Base.eltype(::DeviceSparseMatrixCSC{Tv}) where Tv = Tv +Base.IndexStyle(::Type{<:DeviceSparseMatrixCSC}) = IndexCartesian() +function Base.getindex(A::DeviceSparseMatrixCSC{Tv}, i::Integer, j::Integer) where Tv + @boundscheck checkbounds(A, i, j) + # Scalar device indexing; acceptable for rare host-side inspection. + col_start = Int(A.colptr[j]) + col_end = Int(A.colptr[j + 1]) - 1 + for p in col_start:col_end + Int(A.rowval[p]) == i && return A.nzval[p] + end + return zero(Tv) +end +# Memory space follows the nonzero values buffer. +value_memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A.nzval) +memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A) +# Whole-object aliasing (storage may be replaced on write); stamp with device space. +function aliasing(A::DeviceSparseMatrixCSC, _=identity) + space = value_memory_space(A) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(A)), space) + return ObjectAliasing(ptr, sizeof(typeof(A))) +end +aliases_as_whole(::DeviceSparseMatrixCSC) = true +function Base.pointer(::DeviceSparseMatrixCSC) + throw(ArgumentError("`pointer(::DeviceSparseMatrixCSC)` is unsupported; \ + alias as a whole object via `Dagger.aliasing`.")) +end + +# Allocation hooks used by `spzeros`/`sprand` under a GPU compute scope. +# GPU Exts override these for their processor type. +allocate_sparse_zeros(proc::Processor, ::Type{T}, dims::Dims) where T = + allocate_sparse_zeros_default(proc, T, dims) +allocate_sparse_rand(proc::Processor, ::Type{T}, dims::Dims, sparsity::AbstractFloat) where T = + allocate_sparse_rand_default(proc, T, dims, sparsity) +# Defaults are filled in by SparseArraysExt (host SparseMatrixCSC / SparseVector). +function allocate_sparse_zeros_default end +function allocate_sparse_rand_default end +# Upload host CSC onto device vectors of type `Arr` (filled in by SparseArraysExt). +function device_sparse_from_host end \ No newline at end of file diff --git a/test/cudaenv/Project.toml b/test/cudaenv/Project.toml index 6e245a12a..38ea2a004 100644 --- a/test/cudaenv/Project.toml +++ b/test/cudaenv/Project.toml @@ -4,6 +4,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/gpu.jl b/test/gpu.jl index 474a0bd08..613c3442e 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -1,11 +1,45 @@ using Random using LinearAlgebra +using SparseArrays @everywhere begin using Distributed, Dagger import Dagger: Kernel using KernelAbstractions end + +# Sparse DArray SpGEMM / SpMV under a GPU compute scope. `check_tile(chunk)` +# verifies a distributed tile is device-resident (backend-specific). +function test_gpu_sparse_darray(scope; check_tile) + SA = sprand(Float32, 8, 8, 0.35) + SB = sprand(Float32, 8, 8, 0.35) + x = rand(Float32, 8) + Dagger.with_options(;scope) do + DSA = distribute(SA, Blocks(4, 4)) + DSB = distribute(SB, Blocks(4, 4)) + for chunk in DSA.chunks + @test check_tile(fetch(chunk; raw=true)) + end + + @test collect(DSA * DSB) ≈ SA * SB + @test collect(DSA * DSB') ≈ SA * SB' + @test collect(DSA' * DSB) ≈ SA' * SB + + DSC = similar(DSA) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SA * SB + + Dx = distribute(x, Blocks(4)) + @test collect(DSA * Dx) ≈ SA * x + + Z = SparseArrays.spzeros(Blocks(4, 4), Float32, 8, 8) + # `collect` densifies; check emptiness on device tiles / dense gather. + @test iszero(sum(abs, collect(Z))) + for chunk in Z.chunks + @test check_tile(fetch(chunk; raw=true)) + end + end +end @everywhere begin function isongpu(X) return !(X isa Array) @@ -227,6 +261,17 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cuda_gpu=gpu) + CUDAExt = Base.get_extension(Dagger, :CUDAExt) + test_gpu_sparse_darray(scope; check_tile=chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa CUDA.CUSPARSE.CuSparseMatrixCSC && + chunk.space isa CUDAExt.CUDAVRAMMemorySpace + end) + end end end @@ -398,6 +443,17 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, rocm_gpu=gpu) + ROCExt = Base.get_extension(Dagger, :ROCExt) + test_gpu_sparse_darray(scope; check_tile=chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa AMDGPU.rocSPARSE.ROCSparseMatrixCSC && + chunk.space isa ROCExt.ROCVRAMMemorySpace + end) + end end end @@ -569,6 +625,17 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, intel_gpu=gpu) + IntelExt = Base.get_extension(Dagger, :IntelExt) + test_gpu_sparse_darray(scope; check_tile=chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa IntelExt.IntelVRAMMemorySpace + end) + end end end @@ -714,6 +781,17 @@ end @test_broken array[2, 1] == 4.0f0 @test_broken array[2, 2] == 5.0f0 end + + @testset "Sparse DArray" begin + scope = Dagger.scope(worker=1, metal_gpu=1) + MetalExt = Base.get_extension(Dagger, :MetalExt) + test_gpu_sparse_darray(scope; check_tile=chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa MetalExt.MetalVRAMMemorySpace + end) + end end end @@ -820,5 +898,16 @@ end end @test A ≈ ref .+ 1 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cl_device=gpu) + OpenCLExt = Base.get_extension(Dagger, :OpenCLExt) + test_gpu_sparse_darray(scope; check_tile=chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa OpenCLExt.CLMemorySpace + end) + end end end diff --git a/test/metalenv/Project.toml b/test/metalenv/Project.toml index b3d5e9632..903bfcb7e 100644 --- a/test/metalenv/Project.toml +++ b/test/metalenv/Project.toml @@ -4,6 +4,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/oneapienv/Project.toml b/test/oneapienv/Project.toml index 30e7631f3..1c5664279 100644 --- a/test/oneapienv/Project.toml +++ b/test/oneapienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" diff --git a/test/openclenv/Project.toml b/test/openclenv/Project.toml index b7718eca7..f44d490e1 100644 --- a/test/openclenv/Project.toml +++ b/test/openclenv/Project.toml @@ -4,7 +4,9 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" [sources] Dagger = {path = "../.."} diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml index d350b4c4a..3d3d50b69 100644 --- a/test/rocmenv/Project.toml +++ b/test/rocmenv/Project.toml @@ -4,6 +4,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] From 03a084206446eaaad12f5c694e3c9fa2e00ed114 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 22:31:51 -0700 Subject: [PATCH 26/43] DArray/sparse: Add GPU sparse Krylov solver tests Exercise cg/minres/gmres/bicgstab and Jacobi preconditioning on GPU sparse tiles, with host-safe diagonal/factorize hooks for device CSC. Co-authored-by: Cursor --- src/array/iterativesolvers.jl | 11 +++- src/array/sparse.jl | 19 +++++-- test/gpu.jl | 98 +++++++++++++++++++++++++++++++---- test/openclenv/Project.toml | 1 + test/rocmenv/Project.toml | 1 + 5 files changed, 114 insertions(+), 16 deletions(-) diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl index 95e9100ac..f3d2f116c 100644 --- a/src/array/iterativesolvers.jl +++ b/src/array/iterativesolvers.jl @@ -239,7 +239,16 @@ end # Apply one block operator: `y = op⁻¹ x`. Default uses `\`; backends override for # operator types where `\` is unavailable (e.g. ILU/AMG use `ldiv!`). -_block_apply!(y, op, x) = (y .= op \ x; return nothing) +# Host factorizations (UMFPACK/`lu` of gathered CSC) need a host RHS; GPU vector +# chunks are gathered, solved, and written back. CPU `Array` short-circuits. +function _block_apply!(y, op, x) + if x isa Array + y .= op \ x + else + copyto!(y, op \ Adapt.adapt(Array, x)) + end + return nothing +end # Underlying matrix of a tile (overridden for `DSparseArray` in `sparse.jl`). _tile_matrix(A) = A diff --git a/src/array/sparse.jl b/src/array/sparse.jl index 147e73e84..c5ab69b91 100644 --- a/src/array/sparse.jl +++ b/src/array/sparse.jl @@ -170,13 +170,22 @@ function matvecmul!(C, transA::Char, A::DSparseMatrix, B, alpha, beta) return matvecmul!(C, transA, A.mat, B, alpha, beta) end -# Factorize the inner sparse storage (e.g. sparse LU/UMFPACK) for block-Jacobi, -# rather than the `DSparseArray` wrapper (whose generic `lu` would densify). -_factorize_tile(M::DSparseArray) = LinearAlgebra.lu(M.mat) +# Factorize via host CSC so GPU sparse tiles (rocSPARSE / DeviceSparseMatrixCSC) +# can still feed block-Jacobi; CPU SparseMatrixCSC is a no-op collect. +_factorize_tile(M::DSparseArray) = LinearAlgebra.lu(_sparse_collect(M.mat)) # Expose the inner sparse storage so preconditioner backends (ILU, AMG) build on -# the actual `SparseMatrixCSC` rather than the `DSparseArray` wrapper. -_tile_matrix(M::DSparseArray) = M.mat +# the actual `SparseMatrixCSC` rather than the `DSparseArray` wrapper. GPU mats +# are gathered to host CSC (those backends are host-only). +_tile_matrix(M::DSparseArray) = _sparse_collect(M.mat) + +# Jacobi diagonal extract without GPU scalar indexing. `diag(::SparseMatrixCSC)` +# returns a `SparseVector`; densify before `copyto!` into a device vector. +function _set_inv_diag!(out, tile::DSparseArray) + d = Vector(LinearAlgebra.diag(_sparse_collect(tile.mat))) + copyto!(out, inv.(d)) + return nothing +end function transpose_tile end function copytile!(A::DSparseMatrix, B::DSparseMatrix) diff --git a/test/gpu.jl b/test/gpu.jl index 613c3442e..28b908eb6 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -1,6 +1,7 @@ using Random using LinearAlgebra using SparseArrays +using Krylov @everywhere begin using Distributed, Dagger @@ -40,6 +41,73 @@ function test_gpu_sparse_darray(scope; check_tile) end end end + +# Sparse Krylov solvers under a GPU scope (mirrors the sparse backend cases in +# `array/linalg/iterativesolvers.jl`, trimmed for GPU CI time). +function _gpu_sparse_laplacian(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end +function _gpu_sparse_advection(T, n) + return _gpu_sparse_laplacian(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end +function test_gpu_sparse_solvers(scope; check_tile=nothing) + n, k = 32, 8 + A_part, b_part = Blocks(k, k), Blocks(k) + T = Float32 + + Dagger.with_options(;scope) do + Asp = _gpu_sparse_laplacian(T, n) + b = rand(T, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, b_part) + if check_tile !== nothing + @test check_tile(fetch(DA.chunks[1]; raw=true)) + end + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-3 + end + + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = T(1e-6), rtol = T(1e-5)) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + # Diagonal Jacobi is device-friendly. Block-Jacobi uses host UMFPACK + # factors that Datadeps cannot place under a GPU-only compute scope, so + # it is covered on CPU in `array/linalg/iterativesolvers.jl` instead. + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(T(1) / T(4), n) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (T(1) / T(4)) .* b + x, stats = Dagger.cg(DA, Db; M = P, atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + Anonsym = _gpu_sparse_advection(T, n) + bn = rand(T, n) + xrefn = Matrix(Anonsym) \ bn + DAn = distribute(Anonsym, A_part) + Dbn = distribute(bn, b_part) + @testset "$(nameof(solver)) nonsym" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DAn, Dbn; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xrefn rtol = 1e-3 + end + end +end + @everywhere begin function isongpu(X) return !(X isa Array) @@ -265,12 +333,14 @@ end @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs scope = Dagger.scope(worker=1, cuda_gpu=gpu) CUDAExt = Base.get_extension(Dagger, :CUDAExt) - test_gpu_sparse_darray(scope; check_tile=chunk -> begin + check_tile = chunk -> begin v = Dagger.MemPool.poolget(chunk.handle) return v isa Dagger.DSparseArray && v.mat isa CUDA.CUSPARSE.CuSparseMatrixCSC && chunk.space isa CUDAExt.CUDAVRAMMemorySpace - end) + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) end end end @@ -447,12 +517,14 @@ end @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs scope = Dagger.scope(worker=1, rocm_gpu=gpu) ROCExt = Base.get_extension(Dagger, :ROCExt) - test_gpu_sparse_darray(scope; check_tile=chunk -> begin + check_tile = chunk -> begin v = Dagger.MemPool.poolget(chunk.handle) return v isa Dagger.DSparseArray && v.mat isa AMDGPU.rocSPARSE.ROCSparseMatrixCSC && chunk.space isa ROCExt.ROCVRAMMemorySpace - end) + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) end end end @@ -629,12 +701,14 @@ end @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs scope = Dagger.scope(worker=1, intel_gpu=gpu) IntelExt = Base.get_extension(Dagger, :IntelExt) - test_gpu_sparse_darray(scope; check_tile=chunk -> begin + check_tile = chunk -> begin v = Dagger.MemPool.poolget(chunk.handle) return v isa Dagger.DSparseArray && v.mat isa Dagger.DeviceSparseMatrixCSC && chunk.space isa IntelExt.IntelVRAMMemorySpace - end) + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) end end end @@ -785,12 +859,14 @@ end @testset "Sparse DArray" begin scope = Dagger.scope(worker=1, metal_gpu=1) MetalExt = Base.get_extension(Dagger, :MetalExt) - test_gpu_sparse_darray(scope; check_tile=chunk -> begin + check_tile = chunk -> begin v = Dagger.MemPool.poolget(chunk.handle) return v isa Dagger.DSparseArray && v.mat isa Dagger.DeviceSparseMatrixCSC && chunk.space isa MetalExt.MetalVRAMMemorySpace - end) + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) end end end @@ -902,12 +978,14 @@ end @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs scope = Dagger.scope(worker=1, cl_device=gpu) OpenCLExt = Base.get_extension(Dagger, :OpenCLExt) - test_gpu_sparse_darray(scope; check_tile=chunk -> begin + check_tile = chunk -> begin v = Dagger.MemPool.poolget(chunk.handle) return v isa Dagger.DSparseArray && v.mat isa Dagger.DeviceSparseMatrixCSC && chunk.space isa OpenCLExt.CLMemorySpace - end) + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) end end end diff --git a/test/openclenv/Project.toml b/test/openclenv/Project.toml index f44d490e1..8c25145c4 100644 --- a/test/openclenv/Project.toml +++ b/test/openclenv/Project.toml @@ -1,5 +1,6 @@ [deps] Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml index 3d3d50b69..7b8ac67f3 100644 --- a/test/rocmenv/Project.toml +++ b/test/rocmenv/Project.toml @@ -1,6 +1,7 @@ [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" From c30a308c1129468858d16cbceff2365321375e0f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 11 May 2026 15:35:25 -0700 Subject: [PATCH 27/43] datadeps: Add AOT DAG scheduling --- src/Dagger.jl | 1 + src/datadeps/aliasing.jl | 9 +- src/datadeps/queue.jl | 62 +++- src/datadeps/scheduling.jl | 309 ++++++++++++++++++- src/datadeps/types.jl | 30 ++ src/memory-spaces.jl | 47 +++ test/datadeps/scheduling.jl | 577 ++++++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 8 files changed, 1022 insertions(+), 14 deletions(-) create mode 100644 src/datadeps/types.jl create mode 100644 test/datadeps/scheduling.jl diff --git a/src/Dagger.jl b/src/Dagger.jl index 3fe2a54ba..71b15bbf8 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -102,6 +102,7 @@ include("sch/Sch.jl"); using .Sch include("tochunk.jl") # Data dependency task queue +include("datadeps/types.jl") include("datadeps/aliasing.jl") include("datadeps/chunkview.jl") include("datadeps/remainders.jl") diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 22c7ac071..2f90a4d48 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1,5 +1,3 @@ -import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv - export In, Out, InOut, Deps, spawn_datadeps #= @@ -393,6 +391,9 @@ function aliased_object!(f, cache::AliasedObjectCache, x; ainfo=aliasing(cache.a end struct DataDepsState + # The DAG spec for the task + dag_spec::DAGSpec + # The mapping of original raw argument to its Chunk # N.B. Values are Chunks, or raw remote handles (e.g. ChunkView under MPI) raw_arg_to_chunk::IdDict{Any,Any} @@ -462,7 +463,7 @@ struct DataDepsState ainfos_owner::Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}} ainfos_readers::Dict{AliasingWrapper,Vector{Pair{DTask,Int}}} - function DataDepsState() + function DataDepsState(dag_spec::DAGSpec) arg_to_chunk = IdDict{Any,Any}() arg_origin = IdDict{Any,MemorySpace}() remote_args = Dict{MemorySpace,IdDict{Any,Any}}() @@ -487,7 +488,7 @@ struct DataDepsState ainfos_owner = Dict{AliasingWrapper,Union{Pair{DTask,Int},Nothing}}() ainfos_readers = Dict{AliasingWrapper,Vector{Pair{DTask,Int}}}() - return new(arg_to_chunk, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, + return new(dag_spec, arg_to_chunk, arg_origin, remote_args, remote_arg_to_original, remote_arg_w, ainfo_arg, arg_history, arg_owner, arg_current, arg_overlaps, ainfo_backing_chunk, supports_inplace_cache, ainfo_cache, ainfos_lookup, ainfos_overlaps, ainfos_owner, ainfos_readers) end end diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 0ef6c6655..989e223b5 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -133,17 +133,61 @@ function distribute_tasks!(queue::DataDepsTaskQueue) all_scope = UnionScope(map(ExactScope, all_procs)) exec_spaces = unique(vcat(map(proc->collect(memory_spaces(proc)), all_procs)...)) - # Round-robin assign tasks to processors upper_queue = get_options(:task_queue) + # Compute DAG spec + dag_spec = DAGSpec() + state = DataDepsState(dag_spec) + for (spec, task) in queue.seen_tasks + if !dag_add_task!(dag_spec, state, spec, task) + # This task needs to be deferred + break + end + end + + # Attempt to find any matching DAG specs and reuse their schedule + schedule = Dict{DTask, Processor}() + schedule_cache = datadeps_schedule_cache(queue.scheduler) + cache_hit = false + for (other_spec, spec_schedule) in schedule_cache + if datadeps_dag_equivalent(queue.scheduler, dag_spec, other_spec) + @dagdebug nothing :spawn_datadeps "Found matching DAG spec!" + for (id, proc) in spec_schedule.id_to_proc + uid = dag_spec.id_to_uid[id] + task_idx = findfirst(spec_task -> spec_task.task.uid == uid, queue.seen_tasks) + task = queue.seen_tasks[task_idx].task + schedule[task] = proc + end + cache_hit = true + break + end + end + + if !cache_hit && !isempty(dag_spec) + # Compute a fresh AOT schedule (no-op for schedulers that fall back + # to JIT in distribute_task!) + datadeps_schedule_dag_aot!(queue.scheduler, schedule, dag_spec, all_procs, all_scope) + + # Persist the schedule for reuse by future equivalent DAGs + if !isempty(schedule) + spec_schedule = DAGSpecSchedule() + for (task, proc) in schedule + id = dag_spec.uid_to_id[task.uid] + spec_schedule.id_to_proc[id] = proc + end + push!(schedule_cache, dag_spec => spec_schedule) + end + end + # Start launching tasks and necessary copies - state = DataDepsState() + state = DataDepsState(dag_spec) write_num = 1 proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) for pair in queue.seen_tasks spec = pair.spec task = pair.task - write_num = distribute_task!(queue, state, all_procs, all_scope, spec, task, spec.fargs, proc_to_scope_lfu, write_num) + proc = get(schedule, task, nothing) + write_num = distribute_task!(queue, state, all_procs, all_scope, spec, task, spec.fargs, proc_to_scope_lfu, write_num; proc) end # Copy args from remote to local @@ -242,7 +286,7 @@ map_or_ntuple(f, xs::Vector) = map(f, 1:length(xs)) # N.B. Accept any `Tuple` (typed specs produce heterogeneous tuples of # `TypedArgument{T}`, not a homogeneous `NTuple{N,T}`). @inline map_or_ntuple(@specialize(f), xs::Tuple) = ntuple(f, Val(length(xs))) -function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int) where typed +function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int; proc::Union{Processor,Nothing}=nothing) where typed @specialize spec fargs if typed @@ -254,8 +298,14 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr DATADEPS_CURRENT_TASK[] = task task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) - scheduler = queue.scheduler - our_proc = datadeps_schedule_task(scheduler, state, all_procs, all_scope, task_scope, spec, task) + + if proc === nothing + # Schedule and JIT assign tasks to processors + our_proc = datadeps_schedule_task_jit!(queue.scheduler, all_procs, all_scope, task_scope, spec, task) + else + # Use the provided processor from AOT scheduling + our_proc = proc + end @assert our_proc in all_procs our_space = only(memory_spaces(our_proc)) check_uniform(our_proc) diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 62a350c72..cb6ad861c 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -1,10 +1,253 @@ -abstract type DataDepsScheduler end +### DAG Analysis ### + +Base.length(spec::DAGSpec) = nv(spec.g) +Base.isempty(spec::DAGSpec) = length(spec) == 0 + +function dag_add_task!(dspec::DAGSpec, state, tspec::DTaskSpec, task::DTask) + # Check if this task depends on any other tasks within the DAG, + # which we are not yet ready to handle + for (idx, _arg) in enumerate(tspec.fargs) + arg, deps = unwrap_inout(value(_arg)) + for (dep_mod, readdep, writedep) in deps + if arg isa DTask + if arg.uid in keys(dspec.uid_to_id) + # Within-DAG dependency, bail out + return false + end + end + end + end + + add_vertex!(dspec.g) + id = nv(dspec.g) + + # Record function signature + dspec.id_to_functype[id] = chunktype(tspec.fargs[1]) + argtypes = DatadepsArgSpec[] + for (idx, _arg) in enumerate(tspec.fargs) + arg, deps = unwrap_inout(value(_arg)) + pos = raw_position(_arg) + for (dep_mod, readdep, writedep) in deps + if arg isa DTask + #= TODO: Re-enable this when we can handle within-DAG dependencies + if arg.uid in keys(dspec.uid_to_id) + # Within-DAG dependency + arg_id = dspec.uid_to_id[arg.uid] + push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing())) + add_edge!(dspec.g, arg_id, id) + continue + end + =# + + # External DTask, so fetch this and track it as a raw value + arg = fetch(arg; raw=true) + end + ainfo = aliasing(arg, dep_mod) + # FIXME: Generate syncdeps and add edges + push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo)) + end + end + dspec.id_to_argtypes[id] = argtypes + dspec.id_to_scope[id] = @something(tspec.options.compute_scope, + tspec.options.scope, + DefaultScope()) + + # FIXME: Record syncdeps + dspec.id_to_uid[id] = task.uid + dspec.uid_to_id[task.uid] = id + dspec.id_to_spec[id] = tspec + dspec.id_to_task[id] = task + + return true +end +function dag_build_edges!(dag_spec::DAGSpec) + # Naively build edges based on exact argument comparisons (no aliasing) + arg_writes = Dict{Any,Int}() + for idx in 1:nv(dag_spec.g) + task_spec = dag_spec.id_to_spec[idx] + + # Get the raw arguments for this task + task_raw_args = Vector{Any}() + for arg in task_spec.fargs + arg, deps = unwrap_inout(arg) + for (dep_mod, readdep, writedep) in deps + # Get the raw argument + raw_arg = arg isa DTask ? fetch(arg; raw=true) : arg + + # Did any previous task write to this argument? + if haskey(arg_writes, raw_arg) && arg_writes[raw_arg] != idx + prev_task_id = arg_writes[raw_arg] + add_edge!(dag_spec.g, prev_task_id, idx) + end + + if writedep + # Record this write + arg_writes[raw_arg] = idx + end + end + end + end +end +function dag_has_task(dspec::DAGSpec, task::DTask) + return task.uid in keys(dspec.uid_to_id) +end + +### DAGSpec Equivalence (scheduler-dispatched) ### + +""" + datadeps_dag_equivalent(scheduler, dspec1::DAGSpec, dspec2::DAGSpec) -> Bool + +Returns `true` if a schedule cached for `dspec2` may safely be reused for +`dspec1`, as judged by `scheduler`. This is the top-level entry point used by +`distribute_tasks!` when looking up a cached schedule. + +The default implementation requires: +- Same number of vertices and edges in the dependency graph +- Identical edge sets (per-vertex `outneighbors`) +- Per-vertex agreement on function type and on each task's compute scope +- Per-vertex agreement on the multiset of argspecs, compared via + `datadeps_argspec_equivalent` + +Schedulers that want completely custom matching (or want to opt out of caching +entirely) can override this directly. Schedulers that only want to tweak how +individual arguments are compared should instead override +`datadeps_argspec_equivalent` or `datadeps_ainfo_equivalent`. +""" +function datadeps_dag_equivalent(scheduler::DataDepsScheduler, + dspec1::DAGSpec, dspec2::DAGSpec) + # Graph shape + nv(dspec1.g) == nv(dspec2.g) || return false + ne(dspec1.g) == ne(dspec2.g) || return false + + @inbounds for id in 1:nv(dspec1.g) + # outneighbors covers all edges; inneighbors would be redundant + outneighbors(dspec1.g, id) == outneighbors(dspec2.g, id) || return false + + # Function type must match exactly + dspec1.id_to_functype[id] === dspec2.id_to_functype[id] || return false + + # Per-task compute scope must match (different scopes can produce + # different schedules and must not be aliased) + dspec1.id_to_scope[id] == dspec2.id_to_scope[id] || return false + + # Argspecs must match as a multiset (Deps can put multiple argspecs at + # the same position) + _argspecs_equivalent(scheduler, + dspec1.id_to_argtypes[id], + dspec2.id_to_argtypes[id]) || return false + end + + return true +end + +# Backwards-compatible default `==`: use the no-scheduler default (i.e. as if +# all schedulers behaved like the base `DataDepsScheduler`). The runtime path +# in `distribute_tasks!` calls `datadeps_dag_equivalent` directly. +Base.:(==)(dspec1::DAGSpec, dspec2::DAGSpec) = + datadeps_dag_equivalent(_DefaultEquivalenceScheduler(), dspec1, dspec2) + +# A private marker scheduler used to provide a default for `Base.:(==)` on +# `DAGSpec`. Not exported and not intended for direct use. +struct _DefaultEquivalenceScheduler <: DataDepsScheduler end + +""" + datadeps_argspec_equivalent(scheduler, + a1::DatadepsArgSpec, + a2::DatadepsArgSpec) -> Bool + +Returns `true` if argspecs `a1` and `a2` are interchangeable for the purposes +of `scheduler`'s cached-schedule lookup. The default requires equal positions, +equal value types, equal dep_mods, and structurally-equivalent ainfos (per +`datadeps_ainfo_equivalent`). +""" +function datadeps_argspec_equivalent(scheduler::DataDepsScheduler, + a1::DatadepsArgSpec, a2::DatadepsArgSpec) + a1.pos == a2.pos || return false + a1.value_type === a2.value_type || return false + a1.dep_mod === a2.dep_mod || return false + return datadeps_ainfo_equivalent(scheduler, a1.ainfo, a2.ainfo) +end + +""" + datadeps_ainfo_equivalent(scheduler, + a1::AbstractAliasing, + a2::AbstractAliasing) -> Bool + +Returns `true` if aliasings `a1` and `a2` are interchangeable for the purposes +of `scheduler`'s cached-schedule lookup. The default uses +`equivalent_structure`, which compares shape/strides/lengths while ignoring +absolute pointer addresses, enabling reuse across re-allocations. + +Schedulers can override this to choose a different equivalence strategy, e.g. +pointer-identical (strictest), locality-only (memory-space only), or fully +permissive. +""" +datadeps_ainfo_equivalent(::DataDepsScheduler, + a1::AbstractAliasing, a2::AbstractAliasing) = + equivalent_structure(a1, a2) + +# Compare two argspec vectors as multisets, since `Deps` can place multiple +# argspecs at the same position. We pair each argspec in `as1` with a not-yet- +# matched argspec in `as2`; both lists must be exhausted simultaneously. +function _argspecs_equivalent(scheduler::DataDepsScheduler, + as1::Vector{DatadepsArgSpec}, + as2::Vector{DatadepsArgSpec}) + length(as1) == length(as2) || return false + n = length(as1) + n == 0 && return true + matched = falses(n) + @inbounds for a1 in as1 + found = false + for j in 1:n + matched[j] && continue + if datadeps_argspec_equivalent(scheduler, a1, as2[j]) + matched[j] = true + found = true + break + end + end + found || return false + end + return true +end + +### Schedule Cache (scheduler-owned) ### + +struct DAGSpecSchedule + id_to_proc::Dict{Int, Processor} + DAGSpecSchedule() = new(Dict{Int, Processor}()) +end + +# Per-scheduler-type cache. Each entry in the inner Vector is a (DAGSpec => +# DAGSpecSchedule) pair recorded by a prior call. The outer Dict partitions +# the cache by `typeof(scheduler)` so schedulers don't contaminate each other. +const DATADEPS_DAG_SPECS = + TaskLocalValue{Dict{Type, Vector{Pair{DAGSpec, DAGSpecSchedule}}}}( + ()->Dict{Type, Vector{Pair{DAGSpec, DAGSpecSchedule}}}()) + +""" + datadeps_schedule_cache(scheduler) -> Vector{Pair{DAGSpec, DAGSpecSchedule}} + +Returns the schedule cache that `scheduler` should consult for prior schedules +and append newly-computed schedules to. The default implementation returns a +task-local, per-scheduler-type cache. + +Override this to implement custom caching strategies (e.g., bounded LRU, no +cache at all, cross-task shared cache). +""" +function datadeps_schedule_cache(scheduler::DataDepsScheduler) + cache_by_type = DATADEPS_DAG_SPECS[] + return get!(Vector{Pair{DAGSpec, DAGSpecSchedule}}, + cache_by_type, typeof(scheduler)) +end + +### JIT Schedulers ### mutable struct RoundRobinScheduler <: DataDepsScheduler proc_idx::Int RoundRobinScheduler() = new(1) end -function datadeps_schedule_task(sched::RoundRobinScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +function datadeps_schedule_task_jit!(sched::RoundRobinScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) proc_idx = sched.proc_idx our_proc = all_procs[proc_idx] if task_scope == all_scope @@ -24,7 +267,7 @@ function datadeps_schedule_task(sched::RoundRobinScheduler, state::DataDepsState end struct NaiveScheduler <: DataDepsScheduler end -function datadeps_schedule_task(sched::NaiveScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +function datadeps_schedule_task_jit!(sched::NaiveScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) raw_args = map(arg->tochunk(value(arg)), spec.fargs) our_proc = remotecall_fetch(1, all_procs, raw_args) do all_procs, raw_args Sch.init_eager() @@ -59,7 +302,7 @@ struct UltraScheduler <: DataDepsScheduler Dict{MemorySpace,Int}()) end end -function datadeps_schedule_task(sched::UltraScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) +function datadeps_schedule_task_jit!(sched::UltraScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) args = Base.mapany(spec.fargs) do arg pos, data = arg data, _ = unwrap_inout(data) @@ -121,3 +364,61 @@ function datadeps_schedule_task(sched::UltraScheduler, state::DataDepsState, all return our_proc end + +### AOT Schedulers ### + +function datadeps_schedule_dag_aot!(scheduler, schedule, dag_spec, all_procs, all_scope) + # Fallback to JIT scheduling (done in distribute_task!) + return +end + +struct LayeredScheduler <: DataDepsScheduler end +function datadeps_schedule_dag_aot!(scheduler::LayeredScheduler, schedule, dag_spec, all_procs, all_scope) + layer = 1 + layer_data = Vector{Any}() + layers = Vector{Vector{Int}}() + push!(layers, Vector{Int}()) + for idx in 1:nv(dag_spec.g) + spec = dag_spec.id_to_spec[idx] + + # Get the raw arguments for this task + task_raw_args = Vector{Any}() + for arg in spec.fargs + arg, deps = unwrap_inout(arg) + for (dep_mod, readdep, writedep) in deps + # We only care about write dependencies + writedep || continue + + # Get the raw argument + raw_arg = arg isa DTask ? fetch(arg; raw=true) : arg + + push!(task_raw_args, raw_arg) + end + end + + # Determine if this task stays in this layer + if any(raw_arg -> raw_arg in layer_data, task_raw_args) + # This argument is already written by a previous task in this layer + # Generate new layer + layer += 1 + empty!(layer_data) + push!(layers, Vector{Int}()) + end + + # Add our data and this task to the current layer + append!(layer_data, task_raw_args) + push!(layers[layer], idx) + end + + # Perform round-robin scheduling within each layer + for layer in layers + sub_scheduler = RoundRobinScheduler() + for idx in layer + spec = dag_spec.id_to_spec[idx] + task = dag_spec.id_to_task[idx] + task_scope = @something(spec.options.compute_scope, spec.options.scope, DefaultScope()) + our_proc = datadeps_schedule_task_jit!(sub_scheduler, all_procs, all_scope, task_scope, spec, task) + schedule[task] = our_proc + end + end +end \ No newline at end of file diff --git a/src/datadeps/types.jl b/src/datadeps/types.jl new file mode 100644 index 000000000..70212e10e --- /dev/null +++ b/src/datadeps/types.jl @@ -0,0 +1,30 @@ +import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv, ne + +### Scheduling ### + +struct DatadepsArgSpec + pos::Union{Int, Symbol} + value_type::Type + dep_mod::Any + ainfo::AbstractAliasing +end +struct DTaskDAGID{id} end +struct DAGSpec + g::SimpleDiGraph{Int} + id_to_uid::Dict{Int, UInt} + uid_to_id::Dict{UInt, Int} + id_to_functype::Dict{Int, Type} # FIXME: DatadepsArgSpec + id_to_argtypes::Dict{Int, Vector{DatadepsArgSpec}} + id_to_scope::Dict{Int, AbstractScope} + id_to_spec::Dict{Int, DTaskSpec} + id_to_task::Dict{Int, DTask} + DAGSpec() = new(SimpleDiGraph{Int}(), + Dict{Int, UInt}(), Dict{UInt, Int}(), + Dict{Int, Type}(), + Dict{Int, Vector{DatadepsArgSpec}}(), + Dict{Int, AbstractScope}(), + Dict{Int, DTaskSpec}(), + Dict{Int, DTask}()) +end + +abstract type DataDepsScheduler end \ No newline at end of file diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 015ca502c..31611cbc9 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -125,8 +125,27 @@ mutable struct AliasingWrapper <: AbstractAliasing AliasingWrapper(inner::AbstractAliasing) = new(inner, hash(inner)) end memory_spans(x::AliasingWrapper) = memory_spans(x.inner) +""" + equivalent_structure(x::AbstractAliasing, y::AbstractAliasing) -> Bool + +Returns `true` if `x` and `y` describe aliasing regions of the same *structure*, +ignoring absolute memory addresses. Two aliasings that originate from different +allocations of identically-shaped data should compare equal. + +Used as the default `datadeps_ainfo_equivalent` for `DAGSpec` cache lookups, +which lets schedulers reuse a previously-computed schedule across re-runs of +the same algorithm with freshly-allocated inputs. + +Default fallback returns `false` when the two arguments are of differing +concrete subtypes; specialized methods exist for each concrete subtype. +""" +equivalent_structure(x::AbstractAliasing, y::AbstractAliasing) = false equivalent_structure(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash || equivalent_structure(x.inner, y.inner) +equivalent_structure(x::AliasingWrapper, y::AbstractAliasing) = + equivalent_structure(x.inner, y) +equivalent_structure(x::AbstractAliasing, y::AliasingWrapper) = + equivalent_structure(x, y.inner) Base.hash(x::AliasingWrapper, h::UInt64) = hash(x.hash, h) Base.isequal(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash Base.:(==)(x::AliasingWrapper, y::AliasingWrapper) = x.hash == y.hash @@ -331,8 +350,13 @@ end struct NoAliasing <: AbstractAliasing end memory_spans(::NoAliasing) = MemorySpan{CPURAMMemorySpace}[] +equivalent_structure(::NoAliasing, ::NoAliasing) = true struct UnknownAliasing <: AbstractAliasing end memory_spans(::UnknownAliasing) = [MemorySpan{CPURAMMemorySpace}(C_NULL, typemax(UInt))] +# Two UnknownAliasings have no known structure to compare; conservatively +# treat them as equivalent so cache lookups don't permanently miss when +# aliasing analysis falls back to UnknownAliasing. +equivalent_structure(::UnknownAliasing, ::UnknownAliasing) = true error_unknown_aliasing(T) = throw(ConcurrencyViolationError("Cannot resolve aliasing for object of type $T, execution may become sequential")) @@ -357,6 +381,13 @@ Base.:(==)(ca1::CombinedAliasing, ca2::CombinedAliasing) = ca1.sub_ainfos == ca2.sub_ainfos Base.hash(ca1::CombinedAliasing, h::UInt) = hash(ca1.sub_ainfos, hash(CombinedAliasing, h)) +function equivalent_structure(ca1::CombinedAliasing, ca2::CombinedAliasing) + length(ca1.sub_ainfos) == length(ca2.sub_ainfos) || return false + @inbounds for i in eachindex(ca1.sub_ainfos) + equivalent_structure(ca1.sub_ainfos[i], ca2.sub_ainfos[i]) || return false + end + return true +end struct ObjectAliasing{S<:MemorySpace} <: AbstractAliasing ptr::RemotePtr{Cvoid,S} @@ -374,6 +405,8 @@ function memory_spans(oa::ObjectAliasing{S}) where S span = MemorySpan{S}(oa.ptr, oa.sz) return [span] end +equivalent_structure(x::ObjectAliasing{S}, y::ObjectAliasing{S}) where S = + x.sz == y.sz # Acceleration entry point used by Datadeps. Unwrap whole-object containers # (e.g. views of `DSparseArray`) before computing aliasing. SPMD backends @@ -526,6 +559,8 @@ end memory_spans(a::ContiguousAliasing{S}) where S = MemorySpan{S}[a.span] will_alias(x::ContiguousAliasing{S}, y::ContiguousAliasing{S}) where S = will_alias(x.span, y.span) +equivalent_structure(x::ContiguousAliasing{S}, y::ContiguousAliasing{S}) where S = + x.span.len == y.span.len struct IteratedAliasing{T} <: AbstractAliasing x::T end @@ -594,6 +629,14 @@ function aliasing(x::SubArray{T,N}) where {T,N} return UnknownAliasing() end end +function equivalent_structure(x::StridedAliasing{T,N,S}, y::StridedAliasing{T,N,S}) where {T,N,S} + x.base_inds == y.base_inds || return false + x.lengths == y.lengths || return false + x.strides == y.strides || return false + # Compare the offset of the view into its parent, not the absolute pointers, + # so views into different (but identically-shaped) parents match. + return (x.ptr.addr - x.base_ptr.addr) == (y.ptr.addr - y.base_ptr.addr) +end function will_alias(x::StridedAliasing{T1,N1,S1}, y::StridedAliasing{T2,N2,S2}) where {T1,T2,N1,N2,S1,S2} # Check if the base pointers are the same # FIXME: Conservatively incorrect via `unsafe_wrap` and friends @@ -643,6 +686,8 @@ function memory_spans(a::TriangularAliasing{T,S}) where {T,S} end return spans end +equivalent_structure(x::TriangularAliasing{T,S}, y::TriangularAliasing{T,S}) where {T,S} = + x.stride == y.stride && x.isupper == y.isupper && x.diagonal == y.diagonal function aliasing(x::UpperTriangular{T}) where T p = parent(x) space = memory_space(p) @@ -683,6 +728,8 @@ function aliasing(x::AbstractMatrix{T}, ::Type{Diagonal}) where T rptr = RemotePtr{Cvoid}(ptr, S) return DiagonalAliasing{T,typeof(S)}(rptr, size(parent(x), 1)) end +equivalent_structure(x::DiagonalAliasing{T,S}, y::DiagonalAliasing{T,S}) where {T,S} = + x.stride == y.stride # FIXME: Bidiagonal # FIXME: Tridiagonal diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl new file mode 100644 index 000000000..8f5bb1070 --- /dev/null +++ b/test/datadeps/scheduling.jl @@ -0,0 +1,577 @@ +import Dagger: DAGSpec, DatadepsArgSpec, dag_add_task!, + equivalent_structure, + datadeps_dag_equivalent, datadeps_argspec_equivalent, + datadeps_ainfo_equivalent, datadeps_schedule_cache, + LayeredScheduler, RoundRobinScheduler, DataDepsScheduler, + NoAliasing, UnknownAliasing, ContiguousAliasing, + StridedAliasing, TriangularAliasing, DiagonalAliasing, + ObjectAliasing, CombinedAliasing, AliasingWrapper, + DATADEPS_SCHEDULER, In, Out, InOut, Deps +import Dagger +using LinearAlgebra +using Test + +# ---------- equivalent_structure unit tests ---------- + +@testset "equivalent_structure" begin + @testset "NoAliasing/UnknownAliasing" begin + @test equivalent_structure(NoAliasing(), NoAliasing()) + @test equivalent_structure(UnknownAliasing(), UnknownAliasing()) + @test !equivalent_structure(NoAliasing(), UnknownAliasing()) + @test !equivalent_structure(UnknownAliasing(), NoAliasing()) + end + + @testset "ContiguousAliasing" begin + # Same shape, different allocation → equivalent + @test equivalent_structure(Dagger.aliasing(rand(10, 10)), + Dagger.aliasing(rand(10, 10))) + @test equivalent_structure(Dagger.aliasing(zeros(Float64, 32)), + Dagger.aliasing(rand(Float64, 32))) + # Different element counts → not equivalent + @test !equivalent_structure(Dagger.aliasing(rand(10, 10)), + Dagger.aliasing(rand(5, 5))) + # Different element types → not equivalent (S/T params differ) + @test !equivalent_structure(Dagger.aliasing(rand(Float64, 10)), + Dagger.aliasing(rand(Float32, 10))) + end + + @testset "StridedAliasing" begin + A1 = rand(10, 10) + A2 = rand(10, 10) # same parent shape, different allocation + # Same view shape & relative offset → equivalent + @test equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A2, 2:5, 3:6))) + # Same view shape, different relative offset → not equivalent + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A1, 3:6, 4:7))) + # Different parent shapes → strides differ → not equivalent + B = rand(20, 20) + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(B, 2:5, 3:6))) + # Different view dim lengths → not equivalent + @test !equivalent_structure(Dagger.aliasing(view(A1, 2:5, 3:6)), + Dagger.aliasing(view(A1, 2:6, 3:6))) + end + + @testset "TriangularAliasing" begin + A1 = rand(10, 10); A2 = rand(10, 10) + @test equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UpperTriangular(A2))) + @test equivalent_structure(Dagger.aliasing(LowerTriangular(A1)), + Dagger.aliasing(LowerTriangular(A2))) + # Upper vs Lower → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(LowerTriangular(A1))) + # Unit vs non-unit → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UnitUpperTriangular(A1))) + # Different size → not equivalent + @test !equivalent_structure(Dagger.aliasing(UpperTriangular(A1)), + Dagger.aliasing(UpperTriangular(rand(5, 5)))) + end + + @testset "DiagonalAliasing" begin + A1 = rand(10, 10); A2 = rand(10, 10) + @test equivalent_structure(Dagger.aliasing(A1, Diagonal), + Dagger.aliasing(A2, Diagonal)) + @test !equivalent_structure(Dagger.aliasing(A1, Diagonal), + Dagger.aliasing(rand(5, 5), Diagonal)) + end + + @testset "ObjectAliasing" begin + r1 = Ref(rand(10)) + r2 = Ref(rand(10)) # same Ref{Vector{Float64}} shape + a1 = first(Dagger.aliasing(r1).sub_ainfos)::ObjectAliasing + a2 = first(Dagger.aliasing(r2).sub_ainfos)::ObjectAliasing + @test equivalent_structure(a1, a2) + end + + @testset "CombinedAliasing (Ref)" begin + @test equivalent_structure(Dagger.aliasing(Ref(rand(8))), + Dagger.aliasing(Ref(rand(8)))) + @test !equivalent_structure(Dagger.aliasing(Ref(rand(8))), + Dagger.aliasing(Ref(rand(4)))) + end + + @testset "AliasingWrapper delegation" begin + A1 = rand(10, 10); A2 = rand(10, 10) + w1 = AliasingWrapper(Dagger.aliasing(A1)) + w2 = AliasingWrapper(Dagger.aliasing(A2)) + @test equivalent_structure(w1, w2) + # Wrapper vs bare also works + @test equivalent_structure(w1, Dagger.aliasing(A2)) + @test equivalent_structure(Dagger.aliasing(A1), w2) + end + + @testset "Cross-type returns false" begin + @test !equivalent_structure(Dagger.aliasing(rand(10)), + Dagger.aliasing(view(rand(10, 10), 1:5, 1:5))) + @test !equivalent_structure(Dagger.aliasing(rand(10, 10)), + NoAliasing()) + @test !equivalent_structure(NoAliasing(), + Dagger.aliasing(rand(10, 10))) + end +end + +# ---------- Helpers for DAG capture via the schedule cache ---------- + +""" +Run `f()` inside `spawn_datadeps` with `scheduler`, returning the resulting +schedule cache snapshot (a Vector). The cache is cleared before the run so +each call starts from a known state. +""" +function run_with_fresh_cache(f, scheduler::DataDepsScheduler = LayeredScheduler()) + cache = datadeps_schedule_cache(scheduler) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => scheduler) do + Dagger.spawn_datadeps(f) + end + return cache +end + +""" +Run a sequence of `n` invocations of `f()` under `scheduler` and return the +cache. Useful for asserting cache size after repeated runs. +""" +function run_n_times(f, n::Int; scheduler::DataDepsScheduler = LayeredScheduler()) + cache = datadeps_schedule_cache(scheduler) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => scheduler) do + for _ in 1:n + Dagger.spawn_datadeps(f) + end + end + return cache +end + +# Simple test kernels +add!(X, Y) = (X .+= Y; X) +scale!(X, a) = (X .*= a; X) + +# ---------- DAGSpec construction ---------- + +@testset "DAGSpec construction" begin + @testset "Vertices and function types" begin + A = rand(64) + B = rand(64) + cache = run_with_fresh_cache() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + @test length(cache) == 1 + spec = first(cache).first + @test length(spec) == 2 + @test spec.id_to_functype[1] === typeof(add!) + @test spec.id_to_functype[2] === typeof(scale!) + # `tspec.fargs` includes the function as fargs[1], so argtypes also + # includes the function argspec → 3 entries per vertex. + @test length(spec.id_to_argtypes[1]) == 3 + @test length(spec.id_to_argtypes[2]) == 3 + end + + @testset "Argspec positions, types, and dep_mods" begin + A = rand(32, 32) + cache = run_with_fresh_cache() do + Dagger.@spawn add!(InOut(A), In(A)) + end + spec = first(cache).first + argspecs = spec.id_to_argtypes[1] + # 3 argspecs: function (pos=0), A (InOut, pos=1), A (In, pos=2) + @test length(argspecs) == 3 + @test all(a -> a.pos isa Int, argspecs) + # The two data argspecs (positions 1 and 2) refer to a Matrix{Float64} + data_argspecs = filter(a -> a.pos != 0, argspecs) + @test length(data_argspecs) == 2 + @test all(a -> a.value_type === Matrix{Float64}, data_argspecs) + @test all(a -> a.dep_mod === identity, data_argspecs) + @test all(a -> a.ainfo isa ContiguousAliasing, data_argspecs) + # Function argspec is at position 0 with value_type === typeof(add!) + f_argspec = only(filter(a -> a.pos == 0, argspecs)) + @test f_argspec.value_type === typeof(add!) + end + + @testset "Deps with multiple dep_mods at one position" begin + X = Ref(rand(64)) + do_nothing(R) = nothing + cache = run_with_fresh_cache() do + Dagger.@spawn do_nothing(Deps(X, InOut(:x), In(:x))) + end + spec = first(cache).first + argspecs = spec.id_to_argtypes[1] + # Both Deps entries for :x should be present, even though they share pos. + x_deps = filter(a -> a.dep_mod === :x, argspecs) + @test length(x_deps) == 2 + @test all(a -> a.pos == x_deps[1].pos, x_deps) + end + + @testset "Per-task scope recorded" begin + A = rand(64) + my_scope = Dagger.scope(worker=1) + cache = run_with_fresh_cache() do + Dagger.@spawn scope=my_scope add!(InOut(A), In(A)) + end + spec = first(cache).first + @test spec.id_to_scope[1] == my_scope + end +end + +# ---------- DAGSpec equivalence: cache reuse ---------- + +@testset "DAGSpec equivalence (end-to-end cache reuse)" begin + @testset "Identical algorithm with fresh allocations → cache hit" begin + cache = run_n_times(3) do + A = rand(128) + B = rand(128) + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn add!(InOut(A), In(B)) + end + @test length(cache) == 1 + end + + @testset "Different array sizes → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for sz in (64, 128, 256) + A = rand(sz); B = rand(sz) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end + + @testset "Different element types → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for T in (Float64, Float32, Int) + A = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) + B = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end + + @testset "Different task counts → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for ntasks in (1, 2, 4) + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + for _ in 1:ntasks + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + end + @test length(cache) == 3 + end + + @testset "Different per-task scope → cache miss" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + scopes = [Dagger.DefaultScope(), + Dagger.ExactScope(Dagger.ThreadProc(1, 1))] + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for scp in scopes + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn scope=scp add!(InOut(A), In(B)) + end + end + end + @test length(cache) == length(scopes) + end + + @testset "Different functions → cache miss" begin + cache = run_with_fresh_cache() do + A = rand(64); B = rand(64) + Dagger.@spawn add!(InOut(A), In(B)) + end + @test length(cache) == 1 + # Repeat with the other function + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + A = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + @test length(cache) == 2 + end +end + +# ---------- Multi-scheduler cache partitioning ---------- + +@testset "Cache partitioning between schedulers" begin + ls_cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(ls_cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + @test length(ls_cache) == 1 + + # RoundRobinScheduler doesn't AOT-schedule, so it shouldn't add to its + # cache, and it must NOT see LayeredScheduler's entry. + rr_cache = datadeps_schedule_cache(RoundRobinScheduler()) + @test rr_cache !== ls_cache + Base.ScopedValues.with(DATADEPS_SCHEDULER => RoundRobinScheduler()) do + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + @test isempty(rr_cache) + @test length(ls_cache) == 1 +end + +# ---------- Benchmark algorithms ---------- + +# A small mock-Cholesky-like algorithm that doesn't depend on DArray, to keep +# the test fast and self-contained. +function mock_cholesky!(M::Matrix{<:Matrix}) + mt = size(M, 1) + for k in 1:mt + Dagger.@spawn LinearAlgebra.LAPACK.potrf!('L', InOut(M[k, k])) + for m in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.trsm!('R', 'L', 'T', 'N', + 1.0, In(M[k, k]), + InOut(M[m, k])) + end + for n in (k+1):mt + Dagger.@spawn LinearAlgebra.BLAS.syrk!('L', 'N', -1.0, + In(M[n, k]), 1.0, + InOut(M[n, n])) + for m in (n+1):mt + Dagger.@spawn LinearAlgebra.BLAS.gemm!('N', 'T', -1.0, + In(M[m, k]), + In(M[n, k]), 1.0, + InOut(M[m, n])) + end + end + end +end + +make_spd_blocks(sz::Int, nb::Int) = begin + M_dense = rand(sz, sz); M_dense = M_dense * M_dense' + M_dense[diagind(M_dense)] .+= sz + return [M_dense[i:(i+nb-1), j:(j+nb-1)] for i in 1:nb:sz, j in 1:nb:sz] +end + +@testset "Benchmark: mock Cholesky" begin + @testset "Repeated runs, fresh allocations → 1 cache entry" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for _ in 1:3 + M = make_spd_blocks(64, 16) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 1 + spec = first(cache).first + # 4x4 block grid → potrf + trsms + syrks + gemms. + # Count expected: sum_{k=1}^{4} (1 + (4-k) + (4-k) + (4-k)(4-k-1)/2 * 2) + # Or just sanity check >0. + @test length(spec) > 0 + end + + @testset "Different block-grid sizes → distinct cache entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + # 2x2 grid (sz=32, nb=16) vs 4x4 grid (sz=64, nb=16) → different task counts + for (sz, nb) in ((32, 16), (64, 16)) + M = make_spd_blocks(sz, nb) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 2 + end + + @testset "Same block-grid shape, different block size → distinct entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + # Both 2x2 grids; block size 16 vs 32 → same task count, different ContiguousAliasing length + for (sz, nb) in ((32, 16), (64, 32)) + M = make_spd_blocks(sz, nb) + Dagger.spawn_datadeps() do + mock_cholesky!(M) + end + end + end + @test length(cache) == 2 + end +end + +@testset "Benchmark: tree reduce" begin + function tree_reduce!(As) + to_reduce = Vector[] + push!(to_reduce, As) + while !isempty(to_reduce) + xs = pop!(to_reduce) + n = length(xs) + if n == 2 + Dagger.@spawn Base.mapreducedim!(identity, +, InOut(xs[1]), In(xs[2])) + elseif n > 2 + push!(to_reduce, [xs[1], xs[div(n, 2)+1]]) + push!(to_reduce, xs[1:div(n, 2)]) + push!(to_reduce, xs[div(n, 2)+1:end]) + end + end + end + + @testset "Repeated runs → 1 cache entry" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for _ in 1:3 + As = [rand(64) for _ in 1:8] + Dagger.spawn_datadeps() do + tree_reduce!(As) + end + end + end + @test length(cache) == 1 + end + + @testset "Different reduction widths → distinct entries" begin + cache = datadeps_schedule_cache(LayeredScheduler()) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + for n in (4, 8, 16) + As = [rand(64) for _ in 1:n] + Dagger.spawn_datadeps() do + tree_reduce!(As) + end + end + end + @test length(cache) == 3 + end +end + +# ---------- Direct datadeps_dag_equivalent semantics ---------- + +@testset "datadeps_dag_equivalent (direct)" begin + LS = LayeredScheduler() + + # Build two DAGSpecs from equivalent-but-freshly-allocated workloads. + cache = datadeps_schedule_cache(LS) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LS) do + for _ in 1:2 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + end + # Cache hit on second run → still 1 entry + @test length(cache) == 1 + spec_a = first(cache).first + + # Now build a structurally-different spec. + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => LS) do + A = rand(128); B = rand(128) # different size + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + Dagger.@spawn scale!(InOut(A), 2.0) + end + end + spec_b = first(cache).first + + @test datadeps_dag_equivalent(LS, spec_a, spec_a) + @test datadeps_dag_equivalent(LS, spec_b, spec_b) + @test !datadeps_dag_equivalent(LS, spec_a, spec_b) +end + +# ---------- Scheduler overrides ---------- + +# A scheduler that opts out of all caching by always returning false. +struct NoCacheScheduler <: DataDepsScheduler end +Dagger.datadeps_dag_equivalent(::NoCacheScheduler, ::DAGSpec, ::DAGSpec) = false +function Dagger.datadeps_schedule_dag_aot!(::NoCacheScheduler, schedule, dag_spec, all_procs, all_scope) + for idx in 1:Dagger.nv(dag_spec.g) + task = dag_spec.id_to_task[idx] + schedule[task] = first(all_procs) + end +end +Dagger.datadeps_schedule_task_jit!(::NoCacheScheduler, all_procs, all_scope, task_scope, spec, task) = + first(all_procs) + +@testset "Scheduler-overridable equivalence" begin + @testset "Always-false override never hits cache" begin + nc = NoCacheScheduler() + cache = datadeps_schedule_cache(nc) + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => nc) do + for _ in 1:3 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + end +end + +# A scheduler that uses a strict pointer-identical comparison (the original +# intent of `equivalent_structure`). +struct PtrStrictScheduler <: DataDepsScheduler end +Dagger.datadeps_ainfo_equivalent(::PtrStrictScheduler, + a1::Dagger.AbstractAliasing, + a2::Dagger.AbstractAliasing) = + hash(a1) == hash(a2) +function Dagger.datadeps_schedule_dag_aot!(::PtrStrictScheduler, schedule, dag_spec, all_procs, all_scope) + for idx in 1:Dagger.nv(dag_spec.g) + task = dag_spec.id_to_task[idx] + schedule[task] = first(all_procs) + end +end +Dagger.datadeps_schedule_task_jit!(::PtrStrictScheduler, all_procs, all_scope, task_scope, spec, task) = + first(all_procs) + +@testset "PtrStrictScheduler: structural matches miss, pointer-identical hits" begin + ps = PtrStrictScheduler() + cache = datadeps_schedule_cache(ps) + + # Fresh allocations each run → pointer hash differs → all miss. + empty!(cache) + Base.ScopedValues.with(DATADEPS_SCHEDULER => ps) do + for _ in 1:3 + A = rand(64); B = rand(64) + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 3 + + # Same allocation reused each run → pointer hash matches → cache hits. + empty!(cache) + A = rand(64); B = rand(64) + Base.ScopedValues.with(DATADEPS_SCHEDULER => ps) do + for _ in 1:3 + Dagger.spawn_datadeps() do + Dagger.@spawn add!(InOut(A), In(B)) + end + end + end + @test length(cache) == 1 +end diff --git a/test/runtests.jl b/test/runtests.jl index d327ee50a..4aacf1c7b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,7 @@ tests = [ ("Task Queues", "task-queues.jl"), ("Task Affinity", "task-affinity.jl"), ("Datadeps", "datadeps.jl"), + ("Datadeps - Scheduling", "datadeps/scheduling.jl"), ("Streaming", "streaming.jl"), ("Domain Utilities", "domain.jl"), ("Array - Allocation", "array/allocation.jl"), From 85d50bc7513a53f89b1c97c51aba66d2604818b5 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 24 Mar 2026 15:51:48 -0700 Subject: [PATCH 28/43] datadeps: Add hierarchical scheduling --- src/Dagger.jl | 1 + src/datadeps/aliasing.jl | 3 + src/datadeps/hierarchical.jl | 1086 ++++++++++++++++++++++++++++++++++ src/datadeps/queue.jl | 175 ++++-- src/datadeps/scheduling.jl | 10 +- test/datadeps.jl | 2 +- test/datadeps/scheduling.jl | 45 +- 7 files changed, 1254 insertions(+), 68 deletions(-) create mode 100644 src/datadeps/hierarchical.jl diff --git a/src/Dagger.jl b/src/Dagger.jl index 71b15bbf8..ccc3925bd 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -108,6 +108,7 @@ include("datadeps/chunkview.jl") include("datadeps/remainders.jl") include("datadeps/scheduling.jl") include("datadeps/queue.jl") +include("datadeps/hierarchical.jl") # Stencils include("utils/haloarray.jl") diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 2f90a4d48..00f4701da 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1,3 +1,6 @@ +import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv, + weakly_connected_components, topological_sort_by_dfs, vertices + export In, Out, InOut, Deps, spawn_datadeps #= diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl new file mode 100644 index 000000000..b4e2f0734 --- /dev/null +++ b/src/datadeps/hierarchical.jl @@ -0,0 +1,1086 @@ +# Hierarchical scheduling for datadeps +# Spreads scheduling work across multiple threads/workers via a 4-phase pipeline: +# Phase 1: Parallel aliasing info construction +# Phase 2: Sequential DAG construction from aliasing overlaps +# Phase 3: Data-affinity DAG partitioning +# Phase 4: Parallel local scheduling per partition -- each partition runs on its +# own task, prepares/submits its tasks concurrently with the others, +# and computes its *own* (partition-local) AOT schedule over just its +# processors. This maximizes scheduling parallelism and scalability; +# no global synchronization or global AOT pass is imposed. +# +# N.B. Concurrent task submission from many partition tasks used to intermittently +# hang in the core eager scheduler. The root cause was `Distributed.Future`, +# which is unsafe under concurrent same-process access; Dagger now backs task +# futures with `MemPool.DFuture` (a thread-safe, `DEvent`-based future), so +# parallel Phase 4 is safe. The actual task submission (`enqueue!`) is still +# serialized via `LockedEnqueueQueue`, but the (expensive) per-task +# `distribute_task!` preparation runs concurrently across partitions. + +struct HierarchicalTaskInfo + arg_w::ArgumentWrapper + readdep::Bool + writedep::Bool +end + +struct HierarchicalTaskMeta + pair::DTaskPair + arg_chunk::Union{Chunk, Nothing} + may_alias::Bool + inplace_move::Bool + deps::Vector{HierarchicalTaskInfo} + # Indices of same-region producer tasks whose results are passed as + # arguments (e.g. `In(t1)`). These cannot be `fetch`'d during the + # pre-scan because they are not launched yet; they become hard DAG edges. + value_deps::Vector{Int} +end + +### Cross-partition chunk ownership ### +# +# Each partition schedules with its own `DataDepsState`, so `arg_owner` / +# `arg_history` / physical slots are per-partition. When a backing chunk is +# written by tasks in >=2 partitions that live in *different* memory spaces, +# each partition would otherwise copy that chunk from its origin, write its own +# sub-range, and record itself as owner -- the physical copies then diverge and +# the final copy-back keeps only one, silently losing the others' writes. +# +# The registry below carries the single authoritative version ("ownership") of +# each such shared chunk across partition boundaries. It is deadlock-free by +# construction: a single lock (never per-argument locks acquired in different +# orders) plus the global DAG ordering -- a consumer partition always `wait`s on +# its cross-partition predecessor's `task_submitted` event before scheduling, so +# the producer's ownership commit is always visible before the consumer reads it. +# The lock only provides memory-safety for concurrent access to *different* +# chunks; per-chunk correctness comes from the DAG order. +"Authoritative, cross-partition ownership state for a single shared backing chunk." +mutable struct ChunkOwnership + owner_space::MemorySpace # space holding the current version + owner_slot::Any # physical slot chunk in `owner_space` + owner_task::Union{DTask,Nothing} # producer of the current version (nothing => origin data) + owner_state::Union{DataDepsState,Nothing} # owning partition's state (for copy-back) + const origin_space::MemorySpace # the chunk's home space +end + +struct SharedChunkRegistry + entries::IdDict{Any,ChunkOwnership} # backing chunk (identity) => ownership + lock::ReentrantLock +end +SharedChunkRegistry() = SharedChunkRegistry(IdDict{Any,ChunkOwnership}(), ReentrantLock()) + +is_shared_chunk(::Nothing, @nospecialize(chunk)) = false +is_shared_chunk(reg::SharedChunkRegistry, @nospecialize(chunk)) = haskey(reg.entries, chunk) + +""" + build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) -> SharedChunkRegistry or nothing + +Detects backing chunks accessed by partitions spanning >=2 distinct memory +spaces (the only case that can split-brain) and returns a registry seeded with +origin ownership for each. Returns `nothing` when all partitions share a single +space (e.g. single-worker), so the fast path is entirely unchanged there. +""" +function build_shared_chunk_registry(task_metas::Vector{HierarchicalTaskMeta}, + vertex_to_partition::Vector{Int}, + partition_space::Vector{<:MemorySpace}) + length(unique(partition_space)) <= 1 && return nothing + + chunk_spaces = IdDict{Any,Set{MemorySpace}}() + chunk_origin = IdDict{Any,MemorySpace}() + for v in eachindex(task_metas) + pspace = partition_space[vertex_to_partition[v]] + for dep in task_metas[v].deps + chunk = dep.arg_w.arg + push!(get!(()->Set{MemorySpace}(), chunk_spaces, chunk), pspace) + haskey(chunk_origin, chunk) || (chunk_origin[chunk] = memory_space(chunk)) + end + end + + reg = SharedChunkRegistry() + for (chunk, spaces) in chunk_spaces + length(spaces) >= 2 || continue + origin = chunk_origin[chunk] + reg.entries[chunk] = ChunkOwnership(origin, chunk, nothing, nothing, origin) + end + isempty(reg.entries) && return nothing + return reg +end + +""" + _sync_incoming_ownership!(state, registry, our_space, task_arg_ws, write_num) + +Before a task's copy-to phase, for each shared backing-chunk argument whose +globally-current owner (per `registry`) lives in a space other than `our_space`, +seed this partition's `state` so the existing copy-to machinery pulls a fresh +whole-chunk copy from the true owner (with a syncdep on the producing task). A +no-op for private chunks or when we already hold the authoritative version. +""" +function _sync_incoming_ownership!(state::DataDepsState, registry::SharedChunkRegistry, + our_space::MemorySpace, task_arg_ws, write_num::Int) + map_or_ntuple(task_arg_ws) do idx + arg_ws = task_arg_ws[idx] + (arg_ws.may_alias && arg_ws.inplace_move) || return + chunk = arg_ws.arg + entry = get(registry.entries, chunk, nothing) + entry === nothing && return + + owner_space, owner_slot, owner_task = @lock registry.lock begin + (entry.owner_space, entry.owner_slot, entry.owner_task) + end + owner_space == our_space && return # we already hold the authoritative copy + + # Register the owner's physical slot so slot/aliasing lookups in this + # partition reuse it (rather than generating a fresh, stale copy). + dest_args = get!(IdDict{Any,Any}, state.remote_args, owner_space) + if !haskey(dest_args, chunk) + dest_args[chunk] = owner_slot + state.remote_arg_to_original[owner_slot] = chunk + end + + map_or_ntuple(arg_ws.deps) do dep_idx + dep = arg_ws.deps[dep_idx] + arg_w = dep.arg_w + # Point ownership at the owner space and clear any locally-merged + # history, so `compute_remainder_for_arg!` takes the `FullCopy` + # (whole-chunk) path from the owner rather than a partial remainder. + state.arg_owner[arg_w] = owner_space + haskey(state.arg_history, arg_w) && empty!(state.arg_history[arg_w]) + src_ainfo = aliasing!(state, owner_space, arg_w) + if owner_task !== nothing + # Make the ensuing copy-to (via `get_read_deps!`) wait on the + # producer, so we never copy the owner slot before it is written. + state.ainfos_owner[src_ainfo] = owner_task => (write_num - 1) + end + return + end + return + end + return +end + +""" + _commit_ownership!(state, registry, our_space, task, task_arg_ws) + +After a task is recorded as a writer, publish it as the new authoritative owner +of each shared backing chunk it writes, so subsequent cross-partition consumers +pull from here. Ordering (and thus visibility) is guaranteed by the DAG's +`task_submitted` handshake; the lock only guards concurrent commits for other +chunks. +""" +function _commit_ownership!(state::DataDepsState, registry::SharedChunkRegistry, + our_space::MemorySpace, task::DTask, task_arg_ws) + map_or_ntuple(task_arg_ws) do idx + arg_ws = task_arg_ws[idx] + (arg_ws.may_alias && arg_ws.inplace_move) || return + chunk = arg_ws.arg + entry = get(registry.entries, chunk, nothing) + entry === nothing && return + + wrote = false + map_or_ntuple(arg_ws.deps) do dep_idx + arg_ws.deps[dep_idx].writedep && (wrote = true) + return + end + wrote || return + + dest_args = get(state.remote_args, our_space, nothing) + slot = dest_args === nothing ? nothing : get(dest_args, chunk, nothing) + @lock registry.lock begin + entry.owner_space = our_space + slot !== nothing && (entry.owner_slot = slot) + entry.owner_task = task + entry.owner_state = state + end + return + end + return +end + +# Below this many tasks, the fixed costs of spawning threads and merging +# per-thread results outweigh the benefit of parallelizing the pre-scan. +const COLLECT_ALIASED_ARGS_MIN_CHUNK = 256 + +# Thread-safe "get or compute" against a shared `IdDict` cache. The +# expensive computation (`f`) is performed outside of the lock so that +# multiple threads can make progress on distinct keys concurrently; if two +# threads race on the same key, the loser's result is simply discarded (the +# corresponding `Chunk`/Bool is cheap to let the GC reclaim) so that every +# thread agrees on a single canonical value (e.g. `Chunk`) per raw argument. +@inline function _cached_get!(f, cache::IdDict{Any,V}, cache_lock::Union{ReentrantLock,Nothing}, key) where V + if cache_lock === nothing + return get!(f, cache, key) + end + @lock cache_lock begin + haskey(cache, key) && return cache[key]::V + end + result = f()::V + @lock cache_lock begin + return get!(cache, key, result)::V + end +end + +""" + collect_aliased_args(seen_tasks) -> (task_metas, unique_arg_ws) + +Pre-scans all tasks to collect per-task dependency metadata and the set of +unique `ArgumentWrapper`s that need aliasing analysis. This mirrors the logic +in `populate_task_info!` but only inspects arguments without modifying any +scheduling state. + +For large batches of tasks (the common case for e.g. panel-factorization +algorithms which submit many small tasks per `spawn_datadeps` region), the +pre-scan itself (not just the aliasing computation in +`build_aliasing_parallel`) can dominate scheduling time, since it touches +every argument of every task. This is parallelized across threads: each +thread scans a contiguous range of `seen_tasks` into its own disjoint slice +of `task_metas` and its own local `unique_arg_ws` map (merged at the end), +while sharing (lock-protected) caches for `Chunk`-wrapping and +`supports_inplace_move`, ensuring a single canonical `Chunk` identity per +raw argument regardless of which thread first observes it. +""" +function collect_aliased_args(seen_tasks::Vector{DTaskPair}) + n = length(seen_tasks) + task_metas = Vector{HierarchicalTaskMeta}(undef, n) + n == 0 && return task_metas, Dict{ArgumentWrapper,ArgumentWrapper}() + + # Map in-region tasks to vertex indices so we can record value deps + # without fetching unlaunched DTasks during the pre-scan. + task_to_idx = IdDict{DTask,Int}() + for (i, pair) in enumerate(seen_tasks) + task_to_idx[pair.task] = i + end + + supports_cache = IdDict{Any,Bool}() + raw_arg_cache = IdDict{Any,Chunk}() + + nchunks = Threads.nthreads() <= 1 ? 1 : min(Threads.nthreads(), cld(n, COLLECT_ALIASED_ARGS_MIN_CHUNK)) + + if nchunks <= 1 + unique_arg_ws = Dict{ArgumentWrapper, ArgumentWrapper}() + _collect_aliased_args_range!(task_metas, unique_arg_ws, seen_tasks, 1:n, + supports_cache, raw_arg_cache, nothing, task_to_idx) + return task_metas, unique_arg_ws + end + + cache_lock = ReentrantLock() + chunk_size = cld(n, nchunks) + starts = collect(1:chunk_size:n) + per_chunk_arg_ws = Vector{Dict{ArgumentWrapper,ArgumentWrapper}}(undef, length(starts)) + + @sync for (ci, start) in enumerate(starts) + range = start:min(start+chunk_size-1, n) + Threads.@spawn begin + local_arg_ws = Dict{ArgumentWrapper,ArgumentWrapper}() + _collect_aliased_args_range!(task_metas, local_arg_ws, seen_tasks, range, + supports_cache, raw_arg_cache, cache_lock, task_to_idx) + per_chunk_arg_ws[ci] = local_arg_ws + end + end + + unique_arg_ws = per_chunk_arg_ws[1] + for ci in 2:length(per_chunk_arg_ws) + merge!(unique_arg_ws, per_chunk_arg_ws[ci]) + end + + return task_metas, unique_arg_ws +end + +function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, + unique_arg_ws::Dict{ArgumentWrapper,ArgumentWrapper}, + seen_tasks::Vector{DTaskPair}, + range::UnitRange{Int}, + supports_cache::IdDict{Any,Bool}, + raw_arg_cache::IdDict{Any,Chunk}, + cache_lock::Union{ReentrantLock,Nothing}, + task_to_idx::IdDict{DTask,Int}) + for task_idx in range + pair = seen_tasks[task_idx] + spec = pair.spec + task = pair.task + fargs = spec.fargs + + all_deps = HierarchicalTaskInfo[] + value_deps = Int[] + first_chunk = nothing + task_may_alias = false + task_inplace = false + + for arg_idx in (is_typed(spec) ? (1:length(fargs)) : eachindex(fargs)) + _arg = fargs[arg_idx] + _arg_with_deps = value(_arg) + + arg_pre_unwrap, deps = unwrap_inout(_arg_with_deps) + + # Same-region DTask arguments are not launched yet, so we cannot + # `fetch` them here (that is what the sequential `distribute_task!` + # path does *after* launching the producer). Record a value + # dependency instead; aliasing of the result is handled later in + # `distribute_task!` once the producer has been submitted. + if arg_pre_unwrap isa DTask && !istaskstarted(arg_pre_unwrap) + pred_idx = get(task_to_idx, arg_pre_unwrap, 0) + if pred_idx != 0 && pred_idx != task_idx + push!(value_deps, pred_idx) + end + continue + end + + arg = arg_pre_unwrap isa DTask ? fetch(arg_pre_unwrap; raw=true) : arg_pre_unwrap + + may_alias = type_may_alias(typeof(arg)) + inplace_move = may_alias && _cached_get!(supports_cache, cache_lock, arg) do + supports_inplace_move(arg) + end + + if !may_alias || !inplace_move + continue + end + + arg_chunk = _cached_get!(raw_arg_cache, cache_lock, arg) do + arg isa Chunk ? arg : tochunk(arg) + end + + if first_chunk === nothing + first_chunk = arg_chunk + task_may_alias = true + task_inplace = true + end + + for (dep_mod, readdep, writedep) in deps + arg_w = ArgumentWrapper(arg_chunk, dep_mod) + unique_arg_ws[arg_w] = arg_w + push!(all_deps, HierarchicalTaskInfo(arg_w, readdep, writedep)) + end + end + + task_metas[task_idx] = HierarchicalTaskMeta( + pair, first_chunk, task_may_alias, task_inplace, all_deps, value_deps + ) + end +end + +""" + build_aliasing_parallel(unique_arg_ws) -> (lookup, ainfos_overlaps, arg_to_ainfo) + +Phase 1: Computes `AliasingWrapper` for every unique `ArgumentWrapper` in +parallel. On each worker, threads are used to compute aliasing info for local +data. Results are gathered and reduced into a single `AliasingLookup` with +overlap information. +""" +function build_aliasing_parallel(unique_arg_ws::Dict{ArgumentWrapper, ArgumentWrapper}) + arg_ws_vec = collect(values(unique_arg_ws)) + + by_worker = Dict{Int, Vector{ArgumentWrapper}}() + for arg_w in arg_ws_vec + wid = root_worker_id(memory_space(arg_w.arg)) + worker_args = get!(Vector{ArgumentWrapper}, by_worker, wid) + push!(worker_args, arg_w) + end + + arg_to_ainfo = Dict{ArgumentWrapper, AliasingWrapper}() + + if length(by_worker) == 1 + # Common single-worker case: avoid the `@sync`/`Threads.@spawn`/lock + # overhead entirely, since there's nothing to run concurrently with. + # `_compute_aliasing_batch` still uses threads internally when there + # are enough args to make it worthwhile. + wid, worker_args = only(by_worker) + results = wid == myid() ? _compute_aliasing_batch(worker_args) : + remotecall_fetch(_compute_aliasing_batch, wid, worker_args) + # Key by the *local* `arg_w`, not the pair's: for a remote worker the + # returned `ArgumentWrapper` is a deserialized copy that need not be + # identity/hash-equal to the entry in `arg_ws_vec` we later look up + # (which would raise a `KeyError`). `_compute_aliasing_batch` preserves + # input order, so pair by index. + for i in eachindex(worker_args) + arg_to_ainfo[worker_args[i]] = results[i].second + end + else + all_results_lock = ReentrantLock() + @sync for (wid, worker_args) in by_worker + Threads.@spawn begin + results = if wid == myid() + _compute_aliasing_batch(worker_args) + else + remotecall_fetch(_compute_aliasing_batch, wid, worker_args) + end + # Key by the *local* `arg_w` (see single-worker note above): a + # remote worker returns deserialized `ArgumentWrapper` copies + # that may not compare equal to our `arg_ws_vec` lookup keys. + @lock all_results_lock begin + for i in eachindex(worker_args) + arg_to_ainfo[worker_args[i]] = results[i].second + end + end + end + end + end + + lookup = AliasingLookup() + ainfos_overlaps = Dict{AliasingWrapper, Set{AliasingWrapper}}() + + for arg_w in arg_ws_vec + ainfo = arg_to_ainfo[arg_w] + if haskey(ainfos_overlaps, ainfo) + continue + end + + ainfo_idx = push!(lookup, ainfo) + overlaps = Set{AliasingWrapper}() + push!(overlaps, ainfo) + for other_ainfo in intersect(lookup, ainfo; ainfo_idx) + ainfo == other_ainfo && continue + push!(overlaps, other_ainfo) + push!(ainfos_overlaps[other_ainfo], ainfo) + end + ainfos_overlaps[ainfo] = overlaps + end + + return lookup, ainfos_overlaps, arg_to_ainfo +end + +# Below this many args, the fixed cost of forking/joining `Threads.@threads` +# outweighs the benefit of parallelizing the (typically cheap) `aliasing()` calls. +const COMPUTE_ALIASING_BATCH_MIN_PARALLEL = 8 + +function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) + n = length(arg_ws) + results = Vector{Pair{ArgumentWrapper, AliasingWrapper}}(undef, n) + if n >= COMPUTE_ALIASING_BATCH_MIN_PARALLEL && Threads.nthreads() > 1 + Threads.@threads for i in 1:n + arg_w = arg_ws[i] + ainfo = AliasingWrapper(aliasing(arg_w.arg, arg_w.dep_mod)) + results[i] = arg_w => ainfo + end + else + for i in 1:n + arg_w = arg_ws[i] + ainfo = AliasingWrapper(aliasing(arg_w.arg, arg_w.dep_mod)) + results[i] = arg_w => ainfo + end + end + return results +end + +""" + build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + -> SimpleDiGraph + +Phase 2: Walks tasks in submission order and builds a `SimpleDiGraph` encoding +data dependencies based on the pre-computed aliasing overlaps. Uses the same +WAW / RAW / WAR rules as `get_write_deps!` / `get_read_deps!`. +""" +function build_dependency_dag(task_metas::Vector{HierarchicalTaskMeta}, + arg_to_ainfo::Dict{ArgumentWrapper, AliasingWrapper}, + ainfos_overlaps::Dict{AliasingWrapper, Set{AliasingWrapper}}) + n = length(task_metas) + dag = SimpleDiGraph(n) + + ainfos_owner = Dict{AliasingWrapper, Union{Nothing, Pair{Int,Int}}}() + ainfos_readers = Dict{AliasingWrapper, Vector{Pair{Int,Int}}}() + + write_num = 1 + for v in 1:n + meta = task_metas[v] + + # Hard edges from same-region DTask value arguments (producer must + # be launched before we can fetch its result in distribute_task!). + for pred_v in meta.value_deps + if pred_v != v + add_edge!(dag, pred_v, v) + end + end + + # Add dependency edges + for dep in meta.deps + ainfo = get(arg_to_ainfo, dep.arg_w, nothing) + ainfo === nothing && continue + + if !haskey(ainfos_owner, ainfo) + ainfos_owner[ainfo] = nothing + ainfos_readers[ainfo] = Pair{Int,Int}[] + end + + overlaps = get(ainfos_overlaps, ainfo, Set{AliasingWrapper}()) + + if dep.writedep + for other_ainfo in overlaps + owner = get(ainfos_owner, other_ainfo, nothing) + if owner !== nothing + pred_v, pred_wn = owner + if pred_wn != write_num && pred_v != v + add_edge!(dag, pred_v, v) + end + end + for (reader_v, reader_wn) in get(ainfos_readers, other_ainfo, Pair{Int,Int}[]) + if reader_wn != write_num && reader_v != v + add_edge!(dag, reader_v, v) + end + end + end + else + for other_ainfo in overlaps + owner = get(ainfos_owner, other_ainfo, nothing) + if owner !== nothing + pred_v, pred_wn = owner + if pred_wn != write_num && pred_v != v + add_edge!(dag, pred_v, v) + end + end + end + end + end + + # Update ownership tracking + for dep in meta.deps + ainfo = get(arg_to_ainfo, dep.arg_w, nothing) + ainfo === nothing && continue + + if !haskey(ainfos_owner, ainfo) + ainfos_owner[ainfo] = nothing + ainfos_readers[ainfo] = Pair{Int,Int}[] + end + + if dep.writedep + ainfos_owner[ainfo] = v => write_num + empty!(ainfos_readers[ainfo]) + else + push!(ainfos_readers[ainfo], v => write_num) + end + end + + write_num += 1 + end + + return dag +end + +""" + partition_dag(dag, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs) + +Phase 3: Assigns each task vertex to a partition using data-affinity. For +multi-worker setups, tasks are assigned to the worker owning the most argument +data. For single-worker multi-threaded setups, tasks are balanced across +available processors in topological order. +""" +function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMeta}, + all_procs::Vector{<:Processor}) + n = length(task_metas) + workers = unique(root_worker_id.(only.(memory_spaces.(all_procs)))) + n_workers = length(workers) + + procs_by_worker = Dict{Int, Vector{Processor}}() + for proc in all_procs + wid = root_worker_id(only(memory_spaces(proc))) + push!(get!(Vector{Processor}, procs_by_worker, wid), proc) + end + + multi_worker = n_workers > 1 + if multi_worker + n_partitions = n_workers + partition_worker = workers + else + n_partitions = min(length(all_procs), n) + partition_worker = fill(first(workers), n_partitions) + end + + vertex_to_partition = Vector{Int}(undef, n) + + if multi_worker + worker_to_partition = Dict(w => i for (i, w) in enumerate(workers)) + default_scope = DefaultScope() + for v in 1:n + meta = task_metas[v] + task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) + + # Workers whose processors are eligible under this task's scope. + # A non-default scope that spans multiple workers must still spread + # work across those workers (via affinity / round-robin) -- picking + # only the first match pins everything to worker 1 and breaks + # multi-worker execution. + if task_scope == default_scope + matching = collect(1:n_partitions) + else + matching = Int[] + for (pid, wid) in enumerate(workers) + wprocs = procs_by_worker[wid] + if any(proc -> proc_in_scope(proc, task_scope), wprocs) + push!(matching, pid) + end + end + if isempty(matching) + matching = [1] + end + end + + if length(matching) == 1 + vertex_to_partition[v] = only(matching) + continue + end + + affinity = zeros(Int, n_workers) + for dep in meta.deps + arg_space = memory_space(dep.arg_w.arg) + arg_wid = root_worker_id(arg_space) + idx = get(worker_to_partition, arg_wid, 0) + if idx > 0 && idx in matching + affinity[idx] += 1 + end + end + best_pid = matching[1] + best_aff = -1 + for pid in matching + if affinity[pid] > best_aff + best_aff = affinity[pid] + best_pid = pid + end + end + if best_aff <= 0 + vertex_to_partition[v] = matching[mod1(v, length(matching))] + else + vertex_to_partition[v] = best_pid + end + end + else + topo = try + topological_sort_by_dfs(dag) + catch + collect(1:n) + end + + default_scope = DefaultScope() + partition_load = zeros(Int, n_partitions) + for v in topo + meta = task_metas[v] + task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) + + if task_scope != default_scope + assigned = false + for pid in 1:n_partitions + pidx = mod1(pid, length(all_procs)) + if proc_in_scope(all_procs[pidx], task_scope) + vertex_to_partition[v] = pid + partition_load[pid] += 1 + assigned = true + break + end + end + if !assigned + best = argmin(partition_load) + vertex_to_partition[v] = best + partition_load[best] += 1 + end + else + best = argmin(partition_load) + vertex_to_partition[v] = best + partition_load[best] += 1 + end + end + end + + partition_procs = Vector{Vector{Processor}}(undef, n_partitions) + if multi_worker + for pid in 1:n_partitions + wid = partition_worker[pid] + partition_procs[pid] = procs_by_worker[wid] + end + else + for pid in 1:n_partitions + partition_procs[pid] = copy(all_procs) + end + end + + return vertex_to_partition, n_partitions, partition_procs, multi_worker +end + + +""" + _schedule_vertex!(v, partition_id, temp_queue, state, local_procs, + local_scope, dag, seen_tasks, vertex_to_partition, + schedule, proc_to_scope_lfu, write_num) -> write_num + +Schedules a single (already topologically-ordered) task vertex `v` into its +partition's `state` via `distribute_task!`, returning the updated `write_num`. + +Records cross-partition predecessors as explicit `ThunkSyncdep`s (same-partition +deps are derived from `state` by `distribute_task!`), and passes the +partition-local AOT-assigned processor when one is available and usable for this +partition (else falls back to JIT scheduling). + +Callers must guarantee that every predecessor of `v` (in any partition) has +already been submitted before calling this, so the `ThunkSyncdep`s are valid. +""" +function _schedule_vertex!(v::Int, partition_id::Int, + temp_queue::DataDepsTaskQueue, + state::DataDepsState, + local_procs::Vector{<:Processor}, + local_scope, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + vertex_to_partition::Vector{Int}, + schedule::Dict{DTask,Processor}, + proc_to_scope_lfu, + write_num::Int, + registry::Union{SharedChunkRegistry,Nothing}) + pair = seen_tasks[v] + spec = pair.spec + task = pair.task + + if spec.options.syncdeps === nothing + spec.options.syncdeps = Set{ThunkSyncdep}() + end + for pred_v in inneighbors(dag, v) + if vertex_to_partition[pred_v] != partition_id + pred_task = seen_tasks[pred_v].task + push!(spec.options.syncdeps, ThunkSyncdep(pred_task)) + end + end + + # Use the AOT-assigned processor when one is available and it lies within + # this partition's processors; otherwise fall back to JIT scheduling + # (`proc === nothing`) via `temp_queue`'s scheduler. (An AOT proc can lie + # outside `local_procs` only in the multi-worker case, where partitioning + # restricts each partition to one worker's processors; single-worker + # partitions always span all procs, so the AOT assignment is always usable + # there.) + proc = get(schedule, task, nothing) + if proc !== nothing && !(proc in local_procs) + proc = nothing + end + + return distribute_task!(temp_queue, state, local_procs, local_scope, + spec, task, spec.fargs, + proc_to_scope_lfu, write_num; proc, ownership=registry) +end + +""" + schedule_partition_full!(queue, queue_lock, partition_id, partition_verts, + dag, seen_tasks, local_procs, + vertex_to_partition, task_submitted, + region_uids) -> DataDepsState + +Per-partition scheduling for both single-worker and multi-worker hierarchical +paths. Uses existing `distribute_task!` logic with per-partition +`DataDepsState`, `all_procs` limited to this partition's processors, and +cross-partition syncdeps from the precomputed DAG. + +AOT scheduling is run *locally per partition*: a partition-local `DAGSpec` is +built from just this partition's tasks and scheduled over just `local_procs` via +`datadeps_build_schedule!`. `region_uids` (all in-region task uids) lets that +builder bail cleanly when a task depends on a same-region producer in another +partition (which it must not `fetch`). Tasks without an AOT assignment fall back +to JIT scheduling in `distribute_task!`. +""" +function schedule_partition_full!(queue::DataDepsTaskQueue, + queue_lock::ReentrantLock, + partition_id::Int, + partition_verts::Vector{Int}, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + local_procs::Vector{<:Processor}, + vertex_to_partition::Vector{Int}, + task_submitted::Vector{Base.Event}, + region_uids::Set{UInt}, + registry::Union{SharedChunkRegistry,Nothing}) + if isempty(partition_verts) || isempty(local_procs) + return DataDepsState(DAGSpec()) + end + + local_scope = UnionScope(map(ExactScope, local_procs)) + + # N.B. Hierarchical scheduling doesn't build a global `DAGSpec` (that's + # only used by the AOT DAG-caching path in `distribute_tasks!`), so each + # partition just gets an empty one; `distribute_task!` never reads + # `state.dag_spec` directly. + state = DataDepsState(DAGSpec()) + write_num = 1 + proc_to_scope_lfu = BasicLFUCache{Processor,AbstractScope}(1024) + + vert_set = Set{Int}(partition_verts) + topo = try + topological_sort_by_dfs(dag) + catch + collect(vertices(dag)) + end + ordered_verts = filter(v -> v in vert_set, topo) + + locked_queue = LockedEnqueueQueue(get_options(:task_queue), queue_lock) + # N.B. Each partition gets its own fresh scheduler shard via `similar` + # rather than sharing `queue.scheduler` across all partitions. Two + # independent problems would arise from sharing a single scheduler + # instance here: + # 1) Data race: e.g. `RoundRobinScheduler.proc_idx` would be + # concurrently read/written by every partition's `Threads.@spawn` + # task with no synchronization. + # 2) Semantic bug (worse than the race, and *not* fixed by adding a + # lock): each partition schedules only onto its own worker's + # `local_procs`, which generally has a *different length* than + # other partitions' (or the global) processor list. A `proc_idx` + # counter advanced by one partition's `local_procs` is meaningless + # -- and can be out-of-bounds -- when applied to another + # partition's differently-sized `local_procs`. This reliably + # crashes with a `BoundsError` under multi-worker hierarchical + # scheduling. Giving each partition its own scheduler instance, + # scoped to its own `local_procs`, fixes both issues at once. + temp_queue = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) + + # Per-partition (local) AOT scheduling over just this partition's procs. + # `partition_verts` is in ascending vertex (= submission) order. + partition_pairs = DTaskPair[seen_tasks[v] for v in partition_verts] + _pdag, schedule = datadeps_build_schedule!(temp_queue.scheduler, partition_pairs, + local_procs, local_scope; region_uids) + + # N.B. If this partition throws partway through (e.g. from + # `distribute_task!`), any of our vertices that haven't yet been + # `notify`'d will never be, which would leave other partitions blocked + # forever in `wait(task_submitted[pred_v])` below -- turning a normal, + # reportable exception into a silent, permanent hang (since the + # enclosing `@sync` in `distribute_tasks_hierarchical!` can't finish, + # and thus can't propagate our exception, until *every* spawned + # partition task completes, including the ones stuck waiting on us). + # The `finally` ensures every one of our events gets notified no matter + # how we exit, so that sibling partitions can unblock (and themselves + # fail/finish) and our real exception can actually surface. + try + for v in ordered_verts + for pred_v in inneighbors(dag, v) + if vertex_to_partition[pred_v] != partition_id + wait(task_submitted[pred_v]) + end + end + + # N.B. The per-task `distribute_task!` preparation runs concurrently + # across partitions; only the final task submission is serialized + # (via `LockedEnqueueQueue`, `temp_queue`'s upper queue). This is + # safe now that task futures are backed by `MemPool.DFuture` rather + # than the concurrency-unsafe `Distributed.Future`. The + # cross-partition `wait`s above happen before scheduling `v`, so the + # `ThunkSyncdep`s recorded for `v` are valid. + write_num = _schedule_vertex!( + v, partition_id, temp_queue, state, local_procs, local_scope, + dag, seen_tasks, vertex_to_partition, schedule, + proc_to_scope_lfu, write_num, registry) + + notify(task_submitted[v]) + end + finally + for v in ordered_verts + notify(task_submitted[v]) + end + end + + return state +end + +struct LockedEnqueueQueue <: AbstractTaskQueue + inner::AbstractTaskQueue + lock::ReentrantLock +end +function enqueue!(leq::LockedEnqueueQueue, pair::DTaskPair) + @lock leq.lock enqueue!(leq.inner, pair) +end +function enqueue!(leq::LockedEnqueueQueue, pairs::Vector{DTaskPair}) + @lock leq.lock enqueue!(leq.inner, pairs) +end + +# Parallel partition scheduling wraps errors in TaskFailedException / +# CompositeException via `@sync`/`Threads.@spawn`. Unwrap so callers and tests +# see the root `SchedulingException` / scheduler error. +function _unwrap_partition_exception(e) + while true + if e isa CompositeException && !isempty(e.exceptions) + e = e.exceptions[1] + elseif e isa TaskFailedException + e = something(e.task.exception, e) + else + return e + end + end +end + +""" + distribute_tasks_hierarchical!(queue) + +Main entry point for hierarchical scheduling. Runs the 4-phase pipeline: +1. Parallel aliasing construction +2. DAG construction +3. Partitioning (by worker affinity, or across local procs on one worker) +4. Parallel per-partition scheduling via `distribute_task!` + +Each partition runs on its own task and computes its *own* partition-local AOT +schedule over just its processors (see `schedule_partition_full!`); there is no +global AOT pass or global synchronization, which is what allows this to scale. +AOT scheduling here therefore need not match the flat `distribute_tasks!` path +exactly. Both drivers use `distribute_task!` for argument preparation and +`DataDepsScheduler` dispatch; the old single-worker "batch enqueue with DAG +syncdeps only" path is intentionally not used (it skipped `distribute_task!` and +broke `ChunkView` / custom schedulers). +""" +function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) + seen_tasks = queue.seen_tasks + if isempty(seen_tasks) + return + end + + # Get the set of all processors + all_procs = Processor[] + scope = get_compute_scope() + for w in procs() + append!(all_procs, get_processors(OSProc(w))) + end + filter!(proc->proc_in_scope(proc, scope), all_procs) + if isempty(all_procs) + throw(Sch.SchedulingException("No processors available, try widening scope")) + end + all_scope = UnionScope(map(ExactScope, all_procs)) + + # All in-region task uids, so each partition's local AOT-schedule builder + # can tell same-region producers (which it must not `fetch`) apart from + # already-materialized external values. + region_uids = Set{UInt}(pair.task.uid for pair in seen_tasks) + + # Phase 1: Collect arguments and compute aliasing in parallel + task_metas, unique_arg_ws = collect_aliased_args(seen_tasks) + _lookup, ainfos_overlaps, arg_to_ainfo = build_aliasing_parallel(unique_arg_ws) + + # Phase 2: Build dependency DAG + dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) + + # Phase 3: Partition the DAG + vertex_to_partition, n_partitions, partition_procs, _multi_worker = + partition_dag(dag, task_metas, all_procs) + + # Detect backing chunks shared across partitions in different memory spaces. + # These need runtime ownership transfer to avoid split-brain concurrent + # writes; `nothing` when all partitions share one space (single-worker). + partition_space = MemorySpace[only(memory_spaces(first(pp))) for pp in partition_procs] + registry = build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) + + # Group vertices by partition + partitions = [Int[] for _ in 1:n_partitions] + for v in 1:length(seen_tasks) + pid = vertex_to_partition[v] + push!(partitions[pid], v) + end + + queue_lock = ReentrantLock() + task_submitted = [Base.Event() for _ in 1:length(seen_tasks)] + wait_all_queue = get_options(:task_queue) + + # Phase 4: parallel per-partition scheduling. Each partition runs on its own + # task, prepares its tasks concurrently, computes its own partition-local + # AOT schedule, and coordinates cross-partition dependencies via the + # `task_submitted` events. Concurrent submission is safe now that task + # futures are backed by `MemPool.DFuture` (see the header note). + partition_states = Vector{DataDepsState}(undef, n_partitions) + try + @sync for pid in 1:n_partitions + Threads.@spawn begin + locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) + with_options(; task_queue=locked_queue) do + partition_states[pid] = schedule_partition_full!( + queue, queue_lock, pid, partitions[pid], + dag, seen_tasks, + partition_procs[pid], vertex_to_partition, + task_submitted, region_uids, registry + ) + end + end + end + catch e + rethrow(_unwrap_partition_exception(e)) + end + + _hierarchical_copy_from_and_free!(partition_states, n_partitions, registry) +end + +function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) + wn = 0 + if haskey(state.arg_history, arg_w) + for entry in state.arg_history[arg_w] + wn = max(wn, entry.write_num) + end + end + return wn +end + +function _hierarchical_copy_from!(state::DataDepsState, arg_w::ArgumentWrapper, write_num::Int) + haskey(state.arg_origin, arg_w.arg) || return + origin_space = state.arg_origin[arg_w.arg] + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) + if remainder isa MultiRemainderAliasing + origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) + enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) + elseif remainder isa FullCopy + origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) + enqueue_copy_from!(state, origin_space, arg_w, origin_scope, write_num) + end + return +end + +function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsState}, n_partitions::Int, + registry::Union{SharedChunkRegistry,Nothing}) + # 1. Shared chunks: write back from the registry's authoritative owner state + # (whose per-partition history is coherent), not the cross-partition + # max-`write_num` heuristic below (per-partition write_nums are not + # comparable across partitions). + if registry !== nothing + for (chunk, entry) in registry.entries + state = entry.owner_state + state === nothing && continue # never written + entry.owner_space == entry.origin_space && continue # already in-place at origin + for (arg_w, _space) in state.arg_owner + arg_w.arg === chunk || continue + _hierarchical_copy_from!(state, arg_w, _hierarchical_max_write_num(state, arg_w) + 1) + end + end + end + + # 2. Private chunks: the last writer across partitions is the authoritative + # owner (shared chunks are skipped here, handled above). + merged_arg_owner = Dict{ArgumentWrapper, Tuple{MemorySpace, Int, DataDepsState}}() + for pid in 1:n_partitions + state = partition_states[pid] + for (arg_w, space) in state.arg_owner + is_shared_chunk(registry, arg_w.arg) && continue + wn = _hierarchical_max_write_num(state, arg_w) + if !haskey(merged_arg_owner, arg_w) || wn > merged_arg_owner[arg_w][2] + merged_arg_owner[arg_w] = (space, wn, state) + end + end + end + + for arg_w in sort(collect(keys(merged_arg_owner)); by=arg_w->arg_w.hash) + _space, wn, state = merged_arg_owner[arg_w] + _hierarchical_copy_from!(state, arg_w, wn + 1) + end + + # 3. Free Datadeps-allocated slots. For shared chunks, also sync on the final + # global writer: an intermediate owner's slot may still be read by a + # cross-partition boundary copy recorded in *another* partition's state, + # and the final writer transitively depends on all such copies. + for pid in 1:n_partitions + state = partition_states[pid] + obj_cache = unwrap(state.ainfo_backing_chunk) + write_num = typemax(Int) - 1 + for remote_space in keys(obj_cache.values) + for (ainfo, remote_arg) in obj_cache.values[remote_space] + if !(ainfo in obj_cache.originals) + remote_proc = first(processors(remote_space)) + free_scope = ExactScope(remote_proc) + free_syncdeps = Set{ThunkSyncdep}() + if haskey(state.ainfo_arg, ainfo) + get_write_deps!(state, remote_space, ainfo, write_num, free_syncdeps) + end + if registry !== nothing + orig = get(state.remote_arg_to_original, remote_arg, nothing) + if orig !== nothing + entry = get(registry.entries, orig, nothing) + if entry !== nothing && entry.owner_task !== nothing + push!(free_syncdeps, ThunkSyncdep(entry.owner_task)) + end + end + end + Dagger.@spawn scope=free_scope syncdeps=free_syncdeps Dagger.unsafe_free!(remote_arg) + end + end + end + end +end diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 989e223b5..e7dfceafe 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -65,7 +65,8 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, traversal::Symbol=:inorder, scheduler::Union{DataDepsScheduler,Nothing}=nothing, aliasing::Bool=true, - launch_wait::Union{Bool,Nothing}=nothing) + launch_wait::Union{Bool,Nothing}=nothing, + hierarchical::Union{Bool,Nothing}=nothing) if !static throw(ArgumentError("Dynamic scheduling is no longer available")) end @@ -78,22 +79,127 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true, wait_all(; check_errors=true) do scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler()) launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool + hierarchical = something(hierarchical, DATADEPS_HIERARCHICAL[], true)::Bool + run_distribute = queue -> begin + if hierarchical + distribute_tasks_hierarchical!(queue) + else + distribute_tasks!(queue) + end + end if launch_wait result = spawn_bulk() do queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) with_options(f; task_queue=queue) - distribute_tasks!(queue) + run_distribute(queue) end else queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) result = with_options(f; task_queue=queue) - distribute_tasks!(queue) + run_distribute(queue) end return result end end const DATADEPS_SCHEDULER = ScopedValue{Union{DataDepsScheduler,Nothing}}(nothing) const DATADEPS_LAUNCH_WAIT = ScopedValue{Union{Bool,Nothing}}(nothing) +const DATADEPS_HIERARCHICAL = ScopedValue{Union{Bool,Nothing}}(nothing) + +"Returns `true` if any argument of `spec` is a same-region (`region_uids`) `DTask`." +function _spec_has_region_dtask_arg(spec::DTaskSpec, region_uids::Set{UInt}) + for _arg in spec.fargs + arg, _ = unwrap_inout(value(_arg)) + if arg isa DTask && arg.uid in region_uids + return true + end + end + return false +end + +""" + datadeps_build_schedule!(scheduler, pairs, all_procs, all_scope; + region_uids=nothing) -> (dag_spec, schedule) + +Builds a DAG spec from `pairs` (in submission order) and computes an AOT +processor assignment (`schedule::Dict{DTask,Processor}`) over `all_procs` / +`all_scope`, reusing a cached schedule when an equivalent DAG has been seen +before (and otherwise caching the freshly-computed one for future reuse). + +The schedule cache is task-local per scheduler type, so callers running on +distinct tasks (e.g. the per-partition scheduling tasks of the hierarchical +path) each maintain an independent cache. Schedulers that don't implement +`datadeps_schedule_dag_aot!` leave the schedule empty, in which case tasks fall +back to JIT scheduling in `distribute_task!`. + +`region_uids`, when provided, is the set of *all* in-region task uids. It lets +this function bail (fall back to JIT for the remaining tasks) before +`dag_add_task!` would try to `fetch` an unlaunched same-region producer that is +absent from this (sub)set of `pairs` -- which is essential when `pairs` covers +only one partition of the region, since a producer in another partition would +otherwise be treated as an external value and block forever. For the flat path +(where `pairs` is the whole region), `dag_add_task!`'s own in-DAG check already +handles this and `region_uids` can be omitted. + +Used by both the flat (`distribute_tasks!`, whole region) and hierarchical +(`distribute_tasks_hierarchical!`, per partition) paths. +""" +function datadeps_build_schedule!(scheduler::DataDepsScheduler, + pairs::Vector{DTaskPair}, + all_procs, all_scope; + region_uids::Union{Set{UInt},Nothing}=nothing) + # Compute the DAG spec. `dag_add_task!` ignores the state argument (it only + # inspects the task spec), so a throwaway state is fine here. + dag_spec = DAGSpec() + dummy_state = DataDepsState(dag_spec) + for (spec, task) in pairs + if region_uids !== nothing && _spec_has_region_dtask_arg(spec, region_uids) + # Depends on an in-region producer that may live outside `pairs`; + # defer the rest to JIT scheduling. + break + end + if !dag_add_task!(dag_spec, dummy_state, spec, task) + # This task depends on an in-region task's result; defer the rest + # to JIT scheduling (they won't appear in `schedule`). + break + end + end + + # Attempt to find any matching DAG specs and reuse their schedule + schedule = Dict{DTask, Processor}() + schedule_cache = datadeps_schedule_cache(scheduler) + cache_hit = false + for (other_spec, spec_schedule) in schedule_cache + if datadeps_dag_equivalent(scheduler, dag_spec, other_spec) + @dagdebug nothing :spawn_datadeps "Found matching DAG spec!" + for (id, proc) in spec_schedule.id_to_proc + uid = dag_spec.id_to_uid[id] + task_idx = findfirst(spec_task -> spec_task.task.uid == uid, pairs) + task = pairs[task_idx].task + schedule[task] = proc + end + cache_hit = true + break + end + end + + if !cache_hit && !isempty(dag_spec) + # Compute a fresh AOT schedule (no-op for schedulers that fall back + # to JIT in distribute_task!) + datadeps_schedule_dag_aot!(scheduler, schedule, dag_spec, all_procs, all_scope) + + # Persist the schedule for reuse by future equivalent DAGs + if !isempty(schedule) + spec_schedule = DAGSpecSchedule() + for (task, proc) in schedule + id = dag_spec.uid_to_id[task.uid] + spec_schedule.id_to_proc[id] = proc + end + push!(schedule_cache, dag_spec => spec_schedule) + end + end + + return dag_spec, schedule +end # Current task uid, propagated into `tochunk` so uniform-execution backends # (MPIExt) can derive deterministic, rank-agreed handle IDs. Core datadeps sets @@ -135,49 +241,9 @@ function distribute_tasks!(queue::DataDepsTaskQueue) upper_queue = get_options(:task_queue) - # Compute DAG spec - dag_spec = DAGSpec() - state = DataDepsState(dag_spec) - for (spec, task) in queue.seen_tasks - if !dag_add_task!(dag_spec, state, spec, task) - # This task needs to be deferred - break - end - end - - # Attempt to find any matching DAG specs and reuse their schedule - schedule = Dict{DTask, Processor}() - schedule_cache = datadeps_schedule_cache(queue.scheduler) - cache_hit = false - for (other_spec, spec_schedule) in schedule_cache - if datadeps_dag_equivalent(queue.scheduler, dag_spec, other_spec) - @dagdebug nothing :spawn_datadeps "Found matching DAG spec!" - for (id, proc) in spec_schedule.id_to_proc - uid = dag_spec.id_to_uid[id] - task_idx = findfirst(spec_task -> spec_task.task.uid == uid, queue.seen_tasks) - task = queue.seen_tasks[task_idx].task - schedule[task] = proc - end - cache_hit = true - break - end - end - - if !cache_hit && !isempty(dag_spec) - # Compute a fresh AOT schedule (no-op for schedulers that fall back - # to JIT in distribute_task!) - datadeps_schedule_dag_aot!(queue.scheduler, schedule, dag_spec, all_procs, all_scope) - - # Persist the schedule for reuse by future equivalent DAGs - if !isempty(schedule) - spec_schedule = DAGSpecSchedule() - for (task, proc) in schedule - id = dag_spec.uid_to_id[task.uid] - spec_schedule.id_to_proc[id] = proc - end - push!(schedule_cache, dag_spec => spec_schedule) - end - end + # Compute the DAG spec and an AOT processor assignment (reused from cache + # when an equivalent DAG has been seen before). + dag_spec, schedule = datadeps_build_schedule!(queue.scheduler, queue.seen_tasks, all_procs, all_scope) # Start launching tasks and necessary copies state = DataDepsState(dag_spec) @@ -286,7 +352,7 @@ map_or_ntuple(f, xs::Vector) = map(f, 1:length(xs)) # N.B. Accept any `Tuple` (typed specs produce heterogeneous tuples of # `TypedArgument{T}`, not a homogeneous `NTuple{N,T}`). @inline map_or_ntuple(@specialize(f), xs::Tuple) = ntuple(f, Val(length(xs))) -function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int; proc::Union{Processor,Nothing}=nothing) where typed +function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_procs, all_scope, spec::DTaskSpec{typed}, task::DTask, fargs, proc_to_scope_lfu, write_num::Int; proc::Union{Processor,Nothing}=nothing, ownership=nothing) where typed @specialize spec fargs if typed @@ -352,6 +418,14 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr return end + # Hierarchical scheduling only: for shared backing chunks whose current + # version was produced by another partition, seed this partition's state so + # the copy-to below pulls a fresh whole-chunk copy from the true owner (and + # syncs on its producer). No-op on the flat path (`ownership === nothing`). + if ownership !== nothing + _sync_incoming_ownership!(state, ownership, our_space, task_arg_ws, write_num) + end + # Copy args from local to remote remote_args = map_or_ntuple(task_arg_ws) do idx arg_ws = task_arg_ws[idx] @@ -480,6 +554,13 @@ function distribute_task!(queue::DataDepsTaskQueue, state::DataDepsState, all_pr return end + # Hierarchical scheduling only: publish this task as the new authoritative + # owner of each shared backing chunk it writes, so later cross-partition + # consumers pull the up-to-date version from here. + if ownership !== nothing + _commit_ownership!(state, ownership, our_space, task, task_arg_ws) + end + write_num += 1 DATADEPS_CURRENT_TASK[] = nothing diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index cb6ad861c..af72869df 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -243,10 +243,16 @@ end ### JIT Schedulers ### +# Default for user-defined schedulers with a zero-arg constructor. Schedulers +# that carry mutable state should specialize `similar` to return a fresh shard +# (used when hierarchical scheduling clones a scheduler per partition). +Base.similar(s::DataDepsScheduler) = typeof(s)() + mutable struct RoundRobinScheduler <: DataDepsScheduler proc_idx::Int RoundRobinScheduler() = new(1) end +Base.similar(::RoundRobinScheduler) = RoundRobinScheduler() function datadeps_schedule_task_jit!(sched::RoundRobinScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) proc_idx = sched.proc_idx our_proc = all_procs[proc_idx] @@ -267,6 +273,7 @@ function datadeps_schedule_task_jit!(sched::RoundRobinScheduler, all_procs, all_ end struct NaiveScheduler <: DataDepsScheduler end +Base.similar(::NaiveScheduler) = NaiveScheduler() function datadeps_schedule_task_jit!(sched::NaiveScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) raw_args = map(arg->tochunk(value(arg)), spec.fargs) our_proc = remotecall_fetch(1, all_procs, raw_args) do all_procs, raw_args @@ -302,6 +309,7 @@ struct UltraScheduler <: DataDepsScheduler Dict{MemorySpace,Int}()) end end +Base.similar(::UltraScheduler) = UltraScheduler() function datadeps_schedule_task_jit!(sched::UltraScheduler, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) args = Base.mapany(spec.fargs) do arg pos, data = arg @@ -421,4 +429,4 @@ function datadeps_schedule_dag_aot!(scheduler::LayeredScheduler, schedule, dag_s schedule[task] = our_proc end end -end \ No newline at end of file +end diff --git a/test/datadeps.jl b/test/datadeps.jl index 1ad135cac..d0f7b4e67 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -998,7 +998,7 @@ end struct DummyErrorScheduler <: Dagger.DataDepsScheduler end struct DummySchedulerError <: Exception end -function Dagger.datadeps_schedule_task(::DummyErrorScheduler, state, all_procs, all_scope, task_scope, spec, task) +function Dagger.datadeps_schedule_task_jit!(::DummyErrorScheduler, all_procs, all_scope, task_scope, spec, task) throw(DummySchedulerError()) end diff --git a/test/datadeps/scheduling.jl b/test/datadeps/scheduling.jl index 8f5bb1070..2c2003640 100644 --- a/test/datadeps/scheduling.jl +++ b/test/datadeps/scheduling.jl @@ -115,6 +115,13 @@ end # ---------- Helpers for DAG capture via the schedule cache ---------- +# N.B. These tests exercise the *flat* AOT DAG-scheduling/caching path +# (`distribute_tasks!` + `datadeps_build_schedule!`), so every `spawn_datadeps` +# call here runs with `hierarchical=false`. Hierarchical scheduling intentionally +# computes a separate, partition-local AOT schedule on each partition's own task +# (no global schedule cache on the calling task), so the caller-visible cache +# these tests inspect is only populated by the flat path. + """ Run `f()` inside `spawn_datadeps` with `scheduler`, returning the resulting schedule cache snapshot (a Vector). The cache is cleared before the run so @@ -123,7 +130,7 @@ each call starts from a known state. function run_with_fresh_cache(f, scheduler::DataDepsScheduler = LayeredScheduler()) cache = datadeps_schedule_cache(scheduler) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => scheduler) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => scheduler) do Dagger.spawn_datadeps(f) end return cache @@ -136,7 +143,7 @@ cache. Useful for asserting cache size after repeated runs. function run_n_times(f, n::Int; scheduler::DataDepsScheduler = LayeredScheduler()) cache = datadeps_schedule_cache(scheduler) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => scheduler) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => scheduler) do for _ in 1:n Dagger.spawn_datadeps(f) end @@ -231,7 +238,7 @@ end @testset "Different array sizes → cache miss" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for sz in (64, 128, 256) A = rand(sz); B = rand(sz) Dagger.spawn_datadeps() do @@ -245,7 +252,7 @@ end @testset "Different element types → cache miss" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for T in (Float64, Float32, Int) A = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) B = T <: AbstractFloat ? rand(T, 64) : T.(rand(1:100, 64)) @@ -260,7 +267,7 @@ end @testset "Different task counts → cache miss" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for ntasks in (1, 2, 4) A = rand(64); B = rand(64) Dagger.spawn_datadeps() do @@ -278,7 +285,7 @@ end empty!(cache) scopes = [Dagger.DefaultScope(), Dagger.ExactScope(Dagger.ThreadProc(1, 1))] - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for scp in scopes A = rand(64); B = rand(64) Dagger.spawn_datadeps() do @@ -296,7 +303,7 @@ end end @test length(cache) == 1 # Repeat with the other function - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do A = rand(64) Dagger.spawn_datadeps() do Dagger.@spawn scale!(InOut(A), 2.0) @@ -311,7 +318,7 @@ end @testset "Cache partitioning between schedulers" begin ls_cache = datadeps_schedule_cache(LayeredScheduler()) empty!(ls_cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do A = rand(64); B = rand(64) Dagger.spawn_datadeps() do Dagger.@spawn add!(InOut(A), In(B)) @@ -323,7 +330,7 @@ end # cache, and it must NOT see LayeredScheduler's entry. rr_cache = datadeps_schedule_cache(RoundRobinScheduler()) @test rr_cache !== ls_cache - Base.ScopedValues.with(DATADEPS_SCHEDULER => RoundRobinScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => RoundRobinScheduler()) do A = rand(64); B = rand(64) Dagger.spawn_datadeps() do Dagger.@spawn add!(InOut(A), In(B)) @@ -370,7 +377,7 @@ end @testset "Repeated runs, fresh allocations → 1 cache entry" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for _ in 1:3 M = make_spd_blocks(64, 16) Dagger.spawn_datadeps() do @@ -389,7 +396,7 @@ end @testset "Different block-grid sizes → distinct cache entries" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do # 2x2 grid (sz=32, nb=16) vs 4x4 grid (sz=64, nb=16) → different task counts for (sz, nb) in ((32, 16), (64, 16)) M = make_spd_blocks(sz, nb) @@ -404,7 +411,7 @@ end @testset "Same block-grid shape, different block size → distinct entries" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do # Both 2x2 grids; block size 16 vs 32 → same task count, different ContiguousAliasing length for (sz, nb) in ((32, 16), (64, 32)) M = make_spd_blocks(sz, nb) @@ -437,7 +444,7 @@ end @testset "Repeated runs → 1 cache entry" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for _ in 1:3 As = [rand(64) for _ in 1:8] Dagger.spawn_datadeps() do @@ -451,7 +458,7 @@ end @testset "Different reduction widths → distinct entries" begin cache = datadeps_schedule_cache(LayeredScheduler()) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LayeredScheduler()) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LayeredScheduler()) do for n in (4, 8, 16) As = [rand(64) for _ in 1:n] Dagger.spawn_datadeps() do @@ -471,7 +478,7 @@ end # Build two DAGSpecs from equivalent-but-freshly-allocated workloads. cache = datadeps_schedule_cache(LS) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LS) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LS) do for _ in 1:2 A = rand(64); B = rand(64) Dagger.spawn_datadeps() do @@ -486,7 +493,7 @@ end # Now build a structurally-different spec. empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => LS) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => LS) do A = rand(128); B = rand(128) # different size Dagger.spawn_datadeps() do Dagger.@spawn add!(InOut(A), In(B)) @@ -519,7 +526,7 @@ Dagger.datadeps_schedule_task_jit!(::NoCacheScheduler, all_procs, all_scope, tas nc = NoCacheScheduler() cache = datadeps_schedule_cache(nc) empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => nc) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => nc) do for _ in 1:3 A = rand(64); B = rand(64) Dagger.spawn_datadeps() do @@ -553,7 +560,7 @@ Dagger.datadeps_schedule_task_jit!(::PtrStrictScheduler, all_procs, all_scope, t # Fresh allocations each run → pointer hash differs → all miss. empty!(cache) - Base.ScopedValues.with(DATADEPS_SCHEDULER => ps) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => ps) do for _ in 1:3 A = rand(64); B = rand(64) Dagger.spawn_datadeps() do @@ -566,7 +573,7 @@ Dagger.datadeps_schedule_task_jit!(::PtrStrictScheduler, all_procs, all_scope, t # Same allocation reused each run → pointer hash matches → cache hits. empty!(cache) A = rand(64); B = rand(64) - Base.ScopedValues.with(DATADEPS_SCHEDULER => ps) do + Base.ScopedValues.with(Dagger.DATADEPS_HIERARCHICAL => false, DATADEPS_SCHEDULER => ps) do for _ in 1:3 Dagger.spawn_datadeps() do Dagger.@spawn add!(InOut(A), In(B)) From 824b7ebf02ae3324915d3e2a8d783fee0208228b Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Fri, 24 Jul 2026 15:18:49 -0700 Subject: [PATCH 29/43] datadeps: Partition hierarchical scheduling by MPI rank Hierarchical Datadeps previously keyed multi-owner partitions on root_worker_id, which is always myid() under MPI. Use partition_affinity_id (rank under MPIExt) so tasks spread across ranks, run Phase 4 with a shared DataDepsState for SPMD-coherent planning, and skip the unused shared-chunk registry on that path so write-back stays correct for Cholesky/LU. Co-authored-by: Cursor --- ext/MPIExt.jl | 18 +- src/acceleration.jl | 12 ++ src/datadeps/hierarchical.jl | 311 +++++++++++++++++++++++++---------- src/memory-spaces.jl | 11 ++ test/mpi.jl | 7 +- 5 files changed, 270 insertions(+), 89 deletions(-) diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index f7521bd6e..586a3ccd9 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -25,7 +25,7 @@ import Dagger: move_rewrap_child_types, move_rewrap_header_mode, move_rewrap_parts, move_rewrap_result_type, move_type, multi_span_gather!, multi_span_scatter!, next_id, NoAliasing, Options, OSProc, - post_stage_array_chunks!, processors, ProcessScope, proc_in_scope, + partition_affinity_id, post_stage_array_chunks!, processors, ProcessScope, proc_in_scope, remotecall_endpoint_toplevel, RemotePtr, root_worker_id, same_node, schedule_argument_move, sch_handle, select_processors_uniform!, set_key_stored!, set_stored!, short_name, span_len, stage_acquire!, stage_release!, @@ -142,7 +142,15 @@ end MPIAcceleration() = MPIAcceleration(MPI.COMM_WORLD) function aliasing(accel::MPIAcceleration, x::Chunk, T) - handle = x.handle::MPIRef + handle = x.handle + # Chunks created under a temporary DistributedAcceleration (or a worker + # thread that did not inherit MPI TLS) carry a DRef; fall back to the + # Distributed unwrap path rather than hard-failing the typeassert. + if !(handle isa MPIRef) + return _with_default_acceleration() do + aliasing(x, T) + end + end @assert accel.comm == handle.comm "MPIAcceleration comm mismatch" tag = to_tag() check_uniform(tag) @@ -473,6 +481,12 @@ function memory_spaces(proc::MPIProcessor) end root_worker_id(mem_space::MPIMemorySpace) = myid() +# Hierarchical Datadeps partitions by this id (analogous to Distributed worker +# id). MPI always reports `root_worker_id == myid()`, so affinity must key on +# the owning rank instead. +partition_affinity_id(space::MPIMemorySpace) = space.rank +partition_affinity_id(proc::MPIProcessor) = proc.rank +partition_affinity_id(proc::MPIOSProc) = proc.rank function processors(memSpace::MPIMemorySpace) rawProc = Processor[] diff --git a/src/acceleration.jl b/src/acceleration.jl index 99a71e1e2..ea411bb0a 100644 --- a/src/acceleration.jl +++ b/src/acceleration.jl @@ -114,6 +114,18 @@ processor_owner_rank(accel::Acceleration, proc::Processor) = first(fire_order_ke processor_order_key(accel::Acceleration, proc::Processor) = (fire_order_key(proc)..., short_name(proc)) +""" + partition_affinity_id(x) -> Int + +Logical owner identity used by hierarchical Datadeps partitioning. Defaults to +`root_worker_id` (Distributed worker id). Accelerations with a different notion +of ownership (e.g. MPI ranks) override for their spaces/processors so affinity +partitioning spreads work across ranks the same way Distributed spreads across +workers. +""" +partition_affinity_id(space::MemorySpace) = root_worker_id(space) +partition_affinity_id(proc::Processor) = partition_affinity_id(only(memory_spaces(proc))) + "Preferred execution rank for `task`, or `nothing` if no preference." task_processor_preference(accel::Acceleration, task, state) = nothing diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index b4e2f0734..6524b20eb 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -1,13 +1,15 @@ # Hierarchical scheduling for datadeps -# Spreads scheduling work across multiple threads/workers via a 4-phase pipeline: +# Spreads scheduling work across multiple threads/workers/MPI ranks via a +# 4-phase pipeline: # Phase 1: Parallel aliasing info construction # Phase 2: Sequential DAG construction from aliasing overlaps -# Phase 3: Data-affinity DAG partitioning -# Phase 4: Parallel local scheduling per partition -- each partition runs on its -# own task, prepares/submits its tasks concurrently with the others, -# and computes its *own* (partition-local) AOT schedule over just its -# processors. This maximizes scheduling parallelism and scalability; -# no global synchronization or global AOT pass is imposed. +# Phase 3: Data-affinity DAG partitioning (by Distributed worker or MPI rank +# via `partition_affinity_id`) +# Phase 4: Per-partition scheduling -- under Distributed each partition runs on +# its own task concurrently; under MPI (uniform execution) partitions +# are scheduled sequentially in global topological order for SPMD +# safety. Each partition computes its *own* (partition-local) AOT +# schedule over just its processors. # # N.B. Concurrent task submission from many partition tasks used to intermittently # hang in the core eager scheduler. The root cause was `Distributed.Future`, @@ -25,7 +27,9 @@ end struct HierarchicalTaskMeta pair::DTaskPair - arg_chunk::Union{Chunk, Nothing} + # Wrapped argument identity for the task's first aliasing arg. Under MPI + # this may be a raw `ChunkView` (kept unwrapped by `datadeps_arg_wrap`). + arg_chunk::Any may_alias::Bool inplace_move::Bool deps::Vector{HierarchicalTaskInfo} @@ -249,9 +253,18 @@ function collect_aliased_args(seen_tasks::Vector{DTaskPair}) end supports_cache = IdDict{Any,Bool}() - raw_arg_cache = IdDict{Any,Chunk}() - - nchunks = Threads.nthreads() <= 1 ? 1 : min(Threads.nthreads(), cld(n, COLLECT_ALIASED_ARGS_MIN_CHUNK)) + # Values are Chunks, or raw remote handles kept by `datadeps_arg_wrap` + # (e.g. ChunkView under MPI). + raw_arg_cache = IdDict{Any,Any}() + + # Under uniform execution (MPI), `tochunk` / `MPIRefID` allocation must stay + # sequential on the root task, and spawned threads would not inherit the + # TaskLocalValue acceleration (falling back to Distributed → DRef chunks). + nchunks = if uniform_execution() || Threads.nthreads() <= 1 + 1 + else + min(Threads.nthreads(), cld(n, COLLECT_ALIASED_ARGS_MIN_CHUNK)) + end if nchunks <= 1 unique_arg_ws = Dict{ArgumentWrapper, ArgumentWrapper}() @@ -264,10 +277,14 @@ function collect_aliased_args(seen_tasks::Vector{DTaskPair}) chunk_size = cld(n, nchunks) starts = collect(1:chunk_size:n) per_chunk_arg_ws = Vector{Dict{ArgumentWrapper,ArgumentWrapper}}(undef, length(starts)) + # Propagate the caller's acceleration into each worker thread; TaskLocalValue + # does not inherit across `Threads.@spawn`. + parent_accel = current_acceleration() @sync for (ci, start) in enumerate(starts) range = start:min(start+chunk_size-1, n) Threads.@spawn begin + set_task_acceleration!(parent_accel) local_arg_ws = Dict{ArgumentWrapper,ArgumentWrapper}() _collect_aliased_args_range!(task_metas, local_arg_ws, seen_tasks, range, supports_cache, raw_arg_cache, cache_lock, task_to_idx) @@ -288,7 +305,7 @@ function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, seen_tasks::Vector{DTaskPair}, range::UnitRange{Int}, supports_cache::IdDict{Any,Bool}, - raw_arg_cache::IdDict{Any,Chunk}, + raw_arg_cache::IdDict{Any,Any}, cache_lock::Union{ReentrantLock,Nothing}, task_to_idx::IdDict{DTask,Int}) for task_idx in range @@ -334,7 +351,10 @@ function _collect_aliased_args_range!(task_metas::Vector{HierarchicalTaskMeta}, end arg_chunk = _cached_get!(raw_arg_cache, cache_lock, arg) do - arg isa Chunk ? arg : tochunk(arg) + # Match `populate_task_info!`: acceleration-aware wrap (MPIRef + # under MPI) rather than a bare `tochunk` that can produce a + # rank-local DRef when TLS acceleration is unset. + arg isa Chunk ? arg : datadeps_arg_wrap(arg) end if first_chunk === nothing @@ -443,16 +463,24 @@ const COMPUTE_ALIASING_BATCH_MIN_PARALLEL = 8 function _compute_aliasing_batch(arg_ws::Vector{ArgumentWrapper}) n = length(arg_ws) results = Vector{Pair{ArgumentWrapper, AliasingWrapper}}(undef, n) - if n >= COMPUTE_ALIASING_BATCH_MIN_PARALLEL && Threads.nthreads() > 1 + accel = current_acceleration() + # Under uniform execution (MPI), aliasing may perform collectives that must + # run in the same sequential order on every rank -- never Threads.@threads. + # Also dispatch through `aliasing(accel, ...)` so MPI Chunks go through the + # owner-broadcast path rather than a local unwrap. + can_parallel = !uniform_execution(accel) && + n >= COMPUTE_ALIASING_BATCH_MIN_PARALLEL && + Threads.nthreads() > 1 + if can_parallel Threads.@threads for i in 1:n arg_w = arg_ws[i] - ainfo = AliasingWrapper(aliasing(arg_w.arg, arg_w.dep_mod)) + ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) results[i] = arg_w => ainfo end else for i in 1:n arg_w = arg_ws[i] - ainfo = AliasingWrapper(aliasing(arg_w.arg, arg_w.dep_mod)) + ainfo = AliasingWrapper(aliasing(accel, arg_w.arg, arg_w.dep_mod)) results[i] = arg_w => ainfo end end @@ -556,52 +584,55 @@ end partition_dag(dag, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs) Phase 3: Assigns each task vertex to a partition using data-affinity. For -multi-worker setups, tasks are assigned to the worker owning the most argument -data. For single-worker multi-threaded setups, tasks are balanced across -available processors in topological order. +multi-owner setups (Distributed workers, or MPI ranks via +`partition_affinity_id`), tasks are assigned to the owner holding the most +argument data. For single-owner multi-threaded setups, tasks are balanced +across available processors in topological order. """ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMeta}, all_procs::Vector{<:Processor}) n = length(task_metas) - workers = unique(root_worker_id.(only.(memory_spaces.(all_procs)))) - n_workers = length(workers) + # Stable order so SPMD ranks (and Distributed workers) agree on partition + # indexing when affinity ids are collected from an unordered processor set. + owners = sort!(unique(partition_affinity_id.(all_procs))) + n_owners = length(owners) - procs_by_worker = Dict{Int, Vector{Processor}}() + procs_by_owner = Dict{Int, Vector{Processor}}() for proc in all_procs - wid = root_worker_id(only(memory_spaces(proc))) - push!(get!(Vector{Processor}, procs_by_worker, wid), proc) + oid = partition_affinity_id(proc) + push!(get!(Vector{Processor}, procs_by_owner, oid), proc) end - multi_worker = n_workers > 1 - if multi_worker - n_partitions = n_workers - partition_worker = workers + multi_owner = n_owners > 1 + if multi_owner + n_partitions = n_owners + partition_owner = owners else n_partitions = min(length(all_procs), n) - partition_worker = fill(first(workers), n_partitions) + partition_owner = fill(first(owners), n_partitions) end vertex_to_partition = Vector{Int}(undef, n) - if multi_worker - worker_to_partition = Dict(w => i for (i, w) in enumerate(workers)) + if multi_owner + owner_to_partition = Dict(o => i for (i, o) in enumerate(owners)) default_scope = DefaultScope() for v in 1:n meta = task_metas[v] task_scope = @something(meta.pair.spec.options.compute_scope, meta.pair.spec.options.scope, default_scope) - # Workers whose processors are eligible under this task's scope. - # A non-default scope that spans multiple workers must still spread - # work across those workers (via affinity / round-robin) -- picking - # only the first match pins everything to worker 1 and breaks - # multi-worker execution. + # Owners whose processors are eligible under this task's scope. + # A non-default scope that spans multiple owners must still spread + # work across those owners (via affinity / round-robin) -- picking + # only the first match pins everything to owner 1 and breaks + # multi-owner execution. if task_scope == default_scope matching = collect(1:n_partitions) else matching = Int[] - for (pid, wid) in enumerate(workers) - wprocs = procs_by_worker[wid] - if any(proc -> proc_in_scope(proc, task_scope), wprocs) + for (pid, oid) in enumerate(owners) + oprocs = procs_by_owner[oid] + if any(proc -> proc_in_scope(proc, task_scope), oprocs) push!(matching, pid) end end @@ -615,11 +646,11 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe continue end - affinity = zeros(Int, n_workers) + affinity = zeros(Int, n_owners) for dep in meta.deps arg_space = memory_space(dep.arg_w.arg) - arg_wid = root_worker_id(arg_space) - idx = get(worker_to_partition, arg_wid, 0) + arg_oid = partition_affinity_id(arg_space) + idx = get(owner_to_partition, arg_oid, 0) if idx > 0 && idx in matching affinity[idx] += 1 end @@ -676,10 +707,10 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe end partition_procs = Vector{Vector{Processor}}(undef, n_partitions) - if multi_worker + if multi_owner for pid in 1:n_partitions - wid = partition_worker[pid] - partition_procs[pid] = procs_by_worker[wid] + oid = partition_owner[pid] + partition_procs[pid] = procs_by_owner[oid] end else for pid in 1:n_partitions @@ -687,7 +718,7 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe end end - return vertex_to_partition, n_partitions, partition_procs, multi_worker + return vertex_to_partition, n_partitions, partition_procs, multi_owner end @@ -869,6 +900,81 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, return state end +""" + schedule_partitions_sequential!(...) -> Vector{DataDepsState} + +SPMD-safe Phase 4: schedule every partition's tasks on the root task in global +topological order. Processor assignment remains partition-local (MPI-rank / +worker affinity), but a *single* shared `DataDepsState` is used so argument +history, remainders, and final copy-back stay coherent across ranks -- the same +model as flat `distribute_tasks!`. Per-partition states would otherwise +split-brain overlapping writes (e.g. whole-array + view + triangular dep_mods). + +Returns a one-element vector containing the shared state (for the hierarchical +copy-from/free epilogue). +""" +function schedule_partitions_sequential!(queue::DataDepsTaskQueue, + queue_lock::ReentrantLock, + partitions::Vector{Vector{Int}}, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + partition_procs::Vector{<:Vector{<:Processor}}, + vertex_to_partition::Vector{Int}, + region_uids::Set{UInt}, + registry::Union{SharedChunkRegistry,Nothing}, + wait_all_queue) + n_partitions = length(partitions) + temp_queues = Vector{DataDepsTaskQueue}(undef, n_partitions) + local_scopes = Vector{AbstractScope}(undef, n_partitions) + schedules = Vector{Dict{DTask,Processor}}(undef, n_partitions) + proc_to_scope_lfus = [BasicLFUCache{Processor,AbstractScope}(1024) for _ in 1:n_partitions] + shared_state = DataDepsState(DAGSpec()) + write_num = 1 + # Shared state already tracks global ownership/history like flat + # `distribute_tasks!`. Do not pass `registry` as `ownership`: sync/commit + # would fight the single-state history, and the epilogue must then ignore + # the registry too (see `distribute_tasks_hierarchical!`). + ownership = nothing + + for pid in 1:n_partitions + local_procs = partition_procs[pid] + if isempty(partitions[pid]) || isempty(local_procs) + temp_queues[pid] = DataDepsTaskQueue(wait_all_queue; scheduler=similar(queue.scheduler)) + local_scopes[pid] = DefaultScope() + schedules[pid] = Dict{DTask,Processor}() + continue + end + local_scope = UnionScope(map(ExactScope, local_procs)) + local_scopes[pid] = local_scope + locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) + temp_queues[pid] = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) + partition_pairs = DTaskPair[seen_tasks[v] for v in partitions[pid]] + _pdag, schedule = datadeps_build_schedule!(temp_queues[pid].scheduler, partition_pairs, + local_procs, local_scope; region_uids) + schedules[pid] = schedule + end + + topo = try + topological_sort_by_dfs(dag) + catch + collect(vertices(dag)) + end + + with_options(; task_queue=LockedEnqueueQueue(wait_all_queue, queue_lock)) do + for v in topo + pid = vertex_to_partition[v] + local_procs = partition_procs[pid] + isempty(local_procs) && continue + write_num = _schedule_vertex!( + v, pid, temp_queues[pid], shared_state, local_procs, + local_scopes[pid], dag, seen_tasks, vertex_to_partition, + schedules[pid], proc_to_scope_lfus[pid], write_num, ownership) + end + end + + return DataDepsState[shared_state] +end + struct LockedEnqueueQueue <: AbstractTaskQueue inner::AbstractTaskQueue lock::ReentrantLock @@ -901,17 +1007,21 @@ end Main entry point for hierarchical scheduling. Runs the 4-phase pipeline: 1. Parallel aliasing construction 2. DAG construction -3. Partitioning (by worker affinity, or across local procs on one worker) -4. Parallel per-partition scheduling via `distribute_task!` +3. Partitioning (by Distributed-worker / MPI-rank affinity, or across local + procs on one owner) +4. Per-partition scheduling via `distribute_task!` -Each partition runs on its own task and computes its *own* partition-local AOT -schedule over just its processors (see `schedule_partition_full!`); there is no -global AOT pass or global synchronization, which is what allows this to scale. +Each partition computes its *own* partition-local AOT schedule over just its +processors (see `schedule_partition_full!`); there is no global AOT pass or +global synchronization, which is what allows this to scale under Distributed. AOT scheduling here therefore need not match the flat `distribute_tasks!` path exactly. Both drivers use `distribute_task!` for argument preparation and `DataDepsScheduler` dispatch; the old single-worker "batch enqueue with DAG syncdeps only" path is intentionally not used (it skipped `distribute_task!` and broke `ChunkView` / custom schedulers). + +Under uniform execution (MPI), Phase 4 runs sequentially on the root task so +SPMD tag / `MPIRefID` allocation stays deterministic across ranks. """ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) seen_tasks = queue.seen_tasks @@ -919,17 +1029,24 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) return end - # Get the set of all processors - all_procs = Processor[] - scope = get_compute_scope() - for w in procs() - append!(all_procs, get_processors(OSProc(w))) + # Match flat `distribute_tasks!`: acceleration-aware processor enumeration + # (Distributed `OSProc`s, or MPI `MPIOSProc`s → `MPIProcessor`s). + accel = current_acceleration() + accel_procs = filter(procs(Sch.eager_context())) do proc + accel_matches_proc(accel, proc) end + all_procs = unique(vcat([collect(get_processors(gp)) for gp in accel_procs]...)) + select_processors_uniform!(all_procs, accel) + scope = get_compute_scope() filter!(proc->proc_in_scope(proc, scope), all_procs) if isempty(all_procs) throw(Sch.SchedulingException("No processors available, try widening scope")) end - all_scope = UnionScope(map(ExactScope, all_procs)) + if uniform_execution(accel) + for proc in all_procs + check_uniform(proc) + end + end # All in-region task uids, so each partition's local AOT-schedule builder # can tell same-region producers (which it must not `fetch`) apart from @@ -944,12 +1061,12 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) dag = build_dependency_dag(task_metas, arg_to_ainfo, ainfos_overlaps) # Phase 3: Partition the DAG - vertex_to_partition, n_partitions, partition_procs, _multi_worker = + vertex_to_partition, n_partitions, partition_procs, _multi_owner = partition_dag(dag, task_metas, all_procs) # Detect backing chunks shared across partitions in different memory spaces. # These need runtime ownership transfer to avoid split-brain concurrent - # writes; `nothing` when all partitions share one space (single-worker). + # writes; `nothing` when all partitions share one space (single-owner). partition_space = MemorySpace[only(memory_spaces(first(pp))) for pp in partition_procs] registry = build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) @@ -964,31 +1081,48 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) task_submitted = [Base.Event() for _ in 1:length(seen_tasks)] wait_all_queue = get_options(:task_queue) - # Phase 4: parallel per-partition scheduling. Each partition runs on its own - # task, prepares its tasks concurrently, computes its own partition-local - # AOT schedule, and coordinates cross-partition dependencies via the - # `task_submitted` events. Concurrent submission is safe now that task - # futures are backed by `MemPool.DFuture` (see the header note). - partition_states = Vector{DataDepsState}(undef, n_partitions) - try - @sync for pid in 1:n_partitions - Threads.@spawn begin - locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) - with_options(; task_queue=locked_queue) do - partition_states[pid] = schedule_partition_full!( - queue, queue_lock, pid, partitions[pid], - dag, seen_tasks, - partition_procs[pid], vertex_to_partition, - task_submitted, region_uids, registry - ) + # Phase 4: per-partition scheduling. Under Distributed, each partition runs + # on its own task concurrently (see the header note). Under uniform + # execution (MPI), planning stays sequential on the root task in global + # topological order so `to_tag` / generic `MPIRefID` allocation is + # SPMD-deterministic and cross-partition deps cannot deadlock. + mpi_shared_state = uniform_execution(accel) + partition_states = try + if mpi_shared_state + schedule_partitions_sequential!( + queue, queue_lock, partitions, dag, seen_tasks, + partition_procs, vertex_to_partition, region_uids, registry, + wait_all_queue) + else + states = Vector{DataDepsState}(undef, n_partitions) + @sync for pid in 1:n_partitions + Threads.@spawn begin + locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) + with_options(; task_queue=locked_queue) do + states[pid] = schedule_partition_full!( + queue, queue_lock, pid, partitions[pid], + dag, seen_tasks, + partition_procs[pid], vertex_to_partition, + task_submitted, region_uids, registry + ) + end end end + states end catch e rethrow(_unwrap_partition_exception(e)) end - _hierarchical_copy_from_and_free!(partition_states, n_partitions, registry) + # Sequential MPI path returns a one-element shared-state vector and does not + # commit into `registry` (`ownership=nothing` there); copy-back must therefore + # ignore the registry and use the shared state's full history (same as flat). + # Passing the registry would skip every multi-partition chunk: the registry + # path needs `owner_state` (never set), and the private path skips shared keys. + # Distributed keeps per-partition states and needs the registry for coherent + # cross-partition write-back. + epilogue_registry = mpi_shared_state ? nothing : registry + _hierarchical_copy_from_and_free!(partition_states, length(partition_states), epilogue_registry) end function _hierarchical_max_write_num(state::DataDepsState, arg_w::ArgumentWrapper) @@ -1015,6 +1149,11 @@ function _hierarchical_copy_from!(state::DataDepsState, arg_w::ArgumentWrapper, return end +# Deterministic sort key for shared-chunk registry / free-list iteration so +# SPMD ranks enqueue copy-back and free tasks in the same order. +_hierarchical_chunk_sort_key(chunk::Chunk) = (short_name(chunk.space), hash(chunk.handle)) +_hierarchical_chunk_sort_key(chunk) = ("", _identity_hash(chunk)) + function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsState}, n_partitions::Int, registry::Union{SharedChunkRegistry,Nothing}) # 1. Shared chunks: write back from the registry's authoritative owner state @@ -1022,12 +1161,15 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat # max-`write_num` heuristic below (per-partition write_nums are not # comparable across partitions). if registry !== nothing - for (chunk, entry) in registry.entries + shared_chunks = sort!(collect(keys(registry.entries)); by=_hierarchical_chunk_sort_key) + for chunk in shared_chunks + entry = registry.entries[chunk] state = entry.owner_state state === nothing && continue # never written entry.owner_space == entry.origin_space && continue # already in-place at origin - for (arg_w, _space) in state.arg_owner - arg_w.arg === chunk || continue + arg_ws = sort!(ArgumentWrapper[arg_w for (arg_w, _) in state.arg_owner if arg_w.arg === chunk]; + by=arg_w->arg_w.hash) + for arg_w in arg_ws _hierarchical_copy_from!(state, arg_w, _hierarchical_max_write_num(state, arg_w) + 1) end end @@ -1056,12 +1198,15 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat # global writer: an intermediate owner's slot may still be read by a # cross-partition boundary copy recorded in *another* partition's state, # and the final writer transitively depends on all such copies. + # Iteration is sorted for SPMD-uniform free-task tagging/enqueue order. for pid in 1:n_partitions state = partition_states[pid] obj_cache = unwrap(state.ainfo_backing_chunk) write_num = typemax(Int) - 1 - for remote_space in keys(obj_cache.values) - for (ainfo, remote_arg) in obj_cache.values[remote_space] + remote_spaces = sort!(collect(keys(obj_cache.values)); by=short_name) + for remote_space in remote_spaces + space_entries = sort!(collect(obj_cache.values[remote_space]); by=p->hash(p.first)) + for (ainfo, remote_arg) in space_entries if !(ainfo in obj_cache.originals) remote_proc = first(processors(remote_space)) free_scope = ExactScope(remote_proc) @@ -1078,7 +1223,7 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat end end end - Dagger.@spawn scope=free_scope syncdeps=free_syncdeps Dagger.unsafe_free!(remote_arg) + Dagger.@spawn scope=free_scope syncdeps=free_syncdeps tag=datadeps_task_tag() Dagger.unsafe_free!(remote_arg) end end end diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 31611cbc9..0074916ca 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -522,6 +522,13 @@ aliasing_unwrapped(x) = aliasing(aliasing_root(x)) aliasing_unwrapped(x, dep_mod) = aliasing(aliasing_root(x), dep_mod) function aliasing(x::Chunk, T) + # Under uniform execution (MPI), `root_worker_id` is always `myid()` and is + # not a valid owner key -- defer to the acceleration so non-owning ranks + # take the owner-broadcast path instead of a local unwrap. + accel = current_acceleration() + if uniform_execution(accel) + return aliasing(accel, x, T) + end if root_worker_id(x.processor) == myid() return aliasing_unwrapped(unwrap(x), T) end @@ -531,6 +538,10 @@ function aliasing(x::Chunk, T) end end function aliasing(x::Chunk) + accel = current_acceleration() + if uniform_execution(accel) + return aliasing(accel, x, identity) + end if root_worker_id(x.processor) == myid() return aliasing_unwrapped(unwrap(x)) end diff --git a/test/mpi.jl b/test/mpi.jl index d8d25921f..120974761 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -543,10 +543,9 @@ read_only(X) = (sum(X); nothing) r1 = min(1, nranks-1) cross = r1 != 0 # single-rank runs degenerate to zero copies - # Read-only cross-rank arg: copy-in plus region-end write-back. The - # write-back is redundant (origin still holds the bytes) but is retained - # so write-back planning can stay history-based; `arg_current` elision is - # unsound for Krylov `similar` workspaces under Distributed. + # Read-only cross-rank arg: copy-in plus region-end write-back. Write-back + # is not elided via `arg_current`; for a pure `In` the bytes are unchanged, + # but planning still inserts the copy-from. A = ones(4, 4) logs = with_logs() do Dagger.spawn_datadeps() do From f75a4109a54f63ec73f826069ecd0cc0def7a50d Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 19:53:01 -0700 Subject: [PATCH 30/43] DArray/stencil: Fuse halo construction and vectorize the sweep Stencil throughput was dominated by two costs unrelated to the actual arithmetic. `@neighbors` returned a `SubArray` view of a `HaloArray`, so every element access paid `HaloArray`'s region-code dispatch plus the view's index translation, which prevented the sweep loop from vectorizing at all. The sweep now splits each chunk into an interior -- where every neighborhood access provably lands in the center array, so it can run against plain arrays under `@inbounds @simd` with the kernel force-inlined -- and a thin boundary shell that keeps the general path. `@neighbors` returns a direct-offset wrapper rather than a `SubArray`, since the latter's index translation was itself enough to block vectorization. Halos were built by a second task per chunk per expression that copied the whole center chunk into a cached `HaloArray`. For a double-buffered stencil that copy is as much memory traffic as the sweep itself. Halos are now assembled inside the sweeping task from the neighbor chunks passed as read dependencies, wrapping the center in place; for boundary conditions whose halo regions are plain slices of a neighbor (e.g. `Wrap`) they are views rather than materialized copies. Expressions that read and write the same chunks take their neighbors from a snapshot instead. This drops the per-chunk halo cache, its finalizer, and the in-place `load_*_region_into!` fill helpers, which only existed to serve it. Co-authored-by: Cursor --- src/array/stencil.jl | 482 ++++++++++++++++++++--------------------- src/utils/haloarray.jl | 73 ++++++- 2 files changed, 303 insertions(+), 252 deletions(-) diff --git a/src/array/stencil.jl b/src/array/stencil.jl index 0f3af7f22..f99066b56 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -57,18 +57,6 @@ function load_neighbor_region(arr, region_code::NTuple{N,Int}, neigh_dist) where return move(task_processor(), copy(@view arr[start_idx:stop_idx])) end -# In-place variant: load region directly into a pre-allocated destination buffer. -function load_neighbor_region_into!(dest, arr, region_code::NTuple{N,Int}, neigh_dist) where N - validate_neigh_dist(neigh_dist, size(arr)) - start_idx = CartesianIndex(ntuple(N) do i - region_code[i] == -1 ? lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1 : firstindex(arr, i) - end) - stop_idx = CartesianIndex(ntuple(N) do i - region_code[i] == +1 ? firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1 : lastindex(arr, i) - end) - copyto!(dest, @view arr[start_idx:stop_idx]) -end - is_past_boundary(size, idx) = any(ntuple(i -> idx[i] < 1 || idx[i] > size[i], length(size))) ############################################################################# @@ -142,9 +130,6 @@ boundary_transition(::Wrap, idx, size) = load_boundary_region(::Wrap, arr, region_code, neigh_dist, boundary_dims) = load_neighbor_region(arr, region_code, neigh_dist) -load_boundary_region_into!(dest, ::Wrap, arr, region_code, neigh_dist, boundary_dims) = - load_neighbor_region_into!(dest, arr, region_code, neigh_dist) - function boundary_source_index(::Wrap, arr, rc, nd, idx_d, d) if rc == -1 return lastindex(arr, d) - nd + idx_d @@ -180,9 +165,6 @@ function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_d return move(task_processor(), result) end -load_boundary_region_into!(dest, pad::Pad, arr, region_code, neigh_dist, boundary_dims) = - fill!(dest, pad.padval) - # Use edge as source index (value will be overridden by apply_boundary_value) boundary_source_index(::Pad, arr, rc, nd, idx_d, d) = rc == -1 ? firstindex(arr, d) : (rc == +1 ? lastindex(arr, d) : idx_d) @@ -247,10 +229,6 @@ function load_boundary_region(::Clamp, arr, region_code::NTuple{N,Int}, neigh_di return move(task_processor(), result) end -function load_boundary_region_into!(dest, ::Clamp, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where N - Kernel(load_boundary_region_kernel)(Clamp(), dest, arr, region_code, neigh_dist, boundary_dims; ndrange=length(dest)) -end - function boundary_source_index(::Clamp, arr, rc, nd, idx_d, d) if rc == -1 return firstindex(arr, d) @@ -362,18 +340,6 @@ function load_boundary_region(::LinearExtrapolate, arr::AbstractArray{T}, region return move(task_processor(), result) end -function load_boundary_region_into!(dest, ::LinearExtrapolate, arr::AbstractArray{T}, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where {T<:Real,N} - extrap_dim = 0 - for d in 1:N - if boundary_dims[d] && region_code[d] != 0 - extrap_dim = d - break - end - end - nd = get_neigh_dist(neigh_dist, extrap_dim) - Kernel(load_boundary_region_kernel)(LinearExtrapolate(), dest, arr, region_code, neigh_dist, boundary_dims, Val(extrap_dim), Val(nd); ndrange=length(dest)) -end - # Use edge as source index (value will be computed by apply_boundary_value) boundary_source_index(::LinearExtrapolate, arr, rc, nd, idx_d, d) = rc == -1 ? firstindex(arr, d) : (rc == +1 ? lastindex(arr, d) : idx_d) @@ -476,41 +442,6 @@ function load_boundary_region(::Reflect{Symm}, arr, region_code::NTuple{N,Int}, return region end -function load_boundary_region_into!(dest, ::Reflect{Symm}, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where {N, Symm} - flipped_code = ntuple(N) do i - (region_code[i] != 0 && boundary_dims[i]) ? -region_code[i] : region_code[i] - end - skip = Symm ? 0 : 1 - start_idx = CartesianIndex(ntuple(N) do i - needs_skip = boundary_dims[i] && region_code[i] != 0 - actual_skip = needs_skip ? skip : 0 - if flipped_code[i] == -1 - lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1 - actual_skip - elseif flipped_code[i] == +1 - firstindex(arr, i) + actual_skip - else - firstindex(arr, i) - end - end) - stop_idx = CartesianIndex(ntuple(N) do i - needs_skip = boundary_dims[i] && region_code[i] != 0 - actual_skip = needs_skip ? skip : 0 - if flipped_code[i] == +1 - firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1 + actual_skip - elseif flipped_code[i] == -1 - lastindex(arr, i) - actual_skip - else - lastindex(arr, i) - end - end) - copyto!(dest, @view arr[start_idx:stop_idx]) - for i in 1:N - GPUArraysCore.@allowscalar if region_code[i] != 0 && boundary_dims[i] - reverse!(dest, dims=i) - end - end -end - function boundary_source_index(::Reflect{Symm}, arr, rc, nd, idx_d, d) where Symm skip = Symm ? 0 : 1 if rc == -1 @@ -641,10 +572,6 @@ function load_boundary_region(boundary::Tuple, arr, region_code::NTuple{N,Int}, return move(task_processor(), result) end -function load_boundary_region_into!(dest, boundary::Tuple, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where N - Kernel(load_boundary_region_kernel)(boundary, dest, arr, region_code, neigh_dist, boundary_dims; ndrange=length(dest)) -end - ############################################################################# # Chunk Selection and Halo Building ############################################################################# @@ -854,97 +781,107 @@ function select_neighborhood_info(chunks, idx, neigh_dist, boundary) return region_metadata, neighbor_chunks end -# Per-thread cache: WeakKeyDict{DArray, Dict{(chunk_idx, halo_width), HaloArray}}. -# WeakKeyDict is used for the outer level so that the cache does not hold a strong reference -# to the source DArray — allowing its GC finalizer to fire when user code drops its last -# reference (see below). Using chunk_idx as part of the inner key ensures that within one -# DArray, every chunk has its own dedicated buffer — so if a single worker thread processes -# multiple same-shaped chunks in the same iteration sequentially, each gets a distinct -# HaloArray and there is no aliasing with a concurrently running inner-stencil task. -# Filling a cached buffer in-place is safe because spawn_datadeps blocks until all inner -# tasks complete before the next iteration's build_halo_consolidated calls run. -const HALO_ARRAY_CACHE = TaskLocalValue{WeakKeyDict{Any,Dict{Any,Any}}}(()->WeakKeyDict{Any,Dict{Any,Any}}()) - -# Called on the main task (outside any Dagger.@spawn) to get or create the per-DArray inner -# cache dict. Keeping all WeakKeyDict operations here — rather than inside spawned tasks — -# avoids every pitfall of passing a DArray to @spawn (Dagger resolving it to a ROCArray on -# GPU, serialization producing a copy with a different objectid under Distributed, etc.). -# A finalizer registered on first encounter frees all cached HaloArrays when the DArray is -# collected; the WeakKeyDict ensures the cache itself does not prevent that collection. -function get_halo_inner_cache(read_darray) - outer_cache = HALO_ARRAY_CACHE[] - if !haskey(outer_cache, read_darray) - inner_cache = Dict{Any,Any}() - outer_cache[read_darray] = inner_cache - finalizer(read_darray) do _ - # GC finalizers cannot yield (no task switches allowed), but unsafe_free! on - # GPU/distributed Chunks calls remotecall_fetch which acquires locks and yields. - # Defer cleanup to a scheduled task so it runs outside the finalizer context. - errormonitor(Threads.@spawn begin - for halo in values(inner_cache) - try - unsafe_free!(halo) - catch e - # fill_halo_inplace! runs inside spawn_datadeps and may be placed - # on a Distributed worker; if that worker has already exited by - # cleanup time, skip the free. Use nameof to work with both - # Distributed and DistributedNext backends. - string(nameof(typeof(e))) == "ProcessExitedException" || rethrow(e) - end - end - end) - end - end - return outer_cache[read_darray] +function build_halo(neigh_dist, boundary, center, all_halos...; own_center::Bool=false) + N = ndims(center) + expected_halos = 3^N - 1 + @assert length(all_halos) == expected_halos "Halo mismatch: N=$N expected $expected_halos halos, got $(length(all_halos))" + center_data = own_center ? copy(center) : center + return HaloArray(center_data, (all_halos...,), ntuple(i->get_neigh_dist(neigh_dist, i), N); own_center) end -# Cache-miss path: allocate an empty HaloArray using similar() so the buffers are created -# on the same device as the center chunk (GPU or CPU). No filling — that happens inside -# spawn_datadeps via fill_halo_inplace!. -function alloc_halo(neigh_dist, center) - N = ndims(center) - center_size = size(center) - halo_width = ntuple(i -> get_neigh_dist(neigh_dist, i), N) - codes = all_region_codes(Val(N)) - halos = ntuple(length(codes)) do i - region_size = halo_region_size(center_size, halo_width, codes[i]) - similar(center, region_size...) +############################################################################# +# Fused halo construction +############################################################################# +# +# The `@stencil` fast path builds a chunk's `HaloArray` inside the same task that +# sweeps it, from the neighboring chunks passed as `In` dependencies. Compared to +# filling a cached `HaloArray` in a separate task this avoids (a) a second task per +# chunk per expression and (b) copying the whole center chunk, which for a +# double-buffered stencil is as much memory traffic as the sweep itself. + +struct ViewHalos end +struct CopyHalos end + +# Whether a boundary condition's halo regions are plain slices of a neighboring +# chunk, in which case they can be `view`s rather than materialized copies. +@inline halo_build_style(@nospecialize(boundary)) = CopyHalos() +@inline halo_build_style(::Wrap) = ViewHalos() +@inline halo_build_style(boundary::Tuple) = _combine_halo_style(map(halo_build_style, boundary)) +@inline _combine_halo_style(::Tuple{}) = ViewHalos() +@inline _combine_halo_style(styles::Tuple) = + first(styles) isa ViewHalos ? _combine_halo_style(Base.tail(styles)) : CopyHalos() + +@inline function neighbor_region_view(arr, region_code::NTuple{N,Int}, neigh_dist) where N + ranges = ntuple(Val(N)) do i + nd = get_neigh_dist(neigh_dist, i) + start_i = region_code[i] == -1 ? lastindex(arr, i) - nd + 1 : firstindex(arr, i) + stop_i = region_code[i] == +1 ? firstindex(arr, i) + nd - 1 : lastindex(arr, i) + start_i:stop_i end - return HaloArray(similar(center), halos, halo_width; own_center=true) + return view(arr, ranges...) end -# Cache-hit path: fill an existing HaloArray in-place and return it. No cache operations — -# the HaloArray was looked up and passed by the main task before spawning. -function fill_halo_inplace!(halo::HaloArray, neigh_dist, boundary, center, region_metadata, neighbor_chunks...) - expected_halos = length(region_metadata) - @assert length(neighbor_chunks) == expected_halos - copyto!(halo.center, center) - for i in 1:expected_halos +@inline function _build_fused_halos(::ViewHalos, neigh_dist, boundary, region_metadata, + neighbor_chunks::NTuple{NH,Any}) where NH + return ntuple(Val(NH)) do i + region_code, _, _ = region_metadata[i] + neighbor_region_view(neighbor_chunks[i], region_code, neigh_dist) + end +end +@inline function _build_fused_halos(::CopyHalos, neigh_dist, boundary, region_metadata, + neighbor_chunks::NTuple{NH,Any}) where NH + return ntuple(Val(NH)) do i region_code, is_boundary, boundary_dims = region_metadata[i] chunk = neighbor_chunks[i] if is_boundary - load_boundary_region_into!(halo.halos[i], boundary, chunk, region_code, neigh_dist, boundary_dims) + load_boundary_region(boundary, chunk, region_code, neigh_dist, boundary_dims) else - load_neighbor_region_into!(halo.halos[i], chunk, region_code, neigh_dist) + load_neighbor_region(chunk, region_code, neigh_dist) end end - return halo end -function build_halo(neigh_dist, boundary, center, all_halos...; own_center::Bool=false) - N = ndims(center) - expected_halos = 3^N - 1 - @assert length(all_halos) == expected_halos "Halo mismatch: N=$N expected $expected_halos halos, got $(length(all_halos))" - center_data = own_center ? copy(center) : center - return HaloArray(center_data, (all_halos...,), ntuple(i->get_neigh_dist(neigh_dist, i), N); own_center) +""" + build_fused_halo(neigh_dist, boundary, region_metadata, center, neighbor_chunks...) + +Wraps `center` (used in place, not copied) plus halo regions taken from +`neighbor_chunks` into a `HaloArray`. `region_metadata` is the per-region +`(region_code, is_boundary, boundary_dims)` triple precomputed on the submitting +task by `select_neighborhood_info`. +""" +function build_fused_halo(neigh_dist, boundary, region_metadata, + center::AbstractArray{T,N}, neighbor_chunks::Vararg{Any,NH}) where {T,N,NH} + validate_neigh_dist(neigh_dist, size(center)) + halo_width = ntuple(i -> get_neigh_dist(neigh_dist, i), Val(N)) + halos = _build_fused_halos(halo_build_style(boundary), neigh_dist, boundary, + region_metadata, neighbor_chunks) + return HaloArray(center, halos, halo_width; own_center=false) end -function load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N} - start_idx = idx - CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr))) - stop_idx = idx + CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr))) - return @view arr[start_idx:stop_idx] +""" + stencil_source_chunks(read_chunks, write_chunks) -> chunks + +The chunk array a neighborhood read should be taken from. Normally that is +`read_chunks` itself, but when the expression writes back into the same chunks it +reads (`A[idx] = f(@neighbors(A[idx]))`), neighbors must come from a snapshot +taken before any chunk is overwritten. +""" +function stencil_source_chunks(read_chunks, write_chunks) + write_set = IdDict{Any,Nothing}() + for chunk in write_chunks + write_set[chunk] = nothing + end + overlaps = any(chunk -> haskey(write_set, chunk), read_chunks) + overlaps || return read_chunks + snapshot_tasks = map(chunk -> Dagger.@spawn(name="stencil_snapshot", copy(chunk)), read_chunks) + return map(task -> fetch(task; raw=true), snapshot_tasks) end +@inline load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N} = + StencilNeighborhood(arr, idx, arr.halo_width) +@inline load_neighborhood(arr::HaloInterior{T,N}, idx) where {T,N} = + StencilNeighborhood(arr.parent, idx, arr.halo_width) + function inner_stencil!(f, output, read_vars) processor = task_processor() inner_stencil_proc!(processor, f, output, read_vars) @@ -954,8 +891,77 @@ end # Non-KA (for CPUs) function inner_stencil_proc!(::ThreadProc, f, output, read_vars) - for idx in CartesianIndices(output) - f(idx, output, read_vars) + cpu_stencil_sweep!(f, output, read_vars) + return +end + +# Widest halo required by any neighborhood variable, per dimension. Zero in every +# dimension means no variable is accessed through a halo, so the whole chunk can +# take the fast path. +@inline _widen_halo(w::NTuple{N,Int}, ::Any) where N = w +@inline _widen_halo(w::NTuple{N,Int}, A::HaloArray{T,N}) where {T,N} = + ntuple(i -> max(w[i], A.halo_width[i]), Val(N)) +@inline _max_halo_width(vars::Tuple{}, w::NTuple{N,Int}) where N = w +@inline _max_halo_width(vars::Tuple, w::NTuple{N,Int}) where N = + _max_halo_width(Base.tail(vars), _widen_halo(w, first(vars))) + +# Interior stand-ins: HaloArrays become HaloInterior (direct center indexing), +# everything else is passed through untouched. +@inline _interior_var(A::HaloArray) = HaloInterior(A.center, A.halo_width) +@inline _interior_var(A) = A + +""" + cpu_stencil_sweep!(f, output, read_vars) + +Applies `f` at every index of `output`, splitting the chunk into an interior +where all neighborhood accesses land in the center array, and a thin boundary +shell where they may reach into halos. + +The split is what makes stencils fast: the interior (essentially the whole chunk) +runs against plain arrays under `@inbounds @simd` with `f` force-inlined, while +the generic, branch-heavy `HaloArray` path is confined to the shell. +""" +function cpu_stencil_sweep!(f::F, output::AbstractArray{T,N}, read_vars::NamedTuple) where {F,T,N} + w = _max_halo_width(values(read_vars), ntuple(_ -> 0, Val(N))) + ax = axes(output) + + interior_ax = ntuple(Val(N)) do i + (first(ax[i]) + w[i]):(last(ax[i]) - w[i]) + end + if any(isempty, interior_ax) + # Halos are as wide as the chunk itself; there is no interior to split off. + for idx in CartesianIndices(output) + @inline f(idx, output, read_vars) + end + return + end + + interior_vars = map(_interior_var, read_vars) + @inbounds @simd for idx in CartesianIndices(interior_ax) + @inline f(idx, output, interior_vars) + end + + all(iszero, w) && return + + # Boundary shell, decomposed into 2N disjoint slabs: for dimension `d`, the + # low and high slabs span the full extent in dimensions after `d` and only the + # interior extent in dimensions before `d`. + for d in 1:N + for low_side in (true, false) + slab_ax = ntuple(Val(N)) do i + if i < d + interior_ax[i] + elseif i == d + low_side ? (first(ax[i]):(first(ax[i]) + w[i] - 1)) : + ((last(ax[i]) - w[i] + 1):last(ax[i])) + else + first(ax[i]):last(ax[i]) + end + end + for idx in CartesianIndices(slab_ax) + @inline f(idx, output, read_vars) + end + end end return end @@ -1136,130 +1142,110 @@ macro stencil(orig_ex) end new_inner_f = :(($inner_index_var, $inner_write_var, $inner_vars)->$new_inner_ex_body) actual_read_vars = filter(v -> (v != write_var) || (v in keys(neighborhoods)), collect(read_vars)) - new_inner_ex = quote - $inner_vars = (;$(actual_read_vars...)) - $inner_stencil!($new_inner_f, $inner_write_var, $inner_vars) - end - inner_fn = Expr(:->, Expr(:tuple, Expr(:parameters, inner_write_var, actual_read_vars...)), new_inner_ex) - # 2a. For each neighborhood read_var: pre-compute region metadata on the main task - # (to avoid passing DArrays into @spawn), and for cache misses spawn alloc_halo - # outside spawn_datadeps so that HaloArray buffers are allocated on the correct - # device (GPU or CPU) before filling. + # 2a. For each neighborhood read_var, pre-compute on the main task (so no DArray + # is ever passed into @spawn) the chunk array to read neighbors from and, per + # chunk, the region metadata plus the neighboring chunks themselves. # - # fill_halo_inplace! is spawned *inside* spawn_datadeps (step 2c) with explicit - # Read deps on the DArray chunks, so Datadeps automatically enforces the ordering - # between fill tasks and inner stencil tasks — including the cross-chunk case where - # fill_task[C] reads chunk[N] as a neighbor while inner_stencil[N] writes chunk[N]. - # This replaces the explicit wait-barrier that was needed when fills ran outside. - cache_sym_map = Dict{Symbol, NamedTuple}() + # `stencil_source_chunks` substitutes a snapshot when the expression writes back + # into the chunks it reads (`A[idx] = f(@neighbors(A[idx]))`); everywhere else it + # is the identity, and the neighbor chunks are read directly. + neigh_sym_map = Dict{Symbol, NamedTuple}() for read_var in read_vars if read_var in keys(neighborhoods) neigh_dist, boundary = neighborhoods[read_var] - @gensym region_info_table fill_tasks region_meta neighbor_cks halo_width_var inner_cache_var cache_key_var - cache_sym_map[read_var] = (; fill_tasks, halo_width_var, inner_cache_var, cache_key_var, region_info_table) - # Validate and compute halo_width on the main task so errors are immediate. + @gensym region_info_table src_chunks region_meta neighbor_cks + neigh_sym_map[read_var] = (; region_info_table, src_chunks) push!(final_ex.args, :($validate_neigh_dist($neigh_dist, ndims($read_var)))) - push!(final_ex.args, :($halo_width_var = ntuple(i -> $get_neigh_dist($neigh_dist, i), ndims($read_var)))) - push!(final_ex.args, :($inner_cache_var = $get_halo_inner_cache($read_var))) - # Pre-compute region metadata for every chunk on the main task. - push!(final_ex.args, :($region_info_table = Array{Any}(undef, size($chunks($read_var))))) - push!(final_ex.args, :($fill_tasks = Array{$DTask}(undef, size($chunks($read_var))))) + push!(final_ex.args, :($src_chunks = $stencil_source_chunks($chunks($read_var), $chunks($write_var)))) + push!(final_ex.args, :($region_info_table = Array{Any}(undef, size($src_chunks)))) push!(final_ex.args, quote - for $chunk_idx in $CartesianIndices($chunks($read_var)) - ($region_meta, $neighbor_cks) = $select_neighborhood_info($chunks($read_var), $chunk_idx, $neigh_dist, $boundary) + for $chunk_idx in $CartesianIndices($src_chunks) + ($region_meta, $neighbor_cks) = $select_neighborhood_info($src_chunks, $chunk_idx, $neigh_dist, $boundary) $region_info_table[$chunk_idx] = (tuple($region_meta...), $neighbor_cks) - $cache_key_var = ($chunk_idx, $halo_width_var) - if !haskey($inner_cache_var, $cache_key_var) - # Cache miss: allocate empty buffers on the correct device. - $inner_cache_var[$cache_key_var] = Dagger.@spawn name="stencil_alloc_halo" $alloc_halo($neigh_dist, $chunks($read_var)[$chunk_idx]) - end end end) end end - # 2b. Build the inner-stencil @spawn expression. Neighborhood deps reference - # fill_tasks[chunk_idx] populated by the fill loop (see 2c below). - deps_ex = Any[] - if write_var in read_vars - push!(deps_ex, Expr(:kw, inner_write_var, :($ReadWrite($chunks($write_var)[$chunk_idx])))) - else - push!(deps_ex, Expr(:kw, inner_write_var, :($Write($chunks($write_var)[$chunk_idx])))) - end + # 2b. Build the per-chunk task. Neighborhood variables are passed as their own + # chunk followed by their `3^N - 1` neighboring chunks (a runtime-length group, + # hence the positional splat); the task assembles the HaloArray itself. + prologue_exs = Expr[] # runs per chunk on the submitting task + arg_exs = Any[] # positional args of the spawned task + unpack_exs = Expr[] # runs inside the task, binds each read var + local_vars = Any[] # task-local binding holding each read var's chunk data + @gensym task_args arg_offset for read_var in actual_read_vars - if read_var in keys(neighborhoods) - syms = cache_sym_map[read_var] - # Reference fill_tasks[chunk_idx] — populated by the fill loop before the - # stencil loop runs, so it is always assigned when this expression executes. - push!(deps_ex, Expr(:kw, read_var, :($Read($(syms.fill_tasks)[$chunk_idx])))) - else - if read_var != write_var - push!(deps_ex, Expr(:kw, read_var, :($Read($chunks($read_var)[$chunk_idx])))) - end - end - end - inner_spawn_ex = :(Dagger.@spawn name="stencil_inner_fn" $inner_fn(;$(deps_ex...))) - - # Build the fill @spawn expression — one per neighborhood read_var. - fill_spawn_exs = Expr[] - for read_var in read_vars + # N.B. Bind into a gensym rather than `read_var` itself: the task closure is + # nested inside the user's scope, so assigning `read_var` there would rebind + # the user's DArray variable instead of creating a task-local binding. + local_var = gensym(read_var) if read_var in keys(neighborhoods) neigh_dist, boundary = neighborhoods[read_var] - syms = cache_sym_map[read_var] - @gensym _rmt _nck _rn - push!(fill_spawn_exs, quote - ($_rmt, $_nck) = $(syms.region_info_table)[$chunk_idx] - $_rn = map($Read, $_nck) - $(syms.fill_tasks)[$chunk_idx] = Dagger.@spawn name="stencil_fill_halo" $fill_halo_inplace!( - $ReadWrite($(syms.inner_cache_var)[($chunk_idx, $(syms.halo_width_var))]), - $neigh_dist, $boundary, - $Read($chunks($read_var)[$chunk_idx]), - $_rmt, - $_rn... - ) + syms = neigh_sym_map[read_var] + @gensym region_meta neighbor_cks nneighbors + push!(prologue_exs, quote + ($region_meta, $neighbor_cks) = $(syms.region_info_table)[$chunk_idx] + $nneighbors = length($neighbor_cks) + end) + push!(arg_exs, :($Read($(syms.src_chunks)[$chunk_idx]))) + push!(arg_exs, Expr(:..., :(map($Read, $neighbor_cks)))) + push!(unpack_exs, quote + $local_var = $build_fused_halo($neigh_dist, $boundary, $region_meta, + $task_args[$arg_offset + 1], + $task_args[($arg_offset + 2):($arg_offset + 1 + $nneighbors)]...) + $arg_offset += 1 + $nneighbors + end) + elseif read_var != write_var + push!(arg_exs, :($Read($chunks($read_var)[$chunk_idx]))) + push!(unpack_exs, quote + $local_var = $task_args[$arg_offset + 1] + $arg_offset += 1 end) end + push!(local_vars, local_var) + end + write_dep_ex = if write_var in read_vars + :($ReadWrite($chunks($write_var)[$chunk_idx])) + else + :($Write($chunks($write_var)[$chunk_idx])) end - # 2c. Each expression gets its own spawn_datadeps region with TWO separate loops: - # first all fills, then all stencils. This ordering is critical when write_var is - # also a neighborhood read_var: submitting all fill tasks before any stencil tasks - # ensures Datadeps sees fill[C]'s Read(chunk[N]) BEFORE stencil[N]'s - # ReadWrite(chunk[N]), so stencil[N] is correctly ordered after fill[C]. - # (An interleaved single loop would register stencil[C]'s write to chunk[C] before - # fill[C+1] reads chunk[C], inverting the dependency.) - fill_loop_body = Expr(:block, fill_spawn_exs...) + # The kernel body may mention an accessed variable bare (e.g. `length(A)`), which + # would otherwise capture the user's DArray into the spawned closure -- illegal, + # since Datadeps analyzes the closure's captures for aliasing. Shadow each + # accessed name with a parameter bound to that chunk's local data, so the body + # sees the chunk (with halos, where applicable) instead. + shadow_params = copy(actual_read_vars) + shadow_args = copy(local_vars) + if !(write_var in actual_read_vars) + push!(shadow_params, write_var) + push!(shadow_args, inner_write_var) + end + shadow_body = quote + $inner_vars = (;$([Expr(:kw, v, v) for v in actual_read_vars]...)) + $inner_stencil!($new_inner_f, $inner_write_var, $inner_vars) + end + shadow_fn = Expr(:->, Expr(:tuple, shadow_params...), shadow_body) + inner_fn_body = quote + $arg_offset = 0 + $(unpack_exs...) + $shadow_fn($(shadow_args...)) + end + inner_fn = Expr(:->, Expr(:tuple, inner_write_var, Expr(:..., task_args)), inner_fn_body) + inner_spawn_ex = Expr(:block, prologue_exs..., + :(Dagger.@spawn name="stencil_inner_fn" $inner_fn($write_dep_ex, $(arg_exs...)))) + + # 2c. One spawn_datadeps region per expression, one task per chunk. Because the + # region blocks until all of its tasks finish, the next expression's tasks are + # only submitted once this expression has been applied everywhere, which is what + # gives `@stencil` its "all at once" semantics. push!(final_ex.args, :(Dagger.spawn_datadeps() do - for $chunk_idx in $CartesianIndices($chunks($write_var)) - $fill_loop_body - end for $chunk_idx in $CartesianIndices($chunks($write_var)) $inner_spawn_ex end end)) - - # 2d. After spawn_datadeps completes, resolve any alloc DTasks in the cache to - # their origin Chunks. Do *not* store `fetch(fill_task)`: Datadeps may have run - # fill on a libc-backed remote slot copy and then `unsafe_free!`'d that copy at - # region end, so the fill task's return value can dangle. The origin Chunk (from - # alloc_halo / a prior iteration) receives the updated data via remainder sync - # and is what must stay in the cache for the next iteration. - for read_var in read_vars - if read_var in keys(neighborhoods) - syms = cache_sym_map[read_var] - @gensym cache_entry - push!(final_ex.args, quote - for $chunk_idx in $CartesianIndices($chunks($read_var)) - $(syms.cache_key_var) = ($chunk_idx, $(syms.halo_width_var)) - $cache_entry = $(syms.inner_cache_var)[$(syms.cache_key_var)] - if $cache_entry isa $DTask - $(syms.inner_cache_var)[$(syms.cache_key_var)] = fetch($cache_entry; raw=true) - end - end - end) - end - end end # 3. Return last allocated var if applicable diff --git a/src/utils/haloarray.jl b/src/utils/haloarray.jl index a9bb4a0d4..4e5d74c62 100644 --- a/src/utils/haloarray.jl +++ b/src/utils/haloarray.jl @@ -152,12 +152,77 @@ end Base.IndexStyle(::Type{<:HaloArray}) = IndexCartesian() -# GPU-friendly reductions for SubArray views of HaloArray. +############################################################################# +# HaloInterior +############################################################################# + +""" + HaloInterior(parent, halo_width) + +A stand-in for a `HaloArray` used while sweeping the region of a chunk whose +whole neighborhood lies inside the center. Indexing goes straight to `parent`, +skipping the region-code dispatch that `HaloArray` must perform, which is what +lets the interior sweep vectorize. +""" +struct HaloInterior{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N} + parent::A + halo_width::NTuple{N,Int} +end + +Base.size(A::HaloInterior) = size(A.parent) +Base.axes(A::HaloInterior) = axes(A.parent) +Base.IndexStyle(::Type{<:HaloInterior}) = IndexCartesian() +@inline function Base.getindex(A::HaloInterior{T,N}, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(A.parent, I...) + return @inbounds A.parent[I...] +end +@inline function Base.setindex!(A::HaloInterior{T,N}, value, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(A.parent, I...) + return @inbounds A.parent[I...] = value +end + +############################################################################# +# StencilNeighborhood +############################################################################# + +""" + StencilNeighborhood(parent, center, halo_width) + +The value produced by `@neighbors`: a `(2w+1)^N` window of `parent` centered on +`center`, with 1-based indices. This is a plain offset-index wrapper rather than +a `SubArray` because the latter's index translation blocks vectorization of the +sweep loop (roughly a 5x difference on the interior). +""" +struct StencilNeighborhood{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N} + parent::A + center::CartesianIndex{N} + halo_width::NTuple{N,Int} +end + +Base.size(nb::StencilNeighborhood{T,N}) where {T,N} = + ntuple(i -> 2 * nb.halo_width[i] + 1, Val(N)) +Base.axes(nb::StencilNeighborhood{T,N}) where {T,N} = + ntuple(i -> Base.OneTo(2 * nb.halo_width[i] + 1), Val(N)) +Base.IndexStyle(::Type{<:StencilNeighborhood}) = IndexCartesian() + +@inline function _neighborhood_index(nb::StencilNeighborhood{T,N}, I::NTuple{N,Int}) where {T,N} + return nb.center + CartesianIndex(ntuple(i -> I[i] - nb.halo_width[i] - 1, Val(N))) +end +@inline function Base.getindex(nb::StencilNeighborhood{T,N}, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(nb, I...) + return @inbounds nb.parent[_neighborhood_index(nb, I)] +end +@inline function Base.setindex!(nb::StencilNeighborhood{T,N}, value, I::Vararg{Int,N}) where {T,N} + Base.@boundscheck checkbounds(nb, I...) + return @inbounds nb.parent[_neighborhood_index(nb, I)] = value +end + +# GPU-friendly reductions over neighborhoods. # The standard reduction path uses _foldl_impl/iterate which produces Union # types that SPIR-V and other GPU compilers can't handle, and passes through # keyword-argument forwarding that triggers dynamic dispatch on GPU. # These overrides use CartesianIndices iteration which compiles cleanly. -@inline function Base.mapreduce(f::F, op::OP, A::SubArray{T,N,<:HaloArray}) where {F,OP,T,N} +@inline function Base.mapreduce(f::F, op::OP, A::StencilNeighborhood{T,N}) where {F,OP,T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = f(@inbounds A[first_idx]) @inbounds for idx in CartesianIndices(A) @@ -166,7 +231,7 @@ Base.IndexStyle(::Type{<:HaloArray}) = IndexCartesian() end return result end -@inline function Base.sum(A::SubArray{T,N,<:HaloArray}) where {T,N} +@inline function Base.sum(A::StencilNeighborhood{T,N}) where {T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = @inbounds A[first_idx] @inbounds for idx in CartesianIndices(A) @@ -175,7 +240,7 @@ end end return result end -@inline function Base.sum(f::F, A::SubArray{T,N,<:HaloArray}) where {F,T,N} +@inline function Base.sum(f::F, A::StencilNeighborhood{T,N}) where {F,T,N} first_idx = CartesianIndex(ntuple(d -> firstindex(A, d), Val(N))) result = f(@inbounds A[first_idx]) @inbounds for idx in CartesianIndices(A) From 4c3a3e0bccb1a6dcaf65af2bcaba053f3b7183ed Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 19:53:08 -0700 Subject: [PATCH 31/43] Sch: Let reusable pool tasks migrate off the submitting thread Clearing `sticky` is not enough to make a task migratable. A hand-built `@task` has no thread pool assigned, and `Base.enq_work` routes such a task onto the *current* thread's work queue rather than the shared multi-queue, so it cannot be picked up until the thread that scheduled it yields. `Threads.@spawn` sets both; our task pools only cleared `sticky`. The effect was that fire tasks piled up behind whichever thread was submitting work. A submitter that runs a long stretch without yielding -- Datadeps planning a whole region, say -- held off every fire task it had just created, so nothing began executing until planning had finished. Setting the thread pool as well is what allows execution to overlap submission. Co-authored-by: Cursor --- src/sch/Sch.jl | 14 ++++++++++---- src/utils/tasks.jl | 20 ++++++++++++++++++++ 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 6a353e28c..15edd343c 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -1198,15 +1198,21 @@ Base.hash(task::TaskSpec, h::UInt) = hash(task.thunk_id, hash(TaskSpec, h)) sch_handle, state.uid)) end + # N.B. These must be migratable. Otherwise they pile up on whichever thread + # is submitting work: a submitter that runs a long stretch of + # `schedule_one!` without yielding (e.g. Datadeps planning a whole region) + # holds off every fire task it just created, so nothing starts executing + # until planning finishes. Letting them migrate is what allows execution to + # overlap submission. if !isempty(to_send) if root_worker_id(gproc) == myid() - @reusable_tasks :fire_tasks!_task_cache 32 _->nothing "fire_tasks!" FireTaskSpec(proc, state.chan, to_send) + @reusable_tasks :fire_tasks!_task_cache 32 Dagger.set_task_migratable! "fire_tasks!" FireTaskSpec(proc, state.chan, to_send) else # N.B. We don't batch these because we might get a deserialization # error due to something not being defined on the worker, and then we don't # know which task failed. for task_spec in to_send - @reusable_tasks :fire_tasks!_task_cache 32 _->nothing "fire_tasks!" FireTaskSpec(proc, state.chan, task_spec) + @reusable_tasks :fire_tasks!_task_cache 32 Dagger.set_task_migratable! "fire_tasks!" FireTaskSpec(proc, state.chan, task_spec) end end end @@ -1565,7 +1571,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re if tid !== nothing Dagger.set_task_tid!(t, tid) else - t.sticky = false + Dagger.set_task_migratable!(t) end end "thunk $thunk_id" DoTaskSpec(to_proc, return_queue, task, cancel_token) @@ -1579,7 +1585,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re if tid !== nothing Dagger.set_task_tid!(proc_run_task, tid) else - proc_run_task.sticky = false + Dagger.set_task_migratable!(proc_run_task) end return errormonitor_tracked("processor $to_proc", schedule(proc_run_task)) end diff --git a/src/utils/tasks.jl b/src/utils/tasks.jl index ddd8da2ee..beab71cd7 100644 --- a/src/utils/tasks.jl +++ b/src/utils/tasks.jl @@ -1,3 +1,23 @@ +""" + set_task_migratable!(task) -> task + +Allow `task` to run on any thread of the default pool. + +Clearing `sticky` is not sufficient on its own. A hand-built `@task` has no +thread pool assigned, and `Base.enq_work` routes such a task onto the *current* +thread's work queue rather than the shared multi-queue, so it cannot be picked +up until the thread that scheduled it yields. `Threads.@spawn` sets both; task +pools built with `@task` must do the same or every task they spawn is pinned, +in effect, to whichever thread happened to schedule it. +""" +function set_task_migratable!(task::Task) + task.sticky = false + @static if isdefined(Base.Threads, :_spawn_set_thrpool) + Base.Threads._spawn_set_thrpool(task, :default) + end + return task +end + function set_task_tid!(task::Task, tid::Integer) task.sticky = true ctr = 0 From add5950c315e48787a2f55037c2d6538c18bbe60 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 19:53:16 -0700 Subject: [PATCH 32/43] Sch: Estimate transfer costs per worker, not per processor `estimate_task_costs!` scanned every chunk once per candidate processor, but chunks are located per worker, so the result depends only on the processor's parent. With `procs` commonly being every thread of a single worker, that scan (and the transfer-rate lock acquisition alongside it) was repeated identically once per thread. Both are now done once per parent. Also skip the final sort when every processor costs the same, which is the usual single-worker case; the shuffle just above already gives the arbitrary-but-fair order that sorting would leave us with. Co-authored-by: Cursor --- src/sch/util.jl | 58 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 18 deletions(-) diff --git a/src/sch/util.jl b/src/sch/util.jl index 9057f5740..9242c7bbb 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -637,6 +637,7 @@ function estimate_task_costs(state, procs, task; sig=nothing) return sorted_procs, costs end const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) +const EMPTY_TRANSFER_RATES = Dict{Processor,UInt64}() @reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; sig=nothing) # Find all Chunks @@ -654,28 +655,47 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) end est_time_util = lock(state.signature_time_cost) do stc; get(stc, sig, 1000^3); end - # Estimate total cost for executing this task on each candidate processor + # Estimate network transfer cost per *parent* processor. Chunks are located + # per worker, so this depends only on `get_parent(proc)`, and `procs` is + # commonly every thread of a single worker -- computing it per processor + # would redo identical work (a scan of every chunk) once per thread. + # N.B. We treat same-worker transfers as having zero transfer cost + # TODO: For non-Chunk, model cost from scheduler to worker + # TODO: Measure and model processor move overhead + tx_costs = @reusable_dict :estimate_task_costs_tx_costs Processor Float64 OSProc() 0.0 8 + tx_costs_cleanup = @reuse_defer_cleanup empty!(tx_costs) for proc in procs gproc = get_parent(proc) + haskey(tx_costs, gproc) && continue chunks_filt = Iterators.filter(c->get_parent(processor(c)) != gproc, chunks) + tx_costs[gproc] = impute_sum(datasize(chunk) for chunk in chunks_filt) + end + chunks_cleanup() - # Estimate network transfer costs based on data size - # N.B. We treat same-worker transfers as having zero transfer cost - # TODO: For non-Chunk, model cost from scheduler to worker - # TODO: Measure and model processor move overhead - tx_cost = impute_sum(datasize(chunk) for chunk in chunks_filt) - - # Add fixed cost for cross-worker task transfer (esimated at 1ms) - # TODO: Actually estimate/benchmark this - task_xfer_cost = root_worker_id(gproc) != myid() ? 1_000_000 : 0 # 1ms - pid = Dagger.root_worker_id(gproc) - - tx_rate = lock(state.worker_transfer_rate) do wtr - get(get(wtr, pid, Dict{Processor,UInt64}()), proc, DEFAULT_TRANSFER_RATE) + # Estimate total cost for executing this task on each candidate processor. + # The transfer-rate table is taken once rather than once per processor. + all_equal = true + lock(state.worker_transfer_rate) do wtr + local first_cost = 0.0 + for (idx, proc) in enumerate(procs) + gproc = get_parent(proc) + pid = Dagger.root_worker_id(gproc) + + # Add fixed cost for cross-worker task transfer (esimated at 1ms) + # TODO: Actually estimate/benchmark this + task_xfer_cost = pid != myid() ? 1_000_000 : 0 # 1ms + + tx_rate = get(get(wtr, pid, EMPTY_TRANSFER_RATES), proc, DEFAULT_TRANSFER_RATE) + cost = est_time_util + (tx_costs[gproc]/tx_rate) + task_xfer_cost + costs[proc] = cost + if idx == 1 + first_cost = cost + elseif cost != first_cost + all_equal = false + end end - costs[proc] = est_time_util + (tx_cost/tx_rate) + task_xfer_cost end - chunks_cleanup() + tx_costs_cleanup() # Shuffle procs around, so equally-costly procs are equally considered np = length(procs) @@ -688,8 +708,10 @@ const DEFAULT_TRANSFER_RATE = UInt64(1_000_000) end end - # Sort by lowest cost first - sort!(sorted_procs, by=p->costs[p]) + # Sort by lowest cost first. Skipped when every processor costs the same + # (the usual single-worker case), where the shuffle above already gives the + # arbitrary-but-fair order that sorting would leave us with anyway. + all_equal || sort!(sorted_procs, by=p->costs[p]) end """ From e57a8840670870f700d1945ff20f7e6be9f60e1a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 19:53:16 -0700 Subject: [PATCH 33/43] Datadeps: Reuse a chunk as its own slot when it is already in place `generate_slot!` always routed through `move_rewrap`, even for a `Chunk` that already lives in the destination space, allocating a second Chunk and DRef over the very same memory. That is a MemPool round-trip per argument per region for no benefit. Only leaves qualify for the fast path: `move_rewrap` does more than move bytes, rebuilding wrappers and resolving handles so that what reaches the task is a plain destination-space value (a `ChunkView` becomes a `Chunk` over a real `SubArray`, a nested `Chunk` is flattened). Those results differ from the input even when nothing moves. Co-authored-by: Cursor --- src/datadeps/aliasing.jl | 37 +++++++++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index 00f4701da..42760020a 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -902,6 +902,29 @@ end isremotehandle(x) = false isremotehandle(x::DTask) = true isremotehandle(x::Chunk) = true +""" + slot_is_already_in_place(data, orig_space, dest_space) -> Bool + +Whether `data` can serve as its own Datadeps slot in `dest_space`. + +Only a `Chunk` that already lives in `dest_space` and wraps a `move_rewrap` +*leaf* qualifies. `move_rewrap` does more than move bytes: it rebuilds wrappers +and resolves handles (a `ChunkView` becomes a `Chunk` over a real `SubArray`, a +nested `Chunk` is flattened) so that what reaches the task is a plain +destination-space value. Those results differ from the input even when nothing +moves, so only leaves may be passed through untouched. + +Restricted to locally-owned chunks under non-uniform execution because deciding +this requires unwrapping `data`, and because MPI chunk handles are rank-relative. +""" +function slot_is_already_in_place(data, orig_space, dest_space) + data isa Chunk || return false + orig_space == dest_space || return false + uniform_execution() && return false + root_worker_id(data.processor) == myid() || return false + return move_rewrap_parts(unwrap(data)) === nothing +end + function generate_slot!(state::DataDepsState, dest_space, data) # N.B. We do not perform any sync/copy with the current owner of the data, # because all we want here is to make a copy of some version of the data, @@ -919,8 +942,18 @@ function generate_slot!(state::DataDepsState, dest_space, data) id = rand(Int) @maybelog ctx timespan_start(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data)) tid = something(DATADEPS_CURRENT_TASK[], (;uid=0)).uid - data_chunk = with(DATADEPS_THUNK_ID=>tid) do - remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + data_chunk = if slot_is_already_in_place(data, orig_space, dest_space) + # Nothing to move: the slot for data already in `dest_space` is the data + # itself. Going through `move_rewrap` here would allocate a second Chunk + # (and DRef) over the very same memory, which costs a MemPool round-trip + # per argument per region and buys nothing. Still route through + # `aliased_object!` so the cache records this object as the (never-freed) + # original for its aliasing key, exactly as the general path does. + aliased_object!(Returns(data), aliased_object_cache, data)::Chunk + else + with(DATADEPS_THUNK_ID=>tid) do + remotecall_endpoint_toplevel(move_rewrap, current_acceleration(), aliased_object_cache, from_proc, to_proc, orig_space, dest_space, data) + end end @maybelog ctx timespan_finish(ctx, :move, (;thunk_id=0, id, position=ArgPosition(), processor=to_proc), (;f=nothing, data=data_chunk)) @assert memory_space(data_chunk) == dest_space "space mismatch! $dest_space (dest) != $(memory_space(data_chunk)) (actual) ($(typeof(data)) (data) vs. $(typeof(data_chunk)) (chunk)), spaces ($orig_space -> $dest_space)" From ce90e773377bd3a4cafc1a36e568d02c0c75dbca Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 19:53:20 -0700 Subject: [PATCH 34/43] fixup! datadeps: Add hierarchical scheduling Co-authored-by: Cursor --- src/datadeps/hierarchical.jl | 151 ++++++++++++++++++++++++++++++++--- 1 file changed, 139 insertions(+), 12 deletions(-) diff --git a/src/datadeps/hierarchical.jl b/src/datadeps/hierarchical.jl index 6524b20eb..8c00d8b51 100644 --- a/src/datadeps/hierarchical.jl +++ b/src/datadeps/hierarchical.jl @@ -25,6 +25,56 @@ struct HierarchicalTaskInfo writedep::Bool end +""" + BatchedEnqueueQueue(inner, lock; limit=DATADEPS_BATCH_LIMIT[]) + +Buffers tasks locally and submits them to `inner` in batches of at most `limit`. + +Submitting one task at a time costs a scheduler round-trip each; batching pays +that once per group. It also collapses the lock traffic, which matters because +every partition of the hierarchical path shares a single lock over the region's +queue -- with one acquisition per task the partitions spend most of their time +convoyed behind each other rather than preparing tasks in parallel. + +`limit` trades that amortization against latency: buffered tasks cannot start +executing, so an unbounded batch would keep a partition's whole workload idle +until it finished planning, serializing planning against execution. A small +limit keeps the pipeline fed while still amortizing most of the per-submission +cost. + +`flush_batch!` must additionally be called before anything outside this +partition can observe the buffered tasks (see `schedule_partition_full!`), since +a task's syncdeps must already be submitted when a dependent task is prepared. +""" +struct BatchedEnqueueQueue <: AbstractTaskQueue + inner::AbstractTaskQueue + lock::ReentrantLock + pending::Vector{DTaskPair} + limit::Int +end +BatchedEnqueueQueue(inner::AbstractTaskQueue, lock::ReentrantLock; + limit::Int=DATADEPS_BATCH_LIMIT[]) = + BatchedEnqueueQueue(inner, lock, DTaskPair[], limit) +function enqueue!(beq::BatchedEnqueueQueue, pair::DTaskPair) + push!(beq.pending, pair) + length(beq.pending) >= beq.limit && flush_batch!(beq) + return +end +function enqueue!(beq::BatchedEnqueueQueue, pairs::Vector{DTaskPair}) + append!(beq.pending, pairs) + length(beq.pending) >= beq.limit && flush_batch!(beq) + return +end +function flush_batch!(beq::BatchedEnqueueQueue) + isempty(beq.pending) && return + @lock beq.lock enqueue!(beq.inner, beq.pending) + empty!(beq.pending) + return +end + +"Maximum number of tasks a hierarchical partition buffers before submitting." +const DATADEPS_BATCH_LIMIT = Ref(4) + struct HierarchicalTaskMeta pair::DTaskPair # Wrapped argument identity for the task's first aliasing arg. Under MPI @@ -580,6 +630,31 @@ function build_dependency_dag(task_metas::Vector{HierarchicalTaskMeta}, return dag end +""" +Target number of tasks per partition when partitioning across the local +processors of a single owner. + +Partitioning a single owner parallelizes *planning* only -- every partition +still places work on the same processors -- so it has to earn back what it +duplicates. Each partition builds its own `DataDepsState`, so an argument shared +between partitions has its aliasing and slot bookkeeping redone once per +partition; with neighbour-sharing workloads (stencils) that duplication grows +quickly with the partition count. +""" +const HIER_TASKS_PER_PARTITION = 16 + +""" + single_owner_partition_count(ntasks, nprocs) -> Int + +How many partitions to split `ntasks` into when all processors share an owner. + +Capped at half of `nprocs` because partition planning tasks and the processor +runners execute on the same threads: one partition per processor leaves nothing +to run the work being planned, and measures slower than not partitioning at all. +""" +single_owner_partition_count(ntasks::Int, nprocs::Int) = + clamp(ntasks ÷ HIER_TASKS_PER_PARTITION, 1, max(1, nprocs ÷ 2)) + """ partition_dag(dag, task_metas, all_procs) -> (vertex_to_partition, n_partitions, partition_procs) @@ -608,7 +683,7 @@ function partition_dag(dag::SimpleDiGraph, task_metas::Vector{HierarchicalTaskMe n_partitions = n_owners partition_owner = owners else - n_partitions = min(length(all_procs), n) + n_partitions = single_owner_partition_count(n, length(all_procs)) partition_owner = fill(first(owners), n_partitions) end @@ -782,10 +857,10 @@ function _schedule_vertex!(v::Int, partition_id::Int, end """ - schedule_partition_full!(queue, queue_lock, partition_id, partition_verts, + schedule_partition_full!(queue, batch_queue, partition_id, partition_verts, dag, seen_tasks, local_procs, vertex_to_partition, task_submitted, - region_uids) -> DataDepsState + region_uids, value_dep_verts, registry) -> DataDepsState Per-partition scheduling for both single-worker and multi-worker hierarchical paths. Uses existing `distribute_task!` logic with per-partition @@ -800,7 +875,7 @@ partition (which it must not `fetch`). Tasks without an AOT assignment fall back to JIT scheduling in `distribute_task!`. """ function schedule_partition_full!(queue::DataDepsTaskQueue, - queue_lock::ReentrantLock, + batch_queue::BatchedEnqueueQueue, partition_id::Int, partition_verts::Vector{Int}, dag::SimpleDiGraph, @@ -809,6 +884,7 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, vertex_to_partition::Vector{Int}, task_submitted::Vector{Base.Event}, region_uids::Set{UInt}, + value_dep_verts::Set{Int}, registry::Union{SharedChunkRegistry,Nothing}) if isempty(partition_verts) || isempty(local_procs) return DataDepsState(DAGSpec()) @@ -832,7 +908,6 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, end ordered_verts = filter(v -> v in vert_set, topo) - locked_queue = LockedEnqueueQueue(get_options(:task_queue), queue_lock) # N.B. Each partition gets its own fresh scheduler shard via `similar` # rather than sharing `queue.scheduler` across all partitions. Two # independent problems would arise from sharing a single scheduler @@ -850,7 +925,7 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, # crashes with a `BoundsError` under multi-worker hierarchical # scheduling. Giving each partition its own scheduler instance, # scoped to its own `local_procs`, fixes both issues at once. - temp_queue = DataDepsTaskQueue(locked_queue; scheduler=similar(queue.scheduler)) + temp_queue = DataDepsTaskQueue(batch_queue; scheduler=similar(queue.scheduler)) # Per-partition (local) AOT scheduling over just this partition's procs. # `partition_verts` is in ascending vertex (= submission) order. @@ -869,17 +944,42 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, # The `finally` ensures every one of our events gets notified no matter # how we exit, so that sibling partitions can unblock (and themselves # fail/finish) and our real exception can actually surface. + # A vertex whose successors all live in this partition is only ever observed + # by us, so its submission can stay buffered; one with a successor elsewhere + # must be submitted before we notify, because that successor will record it + # as a syncdep. + has_external_successor = Set{Int}() + for v in ordered_verts + for succ_v in outneighbors(dag, v) + if vertex_to_partition[succ_v] != partition_id + push!(has_external_successor, v) + break + end + end + end + try for v in ordered_verts for pred_v in inneighbors(dag, v) if vertex_to_partition[pred_v] != partition_id + # Anything we have buffered may be a dependency of tasks the + # producing partition is about to prepare; get it submitted + # before we block, or two partitions can deadlock waiting on + # each other's unflushed tasks. + flush_batch!(batch_queue) wait(task_submitted[pred_v]) end end + # A task taking another in-region task's *value* as an argument has + # `distribute_task!` `fetch` that producer, which requires it to have + # actually been launched. A same-partition producer is topologically + # earlier, so it is prepared but possibly still sitting in our batch. + v in value_dep_verts && flush_batch!(batch_queue) + # N.B. The per-task `distribute_task!` preparation runs concurrently # across partitions; only the final task submission is serialized - # (via `LockedEnqueueQueue`, `temp_queue`'s upper queue). This is + # (via `BatchedEnqueueQueue`, `temp_queue`'s upper queue). This is # safe now that task futures are backed by `MemPool.DFuture` rather # than the concurrency-unsafe `Distributed.Future`. The # cross-partition `wait`s above happen before scheduling `v`, so the @@ -889,9 +989,12 @@ function schedule_partition_full!(queue::DataDepsTaskQueue, dag, seen_tasks, vertex_to_partition, schedule, proc_to_scope_lfu, write_num, registry) + v in has_external_successor && flush_batch!(batch_queue) notify(task_submitted[v]) end + flush_batch!(batch_queue) finally + flush_batch!(batch_queue) for v in ordered_verts notify(task_submitted[v]) end @@ -1048,6 +1151,18 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) end end + # With a single owner, partitioning buys parallel planning and nothing else + # (see `HIER_TASKS_PER_PARTITION`). A region too small to fill more than one + # partition would pay for the extra prescan/DAG/partition phases and the + # per-partition state without ever running two planners at once, so plan it + # flat instead. Multi-owner regions always partition: there the split is by + # data ownership, not a throughput heuristic. + if !uniform_execution(accel) && + allequal(partition_affinity_id(proc) for proc in all_procs) && + single_owner_partition_count(length(seen_tasks), length(all_procs)) == 1 + return distribute_tasks!(queue) + end + # All in-region task uids, so each partition's local AOT-schedule builder # can tell same-region producers (which it must not `fetch`) apart from # already-materialized external values. @@ -1070,6 +1185,11 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) partition_space = MemorySpace[only(memory_spaces(first(pp))) for pp in partition_procs] registry = build_shared_chunk_registry(task_metas, vertex_to_partition, partition_space) + # Vertices whose `distribute_task!` will `fetch` an in-region producer's + # value, and so cannot run until that producer has actually been submitted. + value_dep_verts = Set{Int}(v for v in 1:length(task_metas) + if !isempty(task_metas[v].value_deps)) + # Group vertices by partition partitions = [Int[] for _ in 1:n_partitions] for v in 1:length(seen_tasks) @@ -1097,14 +1217,21 @@ function distribute_tasks_hierarchical!(queue::DataDepsTaskQueue) states = Vector{DataDepsState}(undef, n_partitions) @sync for pid in 1:n_partitions Threads.@spawn begin - locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) - with_options(; task_queue=locked_queue) do + batch_queue = BatchedEnqueueQueue(wait_all_queue, queue_lock) + # Copy tasks spawned from within `distribute_task!` pick the + # queue up from options, so they batch alongside their task. + with_options(; task_queue=batch_queue) do + try states[pid] = schedule_partition_full!( - queue, queue_lock, pid, partitions[pid], + queue, batch_queue, pid, partitions[pid], dag, seen_tasks, partition_procs[pid], vertex_to_partition, - task_submitted, region_uids, registry + task_submitted, region_uids, value_dep_verts, + registry ) + finally + flush_batch!(batch_queue) + end end end end @@ -1207,7 +1334,7 @@ function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsStat for remote_space in remote_spaces space_entries = sort!(collect(obj_cache.values[remote_space]); by=p->hash(p.first)) for (ainfo, remote_arg) in space_entries - if !(ainfo in obj_cache.originals) + if !is_original(obj_cache, remote_space, ainfo) remote_proc = first(processors(remote_space)) free_scope = ExactScope(remote_proc) free_syncdeps = Set{ThunkSyncdep}() From a4ec59be14f68f13cd1c85aead70a3a0b55769b5 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 20:38:44 -0700 Subject: [PATCH 35/43] docs/stencils: Explain how block size trades against per-task cost `@stencil` spawns a task per chunk per expression, so the block size decides how many times the scheduler's per-task cost is paid. That cost is large enough to dominate: a 4096x4096 Jacobi sweep on 12 threads is ~1.7x slower at `Blocks(512, 512)` than at `AutoBlocks()`, purely from running 64 tasks per sweep instead of 12. Nothing in the docs pointed at this, and partitioning much finer than the processor count is an easy default to reach for. Co-authored-by: Cursor --- docs/src/stencils.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/src/stencils.md b/docs/src/stencils.md index bfb4838df..27da654e0 100644 --- a/docs/src/stencils.md +++ b/docs/src/stencils.md @@ -36,6 +36,32 @@ B = zeros(Blocks(2, 2), Int, 4, 4) In these examples, the stencil operation is executed for each chunk of the target array. The `idx` variable corresponds to the indices within each chunk. +## Choosing a Block Size + +`@stencil` spawns one task per chunk per expression, so the array's block size sets +both the available parallelism and the number of tasks. Scheduling a task is not +free -- it costs on the order of a hundred microseconds -- so blocks want to be +large enough that the work in a chunk dwarfs that, and numerous enough to keep the +processors busy. Roughly one chunk per processor is a good starting point, which is +what `AutoBlocks` picks: + +```julia +A = rand(AutoBlocks(), Float64, 4096, 4096) +``` + +Partitioning far more finely than the processor count is the common mistake, and it +is expensive: a 4096x4096 Jacobi sweep over 12 threads runs about 1.7x slower with +`Blocks(512, 512)` (64 chunks) than with `AutoBlocks()` (12 chunks), because the +per-task cost is paid 64 times per sweep instead of 12 while the arithmetic is +unchanged. Coarsening past one chunk per processor is bounded by the other side: +chunks are the unit of parallelism, so fewer chunks than processors leaves +processors idle. + +Note that this argues only against *needlessly* fine partitioning. Smaller chunks +still pay for themselves when they buy something else: load balance across uneven +processors, overlap of communication with computation, or fitting a working set +into cache or device memory. + ## Allocation and Return Syntax The `@stencil` macro also supports a functional syntax that automatically allocates an output `DArray` and returns it. This is triggered when the stencil expression (or the last expression in a block) is a "naked" expression (not an assignment). From e0212ae7ddacc05b75c3537ce928d3a6e285201b Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sat, 25 Jul 2026 21:23:56 -0700 Subject: [PATCH 36/43] fixup! DArray/stencil: Fuse halo construction and vectorize the sweep Select halo regions without a runtime tuple index Fusing halos out of views made the halo tuple heterogeneous on GPU backends, which return a plain device array for a contiguous `view` and a `SubArray` for a strided one. `halos[idx]` with a runtime `idx` then infers to a union wide enough that the load after it is a dynamic dispatch, and GPU backends cannot compile a dynamic dispatch at all -- `@stencil` on AMDGPU crashed inside GPUCompiler's IR validation, which segfaults trying to name the offending method. Unroll the selection so each branch sees one concrete array type. Homogeneous tuples, which is what the CPU path always builds, keep the plain indexed load. Co-authored-by: Cursor --- src/utils/haloarray.jl | 52 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/src/utils/haloarray.jl b/src/utils/haloarray.jl index 4e5d74c62..3adee486b 100644 --- a/src/utils/haloarray.jl +++ b/src/utils/haloarray.jl @@ -118,6 +118,54 @@ end end end +""" + halo_load(halos, idx, I) + halo_store!(halos, idx, I, value) + +Read or write element `I` of the `idx`-th halo region. + +`halos` is not necessarily homogeneous. GPU backends return a plain device array +for a `view` that happens to be contiguous and a `SubArray` for one that is not, +so a halo tuple built from views of neighboring chunks can hold both. Indexing +such a tuple with a runtime `idx` infers to a union wide enough that the load +that follows becomes a dynamic dispatch, which GPU backends cannot compile at +all. Selecting the region with an unrolled comparison chain instead gives every +branch one concrete array type. + +Homogeneous tuples -- what the CPU path always produces -- keep the plain +indexed load, so they pay nothing for this. +""" +@generated function halo_load(halos::Tuple, idx::Int, I::NTuple{N,Int}) where N + Ts = fieldtypes(halos) + body = if allequal(Ts) + :(@inbounds halos[idx][I...]) + else + # Last region is the fallback: `region_index` never returns out of range, + # and a trailing `throw` would only add unreachable code the GPU backend + # still has to compile. + ex = :(@inbounds halos[$(length(Ts))][I...]) + for i in (length(Ts) - 1):-1:1 + ex = :(idx === $i ? @inbounds(halos[$i][I...]) : $ex) + end + ex + end + return Expr(:block, Expr(:meta, :inline), body) +end + +@generated function halo_store!(halos::Tuple, idx::Int, I::NTuple{N,Int}, value) where N + Ts = fieldtypes(halos) + body = if allequal(Ts) + :(@inbounds halos[idx][I...] = value) + else + ex = :(@inbounds halos[$(length(Ts))][I...] = value) + for i in (length(Ts) - 1):-1:1 + ex = :(idx === $i ? @inbounds(halos[$i][I...] = value) : $ex) + end + ex + end + return Expr(:block, Expr(:meta, :inline), body) +end + # Define getindex for HaloArray @inline function Base.getindex(tile::HaloArray{T,N}, I::Vararg{Int,N}) where {T,N} Base.@boundscheck checkbounds(tile, I...) @@ -130,7 +178,7 @@ end # Halo region idx = region_index(code) local_idx = compute_local_index(tile, I, code) - return @inbounds tile.halos[idx][local_idx...] + return halo_load(tile.halos, idx, local_idx) end end @@ -146,7 +194,7 @@ end # Halo region idx = region_index(code) local_idx = compute_local_index(tile, I, code) - return @inbounds tile.halos[idx][local_idx...] = value + return halo_store!(tile.halos, idx, local_idx, value) end end From 80737e8f51f97d586535c15f7cae77d92b7de123 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 26 Jul 2026 07:35:41 -0700 Subject: [PATCH 37/43] ROCExt: Stop creating a HIP stream per task in with_context `with_context` recorded the stream to restore by calling `AMDGPU.stream()`, which creates a stream when the running task doesn't have one. Every Dagger task that touches a ROC processor gets a fresh Julia task, so each one manufactured a HIP stream purely to be saved and put back. Those streams accumulate until finalized, and a run gets through enough tasks that stream creation itself stalls inside the driver, after which unrelated HIP calls start reporting illegal addresses. Read the task-local slot directly so a task without a stream reports `nothing` instead of being given one. A 4096^2 stencil sweep over a single GPU chunk goes from 20.4 GB/s to 192 GB/s, and multi-chunk sweeps stop faulting on every run. Co-authored-by: Cursor --- ext/ROCExt.jl | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 70d836369..89223c7c5 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -105,10 +105,42 @@ function with_context!(space::ROCVRAMMemorySpace) end Dagger.with_context!(proc::ROCArrayDeviceProc) = with_context!(proc) Dagger.with_context!(space::ROCVRAMMemorySpace) = with_context!(space) + +""" + task_stream_slot() + restore_stream_slot!(old_stream) + +Save and restore the running task's default HIP stream. + +`AMDGPU.stream()` creates a stream when the current task doesn't have one yet, +so using it to record what to restore would allocate a HIP stream for every +Dagger task that touches a ROC processor. Those streams live until they are +finalized, and a run goes through enough tasks that stream creation itself +eventually stalls inside the driver. Reading the task-local slot directly +reports "no stream" as `nothing` instead of manufacturing one. +""" +function task_stream_slot() + state = AMDGPU.task_local_state() + state === nothing && return nothing + return state.streams[AMDGPU.device_id(state.device)] +end +function restore_stream_slot!(old_stream) + if old_stream !== nothing + stream!(old_stream) + return + end + # The task had no stream of its own; put the slot back the way we found it + # rather than leaving it pointing at Dagger's per-device stream. + state = AMDGPU.task_local_state() + state === nothing && return + state.streams[AMDGPU.device_id(state.device)] = nothing + return +end + function with_context(f, x) old_ctx = context() old_device = AMDGPU.device() - old_stream = stream() + old_stream = task_stream_slot() with_context!(x) try @@ -116,7 +148,7 @@ function with_context(f, x) finally context!(old_ctx) AMDGPU.device!(old_device) - stream!(old_stream) + restore_stream_slot!(old_stream) end end From 69816ccc2fce52baac9ea3844b82a899986c157e Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 26 Jul 2026 07:35:41 -0700 Subject: [PATCH 38/43] DArray/stencil: Split the GPU sweep into interior and boundary Every GPU backend launched one kernel over the whole chunk, with all `@neighbors` accesses going through the general `HaloArray` path: a region code recomputed per access, a branch on it, and an index remap. On a GPU that scalar bookkeeping runs for all (2w+1)^N accesses of every element. Apply the same interior/boundary split the CPU sweep already uses. The interior sweeps against `HaloInterior` stand-ins, which index the center array directly, leaving the general path to the 2N boundary slabs. The five backends launched byte-identical kernels, so this lives in core as `gpu_stencil_sweep!` and each extension now delegates to it. A 4096^2 sweep over a single GPU chunk goes from 192 GB/s to 222 GB/s, against 298 GB/s for an equivalent hand-written AMDGPU kernel. Co-authored-by: Cursor --- ext/CUDAExt.jl | 6 +--- ext/IntelExt.jl | 6 +--- ext/MetalExt.jl | 6 +--- ext/OpenCLExt.jl | 6 +--- ext/ROCExt.jl | 6 +--- src/array/stencil.jl | 63 ++++++++++++++++++++++++++++++++++++++++++ src/utils/haloarray.jl | 3 ++ 7 files changed, 71 insertions(+), 25 deletions(-) diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 813763ce1..02cff791f 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -390,13 +390,9 @@ Adapt.adapt_structure(to::CUDA.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::CuArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:CUDA}) = CuArrayDeviceProc Dagger.gpu_can_compute(::Val{:CUDA}) = CUDA.has_cuda() diff --git a/ext/IntelExt.jl b/ext/IntelExt.jl index 87bb7138f..2f25f7e3d 100644 --- a/ext/IntelExt.jl +++ b/ext/IntelExt.jl @@ -366,13 +366,9 @@ Adapt.adapt_structure(to::oneAPI.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::oneArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:oneAPI}) = oneArrayDeviceProc Dagger.gpu_can_compute(::Val{:oneAPI}) = oneAPI.functional() diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 7e175993c..f0b6efda0 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -432,13 +432,9 @@ Adapt.adapt_structure(to::Metal.Adaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::MtlArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end function Base.show(io::IO, proc::MtlArrayDeviceProc) print(io, "MtlArrayDeviceProc(worker $(proc.owner), device $(something(_get_metal_device(proc)).name))") diff --git a/ext/OpenCLExt.jl b/ext/OpenCLExt.jl index eceb9cf6f..a61c0b310 100644 --- a/ext/OpenCLExt.jl +++ b/ext/OpenCLExt.jl @@ -361,13 +361,9 @@ Adapt.adapt_structure(to::OpenCL.KernelAdaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::CLArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:OpenCL}) = CLArrayDeviceProc Dagger.gpu_can_compute(::Val{:OpenCL}) = length(cl.platforms()) > 0 diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl index 89223c7c5..ddd85d203 100644 --- a/ext/ROCExt.jl +++ b/ext/ROCExt.jl @@ -419,13 +419,9 @@ Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) function Dagger.inner_stencil_proc!(::ROCArrayDeviceProc, f, output, read_vars) - Dagger.Kernel(_inner_stencil!)(f, output, read_vars; ndrange=size(output)) + Dagger.gpu_stencil_sweep!(f, output, read_vars) return end -@kernel function _inner_stencil!(f, output, read_vars) - idx = @index(Global, Cartesian) - f(idx, output, read_vars) -end Dagger.gpu_processor(::Val{:ROC}) = ROCArrayDeviceProc Dagger.gpu_can_compute(::Val{:ROC}) = AMDGPU.functional() diff --git a/src/array/stencil.jl b/src/array/stencil.jl index f99066b56..14a5f8801 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -966,6 +966,69 @@ function cpu_stencil_sweep!(f::F, output::AbstractArray{T,N}, read_vars::NamedTu return end +@kernel function _stencil_box_kernel!(f, output, read_vars, offset) + idx = @index(Global, Cartesian) + @inline f(idx + offset, output, read_vars) +end + +# Sweep `f` over one box of `output`. KernelAbstractions indexes an `ndrange` +# from 1, so a sub-box is expressed as a shifted full launch. +@inline function _stencil_box!(f, output, read_vars, ranges::NTuple{N,AbstractUnitRange}) where N + nd = map(length, ranges) + any(iszero, nd) && return + offset = CartesianIndex(ntuple(i -> first(ranges[i]) - 1, Val(N))) + Kernel(_stencil_box_kernel!)(f, output, read_vars, offset; ndrange=nd) + return +end + +""" + gpu_stencil_sweep!(f, output, read_vars) + +Applies `f` at every index of `output` using KernelAbstractions, with the same +interior/boundary split as [`cpu_stencil_sweep!`](@ref). + +The split matters more here than on the CPU. Every `@neighbors` access on the +general path recomputes a region code, branches on it, and remaps the index; a +GPU runs that scalar bookkeeping for all `(2w+1)^N` accesses of every element. +Sweeping the interior against `HaloInterior` stand-ins reduces it to a direct +load from the center array, leaving the general path to the `2N` boundary slabs. +""" +function gpu_stencil_sweep!(f::F, output::AbstractArray{T,N}, read_vars::NamedTuple) where {F,T,N} + w = _max_halo_width(values(read_vars), ntuple(_ -> 0, Val(N))) + ax = axes(output) + full_ax = ntuple(i -> first(ax[i]):last(ax[i]), Val(N)) + + interior_ax = ntuple(Val(N)) do i + (first(ax[i]) + w[i]):(last(ax[i]) - w[i]) + end + if any(isempty, interior_ax) + # Halos are as wide as the chunk itself; there is no interior to split off. + _stencil_box!(f, output, read_vars, full_ax) + return + end + + _stencil_box!(f, output, map(_interior_var, read_vars), interior_ax) + + all(iszero, w) && return + + for d in 1:N + for low_side in (true, false) + slab_ax = ntuple(Val(N)) do i + if i < d + interior_ax[i] + elseif i == d + low_side ? (first(ax[i]):(first(ax[i]) + w[i] - 1)) : + ((last(ax[i]) - w[i] + 1):last(ax[i])) + else + full_ax[i] + end + end + _stencil_box!(f, output, read_vars, slab_ax) + end + end + return +end + ############################################################################# # @stencil Macro ############################################################################# diff --git a/src/utils/haloarray.jl b/src/utils/haloarray.jl index 3adee486b..48e4f39c2 100644 --- a/src/utils/haloarray.jl +++ b/src/utils/haloarray.jl @@ -304,6 +304,9 @@ Adapt.adapt_structure(to, H::Dagger.HaloArray) = H.halo_width; own_center=H.own_center) +Adapt.adapt_structure(to, A::Dagger.HaloInterior) = + HaloInterior(Adapt.adapt(to, A.parent), A.halo_width) + function aliasing(A::HaloArray) return CombinedAliasing([aliasing(A.center), map(aliasing, A.halos)...]) end From ec75eb8c538dc723ca2fa9c8bbf9c7ccd4f75432 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 26 Jul 2026 10:39:09 -0700 Subject: [PATCH 39/43] fixup! DArray/sparse: Improve copy/indexing and aliasing --- src/array/darray.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/array/darray.jl b/src/array/darray.jl index b774d913f..bb0442714 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -183,7 +183,8 @@ _collect_cat(concat, dims::Int, xs...) = concat(xs...; dims) # Finalize a gathered DArray to a dense `Array{T,N}`. Sparse tile cats may yield # a `DSparseArray` / `SparseMatrixCSC`, and `Base.collect` does not densify those. _collect_dense(::Type{T}, ::Val{N}, x::Array{T,N}) where {T,N} = x -_collect_dense(::Type{T}, ::Val{N}, x) where {T,N} = Array{T,N}(x) +_collect_dense(::Type{T}, ::Val{N}, x::AbstractArray) where {T,N} = Array{T,N}(x) +_collect_dense(::Type, ::Val, x) = x function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} a = fetch(d) From 1bc625e1136ada974fccfd7b1f9af92ed15ee55a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 27 Jul 2026 21:31:37 -0700 Subject: [PATCH 40/43] Automatic partitioning selection --- src/Dagger.jl | 1 + src/tapes/README.md | 140 +++++++++ src/tapes/Tapes.jl | 233 +++++++++++++++ src/tapes/api.jl | 532 ++++++++++++++++++++++++++++++++++ src/tapes/cost.jl | 564 ++++++++++++++++++++++++++++++++++++ src/tapes/integration.jl | 263 +++++++++++++++++ src/tapes/plan.jl | 387 +++++++++++++++++++++++++ src/tapes/tape.jl | 608 +++++++++++++++++++++++++++++++++++++++ src/tapes/tapes_tests.jl | 328 +++++++++++++++++++++ 9 files changed, 3056 insertions(+) create mode 100644 src/tapes/README.md create mode 100644 src/tapes/Tapes.jl create mode 100644 src/tapes/api.jl create mode 100644 src/tapes/cost.jl create mode 100644 src/tapes/integration.jl create mode 100644 src/tapes/plan.jl create mode 100644 src/tapes/tape.jl create mode 100644 src/tapes/tapes_tests.jl diff --git a/src/Dagger.jl b/src/Dagger.jl index ccc3925bd..bc6684c55 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -125,6 +125,7 @@ include("file-io.jl") # Array computations include("procgrid.jl") include("array/darray.jl") +include("tapes/Tapes.jl"); using .Tapes include("array/alloc.jl") include("array/map-reduce.jl") include("array/copy.jl") diff --git a/src/tapes/README.md b/src/tapes/README.md new file mode 100644 index 000000000..e17e2e512 --- /dev/null +++ b/src/tapes/README.md @@ -0,0 +1,140 @@ +# Dagger.Tapes — history-driven partitioning selection + +Speculative selection of `DArray` partitionings from recorded operation +history. Opt-in, disabled by default, and inert when disabled. + +## Files + +``` +tapes/Tapes.jl module entry, TapeConfig, enable!/disable! +tapes/tape.jl keys, site identification, the trie, recording, prediction +tapes/cost.jl MachineModel, ArgView, @cost_model, default models +tapes/plan.jl candidate generation, DP planner, regret gating +tapes/api.jl plan_allocation/track!, @record_op, @expect_ops, introspection +tapes/integration.jl resolve_partitioning, @tracked_alloc, patch points +tapes_tests.jl test suite (rename to test/tapes.jl) +``` + +Copy `tapes/` to `src/tapes/`, add to `src/Dagger.jl` immediately after the +`include("array/darray.jl")` line: + +```julia +include("tapes/Tapes.jl"); using .Tapes +``` + +That alone is inert. The allocation and instrumentation patch points are +written out in the `PATCH POINTS` comment block at the bottom of +`tapes/integration.jl`. + +## How it works + +1. **Identify.** Each allocation gets a `SiteKey`: calling context, element + type, rank, and a log-bucketed size. Three context strategies — + `:backtrace` (unwind and hash raw IPs), `:context` (incremental + probabilistic calling context, O(1)), `:lexical` (macro-expansion site + only, free). + +2. **Record.** `@record_op :cholesky! A` appends `(op, argument position, + arity)` to the array's tape, committed into a per-site prefix trie + immediately. Argument position is part of the identity: being the `A` of + `mul!(C, A, B)` implies a different layout preference from being the `C`. + +3. **Predict.** On a later allocation from the same site, walk the trie + greedily and emit `PredictedOp`s carrying the cumulative probability of + reaching each step. + +4. **Plan.** Cost each candidate layout against each predicted step, then + dynamic-program the chain including redistribution cost on the edges. + Linear chain, so exact in O(n·|L|²). + +5. **Gate.** Adopt the plan only if it beats the fallback by + `gate_margin`, and — below `regret_confidence_threshold` — only if no + single predicted operation is more than `max_regret_ratio` worse than its + own best layout. + +Only `steps[1]` is committed; the planner re-runs at each observed operation +against the branch actually taken. + +## Usage + +```julia +Dagger.Tapes.enable!() + +for i in 1:10 + A = rand(AutoBlocks(), Float64, 4096, 4096) + B = A * A' + 4096I + cholesky!(B) +end + +Dagger.Tapes.report() # what was learned +Dagger.Tapes.explain(Float64, (4096, 4096)) # why a layout was chosen +``` + +Ahead-of-time, no warm-up required: + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, n, n) + ... +end +``` + +Without an operation list, `@expect_ops` just establishes a stable lexical +site key — which makes `site_id = :lexical` as precise as `:backtrace` for the +code you care about, at zero cost. + +Declaring a cost model: + +```julia +Dagger.Tapes.@cost_model my_solve(A, B) = begin + nt = nblocks(A, 1) + serial_time(nt * blocksize(A, 1)^2) + + flops_time(size(A, 1)^2 * size(B, 2)) * imbalance(nt) + + task_time(nt^2 / 2) +end +``` + +Argument order must match the `@record_op` call for that operation. Pass only +`DArray`s. + +## Benchmarking + +Performance becomes warm-up-dependent and bimodal once tapes are active, so +`dagger_bench.py` and anything else reporting a distribution across +repetitions needs to control cold vs warm explicitly: + +- `Tapes.clear!()` between repetitions for cold measurements +- `Tapes.pin!(T, dims, Blocks(...), :cyclicrow)` to make a site deterministic + and history-independent +- `Tapes.disable!()` for a true baseline +- `Tapes.stats().hit_rate` to confirm predictions are actually landing + +## Verify in this order + +1. **`@cost_model` expansion.** The one construct I could not test — no Julia + in the build container. It names the method via + `Expr(:., TAPES_MODULE, QuoteNode(:op_cost))`. If it misbehaves, the two + alternatives are a bare unescaped `op_cost` (relies on hygiene resolving a + *definition*) or `GlobalRef(TAPES_MODULE, :op_cost)`. +2. `Tapes.enable!(); Tapes.explain(Float64, (4096,4096))` — exercises site + identification, candidate generation, the planner and all default models in + one call. +3. The test suite, which runs without a cluster except for the last two + groups. +4. `@record_op` overhead with `site_id = :backtrace` against your 100µs task + floor. + +## TODOs in the code + +Every one carries its rationale inline. The load-bearing ones: + +| where | what | +|:--|:--| +| `record_op!` | **Joint planning.** Tapes are per-array but layout decisions are joint — `mul!(C,A,B)` needs mutually compatible layouts. Currently undervalues layouts that only pay off when all operands agree. Wants union-find over co-participating arrays using the Datadeps aliasing analysis. | +| `integration.jl` | **Residency tracker.** Precondition for read-only replication, precise `redistribution_cost`, heterogeneous cost models and OOC. Highest-leverage missing piece. | +| `integration.jl` | **Datadeps lookahead.** Inside `spawn_datadeps` the full dependency structure is already known — exact lookahead, no prediction. Strictly better where it applies; build before broadening `@record_op` coverage. | +| `maybe_repartition!` | **Repartition mechanism.** Planner emits per-step layouts; nothing acts on them. Hard part is arrays captured by already-submitted-but-unscheduled thunks. | +| `current_machine` | **Calibration.** FLOP rate and bandwidth are guesses; should be measured at `enable!` and refined from `TimespanLogging` task durations. Natural autotuner seam. | +| `backtrace_hash` | **Persistence.** Raw IPs are session-local. Symbolicating `(file, line)` with a per-IP cache would make tapes survive restarts — which is what makes the *first* run fast rather than the second. | +| `ArgSpec` | **Assignment provenance.** `DArray` does not store the assignment it was built with, so recorded specs assume `:arbitrary`. | +| `OP_AFFINITY` | **Learn it.** Hand-maintained table that will rot; the tape already produces the triples needed to fit it. | diff --git a/src/tapes/Tapes.jl b/src/tapes/Tapes.jl new file mode 100644 index 000000000..786a69883 --- /dev/null +++ b/src/tapes/Tapes.jl @@ -0,0 +1,233 @@ +""" + Dagger.Tapes + +Speculative, history-driven selection of `DArray` partitionings. + +# Rationale + +Choosing a partitioning at allocation time is a decision made with no knowledge +of what the array will be *used for*. Any fixed default (currently +`auto_blocks`, which splits the last dimension by processor count) is +necessarily wrong for large classes of operations: a Cholesky factorization +wants square tiles laid out 2D-cyclically, a row-wise `mapreduce` wants fat row +blocks, a triangular solve against a vector wants something else again. + +This module exploits the *principle of persistence*: in real applications, +allocations happen at a small number of program points, and the sequence of +operations applied to the resulting array is highly repeatable across +executions of that program point. So: + +1. At each allocation, identify the *site* (call-stack context + element type + + coarse size bucket). +2. Record the ordered sequence of operations subsequently applied to the array + from that site, into a per-site prefix trie (`TapeRoot`). +3. On a later allocation from the same site, walk the trie to predict the + operation sequence that is about to follow, with a probability attached to + each step. +4. Feed that prediction, plus per-operation cost models, into a dynamic program + that picks the layout minimising expected total cost across the whole + predicted chain (including the cost of any mid-chain repartitioning). + +A misprediction costs a suboptimal layout, never a wrong answer. + +# Status / safety + +The whole subsystem is **opt-in** and disabled by default; when disabled every +hook is a single predictable branch on a `const` struct field, which the +compiler hoists out of hot paths. When enabled, it never overrides an +*explicit* `Blocks(...)` supplied by the user unless +`CONFIG.override_explicit_blocks` is set — it only fills in the choice that +`AutoBlocks()` would otherwise have made blindly. + +# Quick start + +```julia +Dagger.Tapes.enable!() # :backtrace site identification + +# Run your workload once to populate tapes, then again to benefit: +for i in 1:10 + A = rand(AutoBlocks(), Float64, 4096, 4096) + A = A * A' + 4096I + cholesky!(A) +end + +Dagger.Tapes.report() # what was learned +Dagger.Tapes.explain(Float64, (4096, 4096)) # why a layout was chosen +``` + +Cheaper site identification, and ahead-of-time declaration, are available via +[`@expect_ops`](@ref): + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, 4096, 4096) + ... +end +``` + +# Public API + +- [`enable!`](@ref), [`disable!`](@ref), [`CONFIG`](@ref) +- [`@record_op`](@ref) / [`record_op!`](@ref) — instrument an operation +- [`@cost_model`](@ref) — declare a cost model for an operation +- [`@expect_ops`](@ref) — declare an operation chain ahead of time +- [`plan_allocation`](@ref) / [`track!`](@ref) — allocation-site hooks +- [`report`](@ref), [`explain`](@ref), [`dump_tapes`](@ref), [`clear!`](@ref), + [`pin!`](@ref), [`unpin!`](@ref) — introspection and manual override +""" +module Tapes + +using ..Dagger +import ..Dagger: DArray, Blocks, AutoBlocks, AbstractBlocks, AssignmentType +import ..Dagger: auto_blocks, num_processors + +using ScopedValues +using MacroTools: @capture + +export @record_op, @cost_model, @expect_ops + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +""" + TapeConfig + +Tunables for the tape subsystem. The live instance is [`CONFIG`](@ref); mutate +its fields directly, or use [`enable!`](@ref) which takes the same names as +keyword arguments. +""" +Base.@kwdef mutable struct TapeConfig + "Master switch. When `false` every hook is a single branch and returns immediately." + enabled::Bool = false + + """ + How allocation sites are identified. One of: + - `:backtrace` — hash the raw instruction pointers of `backtrace()`. Most + precise, costs roughly 100-500us per allocation depending on stack depth. + Keys are only valid within a session (see TODO on symbolication). + - `:context` — probabilistic calling context (Bond & McKinley, OOPSLA'07): + an incrementally-maintained hash mixed at `@record_op` / `@expect_ops` + boundaries. O(1), but only as precise as the instrumentation density. + - `:lexical` — the macro-expansion site of the allocation call only. + Effectively free, but cannot distinguish two callers of the same wrapper. + """ + site_id::Symbol = :backtrace + + "Number of stack frames hashed in `:backtrace` mode (nearest-first, after `backtrace_skip`)." + backtrace_depth::Int = 24 + "Frames to discard from the top of the stack (Dagger's own plumbing)." + backtrace_skip::Int = 2 + + "Powers-of-two subdivisions per octave when bucketing sizes into a site key. `0` = exact sizes." + size_buckets_per_octave::Int = 2 + + "Maximum number of operations recorded per array (trie depth cap)." + max_tape_length::Int = 24 + "How far ahead the planner looks when predicting." + horizon::Int = 12 + "A trie branch is only followed during prediction if it holds at least this share of its parent's count." + min_branch_prob::Float64 = 0.15 + "Minimum cumulative probability for the first predicted op before we will speculate at all." + min_confidence::Float64 = 0.5 + "Minimum times a site must have been observed before its predictions are used." + min_observations::Int = 2 + + """ + The planned layout must beat the fallback layout by at least this factor to + be adopted (`0.9` = must be at least 10% cheaper). Guards against churn from + cost-model noise. + """ + gate_margin::Float64 = 0.9 + + """ + Reject a plan if, for any *individual* predicted operation, it is more than + this many times worse than the best layout for that operation alone. This is + the minimax-regret safety valve: early on, prefer a layout that is decent for + everything over one that is optimal for the modal chain and catastrophic + otherwise. + """ + max_regret_ratio::Float64 = 3.0 + "Below this cumulative confidence, `max_regret_ratio` is enforced strictly." + regret_confidence_threshold::Float64 = 0.85 + + "Upper bound on candidate layouts considered per allocation (planner is O(n*L^2))." + max_candidates::Int = 16 + + "Override a partitioning the user gave explicitly. Off by default: explicit user intent wins." + override_explicit_blocks::Bool = false + + """ + Permit mid-chain repartitioning when the plan calls for it. Currently a + no-op stub; see `maybe_repartition!`. + """ + allow_repartition::Bool = false + + "LRU cap on the number of distinct sites retained." + max_sites::Int = 4096 + "LRU cap on trie nodes across all sites." + max_nodes::Int = 200_000 + + "Emit `@debug`-style commentary about every decision." + verbose::Bool = false +end + +""" + CONFIG::TapeConfig + +The live configuration. See [`TapeConfig`](@ref) for field documentation. +""" +const CONFIG = TapeConfig() + +""" + is_enabled() -> Bool + +Master predicate guarding every hook. Deliberately trivial so that the branch +is cheap and predictable in the disabled case. +""" +@inline is_enabled() = CONFIG.enabled + +""" + enable!(; kwargs...) + +Turn the tape subsystem on. Keyword arguments set the matching +[`TapeConfig`](@ref) fields. + +```julia +Dagger.Tapes.enable!(site_id=:lexical, horizon=8, verbose=true) +``` +""" +function enable!(; kwargs...) + for (k, v) in kwargs + hasfield(TapeConfig, k) || throw(ArgumentError("unknown Tapes config option: $k")) + setfield!(CONFIG, k, convert(fieldtype(TapeConfig, k), v)) + end + CONFIG.site_id in (:backtrace, :context, :lexical) || + throw(ArgumentError("site_id must be :backtrace, :context or :lexical, got $(CONFIG.site_id)")) + CONFIG.enabled = true + return CONFIG +end + +""" + disable!() + +Turn the tape subsystem off. Learned tapes are retained; use [`clear!`](@ref) +to discard them. +""" +function disable!() + CONFIG.enabled = false + return CONFIG +end + +@inline function vlog(args...) + CONFIG.verbose && @info string("[Tapes] ", args...) + nothing +end + +include("tape.jl") +include("cost.jl") +include("plan.jl") +include("api.jl") +include("integration.jl") + +end # module Tapes diff --git a/src/tapes/api.jl b/src/tapes/api.jl new file mode 100644 index 000000000..24db80f36 --- /dev/null +++ b/src/tapes/api.jl @@ -0,0 +1,532 @@ +# =========================================================================== +# Public API: allocation hooks, instrumentation macros, and introspection. +# =========================================================================== + +# --------------------------------------------------------------------------- +# Ahead-of-time declaration +# --------------------------------------------------------------------------- + +""" + Declaration + +An active [`@expect_ops`](@ref) region. `ops === nothing` means the region only +establishes a lexical scope (so allocations inside it are keyed cheaply and +consistently); a non-`nothing` list is a user assertion about what will happen, +adopted with `prob = declared_prob`. +""" +struct Declaration + token::UInt64 + ops::Union{Nothing,Vector{Symbol}} + declared_prob::Float64 +end + +const ACTIVE_DECLARATION = ScopedValue{Union{Nothing,Declaration}}(nothing) + +""" + @expect_ops begin ... end + @expect_ops [:mul!, :cholesky!, :trsm!] begin ... end + @expect_ops [:cholesky!] prob=0.9 begin ... end + +Declare an operation chain ahead of time for allocations made inside the +region. + +Two distinct uses, both feeding the same planner: + +1. **With an operation list.** You are asserting what will happen to arrays + allocated in this region. The list is adopted as a forecast with confidence + `prob` (default `1.0`) and the planner optimises against it *immediately* — + no warm-up run required, which is the main reason to reach for this. + +2. **Without a list.** The region establishes a stable lexical site key and + mixes into the calling-context hash. Allocations inside are keyed by the + region rather than by an unwound stack, which makes + `CONFIG.site_id = :lexical` (free) or `:context` (nearly free) as precise as + `:backtrace` for the code you care about. Recording and prediction then + proceed normally from observation. + +The declared operations are matched positionally against argument position 1 +with arity 1. That is a simplification — it cannot express "this array is the +`B` of a `trsm!`" — but it covers the common case where the declared array is +the primary operand. + +TODO(richer-declarations): accept `:trsm! => 2` or a small tuple syntax to +declare argument position and arity, and accept per-array declarations rather +than region-wide ones (currently *every* allocation in the region adopts the +same chain, which is wrong when a region allocates both a matrix and its +right-hand side). + +# Example + +```julia +Dagger.@expect_ops [:mul!, :cholesky!, :trsm!] begin + A = rand(AutoBlocks(), Float64, n, n) + B = A * A' + n*I + cholesky!(B) + ldiv!(B, rhs) +end +``` +""" +macro expect_ops(args...) + isempty(args) && error("@expect_ops requires a body") + body = args[end] + rest = args[1:end-1] + + ops = nothing + prob = 1.0 + for a in rest + if @capture(a, prob = p_) + prob = p + elseif a isa Expr && (a.head === :vect || a.head === :tuple) + ops = a + else + error("@expect_ops: unexpected argument `$a`; expected an op list and/or `prob=...`") + end + end + + token = lexical_token(__module__, __source__) + opsexpr = ops === nothing ? :(nothing) : :(Symbol[$(map(esc, ops.args)...)]) + + quote + if $is_enabled() + local __parent__ = $(ACTIVE_DECLARATION)[] + local __tok__ = $mix_context(__parent__ === nothing ? $(CONTEXT_HASH)[] : + __parent__.token, $token) + local __decl__ = $Declaration(__tok__, $opsexpr, Float64($(esc(prob)))) + $(ScopedValues.with)($(ACTIVE_DECLARATION) => __decl__, + $(CONTEXT_HASH) => __tok__) do + $(esc(body)) + end + else + $(esc(body)) + end + end +end + +"Build a synthetic forecast from a user declaration." +function declared_prediction(d::Declaration, self::ArgSpec) + d.ops === nothing && return PredictedOp[] + p = clamp(d.declared_prob, 0.0, 1.0) + return [PredictedOp(OpKey(op, 1, 1), p, ArgSpec[self]) for op in d.ops] +end + +""" + site_for(token::UInt64) -> UInt64 + +The site identifier for the current context, accounting for an enclosing +[`@expect_ops`](@ref) region. + +**Every entry point that keys into the tape store must go through this.** +`plan_allocation`, `pin!` and `explain` each independently derive a site key, +and if they disagree then `pin!` silently pins a site nobody allocates from and +`explain` reports on a site that does not exist — both failures being invisible +rather than loud. +""" +@inline function site_for(token::UInt64) + decl = ACTIVE_DECLARATION[] + return decl === nothing ? current_site(token) : mix_context(decl.token, token) +end + +"Build the full store key for an allocation of `T` and shape `dims` at `token`." +@inline site_key(::Type{T}, dims::Tuple, token::UInt64) where {T} = + SiteKey(site_for(token), T, Int32(length(dims)), size_bucket(dims)) + +# --------------------------------------------------------------------------- +# Operation instrumentation +# --------------------------------------------------------------------------- + +""" + @record_op :opname A B C + +Record that operation `:opname` is about to be applied to the given arrays. +Place this immediately before the operation's task submission. + +Expands to a single `is_enabled()` branch plus a call, so it costs essentially +nothing when the subsystem is off. The macro-expansion site is baked in as a +literal and used both as the lexical site token and as the mixing constant for +the calling-context hash. + +**Pass only `DArray` arguments**, in a stable order, and keep that order in +sync with the matching [`@cost_model`](@ref): argument position is part of an +operation's recorded identity, because layout preference depends on role. + +```julia +function LinearAlgebra.mul!(C::DMatrix, A::DMatrix, B::DMatrix) + Dagger.@record_op :mul! C A B + ... +end +``` +""" +macro record_op(op, args...) + token = lexical_token(__module__, __source__) + quote + if $is_enabled() + $_record_op!($(esc(op)), $token, ($(map(esc, args)...),)) + end + nothing + end +end + +# --------------------------------------------------------------------------- +# Allocation hooks +# --------------------------------------------------------------------------- + +""" + AllocationPlan + +What [`plan_allocation`](@ref) decided. `partitioning` and `assignment` are +what the caller should actually allocate with; the rest is bookkeeping handed +back to [`track!`](@ref). +""" +struct AllocationPlan + partitioning::Union{Blocks,AutoBlocks} + assignment::Any + root::Union{Nothing,TapeRoot} + self::Union{Nothing,ArgSpec} + layout::Union{Nothing,LayoutChoice} + predicted::Vector{PredictedOp} + plan::Union{Nothing,LayoutPlan} +end + +"A plan that changes nothing — returned whenever the subsystem declines to act." +passthrough(part, assignment) = + AllocationPlan(part, assignment, nothing, nothing, nothing, PredictedOp[], nothing) + +""" + plan_allocation(::Type{T}, dims; requested=AutoBlocks(), assignment=:arbitrary, + token=UInt64(0)) -> AllocationPlan + +Decide the partitioning for an array about to be allocated. + +Declines to act — returning `requested` unchanged — when: + +- the subsystem is disabled; +- `requested` is a concrete `Blocks` and `CONFIG.override_explicit_blocks` is + false (explicit user intent wins; we only fill in what `AutoBlocks` would + otherwise have guessed); +- the site has been seen fewer than `CONFIG.min_observations` times; +- the forecast's confidence is below `CONFIG.min_confidence`; +- the planner's gate rejects the plan (insufficient margin over the fallback, + or too much regret at low confidence). + +Always registers the site, so a declined allocation still *learns*. + +Pair with [`track!`](@ref): + +```julia +p = Tapes.plan_allocation(T, dims; requested=dist, assignment=assignment) +A = _allocate(T, dims, p.partitioning, p.assignment) +return Tapes.track!(A, p) +``` +""" +function plan_allocation(::Type{T}, dims::Tuple; + requested = AutoBlocks(), + assignment = :arbitrary, + token::UInt64 = UInt64(0)) where {T} + is_enabled() || return passthrough(requested, assignment) + (requested isa Blocks && !CONFIG.override_explicit_blocks) && + return passthrough(requested, assignment) + isempty(dims) && return passthrough(requested, assignment) + + decl = ACTIVE_DECLARATION[] + key = site_key(T, dims, token) + root = get_root!(key) + + fb = fallback_layout(T, dims) + self = ArgSpec(T, dims, blocksize(fb), :arbitrary) + + # Manual override short-circuits everything. + pinned = root.pinned + if pinned !== nothing + return AllocationPlan(to_blocks(pinned), pinned.assignment, + root, ArgSpec(T, dims, blocksize(pinned), pinned.assignment), + pinned, PredictedOp[], nothing) + end + + # Forecast: a user declaration if present, otherwise the learned trie. + pred = PredictedOp[] + if decl !== nothing && decl.ops !== nothing + pred = declared_prediction(decl, self) + elseif root.nobservations >= CONFIG.min_observations + pred = predict(root.root) + end + + if isempty(pred) || confidence(pred) < CONFIG.min_confidence + vlog("no usable forecast for $(key.eltype)$(dims); using fallback $fb") + return AllocationPlan(requested, assignment, root, self, fb, pred, nothing) + end + + m = current_machine() + cands = candidate_layouts(T, dims, m) + lp = plan_chain(pred, cands, self, m) + + if !lp.accepted + vlog("plan rejected ($(lp.reason)); using fallback $fb") + return AllocationPlan(requested, assignment, root, self, fb, pred, lp) + end + + chosen = lp.steps[1] + vlog("chose $chosen for $(key.eltype)$(dims) ", + "(cost $(round(lp.cost; sigdigits=3))s vs fallback ", + "$(round(lp.fallback_cost; sigdigits=3))s, regret $(round(lp.max_regret; digits=2)))") + return AllocationPlan(to_blocks(chosen), chosen.assignment, root, + ArgSpec(T, dims, blocksize(chosen), chosen.assignment), + chosen, pred, lp) +end + +""" + track!(A::DArray, p::AllocationPlan) -> A + +Attach the recording state produced by [`plan_allocation`](@ref) to the +freshly-allocated array, and return the array so this can be used in tail +position. + +The trace is held in a `WeakKeyDict`, so it never extends `A`'s lifetime and +disappears with it. +""" +function track!(A::DArray, p::AllocationPlan) + root = p.root + (root === nothing || !is_enabled()) && return A + self = p.self + layout = p.layout + self === nothing && return A + layout === nothing && return A + + steps = p.plan === nothing ? LayoutChoice[layout] : p.plan.steps + trace = LiveTrace(root, self, layout, p.predicted, steps) + TRACES[A] = trace + lock(STORE_LOCK) do + root.nobservations += 1 + root.root.count += 1 + end + return A +end + +""" + suggest_partitioning(::Type{T}, dims; kwargs...) -> (partitioning, assignment) + +One-shot convenience wrapper for call sites that cannot conveniently thread an +[`AllocationPlan`](@ref) through to a [`track!`](@ref) call. The array is not +tracked, so it contributes nothing to future predictions — prefer +`plan_allocation` + `track!` wherever the allocation site can be edited +properly. +""" +function suggest_partitioning(::Type{T}, dims::Tuple; kwargs...) where {T} + p = plan_allocation(T, dims; kwargs...) + return (p.partitioning, p.assignment) +end + +# --------------------------------------------------------------------------- +# Manual override +# --------------------------------------------------------------------------- + +""" + pin!(::Type{T}, dims, blocks::Blocks, assignment=:arbitrary; token=UInt64(0)) + +Force a layout for a site, bypassing prediction entirely. Intended for +debugging and for benchmark reproducibility: with a pin in place, the +subsystem's behaviour at that site is deterministic and independent of warm-up +history. + +Note that pinning still requires the *site* to be identified, so in +`:backtrace` mode you must call `pin!` from the same context the allocation +happens in. In `:lexical` mode, wrap the allocation in [`@expect_ops`](@ref) +and pin against that region's token. +""" +function pin!(::Type{T}, dims::Tuple, blocks::Blocks, assignment::Symbol = :arbitrary; + token::UInt64 = UInt64(0)) where {T} + root = get_root!(site_key(T, dims, token)) + lock(STORE_LOCK) do + root.pinned = LayoutChoice(blocks.blocksize, assignment, :pinned) + end + return root +end + +"Remove all pins." +function unpin!() + lock(STORE_LOCK) do + for (_, r) in STORE + r.pinned = nothing + end + end + return nothing +end + +""" + clear!() + +Discard all learned tapes. Does not change [`CONFIG`](@ref). + +Essential for benchmarking: with tapes active, performance becomes +warm-up-dependent and bimodal, so any harness that reports a distribution +across repetitions needs to control explicitly whether each repetition starts +cold or warm. +""" +function clear!() + lock(STORE_LOCK) do + empty!(STORE) + TOTAL_NODES[] = 0 + end + empty!(TRACES) + return nothing +end + +# --------------------------------------------------------------------------- +# Introspection +# --------------------------------------------------------------------------- + +""" + stats() -> NamedTuple + +Aggregate counters: sites tracked, trie nodes, total observations, and +forecast hit/miss counts. +""" +function stats() + lock(STORE_LOCK) do + nobs = 0; hits = 0; misses = 0; nodes = 0 + for (_, r) in STORE + nobs += r.nobservations; hits += r.hits; misses += r.misses; nodes += r.nnodes + end + return (sites = length(STORE), nodes = nodes, observations = nobs, + hits = hits, misses = misses, + hit_rate = (hits + misses) > 0 ? hits / (hits + misses) : NaN, + live_traces = length(TRACES)) + end +end + +""" + report([io]; limit=20) + +Human-readable overview of what has been learned: one line per site with its +observation count and modal predicted chain. + +(Named `report` rather than `summary` to avoid shadowing `Base.summary` inside +this module.) +""" +function report(io::IO = stdout; limit::Int = 20) + s = stats() + println(io, "Dagger.Tapes: ", CONFIG.enabled ? "enabled" : "disabled", + " (site_id=:", CONFIG.site_id, ")") + println(io, " sites=", s.sites, " nodes=", s.nodes, + " observations=", s.observations, " live=", s.live_traces) + roots = lock(STORE_LOCK) do + sort!(collect(values(STORE)); by = r -> -r.nobservations) + end + for (i, r) in enumerate(roots) + i > limit && (println(io, " ... ", length(roots) - limit, " more"); break) + pred = predict(r.root) + chain = isempty(pred) ? "(no prediction)" : + join([string(p.key.op, "@", round(p.prob; digits = 2)) for p in pred], " -> ") + println(io, " [", r.nobservations, "x] ", r.key.eltype, "^", r.key.ndims, + " site=", string(r.key.site; base = 16)[1:min(end, 8)], + r.pinned === nothing ? "" : " PINNED=$(r.pinned)") + println(io, " ", chain) + end + return nothing +end + +""" + dump_tapes([io]) + +Print the full trie for every site, with per-branch counts. Verbose; intended +for debugging prediction quality when `report` shows something surprising. +""" +function dump_tapes(io::IO = stdout) + roots = lock(STORE_LOCK) do + sort!(collect(values(STORE)); by = r -> -r.nobservations) + end + for r in roots + println(io, "site ", string(r.key.site; base = 16), " ", r.key.eltype, + " ndims=", r.key.ndims, " obs=", r.nobservations) + _dump_node(io, r.root, 1) + end + return nothing +end + +function _dump_node(io::IO, n::TapeNode, indent::Int) + kids = n.children + kids === nothing && return nothing + for (k, c) in sort!(collect(kids); by = kv -> -kv[2].count) + stops = max(0, c.count - child_total(c)) + println(io, " "^(2 * indent), k, " x", c.count, + stops > 0 ? " (chain ended here x$stops)" : "") + _dump_node(io, c, indent + 1) + end + return nothing +end + +""" + explain([io], ::Type{T}, dims; requested=AutoBlocks(), assignment=:arbitrary) + +Show, for a hypothetical allocation of `T` and shape `dims` at the *current* +call site, what the subsystem predicts, which candidates it considered, what +each costs, and why it accepted or rejected the resulting plan. + +This is the first thing to reach for when a layout choice looks wrong. Call it +from the same context as the real allocation so the site key matches. + +```julia +Dagger.Tapes.explain(Float64, (4096, 4096)) +``` +""" +function explain(io::IO, ::Type{T}, dims::Tuple; + requested = AutoBlocks(), assignment = :arbitrary) where {T} + if !is_enabled() + println(io, "Tapes is disabled; every allocation uses the caller's request.") + return nothing + end + decl = ACTIVE_DECLARATION[] + key = site_key(T, dims, UInt64(0)) + site = key.site + root = get_root(key) + + println(io, "site = 0x", string(site; base = 16)) + println(io, "key = ", T, " ndims=", length(dims), " bucket=", + unpad_dims(size_bucket(dims), length(dims))) + if root === nothing + println(io, "status = unseen; would use fallback ", fallback_layout(T, dims)) + return nothing + end + println(io, "obs = ", root.nobservations, + " (min required ", CONFIG.min_observations, ")") + + pred = decl !== nothing && decl.ops !== nothing ? + declared_prediction(decl, ArgSpec(T, dims, blocksize(fallback_layout(T, dims)), :arbitrary)) : + predict(root.root) + if isempty(pred) + println(io, "forecast = none; would use fallback ", fallback_layout(T, dims)) + return nothing + end + println(io, "forecast =") + for (i, p) in enumerate(pred) + println(io, " ", i, ". ", p.key, " p=", round(p.prob; digits = 3), + has_cost_model(p.key.op) ? " [modeled]" : " [generic]") + end + + m = current_machine() + self = ArgSpec(T, dims, blocksize(fallback_layout(T, dims)), :arbitrary) + cands = candidate_layouts(T, dims, m) + println(io, "candidates (per-step expected cost, seconds):") + for c in cands + costs = [step_cost(p, self, c, m) for p in pred] + println(io, " ", rpad(string(c), 44), " total=", + round(sum(costs); sigdigits = 4), " ", + join([string(round(x; sigdigits = 3)) for x in costs], " ")) + end + + lp = plan_chain(pred, cands, self, m) + println(io, "plan = ", isempty(lp.steps) ? "(none)" : string(lp.steps[1]), + " (committed; receding horizon)") + length(lp.steps) > 1 && println(io, " full: ", join(string.(lp.steps), " | ")) + println(io, "cost = ", round(lp.cost; sigdigits = 4), + " fallback = ", round(lp.fallback_cost; sigdigits = 4), + " ratio = ", round(lp.cost / lp.fallback_cost; digits = 3), + " (must be <= ", CONFIG.gate_margin, ")") + println(io, "regret = ", round(lp.max_regret; digits = 3), + " (limit ", CONFIG.max_regret_ratio, " below confidence ", + CONFIG.regret_confidence_threshold, ")") + println(io, "verdict= ", lp.accepted ? "ACCEPTED" : "REJECTED ($(lp.reason))") + return nothing +end + +explain(::Type{T}, dims::Tuple; kwargs...) where {T} = explain(stdout, T, dims; kwargs...) +explain(::Type{T}, dims::Integer...; kwargs...) where {T} = explain(stdout, T, dims; kwargs...) diff --git a/src/tapes/cost.jl b/src/tapes/cost.jl new file mode 100644 index 000000000..d0e93e9ac --- /dev/null +++ b/src/tapes/cost.jl @@ -0,0 +1,564 @@ +# =========================================================================== +# Cost models. +# +# The tape tells us *which* operations are coming. To turn that into a layout +# decision we need `cost(op, layout, sizes) -> seconds` for combinations we may +# never have executed. A pure measurement table cannot answer counterfactuals +# ("what would Cholesky cost with 512-square tiles, given we have only ever run +# it with row blocks?") and the (op x layout x size) space is far too large to +# explore online. So the primary model is analytic and parametric; measurements +# are intended to *correct* it, not replace it. +# =========================================================================== + +"The module these macros expand references into. Captured so `@cost_model` can be used from anywhere." +const TAPES_MODULE = @__MODULE__ + +""" + MachineModel + +Coarse performance parameters of the current machine. Deliberately crude: the +planner only ever compares layouts against each other, so systematic error in +these constants largely cancels. What matters is that the *ratios* between +compute, communication and per-task overhead are roughly right. +""" +Base.@kwdef mutable struct MachineModel + "Number of compute processors available." + nprocs::Int = 1 + "Aggregate achievable FLOP/s across all processors, for Float64." + flops_per_sec::Float64 = 5.0e10 + "Effective inter-processor bandwidth in bytes/s." + bandwidth::Float64 = 5.0e9 + "Per-transfer latency in seconds." + latency::Float64 = 1.0e-5 + "Scheduler overhead per task in seconds. Dagger tasks bottom out around 100us." + task_overhead::Float64 = 1.0e-4 + "Usable memory per processor in bytes; used to reject layouts that will not fit." + mem_per_proc::Float64 = 8.0e9 +end + +const MACHINE = Ref{Union{Nothing,MachineModel}}(nothing) + +""" + current_machine() -> MachineModel + +The machine model in effect, constructed lazily from `Dagger.num_processors()` +on first use and cached. + +TODO(calibration): `flops_per_sec` and `bandwidth` are guesses. They should be +measured once at `enable!` time with a small GEMM and a small point-to-point +transfer (a few hundred milliseconds, amortised over the session), and refined +online from the scheduler's existing task timings — Dagger already records +per-task durations through `TimespanLogging`, which is exactly the signal +needed. This is also the natural place to hook the autotuner: it can supply +measured `(op, layout, size) -> seconds` points that override the analytic +model where they exist and leave it in place where they do not. + +TODO(heterogeneity): a single scalar FLOP rate is wrong on a CPU+GPU node, +which is the configuration Dagger most cares about. The model should carry a +per-processor-kind rate and the planner should weight by the fraction of tiles +each kind is expected to receive under the candidate assignment. Note that +`build_procgrid` currently filters to `ThreadProc` for symbolic assignments, so +today's assignment vocabulary cannot express a heterogeneous mapping anyway; +these two limitations should be lifted together. +""" +function current_machine() + m = MACHINE[] + m === nothing || return m + np = try + max(1, num_processors()) + catch + 1 + end + m = MachineModel(; nprocs = np) + MACHINE[] = m + return m +end + +""" + set_machine!(m::MachineModel) + set_machine!(; kwargs...) + +Install or update the machine model; keyword form updates fields in place. +""" +set_machine!(m::MachineModel) = (MACHINE[] = m) +function set_machine!(; kwargs...) + m = current_machine() + for (k, v) in kwargs + hasfield(MachineModel, k) || throw(ArgumentError("unknown MachineModel field: $k")) + setfield!(m, k, convert(fieldtype(MachineModel, k), v)) + end + return m +end + +"Reset the cached machine model so it is rebuilt on next use." +reset_machine!() = (MACHINE[] = nothing) + +# --------------------------------------------------------------------------- +# ArgView: what a cost model actually sees +# --------------------------------------------------------------------------- + +""" + ArgView + +An operation argument as presented to a cost model: shape and element type from +the recorded [`ArgSpec`](@ref), paired with the [`LayoutChoice`](@ref) being +evaluated for it. +""" +struct ArgView + spec::ArgSpec + layout::LayoutChoice +end + +Base.eltype(v::ArgView) = v.spec.eltype +Base.ndims(v::ArgView) = ndims(v.spec) +Base.size(v::ArgView) = size(v.spec) +Base.size(v::ArgView, d::Integer) = size(v.spec, d) +Base.length(v::ArgView) = prod(size(v); init = 1) + +_elsize(::Type{Nothing}) = 0 +_elsize(::Type{T}) where {T} = isbitstype(T) ? sizeof(T) : 8 + +"Bytes per element." +elsize(v::ArgView) = _elsize(v.spec.eltype) + +"Total bytes in the (logical) array." +nbytes(v::ArgView) = Float64(length(v)) * elsize(v) + +"Block size along dimension `d`, or the full tuple." +blocksize(v::ArgView) = blocksize(v.layout) +blocksize(v::ArgView, d::Integer) = blocksize(v.layout, d) + +"Number of blocks along dimension `d`, or the full tuple." +nblocks(v::ArgView) = nblocks(v.layout, size(v)) +nblocks(v::ArgView, d::Integer) = + d <= ndims(v) ? max(1, cld(size(v, d), max(1, blocksize(v, d)))) : 1 + +"Total number of blocks." +ntiles(v::ArgView) = prod(nblocks(v); init = 1) + +"Bytes in one block." +tilebytes(v::ArgView) = Float64(prod(blocksize(v); init = 1)) * elsize(v) + +"Processor assignment strategy of the candidate layout." +assignment(v::ArgView) = v.layout.assignment + +"Longest/shortest block edge ratio; `1.0` for a perfect cube." +function aspect(v::ArgView) + n = ndims(v) + n <= 1 && return 1.0 + bs = blocksize(v) + lo, hi = typemax(Int), 0 + for d in 1:n + b = min(bs[d], size(v, d)) + b <= 0 && continue + lo = min(lo, b); hi = max(hi, b) + end + (lo == typemax(Int) || lo <= 0) && return 1.0 + return hi / lo +end + +# --------------------------------------------------------------------------- +# The cost model registry +# --------------------------------------------------------------------------- + +"Operations with a hand-written [`@cost_model`](@ref), as opposed to the generic fallback." +const MODELED_OPS = Set{Symbol}() + +""" + op_cost(::Val{op}, args::Vector{ArgView}, m::MachineModel) -> Float64 + +Estimated wall-clock seconds for one execution of `op` with the given argument +layouts. Define methods with [`@cost_model`](@ref). + +The fallback delegates to [`generic_op_cost`](@ref), so unmodelled operations +still exert sensible (if blunt) pull on the planner rather than being ignored. +""" +op_cost(::Val{Op}, args::Vector{ArgView}, m::MachineModel) where {Op} = + generic_op_cost(args, m, Op) + +"Whether `op` has a hand-written cost model." +has_cost_model(op::Symbol) = op in MODELED_OPS + +"Guard against NaN / negative / infinite model output poisoning the planner." +@inline function clamp_cost(x::Float64) + (isnan(x) || x < 0) && return Inf + return x +end + +""" + @cost_model f(a, b, ...) = expr + @cost_model function f(a, b, ...) ... end + +Declare a cost model for the operation named `f`, returning estimated seconds. +Usable inside Dagger or from user code as `Dagger.Tapes.@cost_model`. + +The declared argument names are bound to [`ArgView`](@ref)s, **in the order the +operation was recorded with** — that is, the order of arguments passed to +[`@record_op`](@ref), not the order in the underlying function signature. By +convention only `DArray`s are recorded, so scalar arguments such as `alpha` or +`uplo` do not appear. + +An `ArgView` supports `size`, `ndims`, `eltype`, `length`, plus these, all of +which are bound as locals inside the body so no imports are needed: +`blocksize`, `nblocks`, `ntiles`, `tilebytes`, `nbytes`, `elsize`, `aspect`, +`assignment`. + +These helpers are also in scope: + +| helper | meaning | +|:-----------------|:------------------------------------------------------------| +| `flops_time(f)` | seconds for `f` FLOPs spread perfectly over all processors | +| `serial_time(f)` | seconds for `f` FLOPs on one processor (i.e. critical path) | +| `bytes_time(b)` | seconds to move `b` bytes, including one latency | +| `task_time(n)` | scheduler overhead for `n` tasks, spread over all processors | +| `imbalance(n)` | load-imbalance multiplier for `n` tiles over `nprocs` | +| `machine` | the [`MachineModel`](@ref) | +| `nprocs_avail()` | processor count | + +If the operation is recorded with fewer arguments than the model declares, the +generic fallback is used rather than erroring. + +# Example + +```julia +Dagger.Tapes.@cost_model my_solve(A, B) = begin + nt = nblocks(A, 1) + serial_time(nt * blocksize(A, 1)^2) + + flops_time(size(A, 1)^2 * size(B, 2)) * imbalance(nt) + + task_time(nt^2 / 2) +end +``` +""" +macro cost_model(ex) + fname = nothing; fargs = nothing; body = nothing + if @capture(ex, function fnm_(fa__) fb__ end) + fname = fnm; fargs = fa; body = Expr(:block, fb...) + elseif @capture(ex, fnm_(fa__) = fb_) + fname = fnm; fargs = fa; body = fb + else + error("@cost_model expects `f(args...) = expr` or `function f(args...) ... end`") + end + fname isa Symbol || error("@cost_model: operation name must be a plain symbol, got `$fname`") + all(a -> a isa Symbol, fargs) || + error("@cost_model: arguments must be plain names (no types, defaults or slurps)") + + nargs = length(fargs) + A = esc(gensym(:args)) + M = esc(gensym(:machine)) + opq = QuoteNode(fname) + + # Bind the argument names, and re-export the ArgView accessor vocabulary as + # locals so the body works regardless of what the caller has imported. + accessors = (:blocksize, :nblocks, :ntiles, :tilebytes, :nbytes, :elsize, + :aspect, :assignment) + accbinds = [:(local $(esc(s)) = $(getfield(TAPES_MODULE, s))) for s in accessors] + argbinds = [:(local $(esc(a)) = $A[$i]) for (i, a) in enumerate(fargs)] + + # Name the method as `.op_cost` with the module interpolated as a + # value. This is the form that works regardless of where the macro is + # expanded from: an unescaped bare `op_cost` relies on hygiene resolving to + # the macro's defining module (fragile for definitions), and an escaped + # `Dagger.Tapes.op_cost` needs `Dagger` to be bound in the caller's scope, + # which it is not while `Tapes` is itself still being defined. + fnexpr = Expr(:., TAPES_MODULE, QuoteNode(:op_cost)) + + quote + function $fnexpr(::$(Val{fname}), + $A::Vector{$ArgView}, + $M::$MachineModel) + if length($A) < $nargs + return $generic_op_cost($A, $M, $opq) + end + local $(esc(:machine)) = $M + local $(esc(:nprocs_avail)) = () -> $M.nprocs + local $(esc(:flops_time)) = f -> f / $M.flops_per_sec + local $(esc(:serial_time)) = f -> f / ($M.flops_per_sec / max(1, $M.nprocs)) + local $(esc(:bytes_time)) = b -> ($M.latency + b / $M.bandwidth) + local $(esc(:task_time)) = n -> n * $M.task_overhead / max(1, $M.nprocs) + local $(esc(:imbalance)) = n -> (n <= 0 ? 1.0 : + let np = max(1, $M.nprocs); ceil(n / np) / (n / np) end) + $(accbinds...) + $(argbinds...) + return $clamp_cost(Float64($(esc(body)))) + end + push!($MODELED_OPS, $opq) + $opq + end +end + +# --------------------------------------------------------------------------- +# Generic fallback +# --------------------------------------------------------------------------- + +""" + OP_AFFINITY + +Coarse block-shape preference per operation, consulted by +[`generic_op_cost`](@ref) when no analytic model exists. This is how an +unmodelled operation still exerts *some* pull on the planner. + +- `:square` — wants near-cubic tiles (dense factorizations, GEMM) +- `:rowwise` — wants whole rows resident per block +- `:colwise` — wants whole columns resident per block +- `:large` — indifferent to shape, wants few large blocks (elementwise, reductions) +- `:any` — genuinely indifferent + +TODO(learning): this table is hand-maintained and will rot as operations are +added. It should be *learned*: track measured cost against the aspect ratio and +tile count actually used per op, and fit the affinity rather than asserting it. +The tape already produces the `(op, layout, outcome)` triples needed, and this +is a much smaller fitting problem than learning a full cost model. +""" +const OP_AFFINITY = Dict{Symbol,Symbol}( + :cholesky => :square, :cholesky! => :square, :potrf! => :square, + :lu => :square, :lu! => :square, :getrf! => :square, + :qr => :square, :qr! => :square, :geqrf! => :square, + :svd => :square, :svd! => :square, + :mul! => :square, :gemm! => :square, :syrk! => :square, + :trsm! => :square, :ldiv! => :square, :rdiv! => :square, + :trsv! => :rowwise, + :transpose => :square, :adjoint => :square, :permutedims => :square, + :map => :large, :map! => :large, :broadcast => :large, + :materialize! => :large, :copyto! => :large, :fill! => :large, + :reduce => :large, :mapreduce => :large, :sum => :large, :norm => :large, + :stencil => :square, + :sort => :rowwise, :sort! => :rowwise, +) + +""" + shape_penalty(v::ArgView, op::Symbol, m::MachineModel) -> Float64 + +Multiplicative penalty in `[1, Inf)` for a layout that is a poor fit for `op`'s +known affinity, or that is unreasonable on its own terms: tiles too small to +amortise task overhead, tiles too large to fit a processor's memory, or too few +tiles to keep every processor busy. + +These three self-evident penalties matter more than the affinity term. They are +what make layout errors *asymmetric* in the model the way they are in reality — +too-small blocks collapse the scheduler, too-large blocks run out of memory — +so the planner will not trade a small predicted win for a catastrophic risk. +""" +function shape_penalty(v::ArgView, op::Symbol, m::MachineModel) + p = 1.0 + tb = tilebytes(v) + nt = ntiles(v) + + # A task must do enough work to be worth launching. + if tb > 0 + min_useful = m.task_overhead * m.bandwidth / max(1, m.nprocs) + tb < min_useful && (p *= min_useful / tb) + end + + # A tile (and the handful a blocked algorithm holds live) must fit. + budget = m.mem_per_proc / 4 + tb > budget && (p *= 1.0 + 8.0 * (tb / budget - 1.0)) + + # There must be enough tiles to occupy every processor. + nt < m.nprocs && (p *= m.nprocs / max(1, nt)) + + aff = get(OP_AFFINITY, op, :any) + if aff === :square + p *= 1.0 + 0.25 * log2(max(1.0, aspect(v))) + elseif aff === :rowwise + nd = ndims(v) + nd >= 2 && nblocks(v, nd) > 1 && (p *= 1.0 + 0.5 * log2(nblocks(v, nd))) + elseif aff === :colwise + ndims(v) >= 1 && nblocks(v, 1) > 1 && (p *= 1.0 + 0.5 * log2(nblocks(v, 1))) + elseif aff === :large + p *= 1.0 + 0.05 * log2(max(1, nt)) + end + + return p +end + +""" + generic_op_cost(args, m, op=:_unknown) -> Float64 + +Shape-agnostic cost estimate used when no [`@cost_model`](@ref) matches. + +Charges for data volume touched once, per-tile scheduler overhead, load +imbalance, per-tile latency, and the [`shape_penalty`](@ref). + +This also doubles as the planner's "unknown future operation" stand-in: the +residual probability mass of a prediction is charged at this rate, so layouts +that are pathological in general are penalised even when they are perfect for +the modal chain. That is what keeps a misprediction merely suboptimal instead +of catastrophic. +""" +function generic_op_cost(args::Vector{ArgView}, m::MachineModel, op::Symbol = :_unknown) + isempty(args) && return 0.0 + np = max(1, m.nprocs) + total = 0.0 + for v in args + eltype(v) === Nothing && continue + nt = ntiles(v) + c = nbytes(v) / m.bandwidth # touch the data once + c += nt * m.task_overhead / np # per-tile scheduling + c += nt * m.latency # per-tile boundary crossing + c *= nt <= 0 ? 1.0 : ceil(nt / np) / (nt / np) # load imbalance + c *= shape_penalty(v, op, m) + total += c + end + return clamp_cost(total) +end + +""" + cost_of(op, args, m=current_machine()) -> Float64 + +Dispatch to the registered model for `op`, falling back to the affinity-aware +generic model. +""" +cost_of(op::Symbol, args::Vector{ArgView}, m::MachineModel = current_machine()) = + clamp_cost(op_cost(Val(op), args, m)) + +""" + redistribution_cost(from, to, spec, m) -> Float64 + +Cost of changing an array's layout mid-chain. Charged as a full all-to-all of +the array's bytes plus per-tile latency and task overhead on both sides, +because in general every element lands on a different processor. + +TODO(precision): this is a deliberate overestimate in the cases where it is +wrong — `Blocks(1024,1024) -> Blocks(512,512)` under the same assignment is a +purely local subdivision that moves nothing. Computing the true volume needs +both proc grids and an intersection of the two `DomainBlocks`; `build_procgrid` +plus `src/datadeps/aliasing.jl` has the pieces. Overestimating is the safe +direction (it suppresses speculative repartitioning), so this is not urgent, +but it does leave real wins on the table for refinement-style changes. +""" +function redistribution_cost(from::LayoutChoice, to::LayoutChoice, spec::ArgSpec, + m::MachineModel = current_machine()) + from == to && return 0.0 + sz = size(spec) + bytes = Float64(prod(sz; init = 1)) * _elsize(spec.eltype) + nt_to = prod(nblocks(to, sz); init = 1) + nt_from = prod(nblocks(from, sz); init = 1) + return bytes / m.bandwidth + + (nt_to + nt_from) * (m.latency + m.task_overhead / max(1, m.nprocs)) +end + +# --------------------------------------------------------------------------- +# Default models for Dagger's built-in operations. +# +# These are first-order: they capture the terms that actually *differ* between +# layouts (tile count, critical-path length, communication volume) and ignore +# constants that do not, since the planner only ever compares layouts against +# each other. They are meant to be replaced by measurement-corrected versions +# once the autotuner is wired in. +# +# The argument order below matches the order these operations are instrumented +# with in `integration.jl` — keep the two in sync. +# --------------------------------------------------------------------------- + +@cost_model cholesky!(A) = begin + n = size(A, 1) + b = max(1, blocksize(A, 1)) + nt = nblocks(A, 1) + # nt POTRF panels sit on the critical path; the rest is parallel SYRK/GEMM/TRSM. + serial_time(nt * b^3 / 3) + + flops_time(Float64(n)^3 / 3) * imbalance(nt * (nt + 1) / 2) + + task_time(nt^3 / 3 + nt^2) + + bytes_time(nt^2 * b^2 * elsize(A)) + + # Rectangular tiles mismatch TRSM against SYRK and inflate the panel path. + serial_time(nt * b^3 / 3) * (aspect(A) - 1.0) +end + +@cost_model lu!(A) = begin + n = size(A, 1) + b = max(1, blocksize(A, 1)) + nt = nblocks(A, 1) + serial_time(nt * b^3) + + flops_time(2 * Float64(n)^3 / 3) * imbalance(nt * nt) + + task_time(nt^3 / 3 + nt^2) + + # Panel pivoting serialises down a block column. + bytes_time(nt * nt * b^2 * elsize(A)) * (1.0 + 0.5 * (aspect(A) - 1.0)) +end + +@cost_model qr!(A) = begin + mm = size(A, 1); nn = size(A, 2) + b = max(1, blocksize(A, 1)) + ntr = nblocks(A, 1); ntc = nblocks(A, 2) + serial_time(ntc * b^3) + + flops_time(2 * Float64(mm) * nn^2 - 2 * Float64(nn)^3 / 3) * imbalance(ntr * ntc) + + task_time(ntr * ntc^2) + + bytes_time(ntr * ntc * b^2 * elsize(A)) * (1.0 + 0.3 * (aspect(A) - 1.0)) +end + +@cost_model mul!(C, A, B) = begin + mm = size(A, 1); kk = size(A, 2); nn = size(B, 2) + ntm = nblocks(C, 1); ntn = nblocks(C, 2); ntk = nblocks(A, 2) + flops_time(2.0 * mm * kk * nn) * imbalance(ntm * ntn) + + task_time(ntm * ntn * ntk) + + # Each C tile pulls a block row of A and a block column of B. + bytes_time(ntm * ntn * ntk * + (prod(blocksize(A); init = 1) + prod(blocksize(B); init = 1)) * elsize(A)) + + # Mismatched inner blocking forces a repack of one operand. + (blocksize(A, 2) == blocksize(B, 1) ? 0.0 : bytes_time(nbytes(B))) +end + +@cost_model syrk!(C, A) = begin + nn = size(C, 1); kk = size(A, 2) + ntn = nblocks(C, 1); ntk = nblocks(A, 2) + flops_time(Float64(nn)^2 * kk) * imbalance(ntn * (ntn + 1) / 2) + + task_time(ntn * (ntn + 1) * ntk / 2) + + bytes_time(ntn * ntn * ntk * prod(blocksize(A); init = 1) * elsize(A)) +end + +@cost_model trsm!(A, B) = begin + nn = size(A, 1) + nd = ndims(B) + nrhs = nd >= 2 ? size(B, 2) : 1 + nta = nblocks(A, 1); ntb = max(1, nblocks(B, nd)) + # Triangular solve is inherently sequential along the block diagonal. + serial_time(nta * blocksize(A, 1)^2 * max(1, blocksize(B, nd))) + + flops_time(Float64(nn)^2 * nrhs) * imbalance(nta * ntb) + + task_time(nta * nta * ntb / 2) + + bytes_time(nta * ntb * prod(blocksize(B); init = 1) * elsize(B)) +end + +@cost_model trsv!(A, B) = begin + nn = size(A, 1) + nta = nblocks(A, 1) + serial_time(nta * blocksize(A, 1)^2) + + flops_time(Float64(nn)^2) * imbalance(nta) + + task_time(nta * (nta + 1) / 2) + + bytes_time(nta * nta * blocksize(A, 1) * elsize(A)) +end + +@cost_model map(A) = flops_time(Float64(length(A))) * imbalance(ntiles(A)) + task_time(ntiles(A)) +@cost_model map!(A) = flops_time(Float64(length(A))) * imbalance(ntiles(A)) + task_time(ntiles(A)) +@cost_model copyto!(A) = bytes_time(nbytes(A)) / max(1, nprocs_avail()) + task_time(ntiles(A)) + +@cost_model reduce(A) = begin + nt = ntiles(A) + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + (nt > 1 ? bytes_time(nt * elsize(A)) * log2(nt) : 0.0) # tree reduction over tiles +end + +@cost_model mapreduce(A) = begin + nt = ntiles(A) + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + (nt > 1 ? bytes_time(nt * elsize(A)) * log2(nt) : 0.0) +end + +@cost_model transpose(A) = begin + # A full transpose is an all-to-all unless the blocking is symmetric. + sym = ndims(A) >= 2 && blocksize(A, 1) == blocksize(A, 2) + bytes_time(nbytes(A)) * (sym ? 0.5 : 1.0) + task_time(ntiles(A)) +end + +@cost_model permutedims(A) = bytes_time(nbytes(A)) + task_time(ntiles(A)) + +@cost_model stencil(A) = begin + nt = ntiles(A) + tile = Float64(prod(blocksize(A); init = 1)) + halo = 0.0 + for d in 1:ndims(A) + halo += tile / max(1, blocksize(A, d)) # one face per dimension + end + flops_time(Float64(length(A))) * imbalance(nt) + + task_time(nt) + + bytes_time(2 * nt * halo * elsize(A)) +end diff --git a/src/tapes/integration.jl b/src/tapes/integration.jl new file mode 100644 index 000000000..1f12cc5bc --- /dev/null +++ b/src/tapes/integration.jl @@ -0,0 +1,263 @@ +# =========================================================================== +# Integration helpers. +# +# This file holds the glue that touches Dagger's existing code paths. It is +# kept separate from the rest of the subsystem so that the blast radius of the +# feature is obvious: everything else in `tapes/` is self-contained and has no +# effect until something here calls into it. +# +# It is included from `Tapes.jl`, so everything here lives in `Dagger.Tapes` +# and is referenced from Dagger as `Tapes.resolve_partitioning(...)` etc. +# =========================================================================== + +""" + resolve_partitioning(::Type{T}, dims, requested, assignment) -> (part, assignment, plan) + +The single entry point for allocators. Returns the partitioning and assignment +to allocate with, plus the [`AllocationPlan`](@ref) to hand to +[`track!`](@ref). + +When the subsystem is disabled this resolves `AutoBlocks()` exactly as before +(`auto_blocks(dims)`), so behaviour is bit-identical to the status quo. +""" +function resolve_partitioning(::Type{T}, dims::Tuple, requested, assignment) where {T} + if !is_enabled() + part = requested isa AutoBlocks ? auto_blocks(map(Int, dims)::Dims) : requested + return (part, assignment, passthrough(part, assignment)) + end + p = plan_allocation(T, dims; requested = requested, assignment = assignment) + part = p.partitioning isa AutoBlocks ? auto_blocks(map(Int, dims)::Dims) : p.partitioning + return (part, p.assignment, p) +end + +""" + @tracked_alloc T dims requested assignment expr + +Convenience wrapper for an allocation site. Binds `part` and `assign` for use +inside `expr` (which must evaluate to the new `DArray`), and tracks the result. + +```julia +function Base.zeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment = :arbitrary) + Dagger.Tapes.@tracked_alloc T dims p assignment begin + _zeros_impl(part, T, dims, assign) + end +end +``` +""" +macro tracked_alloc(T, dims, requested, assignment, expr) + quote + local __T__ = $(esc(T)) + local __d__ = $(esc(dims)) + local __part__, __assign__, __plan__ = + $resolve_partitioning(__T__, __d__, $(esc(requested)), $(esc(assignment))) + local $(esc(:part)) = __part__ + local $(esc(:assign)) = __assign__ + local __A__ = $(esc(expr)) + __A__ isa $DArray ? $track!(__A__, __plan__) : __A__ + end +end + +# =========================================================================== +# PATCH POINTS +# +# The changes below are written out rather than applied, because they touch +# files this subsystem does not own. Each is small and independently +# revertible. Nothing here changes behaviour while `CONFIG.enabled == false`. +# =========================================================================== + +#= + +## 1. `src/Dagger.jl` + +Insert after the `array/darray.jl` include (currently line 126), before +`array/alloc.jl`: + +```julia +include("tapes/Tapes.jl"); using .Tapes +include("tapes/integration.jl") +``` + +`using .Tapes` brings `@record_op`, `@cost_model` and `@expect_ops` into +`Dagger`, so users write `Dagger.@record_op` and `Dagger.@expect_ops`. + +This must come *after* `array/darray.jl` (which defines `DArray`, `Blocks`, +`AutoBlocks` and `auto_blocks`) and *before* `array/alloc.jl` and the +linear-algebra files, which call into it. + +## 2. `src/array/alloc.jl` — the allocation sites + +Every constructor that accepts `BlocksOrAuto` is a decision point. The pattern +is the same for `rand`, `randn`, `ones`, `zeros`, `sprand` and the +`DArray{T,N}(undef, ...)` family. Taking `zeros` as the example: + +```julia +function Base.zeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) + d = ArrayDomain(map(x -> 1:x, dims)) + s = reduce(vcat, partition(part, d)) + procgrid = build_procgrid(assign, dims, part.blocksize, current_acceleration()) + a = AllocateArray(T, AllocateZeros{T}(), false, d, s, part, procgrid) + return Tapes.track!(_to_darray(a), plan) +end + +# The AutoBlocks method must stop short-circuiting to `auto_blocks` so the +# planner gets a say: +Base.zeros(p::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + _zeros_planned(T, dims, p, assignment) +``` + +The important structural point: **`AutoBlocks` must reach +`resolve_partitioning` un-lowered**. At present `AutoBlocks` methods call +`auto_blocks(A)` immediately and forward to the `Blocks` method, which erases +the distinction between "the user asked for automatic" and "the user asked for +this specific blocking". The tape system needs that distinction — it declines +to override explicit `Blocks` by default. Threading `requested` through as +shown preserves it. + +## 3. `src/array/darray.jl` — `distribute` + +`distribute(A, AutoBlocks(), assignment)` is the other main entry point: + +```julia +function distribute(A::AbstractArray{T,N}, ::AutoBlocks, + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, size(A), AutoBlocks(), assignment) + return Tapes.track!(distribute(A, part, assign), plan) +end +``` + +## 4. Operation instrumentation + +One `@record_op` per operation, immediately before task submission. Argument +order must match the corresponding `@cost_model` in `tapes/cost.jl` — pass only +`DArray`s. + +- `src/array/cholesky.jl`, in `LinearAlgebra._chol!(A::DArray{T,2}, ...)`: + ```julia + Dagger.@record_op :cholesky! A + ``` + +- `src/array/lu.jl`, in the `lu!` entry point: + ```julia + Dagger.@record_op :lu! A + ``` + +- `src/array/qr.jl`, in the `qr!` entry point: + ```julia + Dagger.@record_op :qr! A + ``` + +- `src/array/mul.jl`, in `LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta)`: + ```julia + Dagger.@record_op :mul! C A B + ``` + and in `syrk_dagger!(C, transA, A, alpha, beta)`: + ```julia + Dagger.@record_op :syrk! C A + ``` + Note `generic_matvecmul!` should record as `:mul!` too, with `B` a `DVector`; + the cost model handles `ndims(B) == 1` via its `nd` branch. + +- `src/array/trsm.jl`: + ```julia + Dagger.@record_op :trsm! A B # in trsm! + Dagger.@record_op :trsv! A B # in trsv! + ``` + +- `src/array/map-reduce.jl`: `:map`, `:map!`, `:reduce`, `:mapreduce`. +- `src/array/copy.jl`: `:copyto!`. +- `src/array/permute.jl`: `:transpose`, `:permutedims`. +- `src/array/stencil.jl`: `:stencil`. +- `src/array/sort.jl`: `:sort!`. + +TODO(coverage): `src/array/operators.jl` broadcasting is the highest-value +missing instrumentation — broadcast fusion is where most real user code spends +its operations, and a fused broadcast's layout preference (`:large`, shape +indifferent) is a genuinely useful signal that the current chain misses. It +needs care because a single `materialize!` covers an arbitrary expression tree; +recording it as one `:broadcast` op loses which arrays participated. + +TODO(datadeps): operations submitted inside `spawn_datadeps` are the case where +this system should shine, because Datadeps already knows the full dependency +structure of the region. Rather than instrumenting each operation, the Datadeps +queue itself could record the op sequence per aliased region at +`distribute_tasks!` time — exact lookahead for the region, no prediction +required, which subsumes the tape for that scope. That is strictly better where +it applies and should be built before broadening `@record_op` coverage much +further. + +## 5. Precompilation + +`src/precompile.jl` should run at least one workload with the subsystem enabled +so the planner and macro-generated cost models are precompiled; otherwise the +first *enabled* allocation pays a large latency spike that will be mistaken for +tape overhead. + +=# + +# =========================================================================== +# DESIGN TODOs deferred by explicit decision +# =========================================================================== + +#= + +## Deferred materialization + +The cleanest solution to the whole problem is not to predict at all: keep an +array as a *logical description* until its first consumer is known, then +materialise it with a layout chosen from actual knowledge rather than a +forecast. That gets the first operation's layout exactly right with zero +prediction risk, and reduces the tape's job to optimising operations 2..n. + +Deferred: this requires restructuring how Datadeps tracks and materialises +arrays. Specifically, `DArray` would need an unmaterialised state that +`chunks` accessors can trigger materialisation from, every internal +`A.chunks` access would need to route through that, and the Datadeps aliasing +analysis would need to reason about a domain decomposition that does not exist +yet. `AllocateArray`/`stage` is the natural seam but the change reaches much +further than that. + +Note the interaction: once deferred materialization exists, the tape's role +shifts from "choose the layout" to "choose the layout for the operations after +the first", which makes the confidence gating *more* important, not less — +the easy win is gone and only the speculative part remains. + +## Multiple materialised layouts for read-only data + +For data that is read but not written, maintaining two physical layouts costs +memory but requires no coherence protocol at all. This is exactly what column +stores do with multiple sort-order projections (C-Store / Vertica). For a chain +where operation A wants square tiles and operations B..F want row blocks, +replicating can easily dominate repartitioning. + +Deferred: doing this without introducing OOM failures requires a memory +accounting system that Dagger does not have — something that tracks per-space +residency, knows the high-water mark of the current task graph, and can evict +or refuse a replica under pressure. MemPool has pieces of this +(`MemPool.approx_size`, the storage device abstraction) but there is no +global admission control. Building replication on top of the current +best-effort memory handling would turn a performance feature into a +reliability regression. + +The precondition is a data-residency tracker: (array, layout) -> space, bytes, +last-touch, read-only flag. That is the same structure needed for a proper +`redistribution_cost`, for heterogeneous cost modelling, and for out-of-core +support, so it is probably the highest-leverage missing piece overall. + +## Verified lookahead vs prediction + +A submission-window scheme — buffer N submitted tasks or T microseconds before +scheduling — gives *exact* lookahead with zero prediction risk, and beats this +system on any workload whose operation chain fits in one scope. The tape earns +its keep only where a window cannot reach: operations scattered across user +code, deep call stacks, library boundaries, and loop iterations that a bounded +buffer cannot span. + +These are complements, not alternatives, and the strongest configuration is +both: exact layouts for the verified near-term operations, predicted layouts +for the tail. `plan_chain` already accepts an arbitrary `Vector{PredictedOp}`, +so a window would simply contribute entries with `prob = 1.0` ahead of the +predicted ones. Wiring that up is cheap once a window exists; `@expect_ops` +with an explicit list is the manual version of the same thing. + +=# diff --git a/src/tapes/plan.jl b/src/tapes/plan.jl new file mode 100644 index 000000000..4d478ff9a --- /dev/null +++ b/src/tapes/plan.jl @@ -0,0 +1,387 @@ +# =========================================================================== +# Candidate generation and layout planning. +# +# Given a predicted operation chain and a set of candidate layouts, choosing +# the layout sequence that minimises total cost (including repartitioning on +# the edges) is exactly the classical automatic-data-layout problem of Kennedy +# & Kremer (TOPLAS 1998). They needed 0-1 integer programming because their +# phase graph was general. Ours is a *linear chain*, so it is a shortest path +# through a layered DAG and an exact O(n * |L|^2) dynamic program suffices. +# =========================================================================== + +const ASSIGNMENTS_ALL = (:arbitrary, :blockrow, :blockcol, :cyclicrow, :cycliccol) + +""" + fallback_layout(::Type{T}, dims) -> LayoutChoice + +The blind default: whatever `Dagger.auto_blocks` would pick, with `:arbitrary` +assignment (which spreads by compute capacity). Not good for much in +particular, but acceptable for most things, which is exactly what a fallback +should be. Every speculative plan is measured against this. +""" +function fallback_layout(::Type{T}, dims::Tuple) where {T} + ab = try + auto_blocks(map(Int, dims)::Dims) + catch + Blocks(ntuple(i -> max(1, Int(dims[i])), length(dims))) + end + return LayoutChoice(ab.blocksize, :arbitrary, :auto) +end + +""" + candidate_layouts(::Type{T}, dims, m) -> Vector{LayoutChoice} + +Enumerate plausible layouts for an array of element type `T` and shape `dims`. + +The set is deliberately small and structured rather than a dense sweep: the +planner is quadratic in candidate count, and the interesting decisions are +between *families* (square tiles vs row blocks vs column blocks) and between a +few block sizes within a family, not between 511 and 512. + +Block sizes are drawn from powers of two bounded below by "big enough that a +task amortises `task_overhead`" and above by "small enough that there are at +least a few tiles per processor". + +TODO(rank>2): the N-d cases here are a reasonable guess (split leading dim, +split trailing dim, near-cubic) but untested against real workloads. Tensor +contractions in particular want layouts derived from the contraction indices, +which the tape does not currently record. +""" +function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_machine()) + N = length(dims) + N == 0 && return LayoutChoice[] + esz = max(1, _elsize(T)) + np = max(1, m.nprocs) + out = LayoutChoice[] + + # Smallest tile worth scheduling, and largest that leaves >= 2 tiles/proc. + min_tile_bytes = m.task_overhead * m.bandwidth / np + min_edge = max(32, ceil(Int, (min_tile_bytes / esz)^(1 / N))) + total_bytes = Float64(prod(dims; init = 1)) * esz + max_tile_bytes = max(min_tile_bytes, total_bytes / (2 * np)) + max_edge = max(min_edge, floor(Int, (max_tile_bytes / esz)^(1 / N))) + + function push_unique!(l::LayoutChoice) + l in out || push!(out, l) + return nothing + end + + # --- Family 1: near-cubic tiles at a few sizes ------------------------- + edges = Int[] + e = 64 + while e <= 8192 + (e >= min_edge ÷ 2 && e <= max_edge * 2) && push!(edges, e) + e *= 2 + end + # Also the edge that yields exactly one tile per processor per dimension. + push!(edges, max(1, cld(maximum(dims), max(1, round(Int, np^(1 / N)))))) + for edge in unique!(sort!(edges)) + bs = ntuple(i -> clamp(edge, 1, max(1, Int(dims[i]))), N) + for a in (N >= 2 ? (:cyclicrow, :cycliccol, :arbitrary) : (:arbitrary,)) + push_unique!(LayoutChoice(bs, a, :square)) + end + end + + # --- Family 2: block rows (whole trailing dims resident) --------------- + for div in (np, 2np, 4np) + bs = ntuple(i -> i == 1 ? max(1, cld(Int(dims[1]), div)) : Int(dims[i]), N) + push_unique!(LayoutChoice(bs, :blockrow, :rowblock)) + push_unique!(LayoutChoice(bs, :arbitrary, :rowblock)) + end + + # --- Family 3: block columns (whole leading dims resident) ------------- + if N >= 2 + for div in (np, 2np, 4np) + bs = ntuple(i -> i == N ? max(1, cld(Int(dims[N]), div)) : Int(dims[i]), N) + push_unique!(LayoutChoice(bs, :blockcol, :colblock)) + push_unique!(LayoutChoice(bs, :arbitrary, :colblock)) + end + end + + # --- Family 4: whatever the blind default would have been -------------- + fb = fallback_layout(T, dims) + push_unique!(fb) + + # Trim to budget, spreading across families. The fallback is seeded first + # so it always survives: the planner needs it present as the baseline. + if length(out) > CONFIG.max_candidates + keep = LayoutChoice[fb] + per_family = max(1, cld(CONFIG.max_candidates - 1, 3)) + for fam in (:square, :rowblock, :colblock) + fams = filter(l -> l.label === fam, out) + isempty(fams) && continue + stride = max(1, cld(length(fams), per_family)) + for l in fams[1:stride:end] + l in keep || push!(keep, l) + end + end + out = length(keep) > CONFIG.max_candidates ? keep[1:CONFIG.max_candidates] : keep + end + return out +end + +# --------------------------------------------------------------------------- +# Turning a predicted operation into cost-model input +# --------------------------------------------------------------------------- + +""" + views_for(pred::PredictedOp, self::ArgSpec, layout) -> Vector{ArgView} + +Build the [`ArgView`](@ref) list a cost model expects for one predicted +operation: the tracked array at position `pred.key.pos` takes the candidate +`layout`; the other arguments take whatever layout they were last observed +with. + +TODO(co-argument layouts): using the *last observed* layout for co-arguments is +the weakest link in the whole chain. It means the planner evaluates "what if I +change only this array" rather than solving for a jointly consistent +assignment, so it systematically undervalues layouts that are only good when +all operands agree — precisely the layouts that matter for GEMM and TRSM. Fix +this together with the joint-planning TODO in `record_op!`: once arrays are +grouped into connected components, `views_for` should take a component-wide +assignment rather than a single layout. + +TODO(distributional specs): co-argument specs are the last observation, not a +distribution. If the same site alternates between a 1-column and 512-column +right-hand side, we silently plan for whichever ran last. +""" +function views_for(pred::PredictedOp, self::ArgSpec, layout::LayoutChoice) + pos = Int(pred.key.pos) + arity = Int(pred.key.arity) + specs = pred.argspecs + n = max(arity, length(specs), pos) + views = Vector{ArgView}(undef, n) + @inbounds for i in 1:n + spec = i <= length(specs) ? specs[i] : self + if i == pos + views[i] = ArgView(self, layout) + else + lay = LayoutChoice(unpad_dims(spec.blocksize, spec.ndims), + spec.assignment, :observed) + views[i] = ArgView(spec, lay) + end + end + return views +end + +""" + step_cost(pred, self, layout, m) -> Float64 + +Expected cost contribution of one predicted operation under `layout`. + +Weighted two ways: +- with probability `pred.prob` the predicted operation happens, and costs what + its model says; +- with the residual probability `1 - pred.prob` *something else* happens, which + we charge at the generic rate. That residual term is what makes the objective + risk-aware: a layout that is superb for the predicted chain but absurd in + general (one giant tile, or a million tiny ones) pays for that here. +""" +function step_cost(pred::PredictedOp, self::ArgSpec, layout::LayoutChoice, m::MachineModel) + views = views_for(pred, self, layout) + c_hit = cost_of(pred.key.op, views, m) + p = clamp(pred.prob, 0.0, 1.0) + p >= 1.0 && return c_hit + c_miss = generic_op_cost(ArgView[ArgView(self, layout)], m) + return p * c_hit + (1 - p) * c_miss +end + +# --------------------------------------------------------------------------- +# The planner +# --------------------------------------------------------------------------- + +""" + LayoutPlan + +Result of planning. `steps[i]` is the layout the planner wants for predicted +operation `i`; `steps[1]` is the only one actually committed under the +receding-horizon policy (see [`plan_chain`](@ref)). +""" +struct LayoutPlan + steps::Vector{LayoutChoice} + cost::Float64 + fallback_cost::Float64 + "Worst per-operation ratio of chosen layout to best-possible layout." + max_regret::Float64 + accepted::Bool + reason::Symbol +end + +""" + plan_chain(pred, cands, self, m; start=nothing) -> LayoutPlan + +Exact dynamic program over the predicted chain. + +`D[i, l]` is the minimum expected cost of executing predicted operations +`1..i` with operation `i` running under candidate `l`, including any +repartitioning charged on the edges. Backtracking the argmin over the final +layer gives the optimal layout sequence. + +If `start` is given it is the array's *current* layout, and the first step is +charged the cost of moving from it (used when re-planning mid-chain). + +# Policy note + +The plan covers the whole horizon, but callers commit only `steps[1]` and +re-plan at each subsequent observed operation. That is deliberate: a +speculative repartition is O(data) network traffic, so plan-and-commit exposes +the whole chain's worth of data movement to a single misprediction, whereas +receding-horizon control keeps the multi-operation lookahead benefit while +bounding exposure to one decision at a time. +""" +function plan_chain(pred::Vector{PredictedOp}, cands::Vector{LayoutChoice}, + self::ArgSpec, m::MachineModel; + start::Union{Nothing,LayoutChoice} = nothing) + n = length(pred) + L = length(cands) + fb = fallback_layout(self.eltype, size(self)) + + (n == 0 || L == 0) && + return LayoutPlan([fb], Inf, Inf, 1.0, false, :no_prediction) + + D = fill(Inf, n, L) + P = zeros(Int, n, L) + + # Per-operation cost of every candidate, reused for the regret bound. + C = Matrix{Float64}(undef, n, L) + @inbounds for i in 1:n, l in 1:L + C[i, l] = step_cost(pred[i], self, cands[l], m) + end + + @inbounds for l in 1:L + entry = start === nothing ? 0.0 : redistribution_cost(start, cands[l], self, m) + D[1, l] = C[1, l] + entry + end + @inbounds for i in 2:n + for l in 1:L + best = Inf; bi = 0 + for lp in 1:L + c = D[i-1, lp] + redistribution_cost(cands[lp], cands[l], self, m) + if c < best + best = c; bi = lp + end + end + D[i, l] = best + C[i, l] + P[i, l] = bi + end + end + + # Backtrack. + bestl = 1; bestc = Inf + @inbounds for l in 1:L + D[n, l] < bestc && (bestc = D[n, l]; bestl = l) + end + steps = Vector{LayoutChoice}(undef, n) + l = bestl + for i in n:-1:1 + steps[i] = cands[l] + i > 1 && (l = P[i, l]) + end + + # Same DP, pinned to the fallback layout throughout: the baseline to beat. + fb_idx = findfirst(==(fb), cands) + fb_cost = 0.0 + if fb_idx === nothing + fbviews_cost = 0.0 + for i in 1:n + fbviews_cost += step_cost(pred[i], self, fb, m) + end + fb_cost = fbviews_cost + (start === nothing ? 0.0 : + redistribution_cost(start, fb, self, m)) + else + for i in 1:n + fb_cost += C[i, fb_idx] + end + fb_cost += start === nothing ? 0.0 : redistribution_cost(start, cands[fb_idx], self, m) + end + + # Regret bound: how much worse is the committed layout than the best + # possible layout, for the single operation where it does worst? + committed = steps[1] + ci = findfirst(==(committed), cands) + max_regret = 1.0 + if ci !== nothing + @inbounds for i in 1:n + best_i = minimum(@view C[i, :]) + best_i > 0 && isfinite(best_i) && + (max_regret = max(max_regret, C[i, ci] / best_i)) + end + end + + conf = confidence(pred) + accepted = true + reason = :ok + if !isfinite(bestc) + accepted = false; reason = :nonfinite_cost + elseif bestc > CONFIG.gate_margin * fb_cost + # Not enough predicted win to justify departing from the known-adequate default. + accepted = false; reason = :insufficient_margin + elseif conf < CONFIG.regret_confidence_threshold && max_regret > CONFIG.max_regret_ratio + # Low confidence: refuse layouts that are catastrophic for any single + # predicted operation, even if they win on the weighted sum. + accepted = false; reason = :regret_too_high + end + + return LayoutPlan(steps, bestc, fb_cost, max_regret, accepted, reason) +end + +# --------------------------------------------------------------------------- +# Receding-horizon replanning +# --------------------------------------------------------------------------- + +""" + maybe_repartition!(A::DArray, trace::LiveTrace) + +Hook called after each recorded operation when `CONFIG.allow_repartition` is +set: re-plan from the array's current position in the trie and, if the plan +calls for a different layout and the predicted win exceeds the redistribution +cost, physically repartition `A`. + +**Currently a no-op stub.** The planning half is real (`replan` below); the +acting half is not, and deliberately so — see the TODO. + +TODO(repartition-mechanism): performing this safely requires a real +redistribution primitive that (a) is expressed as Datadeps tasks so it +interleaves with in-flight work rather than serialising behind it, (b) mutates +the `DArray` in place — `domain`, `subdomains`, `chunks` and `partitioning` all +have to change atomically from the perspective of concurrent readers — and +(c) is safe against the array being captured by tasks already submitted but not +yet scheduled. (c) is the hard part: Dagger's eager submission means there may +be queued thunks holding `Chunk` references into the old block structure. The +plausible approach is to route repartitioning through the Datadeps queue as a +normal `InOut` task on the whole array, which makes the dependency explicit and +lets the existing scheduler serialise it correctly, at the cost of a +synchronisation point. + +Until that exists, the planner's multi-step output is informational: it shows +up in `explain` and can guide manual `pin!` decisions, but only `steps[1]` is +ever acted on. +""" +function maybe_repartition!(A::DArray, trace::LiveTrace) + CONFIG.allow_repartition || return nothing + plan = replan(trace) + plan === nothing && return nothing + isempty(plan.steps) && return nothing + want = plan.steps[1] + want == trace.layout && return nothing + plan.accepted || return nothing + vlog("would repartition from $(trace.layout) to $want (not implemented)") + return nothing +end + +""" + replan(trace::LiveTrace) -> Union{Nothing,LayoutPlan} + +Re-run the planner from the array's current position in the trie, taking its +present layout as the starting state. This is the receding-horizon step: the +prediction is now conditioned on the prefix actually observed, so a chain that +diverged from the forecast is re-forecast against the branch it really took. +""" +function replan(trace::LiveTrace) + pred = predict(trace.node) + isempty(pred) && return nothing + self = trace.self + m = current_machine() + cands = candidate_layouts(self.eltype, size(self), m) + return plan_chain(pred, cands, self, m; start = trace.layout) +end diff --git a/src/tapes/tape.jl b/src/tapes/tape.jl new file mode 100644 index 000000000..8e88ffbbf --- /dev/null +++ b/src/tapes/tape.jl @@ -0,0 +1,608 @@ +# =========================================================================== +# Core types, allocation-site identification, and the operation-sequence trie. +# =========================================================================== + +""" +Maximum array rank tracked in fixed-width tuples. Ranks above this are still +handled, but their trailing dimensions collapse into the last slot, which only +costs precision in the site key (never correctness). +""" +const MAX_TRACKED_DIMS = 8 + +const DimTuple = NTuple{MAX_TRACKED_DIMS,Int} + +@inline function pad_dims(t::Tuple)::DimTuple + n = length(t) + ntuple(i -> i <= n ? Int(t[i]) : 0, MAX_TRACKED_DIMS) +end + +@inline unpad_dims(t::DimTuple, n::Integer) = ntuple(i -> t[i], Int(n)) + +# --------------------------------------------------------------------------- +# Keys +# --------------------------------------------------------------------------- + +""" + SiteKey + +Identifies an allocation *site* together with the coarse shape of what it +allocates. Two allocations share a key only if they came from the same program +context, have the same element type and rank, and fall into the same size +bucket — the last of these matters because the best block size for a 1024^2 +matrix is not the best block size for a 100000^2 matrix even at an identical +call site. +""" +struct SiteKey + site::UInt64 + eltype::DataType + ndims::Int32 + szbucket::DimTuple +end + +Base.hash(k::SiteKey, h::UInt) = + hash(k.szbucket, hash(k.ndims, hash(k.eltype, hash(k.site, h)))) +Base.:(==)(a::SiteKey, b::SiteKey) = + a.site == b.site && a.eltype === b.eltype && a.ndims == b.ndims && a.szbucket == b.szbucket + +""" + size_bucket(dims) -> DimTuple + +Quantise `dims` logarithmically so that nearby problem sizes share a tape. +`CONFIG.size_buckets_per_octave == 0` disables bucketing (exact sizes). +""" +function size_bucket(dims::Tuple)::DimTuple + bpo = CONFIG.size_buckets_per_octave + if bpo <= 0 + return pad_dims(dims) + end + n = length(dims) + ntuple(MAX_TRACKED_DIMS) do i + i > n && return 0 + d = Int(dims[i]) + d <= 0 ? 0 : round(Int, log2(d) * bpo) + end +end + +""" + OpKey + +An operation as seen *from the perspective of one of its arguments*. The +argument position is part of the identity because layout preference depends on +role: being the `A` of `mul!(C, A, B)` implies a different preference from +being the `C`. +""" +struct OpKey + op::Symbol + pos::Int16 + arity::Int16 +end +OpKey(op::Symbol, pos::Integer, arity::Integer) = OpKey(op, Int16(pos), Int16(arity)) + +Base.hash(k::OpKey, h::UInt) = hash(k.arity, hash(k.pos, hash(k.op, h))) +Base.:(==)(a::OpKey, b::OpKey) = a.op === b.op && a.pos == b.pos && a.arity == b.arity +Base.show(io::IO, k::OpKey) = print(io, k.op, "[", k.pos, "/", k.arity, "]") + +const ROOT_OPKEY = OpKey(Symbol("#root"), 0, 0) + +# --------------------------------------------------------------------------- +# Argument and layout descriptions +# --------------------------------------------------------------------------- + +""" + ArgSpec + +Metadata-only snapshot of a `DArray` argument. Deliberately holds no reference +to the array itself so that recorded tapes never keep data alive. +""" +struct ArgSpec + eltype::DataType + ndims::Int32 + size::DimTuple + blocksize::DimTuple + assignment::Symbol +end + +function ArgSpec(A::DArray{T,N}) where {T,N} + bs = A.partitioning isa Blocks ? A.partitioning.blocksize : size(A) + # N.B. `DArray` does not persist the assignment used to build it, so we + # cannot recover it here. + # TODO(assignment-provenance): add an `assignment` (or `procgrid`) field to + # `DArray` so recorded specs and cost models can reason about the actual + # processor mapping rather than assuming `:arbitrary`. Without it the + # planner can pick a good *block shape* but has to guess at the mapping + # quality of the arrays it did not itself allocate. + ArgSpec(T, Int32(N), pad_dims(size(A)), pad_dims(bs), :arbitrary) +end + +ArgSpec(::Type{T}, dims::Tuple, bs::Tuple, assignment::Symbol) where {T} = + ArgSpec(T, Int32(length(dims)), pad_dims(dims), pad_dims(bs), assignment) + +Base.eltype(s::ArgSpec) = s.eltype +Base.ndims(s::ArgSpec) = Int(s.ndims) +Base.size(s::ArgSpec) = unpad_dims(s.size, s.ndims) +Base.size(s::ArgSpec, d::Integer) = d <= s.ndims ? s.size[d] : 1 + +""" + LayoutChoice + +A candidate partitioning: a block shape plus a processor assignment strategy. +`label` records the family it came from, purely for explanation output. +""" +struct LayoutChoice + blocksize::DimTuple + ndims::Int32 + assignment::Symbol + label::Symbol +end + +LayoutChoice(bs::Tuple, assignment::Symbol, label::Symbol=:custom) = + LayoutChoice(pad_dims(bs), Int32(length(bs)), assignment, label) + +blocksize(l::LayoutChoice) = unpad_dims(l.blocksize, l.ndims) +blocksize(l::LayoutChoice, d::Integer) = d <= l.ndims ? l.blocksize[d] : 1 + +"Convert to a Dagger `Blocks{N}` suitable for passing to an allocator." +to_blocks(l::LayoutChoice) = Blocks(blocksize(l)) + +Base.:(==)(a::LayoutChoice, b::LayoutChoice) = + a.blocksize == b.blocksize && a.ndims == b.ndims && a.assignment === b.assignment +Base.hash(l::LayoutChoice, h::UInt) = hash(l.assignment, hash(l.ndims, hash(l.blocksize, h))) + +function Base.show(io::IO, l::LayoutChoice) + print(io, "Blocks", blocksize(l), " / :", l.assignment) + l.label === :custom || print(io, " (", l.label, ")") +end + +"Number of blocks along each dimension for an array of `dims` under `l`." +nblocks(l::LayoutChoice, dims::Tuple) = + ntuple(i -> max(1, cld(Int(dims[i]), max(1, blocksize(l, i)))), length(dims)) + +# --------------------------------------------------------------------------- +# Allocation-site identification +# --------------------------------------------------------------------------- + +""" + lexical_token(mod, source) -> UInt64 + +A compile-time-stable identifier for a macro expansion site. Computed at +macro-expansion time and baked into the expansion as a literal, so it costs +nothing at runtime. +""" +function lexical_token(mod::Module, source::LineNumberNode) + h = hash(nameof(mod), UInt64(0x9e3779b97f4a7c15)) + h = hash(something(source.file, :unknown), h) + h = hash(source.line, h) + return h % UInt64 +end + +""" + CONTEXT_HASH + +Probabilistic calling context, in the sense of Bond & McKinley (OOPSLA 2007). +Maintained incrementally as `V := 3V + site` at instrumentation boundaries +rather than by unwinding, giving context sensitivity for the price of a +multiply-add. Propagates into child tasks automatically because +`ScopedValue`s do. +""" +const CONTEXT_HASH = ScopedValue{UInt64}(UInt64(0)) + +@inline mix_context(v::UInt64, token::UInt64) = 3 * v + token + +""" + backtrace_hash() -> UInt64 + +Hash the raw instruction pointers of the current stack, skipping +`CONFIG.backtrace_skip` frames and taking at most `CONFIG.backtrace_depth`. + +Raw pointers are hashed rather than symbolicated frames because symbolication +dominates the cost by orders of magnitude. The consequence is that keys are +only meaningful within a single session and can shift if a frame is +recompiled at a different specialisation. + +TODO(persistence): to persist tapes across sessions — which is what would make +the very first run of a program fast rather than the second — we need a stable +key. Options, in increasing order of cost: (a) symbolicate lazily only when +serialising and re-resolve on load; (b) hash `(file, line)` pairs obtained from +`StackTraces.lookup`, cached per instruction pointer in an `IdDict` so each +unique frame is symbolicated once; (c) require `:lexical`/`:context` mode for +persistable tapes. (b) is probably the right default — the cache makes steady- +state cost comparable to raw pointer hashing. + +TODO(cost): even unsymbolicated, `backtrace()` is the single most expensive +thing this subsystem does. Consider (i) caching the hash in a task-local slot +keyed by a cheap discriminator so repeated allocations in a loop body unwind +once, or (ii) `jl_unw_stepn`-style bounded unwinding via a small C shim rather +than the full `backtrace()` allocation of a `Vector{Ptr}`. +""" +function backtrace_hash() + bt = backtrace() + n = length(bt) + lo = min(CONFIG.backtrace_skip + 1, n + 1) + hi = min(lo + CONFIG.backtrace_depth - 1, n) + h = UInt64(0xcbf29ce484222325) + @inbounds for i in lo:hi + h = (h ⊻ (reinterpret(UInt64, bt[i]) % UInt64)) * UInt64(0x100000001b3) + end + return h +end + +""" + current_site(token::UInt64) -> UInt64 + +Resolve the site identifier for the active `CONFIG.site_id` strategy. `token` +is the lexical token of the calling macro expansion. +""" +@inline function current_site(token::UInt64) + mode = CONFIG.site_id + if mode === :lexical + return token + elseif mode === :context + return mix_context(CONTEXT_HASH[], token) + else # :backtrace + # Fold in the ambient context too: a backtrace cannot see through a + # `Threads.@spawn` or `Dagger.@spawn` boundary, but the scoped context + # hash can, and `@expect_ops` regions set it deliberately. + return mix_context(CONTEXT_HASH[], backtrace_hash()) + end +end + +# --------------------------------------------------------------------------- +# The tape: a prefix trie over operation sequences +# --------------------------------------------------------------------------- + +""" + TapeNode + +One node of a per-site prefix trie. The path from the root to a node is an +observed operation sequence; `count` is how many times that exact prefix was +observed. + +A trie rather than a flat list of recorded sequences because it gives, for +free: incremental commit (no need to wait for an array to die before learning +from it), prefix-conditional prediction (re-planning mid-chain naturally +restricts to the branch actually taken), and per-step branch probabilities +that feed straight into the expected-cost objective. +""" +mutable struct TapeNode + const op::OpKey + const depth::Int + """ + Number of times this exact prefix was observed. For the root this is the + number of arrays allocated at the site (bumped by `track!`), which makes it + the correct denominator for `P(next op = k) = child.count / node.count`. + The residual `node.count - sum(children counts)` is the probability that + the chain *stopped* here — an array that is allocated and never used must + reduce our confidence, or a site that is 90% dead weight would look + perfectly predictable. + """ + count::Int + children::Union{Nothing,Dict{OpKey,TapeNode}} + "Most recently observed specs of *all* arguments of this operation." + argspecs::Vector{ArgSpec} +end + +TapeNode(op::OpKey, depth::Int) = TapeNode(op, depth, 0, nothing, ArgSpec[]) + +"Sum of child counts; `count - child_total` is the number of chains that stopped here." +function child_total(n::TapeNode) + kids = n.children + kids === nothing && return 0 + t = 0 + for (_, c) in kids + t += c.count + end + return t +end + +@inline function children!(n::TapeNode) + c = n.children + c === nothing || return c + d = Dict{OpKey,TapeNode}() + n.children = d + return d +end + +nchildren(n::TapeNode) = n.children === nothing ? 0 : length(n.children) + +""" + TapeRoot + +All history for one [`SiteKey`](@ref). +""" +mutable struct TapeRoot + const key::SiteKey + const root::TapeNode + nobservations::Int + nnodes::Int + hits::Int + misses::Int + last_used::Float64 + "Manual override installed by [`pin!`](@ref); bypasses prediction entirely." + pinned::Union{Nothing,LayoutChoice} +end + +TapeRoot(key::SiteKey) = + TapeRoot(key, TapeNode(ROOT_OPKEY, 0), 0, 1, 0, 0, time(), nothing) + +const STORE = Dict{SiteKey,TapeRoot}() +const STORE_LOCK = ReentrantLock() +const TOTAL_NODES = Ref(0) + +function get_root!(key::SiteKey) + lock(STORE_LOCK) do + r = get(STORE, key, nothing) + if r === nothing + maybe_evict!() + r = TapeRoot(key) + STORE[key] = r + TOTAL_NODES[] += 1 + end + r.last_used = time() + return r + end +end + +get_root(key::SiteKey) = lock(STORE_LOCK) do + get(STORE, key, nothing) +end + +""" +Evict least-recently-used sites when over budget. Called with `STORE_LOCK` held. + +TODO(eviction): LRU on sites is crude — a site observed 10000 times and not +touched for a minute is more valuable than one observed twice a second ago. +Weight by `nobservations` and by realised benefit (`hits`), i.e. evict by +`last_used - w*log(1+hits)`. +""" +function maybe_evict!() + (length(STORE) < CONFIG.max_sites && TOTAL_NODES[] < CONFIG.max_nodes) && return nothing + victims = sort!(collect(STORE); by = kv -> kv[2].last_used) + ndrop = max(1, length(victims) ÷ 4) + for i in 1:ndrop + k, r = victims[i] + TOTAL_NODES[] -= r.nnodes + delete!(STORE, k) + end + vlog("evicted $ndrop site(s); $(length(STORE)) remain") + return nothing +end + +# --------------------------------------------------------------------------- +# Live traces: the association between an allocated array and its tape +# --------------------------------------------------------------------------- + +""" + PredictedOp + +One step of a forecast: the operation, the cumulative probability of reaching +it, and a representative snapshot of the arguments it was last seen with. +""" +struct PredictedOp + key::OpKey + prob::Float64 + argspecs::Vector{ArgSpec} +end + +""" + LiveTrace + +Per-array recording state. Attached to a `DArray` in a `WeakKeyDict`, so it +disappears with the array and never extends its lifetime. +""" +mutable struct LiveTrace + const root::TapeRoot + node::TapeNode + nops::Int + truncated::Bool + """ + Spec of the array itself. Held separately from `root.key` because the key + stores a *bucketed* size, which is right for tape lookup but wrong for + costing — a plan computed against `2^round(log2(n))` would be off by up to + 41% in each dimension. + """ + const self::ArgSpec + "Layout actually chosen at allocation time." + const layout::LayoutChoice + "Forecast made at allocation time (for hit/miss accounting and re-planning)." + predicted::Vector{PredictedOp} + "Per-step layouts the planner would like; `plan[1]` is what was committed." + plan::Vector{LayoutChoice} + "How many recorded operations matched the forecast." + matched::Int + diverged::Bool +end + +LiveTrace(root::TapeRoot, self::ArgSpec, layout::LayoutChoice, + predicted::Vector{PredictedOp}, plan::Vector{LayoutChoice}) = + LiveTrace(root, root.root, 0, false, self, layout, predicted, plan, 0, false) + +const TRACES = WeakKeyDict{DArray,LiveTrace}() + +get_trace(A::DArray) = get(TRACES, A, nothing) +get_trace(@nospecialize(_)) = nothing + +# --------------------------------------------------------------------------- +# Recording +# --------------------------------------------------------------------------- + +""" + advance!(trace, opkey, specs) + +Extend `trace` by one observed operation, committing it into the trie +immediately. +""" +function advance!(trace::LiveTrace, opkey::OpKey, specs::Vector{ArgSpec}) + if trace.nops >= CONFIG.max_tape_length + trace.truncated = true + return nothing + end + + # Divergence accounting against the forecast made at allocation time. + idx = trace.nops + 1 + if idx <= length(trace.predicted) + # Unlocked on purpose: these are diagnostic counters on the hot path, + # and a torn increment costs us an inaccurate hit-rate readout, not + # correctness. Everything the planner actually reads is either + # trace-local or guarded by STORE_LOCK. + if trace.predicted[idx].key == opkey + trace.matched += 1 + trace.root.hits += 1 + else + trace.diverged = true + trace.root.misses += 1 + end + end + + node = trace.node + kids = children!(node) + child = get(kids, opkey, nothing) + if child === nothing + child = TapeNode(opkey, node.depth + 1) + kids[opkey] = child + lock(STORE_LOCK) do + trace.root.nnodes += 1 + TOTAL_NODES[] += 1 + end + end + child.count += 1 + child.argspecs = specs + trace.node = child + trace.nops += 1 + return nothing +end + +""" + record_op!(op::Symbol, token::UInt64, args::Tuple) + +Record that operation `op` is about to be applied to `args`. Called by +[`@record_op`](@ref); the macro supplies `token`. + +Every `DArray` in `args` that carries a live trace has the operation appended, +tagged with its own argument position. Arrays with no trace are ignored. + +TODO(joint-planning): this treats each argument's tape independently, but +layout decisions are *joint* — `mul!(C, A, B)` needs the three layouts to be +mutually compatible, not merely individually good. That is the alignment half +of the classical automatic-data-layout problem (Kennedy & Kremer, TOPLAS 1998) +and it is the harder half. The right structure is to record the co-participating +arrays' identities here, union-find them into connected components, and solve +layout over the component rather than per array. Dagger's Datadeps aliasing +machinery (`src/datadeps/aliasing.jl`) already computes most of the required +overlap information. + +TODO(adoption): when a traced array meets an untraced one (e.g. a user-supplied +`DArray` from `distribute`), we currently learn nothing about the untraced one. +It should be adopted into the traced one's component so the planner at least +knows a redistribution may be needed. +""" +function _record_op!(op::Symbol, token::UInt64, args::Tuple) + is_enabled() || return nothing + arity = length(args) + + # Fast path: if nothing here is traced, we only pay for the scan. + any_traced = false + @inbounds for a in args + if a isa DArray && haskey(TRACES, a) + any_traced = true + break + end + end + any_traced || return nothing + + specs = build_specs(args) + @inbounds for i in 1:arity + a = args[i] + a isa DArray || continue + trace = get_trace(a) + trace === nothing && continue + advance!(trace, OpKey(op, i, arity), specs) + CONFIG.allow_repartition && maybe_repartition!(a, trace) + end + return nothing +end + +""" + record_op!(op::Symbol, args...) + +Function form of [`@record_op`](@ref), for callers that build their argument +list dynamically. Carries no lexical token, so in `:context` site-identification +mode it contributes nothing to the context hash; prefer the macro where the +call site is statically known. +""" +record_op!(op::Symbol, args...) = _record_op!(op, UInt64(0), args) + +""" + build_specs(args) -> Vector{ArgSpec} + +Snapshot every `DArray` argument. Non-`DArray` arguments get a degenerate spec +so that positions line up with the caller's argument list. +""" +function build_specs(args::Tuple) + specs = Vector{ArgSpec}(undef, length(args)) + @inbounds for i in eachindex(args) + a = args[i] + specs[i] = a isa DArray ? ArgSpec(a) : ArgSpec(Nothing, (), (), :none) + end + return specs +end + +# --------------------------------------------------------------------------- +# Prediction +# --------------------------------------------------------------------------- + +""" + predict(node::TapeNode; horizon, min_prob) -> Vector{PredictedOp} + +Walk forward from `node`, greedily following the most-taken branch, and stop +when the branch probability or the cumulative probability falls below +threshold, when the chain is more likely than not to have ended, or when +`horizon` steps have been produced. + +This is a first-order Markov / prediction-by-partial-match walk over the trie. +Each returned step carries the *cumulative* probability of reaching it, which +the planner uses directly as the weight on that operation's cost — an operation +we are 40% sure about should contribute 40% of its cost to the objective. +""" +function predict(node::TapeNode; + horizon::Int = CONFIG.horizon, + min_prob::Float64 = CONFIG.min_branch_prob) + out = PredictedOp[] + cur = node + cum = 1.0 + for _ in 1:horizon + kids = cur.children + (kids === nothing || isempty(kids)) && break + + # Denominator is how many times we were *here*, not how many times we + # left: chains that stopped at this node count against the forecast. + total = max(cur.count, child_total(cur)) + total <= 0 && break + + best = nothing + bestcount = 0 + for (_, c) in kids + if c.count > bestcount + bestcount = c.count + best = c + end + end + best === nothing && break + + p = bestcount / total + p < min_prob && break + cum *= p + cum < min_prob && break + + push!(out, PredictedOp(best.op, cum, best.argspecs)) + cur = best + end + return out +end + +""" + confidence(pred) -> Float64 + +Cumulative probability of the first predicted step, i.e. how much we believe +anything at all is about to happen. +""" +confidence(pred::Vector{PredictedOp}) = isempty(pred) ? 0.0 : pred[1].prob diff --git a/src/tapes/tapes_tests.jl b/src/tapes/tapes_tests.jl new file mode 100644 index 000000000..15efcdf40 --- /dev/null +++ b/src/tapes/tapes_tests.jl @@ -0,0 +1,328 @@ +# Tests for the operation-tape layout predictor. +# +# Drop into `test/` and add `include("tapes.jl")` to `test/runtests.jl`. +# +# Most of these exercise the pure logic (site keys, trie, prediction, cost +# models, planner) with synthetic specs, so they run without a cluster and in +# milliseconds. The end-to-end group at the bottom needs real DArrays. + +using Test +using LinearAlgebra +using Dagger +using Dagger: Tapes +using Dagger.Tapes: SiteKey, OpKey, ArgSpec, ArgView, LayoutChoice, PredictedOp, + TapeNode, TapeRoot, MachineModel, + size_bucket, predict, advance!, plan_chain, candidate_layouts, + fallback_layout, cost_of, generic_op_cost, step_cost, + op_cost, has_cost_model, blocksize, nblocks, ntiles, + to_blocks, confidence, child_total + +# A deterministic machine so cost comparisons are reproducible. +const TESTMACHINE = MachineModel(nprocs = 16, + flops_per_sec = 1.0e12, + bandwidth = 1.0e10, + latency = 1.0e-5, + task_overhead = 1.0e-4, + mem_per_proc = 8.0e9) + +spec(T, dims, bs; assign = :arbitrary) = ArgSpec(T, dims, bs, assign) + +function fresh!() + Tapes.clear!() + Tapes.set_machine!(TESTMACHINE) + return nothing +end + +@testset "Tapes" begin + +@testset "site keys" begin + fresh!() + Tapes.CONFIG.size_buckets_per_octave = 2 + + # Nearby sizes share a bucket; distant ones do not. + @test size_bucket((1024, 1024)) == size_bucket((1100, 1100)) + @test size_bucket((1024, 1024)) != size_bucket((8192, 8192)) + + k1 = SiteKey(UInt64(7), Float64, Int32(2), size_bucket((1024, 1024))) + k2 = SiteKey(UInt64(7), Float64, Int32(2), size_bucket((1024, 1024))) + k3 = SiteKey(UInt64(7), Float32, Int32(2), size_bucket((1024, 1024))) + @test k1 == k2 + @test hash(k1) == hash(k2) + @test k1 != k3 # element type is part of identity + + Tapes.CONFIG.size_buckets_per_octave = 0 + @test size_bucket((1024, 1024)) != size_bucket((1100, 1100)) + Tapes.CONFIG.size_buckets_per_octave = 2 +end + +@testset "trie recording and prediction" begin + fresh!() + key = SiteKey(UInt64(1), Float64, Int32(2), size_bucket((1024, 1024))) + root = Tapes.get_root!(key) + self = spec(Float64, (1024, 1024), (256, 256)) + layout = LayoutChoice((256, 256), :arbitrary, :square) + + chain = [OpKey(:mul!, 1, 3), OpKey(:cholesky!, 1, 1), OpKey(:trsm!, 1, 2)] + + # Ten identical executions of the same chain. + for _ in 1:10 + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1 + root.root.count += 1 + for op in chain + advance!(tr, op, ArgSpec[self]) + end + end + + pred = predict(root.root) + @test length(pred) == 3 + @test [p.key.op for p in pred] == [:mul!, :cholesky!, :trsm!] + @test confidence(pred) ≈ 1.0 + + # Divergence: five runs take a different second step. + for _ in 1:5 + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1 + root.root.count += 1 + advance!(tr, OpKey(:mul!, 1, 3), ArgSpec[self]) + advance!(tr, OpKey(:lu!, 1, 1), ArgSpec[self]) + end + + pred = predict(root.root) + @test pred[1].key.op == :mul! + @test pred[1].prob ≈ 1.0 # first step is still certain + @test pred[2].key.op == :cholesky! # modal branch wins + @test 0.6 < pred[2].prob < 0.7 # 10/15 + + # Prefix conditioning: from the :mul! node the answer is the same, but + # from a node we have never reached there is nothing to say. + mul_node = root.root.children[OpKey(:mul!, 1, 3)] + @test predict(mul_node)[1].key.op == :cholesky! + @test child_total(mul_node) == 15 +end + +@testset "allocated-but-unused reduces confidence" begin + fresh!() + key = SiteKey(UInt64(2), Float64, Int32(2), size_bucket((512, 512))) + root = Tapes.get_root!(key) + self = spec(Float64, (512, 512), (128, 128)) + layout = LayoutChoice((128, 128), :arbitrary, :square) + + # One array gets used, nine are allocated and abandoned. + tr = Tapes.LiveTrace(root, self, layout, PredictedOp[], LayoutChoice[layout]) + root.nobservations += 1; root.root.count += 1 + advance!(tr, OpKey(:cholesky!, 1, 1), ArgSpec[self]) + for _ in 1:9 + root.nobservations += 1; root.root.count += 1 + end + + pred = predict(root.root) + # 1/10 is below min_branch_prob, so we should predict nothing at all + # rather than confidently predicting :cholesky!. + @test isempty(pred) +end + +@testset "cost models" begin + fresh!() + @test has_cost_model(:cholesky!) + @test has_cost_model(:mul!) + @test !has_cost_model(:definitely_not_an_op) + + s = spec(Float64, (8192, 8192), (512, 512)) + + square = LayoutChoice((512, 512), :cyclicrow, :square) + rowblk = LayoutChoice((512, 8192), :blockrow, :rowblock) + views(l) = ArgView[ArgView(s, l)] + + # Cholesky must prefer square tiles to fat row blocks. + @test cost_of(:cholesky!, views(square), TESTMACHINE) < + cost_of(:cholesky!, views(rowblk), TESTMACHINE) + + # An unmodelled op falls through to the generic model without erroring. + c = cost_of(:some_unknown_op, views(square), TESTMACHINE) + @test isfinite(c) && c > 0 + + # Degenerate layouts must be punished, not merely disfavoured. + tiny = LayoutChoice((1, 1), :arbitrary, :square) + huge = LayoutChoice((8192, 8192), :arbitrary, :square) + base = generic_op_cost(views(square), TESTMACHINE) + @test generic_op_cost(views(tiny), TESTMACHINE) > 10 * base + @test generic_op_cost(views(huge), TESTMACHINE) > base +end + +@testset "@cost_model defines a usable model" begin + fresh!() + Tapes.@cost_model __test_op(A, B) = flops_time(Float64(length(A)) * size(B, 1)) + + task_time(ntiles(A)) + @test has_cost_model(:__test_op) + sA = spec(Float64, (100, 100), (50, 50)) + sB = spec(Float64, (100, 10), (50, 10)) + lA = LayoutChoice((50, 50), :arbitrary, :square) + lB = LayoutChoice((50, 10), :arbitrary, :colblock) + v = ArgView[ArgView(sA, lA), ArgView(sB, lB)] + @test cost_of(:__test_op, v, TESTMACHINE) > 0 + + # Too few recorded arguments must degrade, not throw. + @test cost_of(:__test_op, ArgView[ArgView(sA, lA)], TESTMACHINE) > 0 +end + +@testset "candidate generation" begin + fresh!() + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + @test !isempty(cands) + @test length(cands) <= Tapes.CONFIG.max_candidates + @test fallback_layout(Float64, (8192, 8192)) in cands # baseline always present + @test any(c -> c.label === :square, cands) + @test any(c -> c.label === :rowblock, cands) + @test all(c -> all(>(0), blocksize(c)), cands) + + # Vectors must not produce degenerate 2D candidates. + cv = candidate_layouts(Float64, (100_000,), TESTMACHINE) + @test !isempty(cv) + @test all(c -> c.ndims == 1, cv) +end + +@testset "planner" begin + fresh!() + s = spec(Float64, (8192, 8192), (512, 512)) + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + + # A confident, homogeneous Cholesky chain should land on square tiles. + pred = [PredictedOp(OpKey(:cholesky!, 1, 1), 1.0, ArgSpec[s]) for _ in 1:3] + lp = plan_chain(pred, cands, s, TESTMACHINE) + @test lp.accepted + @test lp.cost <= lp.fallback_cost + @test length(lp.steps) == 3 + @test Dagger.Tapes.aspect(ArgView(s, lp.steps[1])) < 4.0 + + # An empty forecast must decline rather than guess. + lp0 = plan_chain(PredictedOp[], cands, s, TESTMACHINE) + @test !lp0.accepted + @test lp0.reason === :no_prediction + + # A low-confidence forecast is subject to the regret bound. + predlow = [PredictedOp(OpKey(:cholesky!, 1, 1), 0.3, ArgSpec[s])] + lplow = plan_chain(predlow, cands, s, TESTMACHINE) + @test lplow.accepted == false || lplow.max_regret <= Tapes.CONFIG.max_regret_ratio + + # The gate must reject a plan that does not beat the fallback by the margin. + old = Tapes.CONFIG.gate_margin + Tapes.CONFIG.gate_margin = 0.0 # nothing can beat this + @test !plan_chain(pred, cands, s, TESTMACHINE).accepted + Tapes.CONFIG.gate_margin = old +end + +@testset "planner accounts for repartitioning" begin + fresh!() + s = spec(Float64, (8192, 8192), (512, 512)) + cands = candidate_layouts(Float64, (8192, 8192), TESTMACHINE) + + # A chain that alternates between two operations with opposed preferences. + # Whatever the planner decides, it must not thrash: the number of distinct + # layouts in the plan should be small relative to the chain length, because + # each change is charged a full redistribution. + pred = PredictedOp[] + for i in 1:8 + op = isodd(i) ? :cholesky! : :reduce + push!(pred, PredictedOp(OpKey(op, 1, 1), 1.0, ArgSpec[s])) + end + lp = plan_chain(pred, cands, s, TESTMACHINE) + @test length(unique(lp.steps)) <= 3 +end + +@testset "enable/disable is inert by default" begin + fresh!() + Tapes.disable!() + @test !Tapes.is_enabled() + p = Tapes.plan_allocation(Float64, (1024, 1024)) + @test p.partitioning isa AutoBlocks # request passed through untouched + @test p.root === nothing + + Tapes.enable!(site_id = :lexical) + @test Tapes.is_enabled() + # An explicit Blocks must survive untouched unless explicitly overridden. + p = Tapes.plan_allocation(Float64, (1024, 1024); requested = Blocks(64, 64)) + @test p.partitioning == Blocks(64, 64) + Tapes.disable!() +end + +@testset "pinning overrides prediction" begin + fresh!() + Tapes.enable!(site_id = :lexical) + Tapes.pin!(Float64, (2048, 2048), Blocks(256, 256), :cyclicrow) + p = Tapes.plan_allocation(Float64, (2048, 2048)) + @test p.partitioning == Blocks(256, 256) + @test p.assignment === :cyclicrow + Tapes.unpin!() + Tapes.disable!() +end + +@testset "introspection does not throw" begin + fresh!() + Tapes.enable!(site_id = :lexical) + Tapes.plan_allocation(Float64, (1024, 1024)) + io = IOBuffer() + @test Tapes.report(io) === nothing + @test Tapes.dump_tapes(io) === nothing + @test Tapes.explain(io, Float64, (1024, 1024)) === nothing + @test Tapes.stats().sites >= 1 + Tapes.disable!() +end + +# --------------------------------------------------------------------------- +# End-to-end. Requires the `@record_op` / `plan_allocation` patch points from +# `tapes/integration.jl` to actually be applied to Dagger's allocators and +# linear-algebra routines; skipped otherwise. +# --------------------------------------------------------------------------- + +@testset "end-to-end learning" begin + fresh!() + Tapes.enable!(site_id = :backtrace, min_observations = 1, verbose = false) + try + n = 512 + function workload() + A = rand(AutoBlocks(), Float64, n, n) + B = A * A' + n * I + return cholesky!(B) + end + + workload() # cold: learns nothing yet + s1 = Tapes.stats() + workload() # warm: should now predict + s2 = Tapes.stats() + + @test s2.observations >= s1.observations + if s2.observations == 0 + @info "Tapes end-to-end: no allocations tracked; integration patches not applied" + else + @test s2.sites >= 1 + end + finally + Tapes.disable!() + Tapes.clear!() + Tapes.reset_machine!() + end +end + +@testset "@expect_ops declares ahead of time" begin + fresh!() + Tapes.enable!(site_id = :lexical, min_observations = 1) + try + n = 512 + Dagger.@expect_ops [:cholesky!, :trsm!] begin + # With a declaration in scope, the very first allocation should be + # planned rather than falling back, since confidence is asserted. + p = Tapes.plan_allocation(Float64, (n, n)) + @test p.root !== nothing + @test !isempty(p.predicted) + @test p.predicted[1].key.op === :cholesky! + @test p.predicted[1].prob ≈ 1.0 + end + finally + Tapes.disable!() + Tapes.clear!() + Tapes.reset_machine!() + end +end + +end # @testset "Tapes" From 09cf77b788029f7903a03018e45b9e471cabda45 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 28 Jul 2026 12:13:36 -0700 Subject: [PATCH 41/43] fixup! Automatic partitioning selection --- ext/FinchExt.jl | 18 +++++----- ext/SparseArraysExt.jl | 18 +++++----- src/Dagger.jl | 4 +++ src/array/alloc.jl | 40 ++++++++++------------- src/array/cholesky.jl | 2 ++ src/array/copy.jl | 18 ++++++---- src/array/darray.jl | 36 ++++++++++++-------- src/array/lu.jl | 4 +++ src/array/map-reduce.jl | 16 +++++++-- src/array/matrix.jl | 1 + src/array/mul.jl | 3 ++ src/array/operators.jl | 2 ++ src/array/permute.jl | 1 + src/array/qr.jl | 1 + src/array/sort.jl | 1 + src/array/stencil.jl | 1 + src/array/trsm.jl | 2 ++ src/precompile.jl | 20 ++++++++++++ src/tapes/api.jl | 1 + src/tapes/plan.jl | 2 +- src/tapes/tape.jl | 6 +++- test/runtests.jl | 1 + src/tapes/tapes_tests.jl => test/tapes.jl | 12 +++---- 23 files changed, 138 insertions(+), 72 deletions(-) rename src/tapes/tapes_tests.jl => test/tapes.jl (98%) diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl index be2f0e356..68bf8f3d8 100644 --- a/ext/FinchExt.jl +++ b/ext/FinchExt.jl @@ -13,12 +13,13 @@ Dagger._sparse_similar(::Finch.Tensor, ::Type{T}, dims::Dims) where {T} = Finch.fspzeros(T, dims...) Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseArray(x) -function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment; + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(part, d), part, assign; return_type=DSparseArray{T,N}) - return Dagger._to_darray(a) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) end Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = Finch.fspzeros(p, T, dims; assignment) @@ -26,15 +27,14 @@ Finch.fspzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = : Finch.fspzeros(p, Float64, dims; assignment) Finch.fspzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = Finch.fspzeros(p, Float64, dims; assignment) -Finch.fspzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - Finch.fspzeros(Dagger.auto_blocks(dims), T, dims; assignment) -function Finch.fsprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) +function Finch.fsprand(p::BlocksOrAuto, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment; + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(part, d), part, assign; return_type=DSparseArray{T,N}) - return Dagger._to_darray(a) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) end Finch.fsprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = Finch.fsprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) @@ -42,8 +42,6 @@ Finch.fsprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::Assignmen Finch.fsprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = Finch.fsprand(p, Float64, dims, sparsity; assignment) -Finch.fsprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - Finch.fsprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) # Materialize a (possibly lazy/`SwizzleArray`) result into a concrete sparse # `Tensor`. `@einsum` returns a lazy `SwizzleArray` for transposed-output diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl index 8820bd45a..757660d1c 100644 --- a/ext/SparseArraysExt.jl +++ b/ext/SparseArraysExt.jl @@ -84,14 +84,15 @@ function _to_device_sparse(like::Dagger.DeviceSparseMatrixCSC, S::SparseMatrixCS return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) end -function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) # Route through `allocate_sparse_zeros` so a GPU compute scope yields # device-resident sparse tiles (vendor sparse or DeviceSparseMatrixCSC). - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_zeros(Dagger.task_processor(), T, _dims)), false, d, Dagger.partition(p, d), p, assignment; + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_zeros(Dagger.task_processor(), T, _dims)), false, d, Dagger.partition(part, d), part, assign; return_type=DSparseArray{T,N}) - return Dagger._to_darray(a) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) end SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = SparseArrays.spzeros(p, T, dims; assignment) @@ -99,15 +100,14 @@ SparseArrays.spzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentTy SparseArrays.spzeros(p, Float64, dims; assignment) SparseArrays.spzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = SparseArrays.spzeros(p, Float64, dims; assignment) -SparseArrays.spzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - SparseArrays.spzeros(Dagger.auto_blocks(dims), T, dims; assignment) -function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) +function SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + part, assign, plan = Dagger.Tapes.resolve_partitioning(T, dims, p, assignment) d = Dagger.ArrayDomain(map(x->1:x, dims)) N = length(dims) - a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_rand(Dagger.task_processor(), T, _dims, sparsity)), false, d, Dagger.partition(p, d), p, assignment; + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_rand(Dagger.task_processor(), T, _dims, sparsity)), false, d, Dagger.partition(part, d), part, assign; return_type=DSparseArray{T,N}) - return Dagger._to_darray(a) + return Dagger.Tapes.track!(Dagger._to_darray(a), plan) end SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = SparseArrays.sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) @@ -115,8 +115,6 @@ SparseArrays.sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::Ass SparseArrays.sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) SparseArrays.sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = SparseArrays.sprand(p, Float64, dims, sparsity; assignment) -SparseArrays.sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - SparseArrays.sprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) _apply_trans(X, t::Char) = t == 'N' ? X : diff --git a/src/Dagger.jl b/src/Dagger.jl index bc6684c55..059f0b465 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -243,6 +243,10 @@ function __init__() catch err @warn "Error parsing JULIA_DAGGER_DEBUG" exception=err end + + # Enable tapes + Tapes.clear!() + Tapes.enable!() end end # module diff --git a/src/array/alloc.jl b/src/array/alloc.jl index 6ffbc8f50..ef926816e 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -75,10 +75,11 @@ end const BlocksOrAuto = Union{Blocks{N} where N, AutoBlocks} -function Base.rand(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.rand(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, rand, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, rand, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.rand(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = rand(p, T, dims; assignment) @@ -86,13 +87,12 @@ Base.rand(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbit rand(p, Float64, dims; assignment) Base.rand(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = rand(p, Float64, dims; assignment) -Base.rand(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - rand(auto_blocks(dims), T, dims; assignment) -function Base.randn(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.randn(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, randn, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, randn, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.randn(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = randn(p, T, dims; assignment) @@ -100,13 +100,12 @@ Base.randn(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbi randn(p, Float64, dims; assignment) Base.randn(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = randn(p, Float64, dims; assignment) -Base.randn(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - randn(auto_blocks(dims), T, dims; assignment) -function Base.ones(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.ones(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, ones, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, ones, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.ones(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = ones(p, T, dims; assignment) @@ -114,13 +113,12 @@ Base.ones(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbit ones(p, Float64, dims; assignment) Base.ones(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = ones(p, Float64, dims; assignment) -Base.ones(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - ones(auto_blocks(dims), T, dims; assignment) -function Base.zeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) +function Base.zeros(p::BlocksOrAuto, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + part, assign, plan = Tapes.resolve_partitioning(T, dims, p, assignment) d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, zeros, false, d, partition(p, d), p, assignment) - return _to_darray(a) + a = AllocateArray(T, zeros, false, d, partition(part, d), part, assign) + return Tapes.track!(_to_darray(a), plan) end Base.zeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = zeros(p, T, dims; assignment) @@ -128,8 +126,6 @@ Base.zeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbi zeros(p, Float64, dims; assignment) Base.zeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = zeros(p, Float64, dims; assignment) -Base.zeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = - zeros(auto_blocks(dims), T, dims; assignment) function Base.zero(x::DArray{T,N}) where {T,N} dims = ntuple(i->x.domain.indexes[i].stop, N) @@ -143,12 +139,12 @@ end function LinearAlgebra._zeros(::Type{T}, B::DVector, n::Integer) where T m = max(size(B, 1), n) sz = (m,) - return zeros(auto_blocks(sz), T, sz) + return zeros(AutoBlocks(), T, sz) end function LinearAlgebra._zeros(::Type{T}, B::DMatrix, n::Integer) where T m = max(size(B, 1), n) sz = (m, size(B, 2)) - return zeros(auto_blocks(sz), T, sz) + return zeros(AutoBlocks(), T, sz) end function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N} diff --git a/src/array/cholesky.jl b/src/array/cholesky.jl index 45ba39185..da780714e 100644 --- a/src/array/cholesky.jl +++ b/src/array/cholesky.jl @@ -17,6 +17,7 @@ function potrf_checked!(uplo, A, info_arr) return _A, info end function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T + Dagger.@record_op :cholesky! A LinearAlgebra.checksquare(A) zone = one(T) @@ -62,6 +63,7 @@ function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T return UpperTriangular(A), info[1] end function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{LowerTriangular}) where T + Dagger.@record_op :cholesky! A LinearAlgebra.checksquare(A) zone = one(T) diff --git a/src/array/copy.jl b/src/array/copy.jl index d703f49d1..28d68b835 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -143,17 +143,23 @@ function copyto_view!(Bpart, Brange, Apart, Arange) return end -Base.copyto!(B::DArray{T,N}, A::DArray{T,N}) where {T,N} = - darray_copyto!(B, A) -Base.copyto!(B::DArray{T,N}, A::Array{T,N}) where {T,N} = - darray_copyto!(B, view(A, B.partitioning)) +function Base.copyto!(B::DArray{T,N}, A::DArray{T,N}) where {T,N} + Dagger.@record_op :copyto! B + return darray_copyto!(B, A) +end +function Base.copyto!(B::DArray{T,N}, A::Array{T,N}) where {T,N} + Dagger.@record_op :copyto! B + return darray_copyto!(B, view(A, B.partitioning)) +end Base.copyto!(B::Array{T,N}, A::DArray{T,N}) where {T,N} = darray_copyto!(view(B, A.partitioning), A) StridedDArray{T,N} = Union{<:DArray{T,N}, SubArray{T,N,<:DArray{T,NP}} where NP} -Base.copyto!(B::StridedDArray, A::StridedDArray) = - darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +function Base.copyto!(B::StridedDArray, A::StridedDArray) + Dagger.@record_op :copyto! parent(B) + return darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +end function Base.copyto!(B::Array, A::StridedDArray) DB = view(B, AutoBlocks()) darray_copyto!(DB, parent(A), parentindices(DB), parentindices(A)) diff --git a/src/array/darray.jl b/src/array/darray.jl index bb0442714..6467a2ea5 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -546,7 +546,11 @@ function distribute(A::AbstractArray{T,N}, dist::Blocks{N}, assignment::Assignme return _to_darray(Distribute(dist, A, procgrid)) end -distribute(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = distribute(A, auto_blocks(A), assignment) +function distribute(A::AbstractArray{T,N}, ::AutoBlocks, + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, size(A), AutoBlocks(), assignment) + return Tapes.track!(distribute(A, part, assign), plan) +end function distribute(x::AbstractArray{T,N}, n::NTuple{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N} p = map((d, dn)->ceil(Int, d / dn), size(x), n) distribute(x, Blocks(p), assignment) @@ -562,33 +566,39 @@ DVector(A::AbstractVector{T}, assignment::AssignmentType{1} = :arbitrary) where DMatrix(A::AbstractMatrix{T}, assignment::AssignmentType{2} = :arbitrary) where T = DMatrix(A, AutoBlocks(), assignment) DArray(A::AbstractArray, assignment::AssignmentType = :arbitrary) = DArray(A, AutoBlocks(), assignment) -DVector(A::AbstractVector{T}, ::AutoBlocks, assignment::AssignmentType{1} = :arbitrary) where T = DVector(A, auto_blocks(A), assignment) -DMatrix(A::AbstractMatrix{T}, ::AutoBlocks, assignment::AssignmentType{2} = :arbitrary) where T = DMatrix(A, auto_blocks(A), assignment) -DArray(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = DArray(A, auto_blocks(A), assignment) +DVector(A::AbstractVector{T}, ::AutoBlocks, assignment::AssignmentType{1} = :arbitrary) where T = + distribute(A, AutoBlocks(), assignment) +DMatrix(A::AbstractMatrix{T}, ::AutoBlocks, assignment::AssignmentType{2} = :arbitrary) where T = + distribute(A, AutoBlocks(), assignment) +DArray(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = + distribute(A, AutoBlocks(), assignment) struct AllocateUndef{S} end (::AllocateUndef{S})(T, dims::Dims{N}) where {S,N} = Array{S,N}(undef, dims) -function DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} +function DArray{T,N}(::UndefInitializer, dist::Union{Blocks{N},AutoBlocks}, dims::NTuple{N,Int}; + assignment::AssignmentType = :arbitrary) where {T,N} + part, assign, plan = Tapes.resolve_partitioning(T, dims, dist, assignment) domain = ArrayDomain(map(x->1:x, dims)) - subdomains = partition(dist, domain) - a = AllocateArray(T, AllocateUndef{T}(), false, domain, subdomains, dist, assignment) - return _to_darray(a) + subdomains = partition(part, domain) + a = AllocateArray(T, AllocateUndef{T}(), false, domain, subdomains, part, assign) + return Tapes.track!(_to_darray(a), plan) end DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, (dims...,); assignment) DArray{T,N}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks(dims), dims; assignment) + DArray{T,N}(undef, AutoBlocks(), dims; assignment) DArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment) + DArray{T,N}(undef, AutoBlocks(), (dims...,); assignment) -DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = +DArray{T}(::UndefInitializer, dist::Union{Blocks{N},AutoBlocks}, dims::NTuple{N,Int}; + assignment::AssignmentType = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, dims; assignment) DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = DArray{T,N}(undef, dist, (dims...,); assignment) DArray{T}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks(dims), dims; assignment) + DArray{T,N}(undef, AutoBlocks(), dims; assignment) DArray{T}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} = - DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment) + DArray{T,N}(undef, AutoBlocks(), (dims...,); assignment) function DArray(A::WrappedDArray{T}; assignment::AssignmentType = :arbitrary) where T B = DArray{T}(undef, size(A); assignment) diff --git a/src/array/lu.jl b/src/array/lu.jl index 45d0c2c38..52575155e 100644 --- a/src/array/lu.jl +++ b/src/array/lu.jl @@ -3,10 +3,12 @@ LinearAlgebra.lu(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlge LinearAlgebra.lu!(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlgebra.NoPivot} = LinearAlgebra.RowMaximum(); check::Bool=true, allowsingular::Bool=false) where {T<:LinearAlgebra.BlasFloat} = LinearAlgebra.lu(A, pivot; check=check, allowsingular=allowsingular) function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.NoPivot(); check) end function LinearAlgebra.lu!(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A check && LinearAlgebra.LAPACK.chkfinite(A) zone = one(T) @@ -260,10 +262,12 @@ end # Implementation of https://inria.hal.science/hal-04984070v1/file/ipdps_paper.pdf function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.RowMaximum; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.RowMaximum(); check, allowsingular) end function LinearAlgebra.lu!(A::DMatrix{T}, ::LinearAlgebra.RowMaximum; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + Dagger.@record_op :lu! A check && LinearAlgebra.LAPACK.chkfinite(A) zone = one(T) diff --git a/src/array/map-reduce.jl b/src/array/map-reduce.jl index 6e16a912b..b0d373943 100644 --- a/src/array/map-reduce.jl +++ b/src/array/map-reduce.jl @@ -41,7 +41,10 @@ function stage(ctx::Context, node::Map) return DArray(RT, domain(primary), domainchunks(primary), thunks, partitioning, concat) end -Base.map(f, x::DArray, xs::DArray...) = _to_darray(Map(f, (x, xs...))) +function Base.map(f, x::DArray, xs::DArray...) + Dagger.@record_op :map x + return _to_darray(Map(f, (x, xs...))) +end #### MapReduce @@ -139,8 +142,15 @@ end import Statistics: mean, var, std import OnlineStats -Base.mapreduce(f::Function, op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) = - _mapreduce_maybesync(f, nothing, op, x, dims, init) +function Base.mapreduce(f::Function, op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) + Dagger.@record_op :mapreduce x + return _mapreduce_maybesync(f, nothing, op, x, dims, init) +end + +function Base.reduce(op::Function, x::DArray; dims=nothing, init=Base._InitialValue()) + Dagger.@record_op :reduce x + return _mapreduce_maybesync(identity, nothing, op, x, dims, init) +end Base.sum(x::DArray; dims=nothing, init=Base._InitialValue()) = sum(identity, x; dims, init) diff --git a/src/array/matrix.jl b/src/array/matrix.jl index b9935f662..9b7e77c1d 100644 --- a/src/array/matrix.jl +++ b/src/array/matrix.jl @@ -1,6 +1,7 @@ # Transpose/Adjoint function copydiag(f, A::DArray{T, 2}) where T + Dagger.@record_op (f === Transpose ? :transpose : :adjoint) A Ac = A.chunks Ac_copy = Matrix{Any}(undef, size(Ac, 2), size(Ac, 1)) _copytile(f, Ac) = copy(f(Ac)) diff --git a/src/array/mul.jl b/src/array/mul.jl index 322b27760..4d07c9f0e 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -88,6 +88,7 @@ function LinearAlgebra.generic_matmatmul!( alpha::Number, beta::Number, ) where {T} + Dagger.@record_op :mul! C A B partC, partA, partB = _repartition_matmatmul(C, A, B, transA, transB) if all(in(('N', 'T', 'C')), (transA, transB)) @@ -296,6 +297,7 @@ function syrk_dagger!( _alpha, _beta, ) where {T} + Dagger.@record_op :syrk! C A Ac = A.chunks Cc = C.chunks @@ -462,6 +464,7 @@ function LinearAlgebra.generic_matvecmul!( _alpha::Number, _beta::Number, ) where {T} + Dagger.@record_op :mul! C A B partC, partA, partB = _repartition_matvecmul(C, A, B, transA) return maybe_copy_buffered(C=>partC, A=>partA, B=>partB) do C, A, B return gemv_dagger!(C, transA, A, B, _alpha, _beta) diff --git a/src/array/operators.jl b/src/array/operators.jl index f93f09a75..837f89fe0 100644 --- a/src/array/operators.jl +++ b/src/array/operators.jl @@ -177,6 +177,7 @@ function imap!(f, A) end function Base.map!(f, a::DArray{T}) where T + Dagger.@record_op :map! a Dagger.spawn_datadeps() do for ca in chunks(a) Dagger.@spawn imap!(f, InOut(ca)) @@ -186,6 +187,7 @@ function Base.map!(f, a::DArray{T}) where T end function Base.map!(f, a::DArray{T}, b::AbstractArray{U}) where {T, U} + Dagger.@record_op :map! a b2 = view(b, a.partitioning) Dagger.spawn_datadeps() do for (c_a, c_b2) in zip(chunks(a), chunks(b2)) diff --git a/src/array/permute.jl b/src/array/permute.jl index 9da51e70b..08b204ca8 100644 --- a/src/array/permute.jl +++ b/src/array/permute.jl @@ -1,4 +1,5 @@ function Base.permutedims(A::DArray{T,N}, perm) where {T,N} + Dagger.@record_op :permutedims A dc = domainchunks(A) new_dc = DomainBlocks( ntuple(i -> dc.start[perm[i]], N), diff --git a/src/array/qr.jl b/src/array/qr.jl index 4409cc1c5..e35a7fcd7 100644 --- a/src/array/qr.jl +++ b/src/array/qr.jl @@ -833,6 +833,7 @@ In-place tiled Compact-WY QR factorization of a distributed matrix. than `p=1`. """ function LinearAlgebra.qr!(A::DMatrix{T}; ib::Union{Int,Nothing}=nothing, p::Int=1) where {T<:Number} + Dagger.@record_op :qr! A p >= 1 || throw(ArgumentError("p must be >= 1, got $p")) ib === nothing || ib >= 1 || throw(ArgumentError("ib must be >= 1, got $ib")) diff --git a/src/array/sort.jl b/src/array/sort.jl index 7abbe373f..109647c46 100644 --- a/src/array/sort.jl +++ b/src/array/sort.jl @@ -301,6 +301,7 @@ function Base.sort(v::ArrayOp; nsamples=2000, order::Ordering=default_ord) v1 = fetch(v) + Dagger.@record_op :sort! v1 ord = Base.Sort.ord(lt,by,rev,order) nchunks = nchunks === nothing ? length(v1.chunks) : nchunks cs = dsort_chunks(v1.chunks, nchunks, nsamples, diff --git a/src/array/stencil.jl b/src/array/stencil.jl index 14a5f8801..c0c9e272b 100644 --- a/src/array/stencil.jl +++ b/src/array/stencil.jl @@ -1304,6 +1304,7 @@ macro stencil(orig_ex) # region blocks until all of its tasks finish, the next expression's tasks are # only submitted once this expression has been applied everywhere, which is what # gives `@stencil` its "all at once" semantics. + push!(final_ex.args, :(Dagger.@record_op :stencil $write_var)) push!(final_ex.args, :(Dagger.spawn_datadeps() do for $chunk_idx in $CartesianIndices($chunks($write_var)) $inner_spawn_ex diff --git a/src/array/trsm.jl b/src/array/trsm.jl index c0c025468..81e098e7a 100644 --- a/src/array/trsm.jl +++ b/src/array/trsm.jl @@ -1,4 +1,5 @@ function trsv!(uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B::DVector{T}) where T + Dagger.@record_op :trsv! A B zone = one(T) mzone = -one(T) @@ -57,6 +58,7 @@ function trsv!(uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B:: end function trsm!(side::Char, uplo::Char, trans::Char, diag::Char, alpha::T, A::DMatrix{T}, B::DVecOrMat{T}) where T + Dagger.@record_op :trsm! A B zone = one(T) mzone = -one(T) diff --git a/src/precompile.jl b/src/precompile.jl index 34b14c503..c3031b224 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -9,6 +9,26 @@ t2 = spawn(+, 1, t1) fetch(t2) + # Exercise the tape planner and default cost models while enabled, so the + # first user-facing enabled allocation does not pay a cold-compile spike. + # Use `:lexical` here: `:backtrace` site IDs are session-local and the + # precompile process's stack is not a useful key to bake into the image. + # Avoid allocating real DArrays here — that would spawn tasks the cleanup + # below is not sized to drain. + let old_site_id = Tapes.CONFIG.site_id + Tapes.enable!(site_id=:lexical) + try + Tapes.backtrace_hash() + Tapes.explain(devnull, Float64, (128, 128)) + Tapes.plan_allocation(Float64, (64, 64); requested = AutoBlocks()) + Tapes.resolve_partitioning(Float64, (64, 64), AutoBlocks(), :arbitrary) + finally + Tapes.disable!() + Tapes.clear!() + Tapes.CONFIG.site_id = old_site_id + end + end + # Clean up refs t1 = nothing; t2 = nothing state = Sch.EAGER_STATE[] diff --git a/src/tapes/api.jl b/src/tapes/api.jl index 24db80f36..512ee29ca 100644 --- a/src/tapes/api.jl +++ b/src/tapes/api.jl @@ -367,6 +367,7 @@ function clear!() TOTAL_NODES[] = 0 end empty!(TRACES) + reset_machine!() return nothing end diff --git a/src/tapes/plan.jl b/src/tapes/plan.jl index 4d478ff9a..a8f10eac0 100644 --- a/src/tapes/plan.jl +++ b/src/tapes/plan.jl @@ -47,7 +47,7 @@ split trailing dim, near-cubic) but untested against real workloads. Tensor contractions in particular want layouts derived from the contraction indices, which the tape does not currently record. """ -function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_machine()) +function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_machine()) where {T} N = length(dims) N == 0 && return LayoutChoice[] esz = max(1, _elsize(T)) diff --git a/src/tapes/tape.jl b/src/tapes/tape.jl index 8e88ffbbf..0a9ced602 100644 --- a/src/tapes/tape.jl +++ b/src/tapes/tape.jl @@ -221,7 +221,11 @@ function backtrace_hash() hi = min(lo + CONFIG.backtrace_depth - 1, n) h = UInt64(0xcbf29ce484222325) @inbounds for i in lo:hi - h = (h ⊻ (reinterpret(UInt64, bt[i]) % UInt64)) * UInt64(0x100000001b3) + # Julia 1.12+ backtraces can mix `Ptr` frames with non-isbits + # `Base.InterpreterIP` entries; only pointers can be reinterpreted. + frame = bt[i] + ip = frame isa Ptr ? (UInt(frame) % UInt64) : (hash(frame) % UInt64) + h = (h ⊻ ip) * UInt64(0x100000001b3) end return h end diff --git a/test/runtests.jl b/test/runtests.jl index 4aacf1c7b..2ac87b2cd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -29,6 +29,7 @@ tests = [ ("Datadeps - Scheduling", "datadeps/scheduling.jl"), ("Streaming", "streaming.jl"), ("Domain Utilities", "domain.jl"), + ("Tapes", "tapes.jl"), ("Array - Allocation", "array/allocation.jl"), ("Array - Indexing", "array/indexing.jl"), ("Array - Core", "array/core.jl"), diff --git a/src/tapes/tapes_tests.jl b/test/tapes.jl similarity index 98% rename from src/tapes/tapes_tests.jl rename to test/tapes.jl index 15efcdf40..63f8561a3 100644 --- a/src/tapes/tapes_tests.jl +++ b/test/tapes.jl @@ -6,9 +6,6 @@ # models, planner) with synthetic specs, so they run without a cluster and in # milliseconds. The end-to-end group at the bottom needs real DArrays. -using Test -using LinearAlgebra -using Dagger using Dagger: Tapes using Dagger.Tapes: SiteKey, OpKey, ArgSpec, ArgView, LayoutChoice, PredictedOp, TapeNode, TapeRoot, MachineModel, @@ -23,7 +20,8 @@ const TESTMACHINE = MachineModel(nprocs = 16, bandwidth = 1.0e10, latency = 1.0e-5, task_overhead = 1.0e-4, - mem_per_proc = 8.0e9) + mem_per_proc = 8.0e9, + calibrated = true) spec(T, dims, bs; assign = :arbitrary) = ArgSpec(T, dims, bs, assign) @@ -33,7 +31,7 @@ function fresh!() return nothing end -@testset "Tapes" begin +was_enabled = Tapes.is_enabled() @testset "site keys" begin fresh!() @@ -325,4 +323,6 @@ end end end -end # @testset "Tapes" +if was_enabled + Tapes.enable!() +end \ No newline at end of file From 009957644b4855439044b006932fa98fbfe943c7 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 28 Jul 2026 13:45:28 -0700 Subject: [PATCH 42/43] fixup! fixup! Automatic partitioning selection --- src/array/lu.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/array/lu.jl b/src/array/lu.jl index 52575155e..8ede04e6f 100644 --- a/src/array/lu.jl +++ b/src/array/lu.jl @@ -3,6 +3,9 @@ LinearAlgebra.lu(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlge LinearAlgebra.lu!(A::DMatrix{T}, pivot::Union{LinearAlgebra.RowMaximum,LinearAlgebra.NoPivot} = LinearAlgebra.RowMaximum(); check::Bool=true, allowsingular::Bool=false) where {T<:LinearAlgebra.BlasFloat} = LinearAlgebra.lu(A, pivot; check=check, allowsingular=allowsingular) function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + # Record against the caller's array, not the internal copy: otherwise the + # allocation site that produced `A` never sees `:lu!` and the tape cannot + # learn from the usual `F = lu(A)` pattern. Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.NoPivot(); check) @@ -262,6 +265,7 @@ end # Implementation of https://inria.hal.science/hal-04984070v1/file/ipdps_paper.pdf function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.RowMaximum; check::Bool = true, allowsingular::Bool = false) where {T<:LinearAlgebra.BlasFloat} + # See NoPivot method: attribute the op to the source array, not the copy. Dagger.@record_op :lu! A A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T)) return LinearAlgebra.lu!(A_copy, LinearAlgebra.RowMaximum(); check, allowsingular) From ab9a1c041d7829e4c0dd19d99a82ef237bf8436e Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 28 Jul 2026 13:47:26 -0700 Subject: [PATCH 43/43] tapes: Add automatic calibration --- src/precompile.jl | 4 +- src/tapes/Tapes.jl | 15 +- src/tapes/cost.jl | 336 ++++++++++++++++++++++++++++++++++++++------- src/tapes/plan.jl | 10 +- 4 files changed, 310 insertions(+), 55 deletions(-) diff --git a/src/precompile.jl b/src/precompile.jl index c3031b224..639474459 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -16,7 +16,9 @@ # Avoid allocating real DArrays here — that would spawn tasks the cleanup # below is not sized to drain. let old_site_id = Tapes.CONFIG.site_id - Tapes.enable!(site_id=:lexical) + # `calibrate=false`: do not bake build-host FLOP/bandwidth rates into + # the image; runtime `enable!` will measure on first use. + Tapes.enable!(site_id=:lexical, calibrate=false) try Tapes.backtrace_hash() Tapes.explain(devnull, Float64, (128, 128)) diff --git a/src/tapes/Tapes.jl b/src/tapes/Tapes.jl index 786a69883..8b15e9d72 100644 --- a/src/tapes/Tapes.jl +++ b/src/tapes/Tapes.jl @@ -188,16 +188,21 @@ is cheap and predictable in the disabled case. @inline is_enabled() = CONFIG.enabled """ - enable!(; kwargs...) + enable!(; calibrate=true, kwargs...) Turn the tape subsystem on. Keyword arguments set the matching [`TapeConfig`](@ref) fields. +By default this also runs [`calibrate_machine!`](@ref) once, so layout ranking +uses measured FLOP/s and memory bandwidth rather than the pessimistic +placeholders. Pass `calibrate=false` to skip (e.g. after an explicit +[`set_machine!`](@ref), or in tests that need a fixed model). + ```julia Dagger.Tapes.enable!(site_id=:lexical, horizon=8, verbose=true) ``` """ -function enable!(; kwargs...) +function enable!(; calibrate::Bool = true, kwargs...) for (k, v) in kwargs hasfield(TapeConfig, k) || throw(ArgumentError("unknown Tapes config option: $k")) setfield!(CONFIG, k, convert(fieldtype(TapeConfig, k), v)) @@ -205,6 +210,12 @@ function enable!(; kwargs...) CONFIG.site_id in (:backtrace, :context, :lexical) || throw(ArgumentError("site_id must be :backtrace, :context or :lexical, got $(CONFIG.site_id)")) CONFIG.enabled = true + # Calibrate once per session unless the user already installed rates via + # `set_machine!` (or passed `calibrate=false` to skip, e.g. in tests). + if calibrate + m = _ensure_machine() + m.calibrated || calibrate_machine!(m) + end return CONFIG end diff --git a/src/tapes/cost.jl b/src/tapes/cost.jl index d0e93e9ac..d526b6fca 100644 --- a/src/tapes/cost.jl +++ b/src/tapes/cost.jl @@ -10,6 +10,8 @@ # are intended to *correct* it, not replace it. # =========================================================================== +import LinearAlgebra + "The module these macros expand references into. Captured so `@cost_model` can be used from anywhere." const TAPES_MODULE = @__MODULE__ @@ -30,28 +32,70 @@ Base.@kwdef mutable struct MachineModel bandwidth::Float64 = 5.0e9 "Per-transfer latency in seconds." latency::Float64 = 1.0e-5 - "Scheduler overhead per task in seconds. Dagger tasks bottom out around 100us." - task_overhead::Float64 = 1.0e-4 + """ + Per-task scheduler cost in seconds, **as serialized wall-clock time**, not as + work to be divided among processors. + + Most of what a Dagger task costs before it runs — Datadeps aliasing + analysis, dependency-edge construction, queue insertion — happens on the + submitting thread, so adding processors barely reduces it. Measured on a + 6-thread node in a `spawn_datadeps` region: ~230us/task for the + three-argument, tile-sharing pattern that blocked linear algebra submits, + and ~150us/task for trivial tasks holding one private chunk each. Dropping + to a single thread only raises the latter to ~360us, so six times the + processors buys a 2.4x reduction: this is far closer to a serial cost than a + parallel one, and cost models must not divide it by `nprocs`. Doing so was + why the planner preferred tile counts an order of magnitude too high. + + The default suits a handful of threads on one node; use + [`calibrate_task_overhead!`](@ref) to measure it. + """ + task_overhead::Float64 = 2.5e-4 "Usable memory per processor in bytes; used to reject layouts that will not fit." mem_per_proc::Float64 = 8.0e9 + """ + Whether [`calibrate_machine!`](@ref) (or the user, via [`set_machine!`](@ref)) + has supplied non-default `flops_per_sec` / `bandwidth`. Untuned defaults + skew candidate generation and layout ranking badly on modern hardware. + """ + calibrated::Bool = false end const MACHINE = Ref{Union{Nothing,MachineModel}}(nothing) +""" + _ensure_machine() -> MachineModel + +Return the cached machine model, constructing an uncalibrated default from +`Dagger.num_processors()` if needed. Does not calibrate (see +[`current_machine`](@ref)). +""" +function _ensure_machine() + m = MACHINE[] + m === nothing || return m + np = try + max(1, num_processors()) + catch + 1 + end + m = MachineModel(; nprocs = np) + MACHINE[] = m + return m +end + """ current_machine() -> MachineModel -The machine model in effect, constructed lazily from `Dagger.num_processors()` -on first use and cached. +The machine model in effect. Constructed lazily from `Dagger.num_processors()` +on first use and cached. When the tape subsystem is enabled and the model has +not yet been calibrated, the first call runs [`calibrate_machine!`](@ref). -TODO(calibration): `flops_per_sec` and `bandwidth` are guesses. They should be -measured once at `enable!` time with a small GEMM and a small point-to-point -transfer (a few hundred milliseconds, amortised over the session), and refined -online from the scheduler's existing task timings — Dagger already records -per-task durations through `TimespanLogging`, which is exactly the signal -needed. This is also the natural place to hook the autotuner: it can supply -measured `(op, layout, size) -> seconds` points that override the analytic -model where they exist and leave it in place where they do not. +TODO(online-refinement): refine these rates from the scheduler's existing task +timings — Dagger already records per-task durations through `TimespanLogging`, +which is exactly the signal needed. This is also the natural place to hook the +autotuner: it can supply measured `(op, layout, size) -> seconds` points that +override the analytic model where they exist and leave it in place where they +do not. TODO(heterogeneity): a single scalar FLOP rate is wrong on a CPU+GPU node, which is the configuration Dagger most cares about. The model should carry a @@ -62,15 +106,143 @@ today's assignment vocabulary cannot express a heterogeneous mapping anyway; these two limitations should be lifted together. """ function current_machine() - m = MACHINE[] - m === nothing || return m - np = try - max(1, num_processors()) + m = _ensure_machine() + if is_enabled() && !m.calibrated + calibrate_machine!(m) + end + return m +end + +""" + calibrate_machine!(m = _ensure_machine(); force=false) -> MachineModel + +Measure aggregate GEMM throughput and streaming memory bandwidth with small +local kernels, and write the results into `m`. Takes a few tens of milliseconds +and is amortised over the session. + +Called automatically on the first [`current_machine`](@ref) after +[`enable!`](@ref), and from `enable!` itself. Pass `force=true` to overwrite a +previous calibration (including one the user installed via [`set_machine!`](@ref)). + +On a multi-threaded BLAS this measures *aggregate* FLOP/s (do not multiply by +`nprocs`). Bandwidth is a single-node DRAM-copy proxy for the "inter-processor" +term the cost models use; on multi-socket / distributed setups it is an +optimistic upper bound. +""" +function calibrate_machine!(m::MachineModel = _ensure_machine(); force::Bool = false) + (!force && m.calibrated) && return m + + try + m.nprocs = max(1, num_processors()) catch - 1 end - m = MachineModel(; nprocs = np) - MACHINE[] = m + + # --- Aggregate GEMM throughput ---------------------------------------- + # BLAS may already be multi-threaded, so the result is aggregate across + # cores, not per-core. Keep this small: `enable!` should stay snappy. + n = 768 + A = rand(Float64, n, n) + B = rand(Float64, n, n) + C = Matrix{Float64}(undef, n, n) + LinearAlgebra.mul!(C, A, B) # warmup / JIT + best = Inf + for _ in 1:3 + best = min(best, @elapsed LinearAlgebra.mul!(C, A, B)) + end + if isfinite(best) && best > 0 + m.flops_per_sec = max(1.0, 2.0 * Float64(n)^3 / best) + end + + # --- Streaming memory bandwidth --------------------------------------- + # Buffer sized to exceed typical LLC so we measure DRAM, not L3. Proxy for + # shared-memory tile movement; optimistic for multi-node. + # + # The destination is freshly allocated on every repetition rather than + # reused. That is deliberate: `bandwidth` is only ever used to price moving + # data *between Dagger chunks* (redistribution, re-blocking, halo + # exchange), and every such move allocates its destination chunks, so it + # always pays first-touch page faults. Reusing a warm destination measures + # something real but irrelevant — on a 6-thread machine it reads 16.7 GB/s + # against 7.8 GB/s cold, and a measured DArray re-block achieves 4.2 GB/s. + # Overstating bandwidth makes every layout change look nearly free, which + # is exactly the error that lets a bad plan clear `CONFIG.gate_margin`. + bytes = 64 * 1024 * 1024 + nelem = bytes ÷ sizeof(Float64) + x = rand(Float64, nelem) + copyto!(Vector{Float64}(undef, nelem), x) # warmup / JIT + bestb = Inf + for _ in 1:3 + bestb = min(bestb, @elapsed copyto!(Vector{Float64}(undef, nelem), x)) + end + if isfinite(bestb) && bestb > 0 + # copy reads x and writes a fresh y. + m.bandwidth = max(1.0, 2.0 * Float64(bytes) / bestb) + end + + # `task_overhead` is left at its default. Probing it needs live Dagger + # spawns, which is unsafe from inside `enable!`/`current_machine`: it can + # deadlock against an in-flight scheduler or BLAS threadpool. Call + # `calibrate_task_overhead!` explicitly from a quiescent point instead, or + # refine it from TimespanLogging task durations (see TODO on + # `current_machine`). + + m.calibrated = true + vlog("calibrated machine: nprocs=", m.nprocs, + " flops_per_sec=", m.flops_per_sec, + " bandwidth=", m.bandwidth) + return m +end + +"An empty task body, used only by `calibrate_task_overhead!`." +_overhead_probe!(c, a, b) = nothing + +""" + calibrate_task_overhead!(m = _ensure_machine(); nt=12) -> MachineModel + +Measure `m.task_overhead` by submitting a GEMM-shaped dependency pattern of +empty tasks over tiny tiles, so the elapsed time is essentially all scheduler +cost. Uses the three-argument, tile-sharing shape that blocked linear algebra +submits, because Datadeps' cost is dominated by aliasing analysis over shared +tiles and a pattern of independent tasks measures ~1.5x too low. + +**Not** called from [`enable!`](@ref) or [`calibrate_machine!`](@ref): it spawns +real Dagger tasks, which can deadlock if a scheduler or BLAS threadpool is +already in flight. Call it once from a quiescent point: + +```julia +Dagger.Tapes.enable!() +Dagger.Tapes.calibrate_task_overhead!() +``` +""" +function calibrate_task_overhead!(m::MachineModel = _ensure_machine(); nt::Int = 12) + tiles = [Dagger.tochunk(zeros(4, 4)) for _ in 1:nt, _ in 1:nt] + ntasks = 0 + submit = function () + n = 0 + Dagger.spawn_datadeps() do + for k in 1:nt, j in (k+1):nt, i in (k+1):nt + Dagger.@spawn _overhead_probe!(Dagger.InOut(tiles[i, j]), + Dagger.In(tiles[i, k]), + Dagger.In(tiles[k, j])) + n += 1 + end + end + return n + end + try + ntasks = submit() # warmup / JIT + ntasks <= 0 && return m + best = Inf + for _ in 1:3 + best = min(best, @elapsed submit()) + end + isfinite(best) && best > 0 && + (m.task_overhead = max(1.0e-6, best / ntasks)) + catch err + vlog("task-overhead calibration failed ($err); keeping ", m.task_overhead) + return m + end + vlog("calibrated task_overhead=", m.task_overhead, " over ", ntasks, " tasks") return m end @@ -79,18 +251,24 @@ end set_machine!(; kwargs...) Install or update the machine model; keyword form updates fields in place. +Supplying `flops_per_sec` or `bandwidth` marks the model as calibrated so +[`enable!`](@ref) will not overwrite it. """ set_machine!(m::MachineModel) = (MACHINE[] = m) function set_machine!(; kwargs...) - m = current_machine() + m = _ensure_machine() for (k, v) in kwargs hasfield(MachineModel, k) || throw(ArgumentError("unknown MachineModel field: $k")) setfield!(m, k, convert(fieldtype(MachineModel, k), v)) end + if any(k -> k === :flops_per_sec || k === :bandwidth || k === :calibrated, keys(kwargs)) + # User-supplied rates (or an explicit flag) count as calibrated. + haskey(kwargs, :calibrated) || (m.calibrated = true) + end return m end -"Reset the cached machine model so it is rebuilt on next use." +"Reset the cached machine model so it is rebuilt (and recalibrated) on next use." reset_machine!() = (MACHINE[] = nothing) # --------------------------------------------------------------------------- @@ -142,6 +320,49 @@ tilebytes(v::ArgView) = Float64(prod(blocksize(v); init = 1)) * elsize(v) "Processor assignment strategy of the candidate layout." assignment(v::ArgView) = v.layout.assignment +""" + square_tile(v::ArgView) -> (edge, reblocked) + +The block edge a blocked dense kernel will *actually* run with, and whether +getting there costs a copy. + +Dagger's factorizations and triangular solves do not honour a non-square +blocking: `lu!`, `cholesky!`, `qr!` and the `ldiv!` family all funnel through +`maybe_copy_buffered(A => Blocks(b, b))` with `b = min(blocksize...)`, which +copies the whole array into square tiles, runs, and copies back. A cost model +that reads `blocksize(A, 1)` directly is therefore costing a layout the runtime +never executes — for `Blocks(4096, 683)` it charges a 4096-tall panel +factorization when what runs is a 683-square one. + +Getting this wrong does not merely misprice one candidate: `AutoBlocks` is +exactly the non-square case, so it inflates the *baseline* every plan is gated +against, which turns `CONFIG.gate_margin` from a safety check into a rubber +stamp. +""" +function square_tile(v::ArgView) + n = ndims(v) + n == 0 && return (1, false) + bs = blocksize(v) + edges = ntuple(d -> max(1, min(bs[d], size(v, d))), n) + edge = minimum(edges) + return (edge, any(!=(edge), edges)) +end + +""" + reblock_cost(v::ArgView, m::MachineModel) -> Float64 + +Cost of the copy-in/copy-out `maybe_copy_buffered` performs when a layout has to +be squared off before a dense kernel will accept it: the array's bytes move +twice, over one task per tile in each of the two blockings. +""" +function reblock_cost(v::ArgView, m::MachineModel) + edge, reblocked = square_tile(v) + reblocked || return 0.0 + nt_sq = prod(ntuple(d -> max(1, cld(size(v, d), edge)), ndims(v)); init = 1) + return 2 * (m.latency + nbytes(v) / m.bandwidth) + + (ntiles(v) + nt_sq) * m.task_overhead +end + "Longest/shortest block edge ratio; `1.0` for a perfect cube." function aspect(v::ArgView) n = ndims(v) @@ -207,13 +428,18 @@ These helpers are also in scope: | helper | meaning | |:-----------------|:------------------------------------------------------------| -| `flops_time(f)` | seconds for `f` FLOPs spread perfectly over all processors | -| `serial_time(f)` | seconds for `f` FLOPs on one processor (i.e. critical path) | -| `bytes_time(b)` | seconds to move `b` bytes, including one latency | -| `task_time(n)` | scheduler overhead for `n` tasks, spread over all processors | -| `imbalance(n)` | load-imbalance multiplier for `n` tiles over `nprocs` | -| `machine` | the [`MachineModel`](@ref) | -| `nprocs_avail()` | processor count | +| `flops_time(f)` | seconds for `f` FLOPs spread perfectly over all processors | +| `serial_time(f)` | seconds for `f` FLOPs on one processor (i.e. critical path) | +| `bytes_time(b)` | seconds to move `b` bytes, including one latency | +| `task_time(n)` | scheduler cost for `n` tasks; serial, see `task_overhead` | +| `reblock_time(v)` | cost of squaring off `v`'s blocking, `0.0` if already square | +| `imbalance(n)` | load-imbalance multiplier for `n` tiles over `nprocs` | +| `machine` | the [`MachineModel`](@ref) | +| `nprocs_avail()` | processor count | + +`square_tile(v)` is also in scope, returning the `(edge, reblocked)` pair a +blocked dense kernel will really run with; prefer it over `blocksize(v, 1)` in +any model for an operation that squares its input. If the operation is recorded with fewer arguments than the model declares, the generic fallback is used rather than erroring. @@ -250,7 +476,7 @@ macro cost_model(ex) # Bind the argument names, and re-export the ArgView accessor vocabulary as # locals so the body works regardless of what the caller has imported. accessors = (:blocksize, :nblocks, :ntiles, :tilebytes, :nbytes, :elsize, - :aspect, :assignment) + :aspect, :assignment, :square_tile) accbinds = [:(local $(esc(s)) = $(getfield(TAPES_MODULE, s))) for s in accessors] argbinds = [:(local $(esc(a)) = $A[$i]) for (i, a) in enumerate(fargs)] @@ -274,7 +500,8 @@ macro cost_model(ex) local $(esc(:flops_time)) = f -> f / $M.flops_per_sec local $(esc(:serial_time)) = f -> f / ($M.flops_per_sec / max(1, $M.nprocs)) local $(esc(:bytes_time)) = b -> ($M.latency + b / $M.bandwidth) - local $(esc(:task_time)) = n -> n * $M.task_overhead / max(1, $M.nprocs) + local $(esc(:task_time)) = n -> n * $M.task_overhead + local $(esc(:reblock_time)) = v -> $reblock_cost(v, $M) local $(esc(:imbalance)) = n -> (n <= 0 ? 1.0 : let np = max(1, $M.nprocs); ceil(n / np) / (n / np) end) $(accbinds...) @@ -345,7 +572,7 @@ function shape_penalty(v::ArgView, op::Symbol, m::MachineModel) # A task must do enough work to be worth launching. if tb > 0 - min_useful = m.task_overhead * m.bandwidth / max(1, m.nprocs) + min_useful = m.task_overhead * m.bandwidth tb < min_useful && (p *= min_useful / tb) end @@ -393,7 +620,7 @@ function generic_op_cost(args::Vector{ArgView}, m::MachineModel, op::Symbol = :_ eltype(v) === Nothing && continue nt = ntiles(v) c = nbytes(v) / m.bandwidth # touch the data once - c += nt * m.task_overhead / np # per-tile scheduling + c += nt * m.task_overhead # per-tile scheduling (serial) c += nt * m.latency # per-tile boundary crossing c *= nt <= 0 ? 1.0 : ceil(nt / np) / (nt / np) # load imbalance c *= shape_penalty(v, op, m) @@ -434,7 +661,7 @@ function redistribution_cost(from::LayoutChoice, to::LayoutChoice, spec::ArgSpec nt_to = prod(nblocks(to, sz); init = 1) nt_from = prod(nblocks(from, sz); init = 1) return bytes / m.bandwidth + - (nt_to + nt_from) * (m.latency + m.task_overhead / max(1, m.nprocs)) + (nt_to + nt_from) * (m.latency + m.task_overhead) end # --------------------------------------------------------------------------- @@ -452,36 +679,38 @@ end @cost_model cholesky!(A) = begin n = size(A, 1) - b = max(1, blocksize(A, 1)) - nt = nblocks(A, 1) + # `cholesky!` squares off its input before factorizing; cost what runs. + b, _ = square_tile(A) + nt = max(1, cld(n, b)) # nt POTRF panels sit on the critical path; the rest is parallel SYRK/GEMM/TRSM. serial_time(nt * b^3 / 3) + flops_time(Float64(n)^3 / 3) * imbalance(nt * (nt + 1) / 2) + task_time(nt^3 / 3 + nt^2) + bytes_time(nt^2 * b^2 * elsize(A)) + - # Rectangular tiles mismatch TRSM against SYRK and inflate the panel path. - serial_time(nt * b^3 / 3) * (aspect(A) - 1.0) + reblock_time(A) end @cost_model lu!(A) = begin n = size(A, 1) - b = max(1, blocksize(A, 1)) - nt = nblocks(A, 1) + b, _ = square_tile(A) + nt = max(1, cld(n, b)) serial_time(nt * b^3) + flops_time(2 * Float64(n)^3 / 3) * imbalance(nt * nt) + - task_time(nt^3 / 3 + nt^2) + - # Panel pivoting serialises down a block column. - bytes_time(nt * nt * b^2 * elsize(A)) * (1.0 + 0.5 * (aspect(A) - 1.0)) + # Row swaps against the trailing submatrix dominate the task count. + task_time(nt^3 / 2 + nt^3 / 3 + nt^2) + + bytes_time(nt * nt * b^2 * elsize(A)) + + reblock_time(A) end @cost_model qr!(A) = begin mm = size(A, 1); nn = size(A, 2) - b = max(1, blocksize(A, 1)) - ntr = nblocks(A, 1); ntc = nblocks(A, 2) + b, _ = square_tile(A) + ntr = max(1, cld(mm, b)); ntc = max(1, cld(nn, b)) serial_time(ntc * b^3) + flops_time(2 * Float64(mm) * nn^2 - 2 * Float64(nn)^3 / 3) * imbalance(ntr * ntc) + task_time(ntr * ntc^2) + - bytes_time(ntr * ntc * b^2 * elsize(A)) * (1.0 + 0.3 * (aspect(A) - 1.0)) + bytes_time(ntr * ntc * b^2 * elsize(A)) + + reblock_time(A) end @cost_model mul!(C, A, B) = begin @@ -508,21 +737,26 @@ end nn = size(A, 1) nd = ndims(B) nrhs = nd >= 2 ? size(B, 2) : 1 - nta = nblocks(A, 1); ntb = max(1, nblocks(B, nd)) + # `ldiv!` squares `A` off (and re-blocks `B` to match) before calling trsm!. + b, _ = square_tile(A) + nta = max(1, cld(nn, b)); ntb = max(1, cld(nrhs, b)) # Triangular solve is inherently sequential along the block diagonal. - serial_time(nta * blocksize(A, 1)^2 * max(1, blocksize(B, nd))) + + serial_time(nta * Float64(b)^2 * b) + flops_time(Float64(nn)^2 * nrhs) * imbalance(nta * ntb) + task_time(nta * nta * ntb / 2) + - bytes_time(nta * ntb * prod(blocksize(B); init = 1) * elsize(B)) + bytes_time(nta * ntb * Float64(b)^2 * elsize(B)) + + reblock_time(A) end @cost_model trsv!(A, B) = begin nn = size(A, 1) - nta = nblocks(A, 1) - serial_time(nta * blocksize(A, 1)^2) + + b, _ = square_tile(A) + nta = max(1, cld(nn, b)) + serial_time(nta * Float64(b)^2) + flops_time(Float64(nn)^2) * imbalance(nta) + task_time(nta * (nta + 1) / 2) + - bytes_time(nta * nta * blocksize(A, 1) * elsize(A)) + bytes_time(nta * nta * Float64(b) * elsize(A)) + + reblock_time(A) end @cost_model map(A) = flops_time(Float64(length(A))) * imbalance(ntiles(A)) + task_time(ntiles(A)) diff --git a/src/tapes/plan.jl b/src/tapes/plan.jl index a8f10eac0..676983f4b 100644 --- a/src/tapes/plan.jl +++ b/src/tapes/plan.jl @@ -55,7 +55,10 @@ function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_mac out = LayoutChoice[] # Smallest tile worth scheduling, and largest that leaves >= 2 tiles/proc. - min_tile_bytes = m.task_overhead * m.bandwidth / np + # `task_overhead` is serialized wall-clock, so it is *not* divided by `np` + # here: a tile has to be big enough to hide the whole submission cost, not + # `1/np` of it. + min_tile_bytes = m.task_overhead * m.bandwidth min_edge = max(32, ceil(Int, (min_tile_bytes / esz)^(1 / N))) total_bytes = Float64(prod(dims; init = 1)) * esz max_tile_bytes = max(min_tile_bytes, total_bytes / (2 * np)) @@ -75,6 +78,11 @@ function candidate_layouts(::Type{T}, dims::Tuple, m::MachineModel = current_mac end # Also the edge that yields exactly one tile per processor per dimension. push!(edges, max(1, cld(maximum(dims), max(1, round(Int, np^(1 / N)))))) + # And the square form of the blind default's block size. Powers of two miss + # it (it is `cld(dims[end], np)`, e.g. 683 for 4096 over 6 procs) and it is + # the layout Dagger's own `maybe_copy_buffered` re-blocks `AutoBlocks` to, + # so it is both a strong candidate and the honest baseline for comparison. + push!(edges, max(1, minimum(blocksize(fallback_layout(T, dims))))) for edge in unique!(sort!(edges)) bs = ntuple(i -> clamp(edge, 1, max(1, Int(dims[i]))), N) for a in (N >= 2 ? (:cyclicrow, :cycliccol, :arbitrary) : (:arbitrary,))