From f4a81fc448c024720786f4e182d68d7ec69421db Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 7 Jul 2026 11:19:42 -0700 Subject: [PATCH 1/9] DArray: Add SVD algorithm --- docs/src/darray.md | 6 +- src/Dagger.jl | 1 + src/array/svd.jl | 670 +++++++++++++++++++++++++++++++++++++++ test/array/linalg/svd.jl | 176 ++++++++++ test/runtests.jl | 1 + 5 files changed, 852 insertions(+), 2 deletions(-) create mode 100644 src/array/svd.jl create mode 100644 test/array/linalg/svd.jl diff --git a/docs/src/darray.md b/docs/src/darray.md index 6419449ff..12c86ec80 100644 --- a/docs/src/darray.md +++ b/docs/src/darray.md @@ -705,8 +705,10 @@ From `LinearAlgebra`: - `mul!` (In-place Matrix-Matrix and Matrix-Vector multiply) - `cholesky`/`cholesky!` (In-place/Out-of-place Cholesky factorization) - `lu`/`lu!` (In-place/Out-of-place LU factorization (`NoPivot` and `RowMaximum`)) -- `\`/`ldiv!` (In-place/Out-of-place Linear solving with LU and Cholesky factorizations) -- `inv` (Out-of-place matrix inversion) +- `qr`/`qr!` (In-place/Out-of-place QR factorization) +- `\`/`ldiv!` (In-place/Out-of-place Linear solving with LU, Cholesky, QR, and SVD factorizations) +- `inv` (Out-of-place matrix inversion, including via SVD) +- `svd`/`svd!`/`svdvals!` (In-place/Out-of-place Singular Value Decomposition) From `AbstractFFTs`: - `fft`/`fft!` diff --git a/src/Dagger.jl b/src/Dagger.jl index 13469536e..593eee4aa 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -129,6 +129,7 @@ include("array/cholesky.jl") include("array/trsm.jl") include("array/lu.jl") include("array/qr.jl") +include("array/svd.jl") # GPU include("gpu.jl") diff --git a/src/array/svd.jl b/src/array/svd.jl new file mode 100644 index 000000000..cf93c4273 --- /dev/null +++ b/src/array/svd.jl @@ -0,0 +1,670 @@ +# Tiled singular value decomposition for `DMatrix`, built on the Datadeps API. +# +# Algorithm: block one-sided Jacobi. The columns of `A` are partitioned into +# block-columns (following `A`'s existing column tiling). For every pair of +# block-columns `(i, j)` we form the small Hermitian Gram matrix +# +# G = [ Aᵢ'Aᵢ Aᵢ'Aⱼ ; Aⱼ'Aᵢ Aⱼ'Aⱼ ] +# +# accumulated over the row tiles, diagonalize it (`G = W Λ W'`), and apply the +# orthogonal rotation `W` to the two block-columns of `A` (and to the +# accumulator `V`). A full pass over all pairs is one "sweep"; sweeps repeat +# until the largest relative off-diagonal coupling drops below `tol`. At +# convergence the columns of `A` are mutually orthogonal: their norms are the +# singular values, and the normalized columns are the left singular vectors. +# +# The design mirrors the other tiled factorizations in this directory: each +# `spawn_datadeps` region emits tile-granular tasks whose `In`/`Out`/`InOut` +# annotations let the scheduler recover parallelism. Independent block-column +# pairs within a sweep run concurrently; pairs sharing a block-column are +# serialized automatically through the aliasing analysis. +# +# Only the thin SVD is produced (U is m×min(m,n), Vᵀ is min(m,n)×n). This +# variant forms Gram matrices of pairs of block-columns, so its accuracy is that +# of standard block Jacobi — excellent for the well- to moderately-conditioned +# case; it is not the bidiagonalization route used for extreme conditioning. + +# ──────────────────────────────────────────────────────────────────────────── +# Tile / host-array kernels (run on chunk data inside Datadeps tasks) +# ──────────────────────────────────────────────────────────────────────────── + +# Accumulate the four sub-blocks of the pair Gram matrix from one row tile. +# `Ai`, `Aj` are the (rows × wᵢ), (rows × wⱼ) tiles of block-columns i and j. +function _svd_gram_acc!(G::AbstractMatrix{T}, Ai::AbstractMatrix{T}, Aj::AbstractMatrix{T}) where {T} + wi = size(Ai, 2) + wj = size(Aj, 2) + n = wi + wj + o = one(T) + @views LinearAlgebra.mul!(G[1:wi, 1:wi], Ai', Ai, o, o) + @views LinearAlgebra.mul!(G[1:wi, wi+1:n], Ai', Aj, o, o) + @views LinearAlgebra.mul!(G[wi+1:n, 1:wi], Aj', Ai, o, o) + @views LinearAlgebra.mul!(G[wi+1:n, wi+1:n], Aj', Aj, o, o) + return G +end + +# Reduce the relative off-diagonal coupling ‖G₁₂‖ / √(‖G₁₁‖‖G₂₂‖) into `off[1]` +# with `max`. This drives the sweep convergence test. +function _svd_offdiag_acc!(off::AbstractVector{R}, G::AbstractMatrix{T}, wi::Int) where {R,T} + n = size(G, 1) + @views num = LinearAlgebra.norm(G[1:wi, wi+1:n]) + @views den = sqrt(LinearAlgebra.norm(G[1:wi, 1:wi]) * LinearAlgebra.norm(G[wi+1:n, wi+1:n])) + val = den == 0 ? zero(R) : R(num / den) + off[1] = max(off[1], val) + return off +end + +# Classical cyclic (two-sided) Jacobi eigenvalue algorithm for a small dense +# Hermitian matrix. +# +# The pair Gram matrices `G` that this file diagonalizes are graded: once a +# block-column pair has nearly converged, `G`'s diagonal spans the squared +# column norms of *all* of `A`'s columns, which can differ by many orders of +# magnitude (e.g. the ratio of the largest to smallest singular value of the +# input). A general-purpose dense eigensolver (`LinearAlgebra.eigen`, backed +# by Hessenberg/QR-type LAPACK routines) is only backward-stable with respect +# to the matrix's *overall* norm: it recovers eigenvectors of `G + E` with +# `‖E‖ = O(eps)‖G‖`, which can dwarf the small eigenvalues themselves and +# destroy the corresponding singular vectors' accuracy after normalization. +# Classical Jacobi rotations, by contrast, are computed from *ratios* of the +# 2×2 sub-block entries and are therefore accurate independent of the overall +# matrix scale (Demmel & Veselic, "Jacobi's method is more accurate than QR", +# 1992) — exactly the property this block algorithm needs to preserve the +# high relative accuracy that one-sided Jacobi SVD is prized for. +function _svd_hermitian_jacobi_eigen(H::LinearAlgebra.Hermitian{T}) where {T<:Number} + n = size(H, 1) + R = real(T) + A = Matrix{T}(H) + V = Matrix{T}(LinearAlgebra.I, n, n) + n <= 1 && return real.(LinearAlgebra.diag(A)), V + maxsweeps = 100 + @inbounds for _s in 1:maxsweeps + rotated = false + for p in 1:n-1, q in p+1:n + apq = A[p, q] + r = abs(apq) + r == 0 && continue + App = real(A[p, p]); Aqq = real(A[q, q]) + # Skip pairs that are already negligible *relative to their own* + # diagonal scale. Using a single tolerance based on the norm of + # the whole matrix (as opposed to this pair's own scale) is + # exactly the kind of absolute test that loses relative accuracy + # for small eigenvalues in a graded matrix — the same failure + # mode this Jacobi solver exists to avoid. + r*r <= eps(R)^2 * App * Aqq && continue + rotated = true + # Rotate out the phase of `apq` so the classical real 2×2 Jacobi + # formulas (scale-invariant in App, Aqq) apply directly; then fold + # the phase back into the accumulated rotation (see e.g. Golub & + # Van Loan, "Matrix Computations", §8.5, for the real case). + phase = apq / r + theta = (Aqq - App) / (2r) + t = theta == 0 ? one(R) : sign(theta) / (abs(theta) + sqrt(1 + theta^2)) + cR = 1 / sqrt(1 + t^2) + sR = t * cR + cphase = conj(phase) + for i in 1:n + Aip = A[i, p]; Aiq = A[i, q] + A[i, p] = cR*Aip - cphase*sR*Aiq + A[i, q] = sR*Aip + cphase*cR*Aiq + end + for j in 1:n + Apj = A[p, j]; Aqj = A[q, j] + A[p, j] = cR*Apj - phase*sR*Aqj + A[q, j] = sR*Apj + phase*cR*Aqj + end + for i in 1:n + Vip = V[i, p]; Viq = V[i, q] + V[i, p] = cR*Vip - cphase*sR*Viq + V[i, q] = sR*Vip + cphase*cR*Viq + end + end + rotated || break + end + return real.(LinearAlgebra.diag(A)), V +end + +# Diagonalize the (symmetrized) Hermitian Gram matrix; store eigenvectors in `W`. +function _svd_rot!(W::AbstractMatrix{T}, G::AbstractMatrix{T}) where {T} + H = LinearAlgebra.Hermitian((G .+ G') ./ 2) + _, vecs = _svd_hermitian_jacobi_eigen(H) + copyto!(W, vecs) + return W +end + +# Apply the block rotation W to a row-tile pair in place: [Cᵢ Cⱼ] ← [Cᵢ Cⱼ] W. +function _svd_apply_rot!(Ci::AbstractMatrix{T}, Cj::AbstractMatrix{T}, W::AbstractMatrix{T}, wi::Int) where {T} + n = size(W, 1) + o = one(T) + @views W11 = W[1:wi, 1:wi] + @views W21 = W[wi+1:n, 1:wi] + @views W12 = W[1:wi, wi+1:n] + @views W22 = W[wi+1:n, wi+1:n] + Ti = Ci * W11 + LinearAlgebra.mul!(Ti, Cj, W21, o, o) + Tj = Ci * W12 + LinearAlgebra.mul!(Tj, Cj, W22, o, o) + copyto!(Ci, Ti) + copyto!(Cj, Tj) + return nothing +end + +# Accumulate squared column norms of one row tile into the block-column's buffer. +function _svd_colnorm_acc!(s::AbstractVector{R}, A::AbstractMatrix{T}) where {R,T} + @inbounds for c in axes(A, 2) + acc = zero(R) + @views for r in axes(A, 1) + acc += abs2(A[r, c]) + end + s[c] += acc + end + return s +end + +# Scale each column of a tile by 1/σ (or zero it when σ = 0), forming U in place. +function _svd_scale_cols!(A::AbstractMatrix{T}, s::AbstractVector{R}) where {R,T} + @inbounds for c in axes(A, 2) + sc = s[c] + if sc > 0 + invs = inv(sc) + @views A[:, c] .*= invs + else + @views A[:, c] .= zero(T) + end + end + return A +end + +# Single-block-column kernels: accumulate A'A and apply A ← A W (no partner). +_svd_gram_self_acc!(G::AbstractMatrix{T}, A::AbstractMatrix{T}) where {T} = + (LinearAlgebra.mul!(G, A', A, one(T), one(T)); G) +_svd_apply_self!(A::AbstractMatrix{T}, W::AbstractMatrix{T}) where {T} = + (Tmp = A * W; copyto!(A, Tmp); nothing) + +_svd_adjoint_tile!(dst, src) = (copyto!(dst, adjoint(src)); dst) +_svd_set_identity_diag!(B) = (@inbounds B[LinearAlgebra.diagind(B)] .= one(eltype(B)); B) + +# ──────────────────────────────────────────────────────────────────────────── +# Small distributed helpers +# ──────────────────────────────────────────────────────────────────────────── + +# Widths of each block-column from a cumulative-length vector. +function _svd_tilewidths(cum::AbstractVector{<:Integer}) + w = Vector{Int}(undef, length(cum)) + prev = 0 + @inbounds for i in eachindex(cum) + w[i] = cum[i] - prev + prev = cum[i] + end + return w +end + +# Locate the block-column and in-tile offset of global column `c`. +@inline function _svd_locate_col(cum::AbstractVector{<:Integer}, c::Int) + k = searchsortedfirst(cum, c) + prev = k == 1 ? 0 : cum[k-1] + return k, c - prev +end + +# Conjugate transpose of a `DMatrix`, tile by tile. +function _svd_adjoint(A::DMatrix{T}) where {T} + m, n = size(A) + mb, nb = A.partitioning.blocksize + At = DArray{T,2}(undef, Blocks(nb, mb), (n, m)) + Ac = A.chunks + Atc = At.chunks + at_mt, at_nt = size(Atc) + Dagger.spawn_datadeps() do + for i in 1:at_mt, j in 1:at_nt + Dagger.@spawn _svd_adjoint_tile!(Out(Atc[i, j]), In(Ac[j, i])) + end + end + return At +end + +# n×n distributed identity with uniform `nb` column/row tiling. +function _svd_identity(::Type{T}, n::Int, nb::Int) where {T} + V = zeros(Blocks(nb, nb), T, n, n) + Vc = V.chunks + nt = size(Vc, 1) + Dagger.spawn_datadeps() do + for k in 1:nt + Dagger.@spawn _svd_set_identity_diag!(InOut(Vc[k, k])) + end + end + return V +end + +# Fill `dst`'s columns from `srcs` (the distinct source tiles this destination +# tile draws from) according to `colmap[ld] = (i, lc)`: column `ld` of `dst` +# comes from column `lc` of `srcs[i]`. +function _svd_permute_tile!(dst::AbstractMatrix{T}, colmap::Vector{Tuple{Int,Int}}, + srcs::Vararg{AbstractMatrix{T}}) where {T} + for (ld, (i, lc)) in enumerate(colmap) + @views dst[:, ld] .= srcs[i][:, lc] + end + return dst +end + +# New `DMatrix` whose column `c` is column `p[c]` of `D` (same partitioning). +# +# Each destination tile is written by exactly one task per row-tile, which +# gathers all the source columns it needs and writes the whole tile at once. +# Spawning one task *per column* (each an `InOut` of a view into a shared +# destination tile) would let the scheduler treat those column writes as +# independent and run them concurrently; that's safe under shared memory +# (`-p 0`) but not once tiles are pulled to separate worker processes, where +# each concurrent task can end up mutating its own copy of the *whole* tile +# and only the last one to write back survives — silently dropping every +# other column's update. Writing the full tile from a single task avoids the +# hazard entirely. +function _svd_permute_columns(D::DMatrix{T}, p::AbstractVector{<:Integer}) where {T} + m, n = size(D) + cum = D.subdomains.cumlength[2] + R = similar(D) + Dc = D.chunks + Rc = R.chunks + rt = size(Dc, 1) + nt = size(Dc, 2) + Dagger.spawn_datadeps() do + for kd in 1:nt + lo = kd == 1 ? 1 : cum[kd-1] + 1 + hi = cum[kd] + srcs_needed = Int[] + colmap = Vector{Tuple{Int,Int}}(undef, hi - lo + 1) + for c in lo:hi + s = Int(p[c]) + ks, ls = _svd_locate_col(cum, s) + i = findfirst(==(ks), srcs_needed) + if i === nothing + push!(srcs_needed, ks) + i = length(srcs_needed) + end + colmap[c - lo + 1] = (i, ls) + end + for k in 1:rt + srcs = ntuple(i -> Dc[k, srcs_needed[i]], length(srcs_needed)) + Dagger.@spawn _svd_permute_tile!(Out(Rc[k, kd]), colmap, srcs...) + end + end + end + return R +end + +# A `sqrt(eps)` threshold (a common rule of thumb, relying on Jacobi's +# asymptotic quadratic convergence to reach full precision one sweep after +# crossing it) is too loose here: the measured relative off-diagonal coupling +# does not always cross `sqrt(eps)` and then immediately collapse to `eps` on +# the very next sweep — for some inputs it lands on a plateau anywhere between +# the two. Stopping as soon as `sqrt(eps)` is reached can therefore leave the +# block-columns short of full convergence, which — since the left singular +# vectors are formed by dividing columns by their own (possibly tiny) norm — +# gets magnified into large errors in `U`. Requiring a full `eps`-level +# relative coupling keeps sweeping until that plateau is actually reached. +_svd_default_tol(::Type{T}) where {T} = eps(real(float(T))) + +# ──────────────────────────────────────────────────────────────────────────── +# Core block one-sided Jacobi (requires m ≥ n) +# ──────────────────────────────────────────────────────────────────────────── + +# Returns `(U, S, V)` where `U` is m×n (or `nothing` when `vectors=false`), +# `S` is the descending singular values, and `V` is n×n (or `nothing`). +function _svd_jacobi!(A::DMatrix{T}; tol::Real=_svd_default_tol(T), + maxsweeps::Integer=30, vectors::Bool=true) where {T<:Number} + m, n = size(A) + m >= n || throw(ArgumentError("_svd_jacobi! requires m ≥ n (got $m×$n)")) + R = real(float(T)) + Ac = A.chunks + mt, nt = size(Ac) + cum = A.subdomains.cumlength[2] + widths = _svd_tilewidths(cum) + nb = A.partitioning.blocksize[2] + local V::DMatrix{T} + + if nt == 1 + # Single block-column: the eigenvectors of A'A orthogonalize the columns + # exactly in one shot, so no sweeps are needed. + G = zeros(T, n, n) + Dagger.spawn_datadeps() do + for k in 1:mt + Dagger.@spawn _svd_gram_self_acc!(InOut(G), In(Ac[k, 1])) + end + end + vals, vecs = _svd_hermitian_jacobi_eigen(LinearAlgebra.Hermitian((G .+ G') ./ 2)) + if !vectors + return nothing, sort(sqrt.(max.(vals, zero(R))); rev=true), nothing + end + W = Matrix{T}(vecs) + V = _svd_identity(T, n, nb) + Vc = V.chunks + Dagger.spawn_datadeps() do + for k in 1:mt + Dagger.@spawn _svd_apply_self!(InOut(Ac[k, 1]), In(W)) + end + for k in 1:size(Vc, 1) + Dagger.@spawn _svd_apply_self!(InOut(Vc[k, 1]), In(W)) + end + end + else + if vectors + V = _svd_identity(T, n, nb) + Vc = V.chunks + vt = size(Vc, 1) + end + # Cyclic Jacobi's off-diagonal coupling converges quadratically at + # first, but — once it nears the floor imposed by the working + # precision — can plateau at a value that never quite dips below a + # fixed `tol` (see `_svd_default_tol`). Without a plateau check, + # every sweep after that point is wasted work, running all the way to + # `maxsweeps` for no further accuracy gain. A single sweep with a + # disappointing improvement ratio isn't a reliable plateau signal on + # its own — with many block-columns some sweeps in the middle of + # otherwise-healthy quadratic convergence can still only manage, + # say, a 2x reduction — so only stop early once two *consecutive* + # sweeps each fail to shrink the coupling by much (a true plateau, + # as opposed to a merely slow sweep, keeps failing to improve). + prev_offv = R(Inf) + stall_count = 0 + for _ in 1:maxsweeps + offv = R[zero(R)] + Dagger.spawn_datadeps() do + for i in 1:nt-1, j in i+1:nt + wi = widths[i] + wj = widths[j] + G = zeros(T, wi + wj, wi + wj) + W = Matrix{T}(undef, wi + wj, wi + wj) + for k in 1:mt + Dagger.@spawn _svd_gram_acc!(InOut(G), In(Ac[k, i]), In(Ac[k, j])) + end + Dagger.@spawn _svd_offdiag_acc!(InOut(offv), In(G), wi) + Dagger.@spawn _svd_rot!(Out(W), In(G)) + for k in 1:mt + Dagger.@spawn _svd_apply_rot!(InOut(Ac[k, i]), InOut(Ac[k, j]), In(W), wi) + end + if vectors + for k in 1:vt + Dagger.@spawn _svd_apply_rot!(InOut(Vc[k, i]), InOut(Vc[k, j]), In(W), wi) + end + end + end + end + offv[1] <= tol && break + stall_count = offv[1] > prev_offv * R(0.9) ? stall_count + 1 : 0 + prev_offv = offv[1] + stall_count >= 2 && break + end + end + + # Singular values from converged column norms. + svbuf = [zeros(R, widths[i]) for i in 1:nt] + Dagger.spawn_datadeps() do + for i in 1:nt, k in 1:mt + Dagger.@spawn _svd_colnorm_acc!(InOut(svbuf[i]), In(Ac[k, i])) + end + end + for i in 1:nt + svbuf[i] .= sqrt.(svbuf[i]) + end + σ = reduce(vcat, svbuf) + + if !vectors + return nothing, sort(σ; rev=true), nothing + end + + # Normalize columns of A in place → left singular vectors, then sort. + Dagger.spawn_datadeps() do + for i in 1:nt, k in 1:mt + Dagger.@spawn _svd_scale_cols!(InOut(Ac[k, i]), In(svbuf[i])) + end + end + p = sortperm(σ; rev=true) + U = _svd_permute_columns(A, p) + Vsorted = _svd_permute_columns(V, p) + unsafe_free!(V) + return U, σ[p], Vsorted +end + +# ──────────────────────────────────────────────────────────────────────────── +# Public LinearAlgebra interface +# ──────────────────────────────────────────────────────────────────────────── + +""" + svd!(A::DMatrix; tol, maxsweeps, full=false) -> SVD + +In-place thin singular value decomposition of a distributed matrix using a +tiled block one-sided Jacobi algorithm (see `src/array/svd.jl`). Overwrites +`A`. Returns a `LinearAlgebra.SVD` holding distributed `U` and `Vt` factors and +a host `Vector` of singular values in descending order. + +`tol` is the relative off-diagonal threshold for sweep convergence and +`maxsweeps` bounds the number of Jacobi sweeps. Only the thin factorization is +supported; `full=true` raises an error. +""" +function LinearAlgebra.svd!(A::DMatrix{T}; tol::Real=_svd_default_tol(T), + maxsweeps::Integer=30, full::Bool=false) where {T<:Number} + full && throw(ArgumentError("full=true SVD is not supported for DMatrix; only the thin SVD is available")) + m, n = size(A) + if m >= n + U, S, V = _svd_jacobi!(A; tol=tol, maxsweeps=maxsweeps) + Vt = _svd_adjoint(V) + unsafe_free!(V) + return LinearAlgebra.SVD(U, S, Vt) + else + # Wide matrix: factor the (tall) adjoint, then swap the roles of U and V. + # A = (Aᴴ)ᴴ = (U₁ Σ V₁ᴴ)ᴴ = V₁ Σ U₁ᴴ ⇒ U = V₁, Vᵀ = U₁ᴴ. + At = _svd_adjoint(A) + U1, S, V1 = _svd_jacobi!(At; tol=tol, maxsweeps=maxsweeps) + unsafe_free!(At) + Vt = _svd_adjoint(U1) + unsafe_free!(U1) + return LinearAlgebra.SVD(V1, S, Vt) + end +end + +""" + svd(A::DMatrix; kwargs...) -> SVD + +Thin singular value decomposition of a distributed matrix. `A` is not modified; +see [`svd!`](@ref) for the in-place variant and the list of keyword arguments. +""" +LinearAlgebra.svd(A::DMatrix; kwargs...) = LinearAlgebra.svd!(copy(A); kwargs...) + +""" + svdvals!(A::DMatrix; tol, maxsweeps) -> Vector + +In-place computation of the singular values of `A` (descending). Overwrites +`A` and skips forming the singular vectors. +""" +function LinearAlgebra.svdvals!(A::DMatrix{T}; tol::Real=_svd_default_tol(T), + maxsweeps::Integer=30) where {T<:Number} + m, n = size(A) + if m >= n + _, S, _ = _svd_jacobi!(A; tol=tol, maxsweeps=maxsweeps, vectors=false) + return S + else + At = _svd_adjoint(A) + _, S, _ = _svd_jacobi!(At; tol=tol, maxsweeps=maxsweeps, vectors=false) + unsafe_free!(At) + return S + end +end + +""" + svdvals(A::DMatrix; kwargs...) -> Vector + +Singular values of `A` in descending order. `A` is not modified. +""" +LinearAlgebra.svdvals(A::DMatrix; kwargs...) = LinearAlgebra.svdvals!(copy(A); kwargs...) + +# ──────────────────────────────────────────────────────────────────────────── +# Solve via SVD: x = V Σ⁺ Uᴴ b (least-squares / min-norm) +# ──────────────────────────────────────────────────────────────────────────── + +# Scale each row of a tile by the reciprocal of the corresponding singular +# value (or zero it when σ is treated as numerically zero). `s` is the full +# descending singular-value vector; `row_offset` is the 0-based global start +# of this tile's rows. +function _svd_scale_rows!(Y::AbstractMatrix{T}, s::AbstractVector{R}, + row_offset::Int, k::Int) where {T,R} + @inbounds for i in axes(Y, 1) + gi = row_offset + i + if gi <= k + sc = s[gi] + invs = iszero(sc) ? zero(T) : T(inv(sc)) + @views Y[i, :] .*= invs + else + @views Y[i, :] .= zero(T) + end + end + return Y +end +function _svd_scale_rows!(y::AbstractVector{T}, s::AbstractVector{R}, + row_offset::Int, k::Int) where {T,R} + @inbounds for i in eachindex(y) + gi = row_offset + i + if gi <= k + sc = s[gi] + y[i] = iszero(sc) ? zero(T) : T(y[i] / sc) + else + y[i] = zero(T) + end + end + return y +end + +# Apply Σ⁺ in place to a distributed vector/matrix whose leading dimension +# matches the singular-value length (or is longer — trailing rows are zeroed). +function _svd_apply_pinv_S!(Y::DVecOrMat{T}, S::AbstractVector{R}, k::Int) where {T,R} + Yc = Y.chunks + if Y isa DVector + cum = Y.subdomains.cumlength[1] + Dagger.spawn_datadeps() do + for i in eachindex(Yc) + off = i == 1 ? 0 : cum[i - 1] + Dagger.@spawn _svd_scale_rows!(InOut(Yc[i]), In(S), off, k) + end + end + else + cum = Y.subdomains.cumlength[1] + mt, nt = size(Yc) + Dagger.spawn_datadeps() do + for i in 1:mt, j in 1:nt + off = i == 1 ? 0 : cum[i - 1] + Dagger.@spawn _svd_scale_rows!(InOut(Yc[i, j]), In(S), off, k) + end + end + end + return Y +end + +""" + ldiv!(F::SVD{<:Any,<:Any,<:DMatrix}, B::DVecOrMat) -> B + +Solve `F.U * Diagonal(F.S) * F.Vt * X ≈ B` in place using the distributed thin +SVD factors. Overwrites `B` with the least-squares / minimum-norm solution +(same contract as `LinearAlgebra.ldiv!(::SVD, ::AbstractVecOrMat)`): for an +`m×n` factorization, `B` must have `max(m, n)` rows and the solution occupies +the leading `n` rows. + +Singular values below `eps(real(T)) * S[1]` are treated as zero (truncated +pseudoinverse), matching the dense `LinearAlgebra` SVD solve. +""" +function LinearAlgebra.ldiv!(F::LinearAlgebra.SVD{T,<:Any,<:DMatrix{T}}, + B::DVecOrMat) where {T} + m, n = size(F) + size(B, 1) == max(m, n) || throw(DimensionMismatch( + "B has $(size(B, 1)) rows but SVD is $m×$n (need max(m,n)=$(max(m, n)) rows)")) + k = searchsortedlast(F.S, eps(real(T)) * F.S[1]; rev=true) + # y ← Uᴴ b (length m for thin U; pad/truncate into B's leading m rows) + # For the thin SVD, U is m×min(m,n). When B is taller than m (wide A), + # only the leading m rows of B hold the RHS; when B is exactly m (tall or + # square), the whole vector is the RHS. + if B isa DVector + b_rhs = size(B, 1) > m ? B[1:m] : B + y = F.U' * b_rhs + _svd_apply_pinv_S!(y, F.S, k) + # x ← V y = Vt' y (length n). Write into leading n entries of B. + x = F.Vt' * y + if size(B, 1) == n + copyto!(B, x) + else + # Tall: B is length m > n; solution occupies 1:n, rest unused. + fill!(B, zero(eltype(B))) + B[1:n] = x + end + unsafe_free!(y) + unsafe_free!(x) + b_rhs !== B && unsafe_free!(b_rhs) + else + B_rhs = size(B, 1) > m ? B[1:m, :] : B + Y = F.U' * B_rhs + _svd_apply_pinv_S!(Y, F.S, k) + X = F.Vt' * Y + if size(B, 1) == n + copyto!(B, X) + else + fill!(B, zero(eltype(B))) + B[1:n, :] = X + end + unsafe_free!(Y) + unsafe_free!(X) + B_rhs !== B && unsafe_free!(B_rhs) + end + return B +end + +function LinearAlgebra.ldiv!(X::DVecOrMat, + F::LinearAlgebra.SVD{T,<:Any,<:DMatrix{T}}, + B::DVecOrMat) where {T} + m, n = size(F) + size(B, 1) == m || throw(DimensionMismatch( + "B has $(size(B, 1)) rows but SVD has $m rows")) + size(X, 1) == n || throw(DimensionMismatch( + "X has $(size(X, 1)) rows but SVD has $n columns (solution length)")) + if ndims(X) == 2 + size(X, 2) == size(B, 2) || throw(DimensionMismatch( + "X and B must have the same number of columns")) + end + # Workspace for `ldiv!(F, ·)` must have max(m, n) rows. + if m >= n + Bc = m == n ? copyto!(X, B) : copy(B) + LinearAlgebra.ldiv!(F, Bc) + if m > n + if X isa DVector + copyto!(X, Bc[1:n]) + else + copyto!(X, Bc[1:n, :]) + end + unsafe_free!(Bc) + end + else + # Underdetermined: pad B into an n-row workspace, solve, result is X. + fill!(X, zero(eltype(X))) + if X isa DVector + X[1:m] = B + else + X[1:m, :] = B + end + LinearAlgebra.ldiv!(F, X) + end + return X +end + +""" + inv(F::SVD{<:Any,<:Any,<:DMatrix}) -> DMatrix + +Pseudoinverse `V Σ⁺ Uᴴ` of a square distributed SVD factorization. Singular +values below `eps(real(T)) * S[1]` are treated as zero. +""" +function LinearAlgebra.inv(F::LinearAlgebra.SVD{T,<:Any,<:DMatrix{T}}) where {T} + LinearAlgebra.checksquare(F) + @inbounds for i in eachindex(F.S) + iszero(F.S[i]) && throw(LinearAlgebra.SingularException(i)) + end + k = searchsortedlast(F.S, eps(real(T)) * F.S[1]; rev=true) + # V Σ⁺ Uᴴ = Vt' * Diagonal(Σ⁺) * U' + Ut = _svd_adjoint(F.U) + _svd_apply_pinv_S!(Ut, F.S, k) + P = F.Vt' * Ut + unsafe_free!(Ut) + return P +end diff --git a/test/array/linalg/svd.jl b/test/array/linalg/svd.jl new file mode 100644 index 000000000..1e7ccb8d6 --- /dev/null +++ b/test/array/linalg/svd.jl @@ -0,0 +1,176 @@ +using LinearAlgebra +# ────────────────────────────────────────────────────────────────────── +# Helper: validate a thin SVD factorization F of the dense matrix A_col +# ────────────────────────────────────────────────────────────────────── +function check_svd(A_col, F; tol=200.0) + T = eltype(A_col) + eps_val = eps(real(T)) + m, n = size(A_col) + k = min(m, n) + + U = collect(F.U) + S = F.S + Vt = collect(F.Vt) + + @test size(U) == (m, k) + @test length(S) == k + @test size(Vt) == (k, n) + + # Reconstruction residual + recon = U * Diagonal(S) * Vt + res_f = opnorm(A_col - recon, 1) / (opnorm(A_col, 1) * max(m, n) * eps_val) + @test res_f < tol + + # Orthonormality of singular vectors + res_u = opnorm(U' * U - I, 1) / (k * eps_val) + res_v = opnorm(Vt * Vt' - I, 1) / (k * eps_val) + @test res_u < tol + @test res_v < tol + + # Singular values: nonnegative, descending, and matching LAPACK + @test all(>=(0), S) + @test issorted(S; rev=true) + @test S ≈ svdvals(A_col) +end + +# ====================================================================== +# 1. svd(DA) — varying shapes, block sizes, element types +# ====================================================================== +@testset "Tile SVD: $T" for T in (Float32, Float64, ComplexF32, ComplexF64) + @testset "Square" begin + @testset "blocks=$bs" for bs in [(32, 32), (16, 32), (16, 16)] + A = rand(T, 128, 128) + DA = distribute(A, Blocks(bs...)) + check_svd(collect(DA), svd(DA)) + # svd must not modify its input + @test collect(DA) ≈ A + end + end + + @testset "Tall" begin + @testset "blocks=$bs" for bs in [(32, 32), (16, 32)] + A = rand(T, 128, 64) + DA = distribute(A, Blocks(bs...)) + check_svd(collect(DA), svd(DA)) + end + end + + @testset "Wide" begin + @testset "blocks=$bs" for bs in [(32, 32), (32, 16)] + A = rand(T, 64, 128) + DA = distribute(A, Blocks(bs...)) + check_svd(collect(DA), svd(DA)) + end + end + + @testset "Irregular tiling" begin + A = rand(T, 100, 60) + DA = distribute(A, Blocks(24, 16)) + check_svd(collect(DA), svd(DA)) + end + + @testset "Single block-column" begin + A = rand(T, 48, 40) + DA = distribute(A, Blocks(48, 64)) # one column tile + check_svd(collect(DA), svd(DA)) + end +end + +# ====================================================================== +# 2. In-place svd! (exercises the in-place path and input destruction) +# ====================================================================== +@testset "In-place svd!: $T" for T in (Float64, ComplexF64) + A = rand(T, 96, 96) + DA = distribute(A, Blocks(32, 32)) + F = svd!(DA) + check_svd(A, F) +end + +# ====================================================================== +# 3. svdvals / svdvals! agree with LAPACK and with svd +# ====================================================================== +@testset "svdvals: $T" for T in (Float64, ComplexF64) + @testset "shape=$sz" for sz in [(128, 128), (128, 48), (48, 128)] + A = rand(T, sz...) + DA = distribute(A, Blocks(32, 32)) + + S = svdvals(DA) + @test S ≈ svdvals(A) + @test collect(DA) ≈ A # svdvals must not modify input + + # svdvals agrees with the singular values from a full svd + @test S ≈ svd(DA).S + + # in-place variant + DA2 = distribute(A, Blocks(32, 32)) + @test svdvals!(DA2) ≈ svdvals(A) + end +end + +# ====================================================================== +# 4. Known spectrum: reconstruct a matrix with prescribed singular values +# ====================================================================== +@testset "Prescribed spectrum: $T" for T in (Float64, ComplexF64) + n = 64 + svals = collect(range(10.0, 1.0; length=n)) + Q1, _ = qr(rand(T, n, n)) + Q2, _ = qr(rand(T, n, n)) + A = Matrix(Q1) * Diagonal(T.(svals)) * Matrix(Q2)' + DA = distribute(A, Blocks(16, 16)) + S = svdvals(DA) + @test S ≈ svals +end + +# ====================================================================== +# 5. full=true is rejected +# ====================================================================== +@testset "full=true unsupported" begin + DA = distribute(rand(Float64, 32, 32), Blocks(16, 16)) + @test_throws ArgumentError svd(DA; full=true) +end + +# ====================================================================== +# 6. Solve via SVD: \, ldiv!, inv +# ====================================================================== +@testset "SVD solve: $T" for T in (Float64, ComplexF64) + tol = T <: Complex ? 1e-8 : 1e-10 + + @testset "shape=$sz" for sz in [(64, 64), (96, 48), (48, 96)] + m, n = sz + A = rand(T, m, n) + DA = distribute(A, Blocks(16, 16)) + F = svd(DA) + Fd = svd(A) + + # Vector RHS + b = rand(T, m) + Db = distribute(b, Blocks(16)) + x = F \ Db + xd = Fd \ b + @test size(x) == (n,) + @test collect(x) ≈ xd rtol=tol + + X = zeros(Blocks(16), T, n) + ldiv!(X, F, Db) + @test collect(X) ≈ xd rtol=tol + + # Matrix RHS + B = rand(T, m, 3) + DB = distribute(B, Blocks(16, 3)) + X = F \ DB + Xd = Fd \ B + @test size(X) == (n, 3) + @test collect(X) ≈ Xd rtol=tol + + X2 = zeros(Blocks(16, 3), T, n, 3) + ldiv!(X2, F, DB) + @test collect(X2) ≈ Xd rtol=tol + end + + @testset "inv (square)" begin + A = rand(T, 64, 64) + DA = distribute(A, Blocks(16, 16)) + P = inv(svd(DA)) + @test collect(P) * A ≈ I rtol=tol + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 5c2e1e3cd..eb50157de 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -35,6 +35,7 @@ tests = [ ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"), ("Array - LinearAlgebra - Solve", "array/linalg/solve.jl"), ("Array - LinearAlgebra - QR", "array/linalg/qr.jl"), + ("Array - LinearAlgebra - SVD", "array/linalg/svd.jl"), ("Array - Permute", "array/permute.jl"), ("Array - Random", "array/random.jl"), ("Array - Stencils", "array/stencil.jl"), From debd1991b0eb47a45a38e63cc1e208e6a4820447 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 7 Jul 2026 11:20:30 -0700 Subject: [PATCH 2/9] CI: Bump CPU job timeouts --- .github/workflows/CI.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c132eac2a..9cd0ffb42 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,7 +19,7 @@ jobs: test: name: Julia ${{ matrix.version }} - ${{ matrix.os }} (${{ matrix.arch }}) runs-on: ${{ matrix.os }} - timeout-minutes: 150 + timeout-minutes: 180 if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} env: JULIA_NUM_THREADS: '1' @@ -55,7 +55,7 @@ jobs: test-multithreaded: name: Julia ${{ matrix.version }} - ${{ matrix.os }} (${{ matrix.arch }}) (multithreaded) runs-on: ${{ matrix.os }} - timeout-minutes: 90 + timeout-minutes: 120 if: ${{ !contains(github.event.head_commit.message, '[skip tests]') }} strategy: fail-fast: false From ea9bdf26bf752b4aa2de963abb3847b0b1074954 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 7 Jul 2026 13:51:44 -0700 Subject: [PATCH 3/9] CI: Upgrade macOS runners to Arm64 v26 --- .github/workflows/CI.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 9cd0ffb42..6841930f3 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -32,7 +32,7 @@ jobs: - {version: '1.12', os: ubuntu-latest, arch: x64} - {version: '1', os: ubuntu-latest, arch: x64} - {version: '1', os: windows-latest, arch: x64} - - {version: '1', os: macos-15, arch: x64} + - {version: '1', os: macos-26, arch: arm64} - {version: 'nightly', os: ubuntu-latest, arch: x64} steps: - uses: actions/checkout@v4 @@ -66,7 +66,7 @@ jobs: - {version: '1.12', os: ubuntu-latest, arch: x64} - {version: '1', os: ubuntu-latest, arch: x64} - {version: '1', os: windows-latest, arch: x64} - - {version: '1', os: macos-15, arch: x64} + - {version: '1', os: macos-26, arch: arm64} - {version: 'nightly', os: ubuntu-latest, arch: x64} steps: - uses: actions/checkout@v4 From 744f8ad25c604829525b53d9ba4526d80c253d78 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 7 Jul 2026 15:58:24 -0700 Subject: [PATCH 4/9] datadeps: Disallow passing DArray to tasks --- src/array/darray.jl | 5 +++++ test/datadeps.jl | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/src/array/darray.jl b/src/array/darray.jl index 32336f95d..fb27adb7a 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -321,6 +321,11 @@ function Base.isequal(x::ArrayOp, y::ArrayOp) x === y end +aliasing(x::DArray) = + throw(ConcurrencyViolationError("DArray aliasing may be mixed and unstable")) +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) Base.similar(D::DArray{T,N1} where T, ::Type{S}, dims::Dims{N2}) where {S,N1,N2} = diff --git a/test/datadeps.jl b/test/datadeps.jl index 84c31be8f..291e10e70 100644 --- a/test/datadeps.jl +++ b/test/datadeps.jl @@ -82,6 +82,13 @@ end end end +@testset "DArray" begin + A = rand(Blocks(2), 4) + @test_throws ConcurrencyViolationError Dagger.spawn_datadeps() do + Dagger.@spawn sum(A) + end +end + function test_move_rewrap_aliasing(obj, dest_space) src_space = Dagger.memory_space(obj) from_proc = first(Dagger.processors(src_space)) From bd5514313d3217493817d86a76c9a66dc54302a0 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 8 Jul 2026 13:46:35 -0700 Subject: [PATCH 5/9] DArray: Fix A' * B --- src/array/mul.jl | 23 +++++++++++++++++------ test/array/linalg/matmul.jl | 11 ++++++++++- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/src/array/mul.jl b/src/array/mul.jl index 3c9d72ff5..3a2637b76 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -487,11 +487,22 @@ function gemv_dagger!( alpha = T(_alpha) beta = T(_beta) - if Ant != Bmt - throw(DimensionMismatch(lazy"A has number of blocks ($Amt,$Ant) but B has number of blocks ($Bmt)")) - end - if Amt != Cmt - throw(DimensionMismatch(lazy"A has number of blocks ($Amt,$Ant) but C has number of blocks ($Cmt)")) + # For op(A)*x: when A is not transposed, x matches A's column-blocks and + # C matches A's row-blocks; when A is [conj-]transposed the roles swap. + if transA == 'N' + if Ant != Bmt + throw(DimensionMismatch(lazy"A has number of blocks ($Amt,$Ant) but B has number of blocks ($Bmt)")) + end + if Amt != Cmt + throw(DimensionMismatch(lazy"A has number of blocks ($Amt,$Ant) but C has number of blocks ($Cmt)")) + end + else + if Amt != Bmt + throw(DimensionMismatch(lazy"A' has number of blocks ($Ant,$Amt) but B has number of blocks ($Bmt)")) + end + if Ant != Cmt + throw(DimensionMismatch(lazy"A' has number of blocks ($Ant,$Amt) but C has number of blocks ($Cmt)")) + end end Dagger.spawn_datadeps() do @@ -510,7 +521,7 @@ function gemv_dagger!( ) end else - # A: [Conj]Trans + # 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!( diff --git a/test/array/linalg/matmul.jl b/test/array/linalg/matmul.jl index 7602fbe42..26fd0afe3 100644 --- a/test/array/linalg/matmul.jl +++ b/test/array/linalg/matmul.jl @@ -163,12 +163,21 @@ function test_gemv!(T, szA, szB, partA, partB) @test collect(DC) ≈ C if szA[1] == szB[1] - # transA + # transA (square / matching inner dim for A') DC = DA' * DB C = A' * B @test collect(DC) ≈ C end + # Tall/wide A': b matches A's row count, result matches A's column count + if szA[1] != szA[2] + B2 = rand(T, szA[1]) + DB2 = distribute(B2, Blocks(partA.blocksize[1])) + DC2 = DA' * DB2 + C2 = A' * B2 + @test collect(DC2) ≈ C2 + end + ## In-place gemm # No transA C = zeros(T, szC...) From c54f857aa4fe51fa7428cce1582cec6b002d2be9 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 24 Mar 2026 15:51:48 -0700 Subject: [PATCH 6/9] datadeps: Add hierarchical scheduling --- src/Dagger.jl | 1 + src/datadeps/aliasing.jl | 3 +- src/datadeps/hierarchical.jl | 789 +++++++++++++++++++++++++++++++++++ src/datadeps/queue.jl | 17 +- src/datadeps/scheduling.jl | 8 + 5 files changed, 814 insertions(+), 4 deletions(-) create mode 100644 src/datadeps/hierarchical.jl diff --git a/src/Dagger.jl b/src/Dagger.jl index 593eee4aa..3298e63a1 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -97,6 +97,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 13432748a..7329c3fe2 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -1,4 +1,5 @@ -import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, nv +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..5d16b7cca --- /dev/null +++ b/src/datadeps/hierarchical.jl @@ -0,0 +1,789 @@ +# 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 + +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 + +# 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) + for (arg_w, ainfo) in results + arg_to_ainfo[arg_w] = ainfo + 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 + @lock all_results_lock begin + for (arg_w, ainfo) in results + arg_to_ainfo[arg_w] = ainfo + 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_partition_full!(queue, queue_lock, partition_id, partition_verts, + dag, seen_tasks, task_metas, local_procs, + vertex_to_partition, task_submitted) -> 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. +""" +function schedule_partition_full!(queue::DataDepsTaskQueue, + queue_lock::ReentrantLock, + partition_id::Int, + partition_verts::Vector{Int}, + dag::SimpleDiGraph, + seen_tasks::Vector{DTaskPair}, + task_metas::Vector{HierarchicalTaskMeta}, + local_procs::Vector{<:Processor}, + vertex_to_partition::Vector{Int}, + task_submitted::Vector{Base.Event}) + if isempty(partition_verts) || isempty(local_procs) + return DataDepsState() + end + + local_scope = UnionScope(map(ExactScope, local_procs)) + + state = DataDepsState() + 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)) + + # 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 + + 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 + + write_num = distribute_task!(temp_queue, state, local_procs, local_scope, + spec, task, spec.fargs, + proc_to_scope_lfu, write_num) + + 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!` + +Both single-worker and multi-worker use `schedule_partition_full!` so argument +preparation and `DataDepsScheduler` dispatch stay correct. 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 + + # 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) + + # 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 (distribute_task! + cross-partition DAG syncdeps) + partition_states = Vector{DataDepsState}(undef, n_partitions) + + if n_partitions == 1 + # Avoid Threads.@spawn overhead when there is nothing to parallelize. + locked_queue = LockedEnqueueQueue(wait_all_queue, queue_lock) + with_options(; task_queue=locked_queue) do + partition_states[1] = schedule_partition_full!( + queue, queue_lock, 1, partitions[1], + dag, seen_tasks, task_metas, + partition_procs[1], vertex_to_partition, + task_submitted + ) + end + else + 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, task_metas, + partition_procs[pid], vertex_to_partition, + task_submitted + ) + end + end + end + catch e + rethrow(_unwrap_partition_exception(e)) + end + end + + _hierarchical_copy_from_and_free!(partition_states, n_partitions) +end + +function _hierarchical_copy_from_and_free!(partition_states::Vector{DataDepsState}, n_partitions::Int) + 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 + 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 + 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] + arg = arg_w.arg + haskey(state.arg_origin, arg) || continue + origin_space = state.arg_origin[arg] + write_num = wn + 1 + 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 + end + + 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 + 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 13d0c6b63..927a026a3 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -58,7 +58,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 @@ -71,22 +72,32 @@ 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 if launch_wait result = spawn_bulk() do queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) with_options(f; task_queue=queue) - distribute_tasks!(queue) + if hierarchical + distribute_tasks_hierarchical!(queue) + else + distribute_tasks!(queue) + end end else queue = DataDepsTaskQueue(get_options(:task_queue); scheduler) result = with_options(f; task_queue=queue) - distribute_tasks!(queue) + if hierarchical + distribute_tasks_hierarchical!(queue) + else + distribute_tasks!(queue) + end 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) function distribute_tasks!(queue::DataDepsTaskQueue) #= TODO: Improvements to be made: diff --git a/src/datadeps/scheduling.jl b/src/datadeps/scheduling.jl index 62a350c72..830c50353 100644 --- a/src/datadeps/scheduling.jl +++ b/src/datadeps/scheduling.jl @@ -1,9 +1,15 @@ abstract type DataDepsScheduler end +# 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(sched::RoundRobinScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) proc_idx = sched.proc_idx our_proc = all_procs[proc_idx] @@ -24,6 +30,7 @@ function datadeps_schedule_task(sched::RoundRobinScheduler, state::DataDepsState end struct NaiveScheduler <: DataDepsScheduler end +Base.similar(::NaiveScheduler) = NaiveScheduler() function datadeps_schedule_task(sched::NaiveScheduler, state::DataDepsState, 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 @@ -59,6 +66,7 @@ struct UltraScheduler <: DataDepsScheduler Dict{MemorySpace,Int}()) end end +Base.similar(::UltraScheduler) = UltraScheduler() function datadeps_schedule_task(sched::UltraScheduler, state::DataDepsState, all_procs, all_scope, task_scope, spec::DTaskSpec, task::DTask) args = Base.mapany(spec.fargs) do arg pos, data = arg From 1ca0e80020a467d0a2d802304aee7aef6d499f9f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Wed, 1 Jul 2026 09:58:10 -0700 Subject: [PATCH 7/9] dagdebug: Disable by default --- src/utils/dagdebug.jl | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/utils/dagdebug.jl b/src/utils/dagdebug.jl index 3812ba271..fe9c1d2b9 100644 --- a/src/utils/dagdebug.jl +++ b/src/utils/dagdebug.jl @@ -2,9 +2,20 @@ function istask end function task_id end # Use a Set for O(1) membership checks (vs O(n) for Vector). -const DAGDEBUG_CATEGORIES = Set{Symbol}([:global, :submit, :schedule, :scope, - :take, :execute, :move, :processor, :finish, - :cancel, :stream]) +# +# N.B. This is empty by default (debug tracing is opt-in only, via the +# `JULIA_DAGGER_DEBUG` environment variable or by `push!`-ing a category +# directly), because `@dagdebug` eagerly formats its message string (see +# `_dagdebug_emit` below) whenever its category is in this set, *before* +# Julia's own `@debug`/logger level check gets a chance to skip printing it. +# For hot-path categories like `:execute`, `:move`, `:schedule`, and +# `:processor`, that formatting (which often stringifies types and values) +# is prohibitively expensive to pay on every task, so it must not run unless +# a user has explicitly asked for tracing. +const DAGDEBUG_VALID_CATEGORIES = (:all, :global, :submit, :schedule, :scope, + :take, :execute, :move, :processor, :finish, + :cancel, :stream, :validate) +const DAGDEBUG_CATEGORIES = Set{Symbol}() # Out-of-line emission keeps call-site IR minimal: just one `in` check + one # function call per @dagdebug site, regardless of how complex the message is. From 9eed252477e8b8b1cf9caadd875a86fa7a6e7165 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 25 Jun 2026 11:39:54 -0700 Subject: [PATCH 8/9] TimespanLogging: Optimize MultiEventLog init and lock access --- lib/TimespanLogging/src/core.jl | 35 ++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/lib/TimespanLogging/src/core.jl b/lib/TimespanLogging/src/core.jl index 815345671..eb515bd37 100644 --- a/lib/TimespanLogging/src/core.jl +++ b/lib/TimespanLogging/src/core.jl @@ -206,23 +206,44 @@ function Base.setindex!(ml::MultiEventLog, c, name::Symbol) ml.consumers[name] = c end -function get_state(ml::MultiEventLog) - lock(event_log_lock) do - mls = get!(()->MultiEventLogState(), MultiEventLogState_PLS, ml.uid) - max_length = reduce(max, map(length, values(mls.consumer_logs)); init=0) +# Resolve (and lazily initialize) the process-local state for `ml`. The caller +# MUST already hold `event_log_lock`. Split out from `get_state` so the hot +# `write_event` path takes the lock exactly once (was twice: get_state + write). +function _get_state_locked(ml::MultiEventLog) + mls = get!(()->MultiEventLogState(), MultiEventLogState_PLS, ml.uid) + # Fast path: the consumer set is stable, so there is nothing to initialize. + # Consumers are only ever added (never removed -- see FIXME), so a length + # mismatch is a sufficient and cheap "needs init" test, and it lets us skip + # the per-event `map(length, ...)` allocation in steady state. + if length(mls.consumers) != length(ml.consumers) + max_length = 0 + for v in values(mls.consumer_logs) + l = length(v) + if l > max_length + max_length = l + end + end for name in keys(ml.consumers) if !haskey(mls.consumers, name) mls.consumers[name] = init_similar(ml.consumers[name]) mls.consumer_logs[name] = Vector{Any}(fill(nothing, max_length)) end end + end + if length(mls.aggregators) != length(ml.aggregators) for name in keys(ml.aggregators) if !haskey(mls.aggregators, name) mls.aggregators[name] = init_similar(ml.aggregators[name]) end end - # FIXME: Remove deleted consumers and aggregators - mls + end + # FIXME: Remove deleted consumers and aggregators + return mls +end + +function get_state(ml::MultiEventLog) + lock(event_log_lock) do + _get_state_locked(ml) end end @@ -230,8 +251,8 @@ end init_similar(x) = x function write_event(ml::MultiEventLog, event::Event) - mls = get_state(ml) lock(event_log_lock) do + mls = _get_state_locked(ml) for name in keys(mls.consumers) cevent = try mls.consumers[name](event) From 725957320f34b83c6db719a6c891ca85d9e57922 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Thu, 25 Jun 2026 11:41:50 -0700 Subject: [PATCH 9/9] TimespanLogging: Optimize non-profiling code paths --- lib/TimespanLogging/src/core.jl | 30 ++++++++++++++++++++++++++++-- src/utils/logging.jl | 3 +++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/lib/TimespanLogging/src/core.jl b/lib/TimespanLogging/src/core.jl index eb515bd37..d36f4c91b 100644 --- a/lib/TimespanLogging/src/core.jl +++ b/lib/TimespanLogging/src/core.jl @@ -299,16 +299,30 @@ end empty_prof() = ProfilerResult(UInt[], Dict{UInt64, Vector{Base.StackTraces.StackFrame}}(), UInt[]) +# Shared, immutable-in-practice empty profiler result. Every non-profiled event +# carries *this* object instead of allocating a fresh `ProfilerResult` (2 vectors +# + a dict) per event. It is only ever read (e.g. `mix_samples` does a read-only +# `vcat`), never mutated, on the non-profiling path. +const EMPTY_PROF = empty_prof() +const _NO_TASKS = Task[] + +# Profiling is opt-in and rare; this flag lets the (very hot) `timespan_finish` +# and per-compute-task `prof_task_put!` paths skip `prof_lock` entirely when +# profiling is off. `Dagger.enable_logging!(profile=true)` sets it. +const PROFILE_TASKS = Ref{Bool}(false) + const prof_refcount = Ref{Threads.Atomic{Int}}(Threads.Atomic{Int}(0)) const prof_lock = Threads.ReentrantLock() const prof_tasks = IdDict{Any, Vector{Task}}() function prof_task_put!(id, task::Task=Base.current_task()) + PROFILE_TASKS[] || return lock(prof_lock) do push!(get!(()->Task[], prof_tasks, id), task) end end function prof_tasks_take!(id) + PROFILE_TASKS[] || return _NO_TASKS lock(prof_lock) do if haskey(prof_tasks, id) pop!(prof_tasks, id) @@ -338,7 +352,9 @@ function timespan_start(ctx, category::Symbol, @nospecialize(id), @nospecialize( Profile.start_timer() end end - ev = Event(:start, category, id, tl, time_ns(), gc_num(), empty_prof()) + # Start events never carry profiler samples (those are gathered at finish), + # so always reuse the shared empty result rather than allocating. + ev = Event(:start, category, id, tl, time_ns(), gc_num(), EMPTY_PROF) write_event(sink, ev) nothing end @@ -351,12 +367,22 @@ categorized by `category`, and uniquely identified by `id`; these two must be the same as previously passed to `timespan_start`. `tl` is the "timeline" of the event, which is just an arbitrary payload attached to the event. """ -function timespan_finish(ctx, category::Symbol, @nospecialize(id), @nospecialize(tl); tasks=prof_tasks_take!(id)) +function timespan_finish(ctx, category::Symbol, @nospecialize(id), @nospecialize(tl); tasks=nothing) sink = log_sink(ctx) isa(sink, NoOpLog) && return do_profile = profile(ctx, category, id, tl) time = time_ns() gcn = gc_num() + if !do_profile + # Hot path: no profiling. Skip `prof_lock`, `Profile.fetch`, and the + # per-event `ProfilerResult` allocation by reusing the shared empty one. + ev = Event(:finish, category, id, tl, time, gcn, EMPTY_PROF) + write_event(sink, ev) + return nothing + end + if tasks === nothing + tasks = prof_tasks_take!(id) + end prof = UInt[] lidict = Dict{UInt64, Vector{Base.StackTraces.StackFrame}}() GC.@preserve tasks begin diff --git a/src/utils/logging.jl b/src/utils/logging.jl index a431bd9eb..2479a9f86 100644 --- a/src/utils/logging.jl +++ b/src/utils/logging.jl @@ -81,6 +81,9 @@ function enable_logging!(;metrics::Bool=false, if profile ml[:profile] = DaggerWebDash.ProfileMetrics() end + if isdefined(TimespanLogging, :PROFILE_TASKS) + TimespanLogging.PROFILE_TASKS[] = profile + end if metrics ml[:wsat] = Dagger.Events.WorkerSaturation() ml[:loadavg] = TimespanLogging.Events.CPULoadAverages()