From 932c63eb0bfb3f3ad29167b35852924d5862fde7 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:19:38 +0000 Subject: [PATCH 1/4] T17: uniform batch in the numeric and solve phases (WIP) Co-Authored-By: Claude Opus 5.5 --- src/SparseDirectSolver.jl | 1 + src/dense/fallback/potrf.jl | 23 +++- src/dense/interface.jl | 36 +++++- src/generic.jl | 33 +++-- src/numeric/assembly.jl | 95 +++++++++++---- src/numeric/batch.jl | 139 +++++++++++++++++++++ src/numeric/extract.jl | 42 ++++++- src/numeric/factorize.jl | 118 +++++++++++++++--- src/numeric/front.jl | 73 ++++++----- src/numeric/ldlt.jl | 180 +++++++++++++++++----------- src/numeric/storage.jl | 99 ++++++++++++--- src/numeric/subtree.jl | 29 +++-- src/solve/permute.jl | 62 ++++++---- src/solve/refinement.jl | 105 ++++++++++------ src/solve/sweeps.jl | 233 +++++++++++++++++++++++++----------- src/solver.jl | 197 +++++++++++++++++++++--------- 16 files changed, 1091 insertions(+), 374 deletions(-) create mode 100644 src/numeric/batch.jl diff --git a/src/SparseDirectSolver.jl b/src/SparseDirectSolver.jl index 16ab9ad..ff7906c 100644 --- a/src/SparseDirectSolver.jl +++ b/src/SparseDirectSolver.jl @@ -52,6 +52,7 @@ include("symbolic/layout.jl") include("symbolic/maps.jl") # numeric storage, CPU reference multifrontal factorization (PLAN §7, the oracle) +include("numeric/batch.jl") include("numeric/storage.jl") include("reference/cholesky.jl") include("reference/ldlt.jl") diff --git a/src/dense/fallback/potrf.jl b/src/dense/fallback/potrf.jl index 077299b..71a2bb1 100644 --- a/src/dense/fallback/potrf.jl +++ b/src/dense/fallback/potrf.jl @@ -115,17 +115,18 @@ end @kernel function _ka_chol_check_kernel!(info, A, n, idx) li = @index(Local, Linear) + b = @index(Group, Linear) if li == 1 - @inbounds if info[idx] == 0 + @inbounds if info[idx + b - 1] == 0 st = Int32(0) for j in 1:n - d = _get(A, j, j, 1) + d = _get(A, j, j, b) if !(isfinite(real(d)) && real(d) > 0 && iszero(imag(d))) st = Int32(j) break end end - info[idx] = st + info[idx + b - 1] = st end end end @@ -145,3 +146,19 @@ function ka_chol_check_info!(info::AbstractVector{Int32}, idx::Integer, L::Abstr kernel!(info, L, size(L, 1), Int(idx); ndrange = 1) return info end + +""" + ka_chol_check_info!(info, offset, L::AbstractArray{T,3}) -> info + +Batched form: the check of `ka_chol_check_info!(info, offset + b, L[:, :, b])` +for every member `b` of the 3-D strided batch `L`, one launch. +""" +function ka_chol_check_info!(info::AbstractVector{Int32}, offset::Integer, L::AbstractArray{<:Any, 3}) + nb = size(L, 3) + 0 <= offset && offset + nb <= length(info) || + throw(DimensionMismatch("info has length $(length(info)) < offset $offset + batch count $nb")) + nb == 0 && return info + kernel! = _ka_chol_check_kernel!(KernelAbstractions.get_backend(L), 1) + kernel!(info, L, size(L, 1), Int(offset) + 1; ndrange = nb) + return info +end diff --git a/src/dense/interface.jl b/src/dense/interface.jl index 63232a9..c53543d 100644 --- a/src/dense/interface.jl +++ b/src/dense/interface.jl @@ -25,6 +25,7 @@ const DENSE_OPS = ( laswp = (generic = nothing, vendor = nothing), gemm_strided_batched = (generic = nothing, vendor = :vendor_gemm_strided_batched), trsm_strided_batched = (generic = nothing, vendor = :vendor_trsm_batched), + potrf_batched = (generic = nothing, vendor = :vendor_potrf_batched), ) # DENSE_OPS as a concretely typed table (`:none` = no such path), for `select_impl` @@ -74,7 +75,8 @@ end dense_impls(op::Symbol, backend, T) -> Vector{Symbol} The implementations of dense op `op` (`:gemm`, `:syrk`, `:herk`, `:trsm`, -`:potrf`, `:getrf`, `:laswp`, `:gemm_strided_batched`, `:trsm_strided_batched`) +`:potrf`, `:getrf`, `:laswp`, `:gemm_strided_batched`, `:trsm_strided_batched`, +`:potrf_batched`) available for element type `T` on `backend` according to [`capabilities`](@ref), in `:auto` preference order (`:vendor`, `:generic`, `:ka`). `:ka` is always present. @@ -382,6 +384,38 @@ function trsm_strided_batched!(side, uplo, trans, diag, α, A::AbstractArray{<:A return B end +""" + potrf_batched_info!(uplo, A, info, offset = 0; impl = :auto) -> info + +Uniform strided batch of [`potrf_info!`](@ref): Cholesky of the `uplo` +triangle of every member `A[:, :, b]` of the 3-D array `A` in place, its +status (0, or the first non-positive pivot column) into `info[offset + b]`. +Implementations: `:vendor` (`vendor_potrf_batched!`, validated on the device +as `potrf_info!`) and `:ka` ([`ka_potrf!`](@ref), one workgroup per member). +Asynchronous. +""" +function potrf_batched_info!(uplo, A::AbstractArray{<:Any, 3}, info::AbstractVector{Int32}, offset::Integer = 0; + impl::Symbol = :auto) + ul = _uplo_char(uplo) + size(A, 1) == size(A, 2) || throw(DimensionMismatch("potrf: members are $(size(A, 1))×$(size(A, 2)), not square")) + 0 <= offset && offset + size(A, 3) <= length(info) || + throw(DimensionMismatch("info has length $(length(info)) < offset $offset + batch count $(size(A, 3))")) + size(A, 3) == 0 && return info + return _potrf_batched_impl!(select_impl(:potrf_batched, A, impl), A, info, offset, ul) +end + +# `potrf_batched_info!` with an already resolved implementation `p` (no capability lookup) +function _potrf_batched_impl!(p::Symbol, A::AbstractArray{<:Any, 3}, info::AbstractVector{Int32}, offset::Integer, + ul::Char = 'L') + if p === :vendor + vendor_potrf_batched!(ul, A, view(info, (offset + 1):(offset + size(A, 3)))) + ka_chol_check_info!(info, offset, A) + else + ka_potrf!(ul, A, info; offset) + end + return info +end + """ strided_batch(buf, offset, m, n, stride, count) -> AbstractArray{T,3} diff --git a/src/generic.jl b/src/generic.jl index 1fa5d8b..f0bab5d 100644 --- a/src/generic.jl +++ b/src/generic.jl @@ -4,6 +4,12 @@ # `Symmetric`/`Hermitian` wrappers, as CUDSS.jl does. A `SparseMatrixCSC` is not # accepted by `cholesky` (that method belongs to CHOLMOD): wrap it with `CSR(A)`. # +# Uniform batches are detected from the values, as in CUDSS.jl: a matrix whose +# `nzval` is longer than `colval` (`length(nzval) ÷ length(colval) > 1` members, +# or an `nnz × nbatch` matrix of values) gives a batched solver; `ldiv!` then +# takes the right-hand sides of every member (`n × nbatch`, `n × nrhs × nbatch` +# or strided vectors). +# # PLAN §3.1 turns iterative refinement on in this layer: the solvers created # by `cholesky`/`ldlt` start with `ir_n_steps = 2` (the handle layer keeps the # cuDSS default 0); `ir_tol` (default 0: no early exit, no synchronization per @@ -29,6 +35,8 @@ change it with [`setparam!`](@ref)). `view` selects the triangle of `A` that is read (`'L'`, `'U'`, `'F'`). As in CUDSS.jl a failed factorization does not throw unless `check = true` (then [`FactorizationError`](@ref)); otherwise check `getparam(solver, "info")`. +A uniform batch (values of several members in `A.nzval`, see +[`DirectSolver`](@ref)) is factorized member by member with one analysis. """ function LinearAlgebra.cholesky(A::CSR{T}, ::NoPivot = NoPivot(); view::Char = 'F', check::Bool = false) where {T} solver = _linear_algebra_defaults!(DirectSolver(A, T <: Real ? "SPD" : "HPD", view)) @@ -39,8 +47,10 @@ function LinearAlgebra.cholesky(A::CSR{T}, ::NoPivot = NoPivot(); view::Char = ' end function _check_info(solver::DirectSolver, check::Bool) - check && solver.info != 0 && - throw(FactorizationError(solver.info, "the matrix is not positive definite (pivot $(solver.info))")) + if check && any(!=(0), solver.info) + info = _info_value(solver) + throw(FactorizationError(info, "the matrix is not positive definite (pivot $info)")) + end return nothing end @@ -97,7 +107,7 @@ end # one method per argument kind: LinearAlgebra has `ldiv!` methods for `Factorization` with # vectors and with matrices, so a `Union` here would be ambiguous -for RHS in (AbstractVector, AbstractMatrix, MatrixDescriptor) +for RHS in (AbstractVector, AbstractMatrix, AbstractArray{<:Any, 3}, MatrixDescriptor) @eval function LinearAlgebra.ldiv!(solver::DirectSolver, B::$RHS) execute!("solve", solver, B, B; asynchronous = false) return B @@ -113,17 +123,19 @@ end ldiv!(X, solver::DirectSolver, B) -> X Solve `A X = B` with the factorization in `solver` (`"solve"`, synchronized), -in place in `B` or into `X`. `B`, `X`: vectors, matrices or -[`MatrixDescriptor`](@ref)s on the solver's backend. +in place in `B` or into `X`. `B`, `X`: vectors, matrices, +`n × nrhs × nbatch` arrays (uniform batch) or [`MatrixDescriptor`](@ref)s on +the solver's backend. """ LinearAlgebra.ldiv!(::DirectSolver, ::AbstractVector) """ solver \\ B -`ldiv!(similar(B), solver, B)` for a vector or matrix `B`. +`ldiv!(similar(B), solver, B)` for a vector, matrix or `n × nrhs × nbatch` array `B`. """ Base.:\(solver::DirectSolver, B::AbstractVector) = ldiv!(similar(B), solver, B) Base.:\(solver::DirectSolver, B::AbstractMatrix) = ldiv!(similar(B), solver, B) +Base.:\(solver::DirectSolver, B::AbstractArray{<:Any, 3}) = ldiv!(similar(B), solver, B) # disambiguation with LinearAlgebra's real-factorization / complex-RHS method: the element types must match Base.:\(solver::DirectSolver{T}, B::Union{Array{Complex{T}, 1}, Array{Complex{T}, 2}}) where {T <: Union{Float32, Float64}} = ldiv!(similar(B), solver, B) @@ -152,11 +164,14 @@ permutations cancel, `det L = 1`), the product of the 1×1 pivots and of the determinants of the 2×2 blocks; `sign` is `±1` (real or Hermitian) or `det D / |det D|` (complex symmetric). With perturbed pivots (`"npivots" > 0`) this is the determinant of `A + E`. Raises [`FactorizationError`](@ref) when the -factorization failed (`"info" ≠ 0`). +factorization failed (`"info" ≠ 0`) and [`NotSupportedError`](@ref) for a +uniform batch (read the members' diagonals with `getparam(solver, "diag")`). """ function LinearAlgebra.logabsdet(solver::DirectSolver{T}) where {T} - solver.stage >= STAGE_FACTORIZED && solver.info == 0 || - throw(FactorizationError(solver.info, "logabsdet needs a successful factorization")) + solver.nbatch == 1 || + throw(NotSupportedError("logabsdet of a uniform batch of $(solver.nbatch) matrices; use getparam(solver, \"diag\")")) + solver.stage >= STAGE_FACTORIZED && solver.info[1] == 0 || + throw(FactorizationError(solver.info[1], "logabsdet needs a successful factorization")) _is_ldlt_structure(solver.structure) && return _logabsdet_ldlt(solver) d = Array(getparam(solver, "diag")) return 2 * sum(x -> log(abs(x)), d; init = zero(real(T))), one(T) diff --git a/src/numeric/assembly.jl b/src/numeric/assembly.jl index 5123da3..cb7589f 100644 --- a/src/numeric/assembly.jl +++ b/src/numeric/assembly.jl @@ -7,7 +7,9 @@ # reproducible. Contribution blocks on the update stack are packed lower # triangles (`_packed`, `m(m+1)/2` entries). On the KA CPU backend, values # derived from the group index are recomputed in every segment between barriers -# (see `src/dense/fallback/common.jl`). +# (see `src/dense/fallback/common.jl`). Uniform batch: one workgroup per (front, +# active member), the kernels see the member's panels (`MemberPanels` in place +# of `front_ptr`), update stack and values (`src/numeric/batch.jl`). "Workgroup size of the assembly kernels." const ASSEMBLY_WORKGROUP = 256 @@ -16,7 +18,8 @@ const ASSEMBLY_WORKGROUP = 256 @inline function _zero_front!(factor, stack, s, li, front_ptr, front_nrows, front_ncols, cb_ptr, ::Val{WG}) where {WG} @inbounds begin z = zero(eltype(factor)) - for q in (front_ptr[s] + li - 1):WG:(front_ptr[s + 1] - 1) + p0 = front_ptr[s] + for q in (p0 + li - 1):WG:(p0 + front_nrows[s] * front_ncols[s] - 1) factor[q] = z end c = cb_ptr[s] @@ -30,12 +33,14 @@ const ASSEMBLY_WORKGROUP = 256 return nothing end -@kernel function _zero_fronts_kernel!(factor, stack, nodes, first, front_ptr, front_nrows, front_ncols, cb_ptr, +@kernel function _zero_fronts_kernel!(factor, stack, nodes, bm, front_ptr, front_nrows, front_ncols, cb_ptr, ::Val{WG}) where {WG} li = @index(Local, Linear) - g = @index(Group, Linear) - @inbounds s = nodes[first + g - 1] - _zero_front!(factor, stack, s, li, front_ptr, front_nrows, front_ncols, cb_ptr, Val(WG)) + G = @index(Group, Linear) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _zero_front!(factor, _mview(stack, k, bm.nbatch), s, li, member_panels(front_ptr, k, bm.nbatch), front_nrows, + front_ncols, cb_ptr, Val(WG)) end """ @@ -48,17 +53,19 @@ Asynchronous. function zero_fronts!(N::Numeric, S::Symbolic, first::Integer, count::Integer) count > 0 || return N WG = ASSEMBLY_WORKGROUP + bm = batch_map(N; first) kernel! = _zero_fronts_kernel!(KernelAbstractions.get_backend(N.factor), WG) - kernel!(N.factor, N.stack, S.group_nodes, Int(first), S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, - Val(WG); ndrange = WG * count) + kernel!(N.factor, N.stack, S.group_nodes, bm, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, + Val(WG); ndrange = WG * count * bm.nact) return N end # value of nzval entry `x` as added at amap offset `off` (negative: conjugate) @inline _amap_value(::Type{T}, x, off) where {T} = off < 0 ? T(conj(x)) : T(x) -# assemble the values of A into the zeroed panel of front `s` (work item `li` of `WG`) -@inline function _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, li, ::Val{WG}) where {WG} +# assemble the values of A into the zeroed panel of front `s` (work item `li` of `WG`); `shift` moves the +# `amap` offsets (single-matrix layout) to the batch member's panel +@inline function _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, shift, li, ::Val{WG}) where {WG} @inbounds begin T = eltype(factor) a = amap_ptr[s] @@ -77,20 +84,26 @@ end v += _amap_value(T, nzval[p], off) kk += 1 end - factor[dest] = v + factor[dest + shift] = v end end end return nothing end -@kernel function _scatter_A_kernel!(factor, nzval, amap, amap_ptr, amap_src, nodes, first, ::Val{WG}) where {WG} +@kernel function _scatter_A_kernel!(factor, nzval, amap, amap_ptr, amap_src, nodes, bm, front_ptr, + ::Val{WG}) where {WG} li = @index(Local, Linear) - g = @index(Group, Linear) - @inbounds s = nodes[first + g - 1] - _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, li, Val(WG)) + G = @index(Group, Linear) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _scatter_front!(factor, _mview(nzval, k, bm.nbatch), amap, amap_ptr, amap_src, s, + _member_shift(front_ptr, s, k, bm.nbatch), li, Val(WG)) end +# shift of the factor offsets of front `s` from the single-matrix layout to batch member `k` +@inline _member_shift(front_ptr, s, k, nb) = @inbounds Int(member_panels(front_ptr, k, nb)[s]) - Int(front_ptr[s]) + """ scatter_A!(numeric, symbolic, nzval, first, count) -> numeric @@ -103,8 +116,10 @@ negative offsets add the conjugate). Asynchronous. function scatter_A!(N::Numeric, S::Symbolic, nzval::AbstractVector, first::Integer, count::Integer) count > 0 || return N WG = ASSEMBLY_WORKGROUP + bm = batch_map(N; first) kernel! = _scatter_A_kernel!(KernelAbstractions.get_backend(N.factor), WG) - kernel!(N.factor, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, Int(first), Val(WG); ndrange = WG * count) + kernel!(N.factor, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, bm, S.front_ptr, Val(WG); + ndrange = WG * count * bm.nact) return N end @@ -147,16 +162,17 @@ end return nothing end -@kernel function _extend_add_kernel!(factor, stack, nodes, first, front_ptr, front_nrows, front_ncols, cb_ptr, +@kernel function _extend_add_kernel!(factor, stack, nodes, bm, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, ::Val{WG}) where {WG} li = @index(Local, Linear) - g = @index(Group, Linear) + G = @index(Group, Linear) # children one after the other (fixed order, a barrier between them): rows # of one child go to distinct destinations, rows of two children may not for k in 1:maxchild - @inbounds s = nodes[first + g - 1] - _extend_add_child!(factor, stack, s, k, li, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, - child_list, relind_ptr, relind, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + mb = _bm_gmember(bm, G) + _extend_add_child!(factor, _mview(stack, mb, bm.nbatch), s, k, li, member_panels(front_ptr, mb, bm.nbatch), + front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, Val(WG)) @synchronize end end @@ -174,9 +190,10 @@ deterministic. Asynchronous. function extend_add!(N::Numeric, S::Symbolic, first::Integer, count::Integer, maxchild::Integer) (count > 0 && maxchild > 0) || return N WG = ASSEMBLY_WORKGROUP + bm = batch_map(N; first) kernel! = _extend_add_kernel!(KernelAbstractions.get_backend(N.factor), WG) - kernel!(N.factor, N.stack, S.group_nodes, Int(first), S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, - S.child_ptr, S.child_list, S.relind_ptr, S.relind, Int(maxchild), Val(WG); ndrange = WG * count) + kernel!(N.factor, N.stack, S.group_nodes, bm, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, + S.child_ptr, S.child_list, S.relind_ptr, S.relind, Int(maxchild), Val(WG); ndrange = WG * count * bm.nact) return N end @@ -191,13 +208,33 @@ end end end +# work item q of block b: the `m×m` block b of `work` (member `members[j0 + b + 1]`) into its member's stack +@kernel function _pack_add_batch_kernel!(stack, c0, work, m, members, j0, nb) + q0 = @index(Global, Linear) + b = (q0 - 1) ÷ (m * m) + q = q0 - b * m * m + @inbounds begin + j = (q - 1) ÷ m + 1 + i = q - (j - 1) * m + if i >= j + k = Int(members[j0 + b + 1]) + v = work[b * m * m + q] + i == j && (v = eltype(stack)(real(v))) # gemm with F21ᴴ: keep the herk real diagonal + _mview(stack, k, nb)[c0 + _packed(i, j, m) - 1] += v + end + end +end + """ pack_add!(stack, c0, work, m) -> stack + pack_add!(stack, c0, work, m, members, j0, count, nbatch) -> stack Add the lower triangle of the column-major `m×m` matrix `work[1:m^2]` to the packed contribution block `stack[c0:(c0 + m(m+1)/2 - 1)]` (the regime-C step -after the vendor `syrk`/`herk` into the workspace). One launch, one work item -per entry. Asynchronous. +after the vendor `syrk`/`herk` into the workspace). Batched form: the `count` +blocks `work[(b-1)m² + 1 : b m²]` go to the stacks of batch members +`members[j0 + b]` (device vector) of an `nbatch`-member update stack. One +launch, one work item per entry. Asynchronous. """ function pack_add!(stack::AbstractVector, c0::Integer, work::AbstractVector, m::Integer) m > 0 || return stack @@ -205,3 +242,11 @@ function pack_add!(stack::AbstractVector, c0::Integer, work::AbstractVector, m:: kernel!(stack, Int(c0), work, Int(m); ndrange = m * m) return stack end + +function pack_add!(stack::AbstractVector, c0::Integer, work::AbstractVector, m::Integer, members::AbstractVector{Int32}, + j0::Integer, count::Integer, nbatch::Integer) + m > 0 && count > 0 || return stack + kernel! = _pack_add_batch_kernel!(KernelAbstractions.get_backend(stack), ASSEMBLY_WORKGROUP) + kernel!(stack, Int(c0), work, Int(m), members, Int(j0), Int(nbatch); ndrange = m * m * count) + return stack +end diff --git a/src/numeric/batch.jl b/src/numeric/batch.jl new file mode 100644 index 0000000..828e655 --- /dev/null +++ b/src/numeric/batch.jl @@ -0,0 +1,139 @@ +# Uniform batch (PLAN §1.6, §3.5): `nbatch` matrices with the pattern of one +# analysis. Every device array of the numeric and solve phases gets a batch +# stride, and every kernel takes the batch index as one more grid dimension, +# flattened into its 1-D ndrange (2-D ndranges allocate on every launch of the +# KA CPU backend): the workgroups (or work items) of a launch over `count` +# fronts become `count × nact`, member slot fastest, for the `nact` active +# members of the batch (`ubatch_index`/`ubatch_mask`). +# +# Storage of member `k` (1-based) of an `nb`-member batch: +# +# * factor panels: interleaved per front, the `nb` panels of supernode `s` one +# after the other, `f×w` each, starting at `nb (panel_ptr[s] - 1) + 1`, so a +# regime-C front of the whole batch is one `f×w×nb` strided batch for the +# batched dense calls ([`member_panels`](@ref)); +# * `info`: interleaved, entry `s` of member `k` at `(s - 1) nb + k`, so the +# statuses of one front are contiguous (batched `potrf`); +# * everything else (`d`, `piv`, `pivot_kind`, the update stack, the regime-C +# workspace per launch, `stats`, `totals`, `aux`, the values `nzval`, the +# right-hand sides): member after member (`(k - 1) len + 1 : k len`). +# +# With `nb = 1` both layouts are the single-matrix layout. + +""" + BatchMap{V <: Union{Nothing, AbstractVector{Int32}}} + +Batch part of a kernel launch (an argument of every numeric and solve +kernel): `members[1:nact]` are the active batch members (1-based, increasing; +`members = nothing`: a single matrix, member 1), +`nbatch` the batch size, `first` the launch's first entry in its node list +(the former `first` argument of the kernels) and `nrhs` the right-hand sides +per member of a solve (1 in the numeric phase). Adapted to the device with its +`members` vector. +""" +struct BatchMap{V <: Union{Nothing, AbstractVector{Int32}}} + members::V + nact::Int + nbatch::Int + first::Int + nrhs::Int +end + +Adapt.@adapt_structure BatchMap + +BatchMap(members::AbstractVector{Int32}, nact::Integer, nbatch::Integer; first::Integer = 1, nrhs::Integer = 1) = + BatchMap{typeof(members)}(members, Int(nact), Int(nbatch), Int(first), Int(nrhs)) + +BatchMap(::Nothing, nact::Integer, nbatch::Integer; first::Integer = 1, nrhs::Integer = 1) = + BatchMap{Nothing}(nothing, Int(nact), Int(nbatch), Int(first), Int(nrhs)) + +# a single matrix (launches without a `Numeric`: tests and benchmarks of the kernels) +single_batch(; first::Integer = 1, nrhs::Integer = 1) = BatchMap(nothing, 1, 1; first, nrhs) + +# workgroup `G` of a launch over `nact` members: (node index g, member slot j), member slot fastest +@inline _bm_node(bm::BatchMap, G) = (G - 1) ÷ bm.nact + 1 +@inline _bm_slot(bm::BatchMap, G) = (G - 1) % bm.nact + 1 +@inline _bm_member(bm::BatchMap, j) = @inbounds Int(bm.members[j]) +@inline _bm_member(::BatchMap{Nothing}, j) = 1 +# member of workgroup `G` +@inline _bm_gmember(bm::BatchMap, G) = _bm_member(bm, _bm_slot(bm, G)) +# member of the compact solve column `r` (columns of member slot j: (j - 1) nrhs + 1 : j nrhs) +@inline _bm_cmember(bm::BatchMap, r) = _bm_member(bm, (r - 1) ÷ bm.nrhs + 1) +# column of the user's right-hand side array (all `nbatch` members) holding the compact column `r` +@inline _bm_ucol(bm::BatchMap, r) = (_bm_cmember(bm, r) - 1) * bm.nrhs + (r - 1) % bm.nrhs + 1 + +# member-major view of member `k` of the `nb` members of `A` (`length(A) ÷ nb` entries each) +@inline function _mview(A, k, nb) + len = length(A) ÷ nb + return @inbounds view(A, ((k - 1) * len + 1):(k * len)) +end + +# interleaved view of member `k`: entry `i` at `(i - 1) nb + k` +@inline _iview(A, k, nb) = @inbounds view(A, k:nb:length(A)) + +""" + MemberPanels{V} + +The panel pointers of member `k` of an `nb`-member batch in the interleaved +factor layout: `P[s] = nb (front_ptr[s] - 1) + (k - 1)(front_ptr[s+1] - front_ptr[s]) + 1` +(the first entry of the `f×w` panel of supernode `s` of member `k`). Passed +to the kernels in place of `front_ptr` for the factor offsets; `P.ptr` is the +analysis's `front_ptr`. Note that `P[s + 1] - P[s]` is not a panel length. +""" +struct MemberPanels{V} + ptr::V + k::Int + nb::Int +end + +@inline function Base.getindex(P::MemberPanels, s::Integer) + @inbounds a = P.ptr[s] + @inbounds b = P.ptr[s + 1] + return oftype(a, P.nb * (a - 1) + (P.k - 1) * (b - a) + 1) +end + +@inline member_panels(front_ptr, k, nb) = MemberPanels(front_ptr, Int(k), Int(nb)) + +""" + panel_offset(panel_ptr, s, k, nb) -> Int + +First factor index of the panel of supernode `s` of batch member `k` in the +interleaved layout (host side, `panel_ptr` of the `Layout`). +""" +panel_offset(panel_ptr::AbstractVector, s::Integer, k::Integer, nb::Integer) = + nb * (panel_ptr[s] - 1) + (k - 1) * (panel_ptr[s + 1] - panel_ptr[s]) + 1 + +""" + active_members(nbatch, ubatch_index, ubatch_mask) -> Vector{Int} + +The batch members a phase processes (1-based, increasing): all of them for +`ubatch_index = -1` and no mask, member `ubatch_index + 1` (`ubatch_index` is +0-based, as in cuDSS), or the members whose `ubatch_mask` flag is 1 (the mask +has `nbatch` entries; with both set, their intersection). Raises +[`InvalidValueError`](@ref) for an index or mask that does not fit the batch. +""" +function active_members(nb::Integer, index::Integer, mask) + (index == -1 || 0 <= index < nb) || + throw(InvalidValueError("ubatch_index = $index is outside the batch (0:$(nb - 1), or -1 for all members)")) + if mask !== nothing + length(mask) == nb || + throw(InvalidValueError("ubatch_mask has $(length(mask)) entries, the batch has $nb members")) + end + members = Int[] + for k in 1:nb + index == -1 || k == index + 1 || continue + mask === nothing || mask[k] != 0 || continue + push!(members, k) + end + return members +end + +# runs of consecutive members: the start positions in `members` (and length(members) + 1 at the end) +function member_runs!(starts::Vector{Int}, members::AbstractVector{<:Integer}, nact::Integer) + empty!(starts) + for j in 1:nact + (j == 1 || members[j] != members[j - 1] + 1) && push!(starts, j) + end + push!(starts, nact + 1) + return starts +end diff --git a/src/numeric/extract.jl b/src/numeric/extract.jl index e3c1db6..1fb3cbe 100644 --- a/src/numeric/extract.jl +++ b/src/numeric/extract.jl @@ -15,17 +15,49 @@ function host_numeric(N::Numeric{T}) where {T} stats, info, piv, pivot_kind = Array(N.stats), Array(N.info), Array(N.piv), Array(N.pivot_kind) return Numeric{T, Vector{T}, Vector{Int64}, Vector{Int32}, Vector{Int8}}(factor, d, stack, work, stats, info, piv, pivot_kind, Array(N.totals), - Array(N.psign), Array(N.aux), N.plan) + Array(N.psign), Array(N.aux), N.plan, + Array(N.members), N.nbatch) end """ - panel(symbolic, numeric, s) -> Matrix + panel(symbolic, numeric, s, k = 1) -> Matrix Host copy of the `f×w` panel of supernode `s` (rows `snrows(s)`, the factor -columns of `s`; the strict upper triangle of its first `w` rows is zero). +columns of `s`; the strict upper triangle of its first `w` rows is zero) of +batch member `k`. """ -function panel(S::Symbolic, N::Numeric{T}, s::Integer) where {T} +function panel(S::Symbolic, N::Numeric{T}, s::Integer, k::Integer = 1) where {T} L, sc = S.layout, S.schedule f, w = sc.rows[s], sc.width[s] - return reshape(Array(view(N.factor, L.panel_ptr[s]:(L.panel_ptr[s + 1] - 1))), f, w) + p0 = panel_offset(L.panel_ptr, s, k, N.nbatch) + return reshape(Array(view(N.factor, p0:(p0 + f * w - 1))), f, w) +end + +""" + member_numeric(numeric, symbolic, k) -> Numeric{T, Vector{T}} + +Host copy of batch member `k` of `numeric` as a single-matrix storage (the +layout of `allocate_numeric(symbolic, T)`), for the reference comparisons and +[`extract_L`](@ref). +""" +function member_numeric(N::Numeric{T}, S::Symbolic, k::Integer) where {T} + nb = N.nbatch + 1 <= k <= nb || throw(InvalidValueError("batch member $k outside 1:$nb")) + H = host_numeric(N) + nb == 1 && return H + L = S.layout + ns = nsupernodes(S) + factor = zeros(T, L.factor_len) + for s in 1:ns + len = L.panel_ptr[s + 1] - L.panel_ptr[s] + p0 = panel_offset(L.panel_ptr, s, k, nb) + copyto!(factor, L.panel_ptr[s], H.factor, p0, len) + end + mv(A) = A[((k - 1) * (length(A) ÷ nb) + 1):(k * (length(A) ÷ nb))] + info = H.info[k:nb:end] + return Numeric{T, Vector{T}, Vector{Int64}, Vector{Int32}, Vector{Int8}}(factor, mv(H.d), mv(H.stack), mv(H.work), + mv(H.stats), info, mv(H.piv), + mv(H.pivot_kind), mv(H.totals), + copy(H.psign), mv(H.aux), + NumericPlan(S), Int32[1], 1) end diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index 58585b6..becce03 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -10,11 +10,14 @@ "Workgroup size of the per-phase reduction of the front statistics." const STATS_WORKGROUP = 256 -@kernel function _cholesky_stats_kernel!(stats, info, super_ptr, front_ncols, ns, ::Val{WG}, ::Val{NF}, +@kernel function _cholesky_stats_kernel!(stats_all, info_all, super_ptr, front_ncols, ns, bm, ::Val{WG}, ::Val{NF}, ::Val{LOG2WG}) where {WG, NF, LOG2WG} li = @index(Local, Linear) + G = @index(Group, Linear) best = @localmem Int64 (WG,) @inbounds begin + stats = _mview(stats_all, _bm_gmember(bm, G), bm.nbatch) + info = _iview(info_all, _bm_gmember(bm, G), bm.nbatch) m = typemax(Int64) for s in li:WG:ns fi = Int64(info[s]) @@ -39,7 +42,8 @@ const STATS_WORKGROUP = 256 @synchronize end if li == 1 - @inbounds info[ns + 1] = best[1] == typemax(Int64) ? Int32(0) : Int32(best[1]) + @inbounds _iview(info_all, _bm_gmember(bm, G), bm.nbatch)[ns + 1] = + best[1] == typemax(Int64) ? Int32(0) : Int32(best[1]) end end @@ -49,14 +53,15 @@ end Fill the per-front statistics of a Cholesky factorization from the per-front `potrf` status `numeric.info[1:ns]` (`npos` = `w`, or the failed local column minus one; `info` = the failed local column) and reduce the smallest failed -factor column into `numeric.info[ns + 1]` (0 = none). One workgroup, no -atomics. Asynchronous. +factor column into `numeric.info[ns + 1]` (0 = none), for every active batch +member. One workgroup per member, no atomics. Asynchronous. """ function cholesky_stats!(N::Numeric, S::Symbolic) WG = STATS_WORKGROUP + bm = batch_map(N) kernel! = _cholesky_stats_kernel!(KernelAbstractions.get_backend(N.factor), WG) - kernel!(N.stats, N.info, S.super_ptr, S.front_ncols, nsupernodes(S), Val(WG), Val(FRONT_STATS_FIELDS), - Val(_ilog2(WG)); ndrange = WG) + kernel!(N.stats, N.info, S.super_ptr, S.front_ncols, nsupernodes(S), bm, Val(WG), Val(FRONT_STATS_FIELDS), + Val(_ilog2(WG)); ndrange = WG * bm.nact) return N end @@ -65,12 +70,14 @@ function _check_numeric(N::Numeric{T}, S::Symbolic, nzval::AbstractVector) where sizeof(T) <= S.elsize || throw(InvalidValueError("the analysis was built for $(S.elsize)-byte elements, the numeric storage has " * "$(sizeof(T))-byte $T")) - length(nzval) == S.nnz || - throw(InvalidValueError("nzval has $(length(nzval)) entries, the analysis expects $(S.nnz)")) - length(N.factor) == S.layout.factor_len && length(N.stack) == S.layout.stack_len && - length(N.info) == nsupernodes(S) + 1 || + nb = N.nbatch + length(nzval) == nb * S.nnz || + throw(InvalidValueError("nzval has $(length(nzval)) entries, the analysis expects $(S.nnz)" * + (nb > 1 ? " per member for $nb batch members" : ""))) + length(N.factor) == nb * S.layout.factor_len && length(N.stack) == nb * S.layout.stack_len && + length(N.info) == nb * (nsupernodes(S) + 1) || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) - length(N.work) == S.layout.work_len && length(N.piv) == S.n && length(N.psign) == S.n || + length(N.work) == nb * S.layout.work_len && length(N.piv) == nb * S.n && length(N.psign) == S.n || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) backend = KernelAbstractions.get_backend(N.factor) for x in (nzval, S.amap, N.info) @@ -84,16 +91,67 @@ end # the dense implementations of the regime-C front kernels, resolved once per phase (backend and `T` are # fixed), so the per-front calls do no capability lookup (its lock and closure allocate); # `factorization_alg = "algo1"` (no vendor calls) turns `:auto` into the KA kernels +# uniform batch: the strided-batched ops (`b*`), `:ka` or `:vendor` from the capabilities, or `:none` (no +# batched path: `impl = :generic`, then every member takes the single-matrix path) function _front_impls(N::Numeric{T}, S::Symbolic, impl::Symbol) where {T} impl === :auto && !S.schedule.vendor_c && (impl = :ka) + batched = N.nbatch > 1 && impl !== :generic + bimpl(op) = batched ? (impl === :ka ? :ka : select_impl(op, N.factor, :auto)) : :none return (potrf = select_impl(:potrf, N.factor, impl), trsm = select_impl(:trsm, N.factor, impl), - herk = select_impl(T <: Real ? :syrk : :herk, N.factor, impl)) + herk = select_impl(T <: Real ? :syrk : :herk, N.factor, impl), bpotrf = bimpl(:potrf_batched), + btrsm = bimpl(:trsm_strided_batched), bgemm = bimpl(:gemm_strided_batched)) end -# regime-C dense kernels of front `s` (assembled panel + contribution block), `p` from `_front_impls` +# regime-C dense kernels of front `s` (assembled panel + contribution block) of every active batch member, +# `p` from `_front_impls` function _factor_front_c!(N::Numeric{T}, S::Symbolic, s::Int, p::NamedTuple) where {T} L, sc = S.layout, S.schedule - _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], L.cb_ptr[s], p) + if N.nbatch == 1 + _factor_panel_c!(N.factor, N.stack, N.work, N.info, s, L.panel_ptr[s], sc.rows[s], sc.width[s], L.cb_ptr[s], + p) + return nothing + end + plan = N.plan + nb = N.nbatch + runs = plan.runs + for r in 1:(length(runs) - 1) + j0, j1 = runs[r], runs[r + 1] - 1 + if p.bpotrf === :none + for j in j0:j1 + k = Int(plan.members_host[j]) + _factor_panel_c!(N.factor, _mview(N.stack, k, nb), N.work, N.info, (s - 1) * nb + k, + panel_offset(L.panel_ptr, s, k, nb), sc.rows[s], sc.width[s], L.cb_ptr[s], p) + end + else + _factor_panels_c_batched!(N, S, s, j0, j1 - j0 + 1, p) + end + end + return nothing +end + +# regime C of front `s` for the `count` consecutive members `members_host[j0:(j0 + count - 1)]`: their panels +# are one `f×w×count` strided batch (interleaved layout), factored by strided-batched `potrf`, `trsm` and `gemm` +# (`F21 F21ᴴ` into `count` blocks of the workspace), then packed into the members' contribution blocks +function _factor_panels_c_batched!(N::Numeric{T}, S::Symbolic, s::Int, j0::Int, count::Int, + p::NamedTuple) where {T} + L, sc = S.layout, S.schedule + nb = N.nbatch + f, w = sc.rows[s], sc.width[s] + m = f - w + k0 = Int(N.plan.members_host[j0]) + off = panel_offset(L.panel_ptr, s, k0, nb) - 1 + P3 = reshape(view(N.factor, (off + 1):(off + f * w * count)), f, w, count) + F11 = view(P3, 1:w, 1:w, :) + ioff = (s - 1) * nb + k0 - 1 # info of front s, member k0 (interleaved) + _potrf_batched_impl!(p.bpotrf, F11, N.info, ioff) + m > 0 || return nothing + F21 = view(P3, (w + 1):f, 1:w, :) + trsm_strided_batched!('R', 'L', 'C', 'N', one(T), F11, F21; impl = p.btrsm) + c0 = L.cb_ptr[s] + c0 > 0 || return nothing # no parent: nothing to update + C = reshape(view(N.work, 1:(m * m * count)), m, m, count) + gemm_strided_batched!(C, F21, F21, -one(T), zero(T); transB = 'C', impl = p.bgemm) + pack_add!(N.stack, c0, N.work, m, N.members, j0 - 1, count, nb) return nothing end @@ -180,9 +238,35 @@ function factorize!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; impl::Sym end end cholesky_stats!(N, S) - copyto!(plan.info_host, 1, N.info, nsupernodes(S) + 1, 1) # the phase's only host synchronization - k = Int(plan.info_host[1]) - return k == 0 ? 0 : S.partition.perm[k] + nb = N.nbatch + copyto!(plan.info_host, 1, N.info, nsupernodes(S) * nb + 1, nb) # the phase's only host synchronization + nb == 1 && return _original_column(S, plan.info_host[1]) + for j in 1:plan.nact[] + k = Int(plan.members_host[j]) + r = _original_column(S, plan.info_host[k]) + r == 0 || return r + end + return 0 +end + +_original_column(S::Symbolic, k::Integer) = k == 0 ? 0 : S.partition.perm[k] + +""" + member_info!(info, numeric, symbolic) -> info + +Write the `info` of the last [`factorize!`](@ref) of every active batch member +`k` into `info[k]` (the original column of the member's first non-positive +pivot, 0 = success; inactive members keep their entries). Host data only: the +statuses were read by `factorize!` (Cholesky) or are 0 (LDLᵀ/LDLᴴ). +""" +function member_info!(info::AbstractVector{<:Integer}, N::Numeric, S::Symbolic) + plan = N.plan + chol = !_is_ldlt_structure(S.structure) + for j in 1:plan.nact[] + k = Int(plan.members_host[j]) + info[k] = chol ? _original_column(S, plan.info_host[k]) : 0 + end + return info end factorize!(N::Numeric, S::Symbolic, A::CSR; impl::Symbol = :auto, opts::Options = Options()) = diff --git a/src/numeric/front.jl b/src/numeric/front.jl index 6b7d99d..9652732 100644 --- a/src/numeric/front.jl +++ b/src/numeric/front.jl @@ -134,13 +134,14 @@ end end """ - front_cholesky_kernel!(backend, WG)(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, first, + front_cholesky_kernel!(backend, WG)(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, bm, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, Val(ASM), Val(W), Val(NL), Val(WG); - ndrange = WG * count) + ndrange = WG * count * bm.nact) -Fused regime-B kernel: workgroup `g` takes front `s = nodes[first + g - 1]` -(width `w ≤ W`), and, when `ASM`, zeroes and assembles it (A through the +Fused regime-B kernel: workgroup `G` takes front `s = nodes[bm.first + g - 1]` +(width `w ≤ W`) of batch member `k` (`g`, `k` from the [`BatchMap`](@ref) `bm`), +and, when `ASM`, zeroes and assembles it (A through the `amap`, then the `maxchild` children's contribution blocks in `child_list` order), then factors `F11 = L11 L11ᴴ` in `@localmem` (`NL = W(W+1)/2` entries), solves `F21 ← F21 L11⁻ᴴ`, updates its contribution block on the @@ -148,59 +149,70 @@ update stack (`cb_ptr[s] > 0`, packed lower triangle) and sets `info[s]` (0, or the first non-positive pivot; the front is then left partially factored). Without `ASM` the panels and contribution blocks must already be assembled. """ -@kernel function front_cholesky_kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, first, +@kernel function front_cholesky_kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, bm, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, ::Val{ASM}, ::Val{W}, ::Val{NL}, ::Val{WG}) where {ASM, W, NL, WG} @uniform TT = eltype(factor) @uniform RT = real(eltype(factor)) li = @index(Local, Linear) - g = @index(Group, Linear) + G = @index(Group, Linear) L11 = @localmem TT (NL,) st = @localmem Int32 (1,) piv = @localmem RT (1,) if ASM - @inbounds s = nodes[first + g - 1] - _zero_front!(factor, stack, s, li, front_ptr, front_nrows, front_ncols, cb_ptr, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _zero_front!(factor, _mview(stack, k, bm.nbatch), s, li, member_panels(front_ptr, k, bm.nbatch), front_nrows, + front_ncols, cb_ptr, Val(WG)) end @synchronize if ASM - @inbounds s = nodes[first + g - 1] - _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, li, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _scatter_front!(factor, _mview(nzval, k, bm.nbatch), amap, amap_ptr, amap_src, s, + _member_shift(front_ptr, s, k, bm.nbatch), li, Val(WG)) end @synchronize - for k in 1:maxchild - @inbounds s = nodes[first + g - 1] - _extend_add_child!(factor, stack, s, k, li, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, - child_list, relind_ptr, relind, Val(WG)) + for kc in 1:maxchild + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _extend_add_child!(factor, _mview(stack, k, bm.nbatch), s, kc, li, member_panels(front_ptr, k, bm.nbatch), + front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, Val(WG)) @synchronize end - @inbounds s = nodes[first + g - 1] - _front_load!(L11, st, factor, s, li, front_ptr, front_nrows, front_ncols, Val(W), Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _front_load!(L11, st, factor, s, li, member_panels(front_ptr, k, bm.nbatch), front_nrows, front_ncols, Val(W), + Val(WG)) @synchronize for j in 1:W - @inbounds s = nodes[first + g - 1] + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] _front_chol_update!(L11, st, piv, j, s, li, front_ncols, Val(W), Val(WG)) @synchronize - @inbounds s = nodes[first + g - 1] + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] _front_chol_scale!(L11, st, piv, j, s, li, front_ncols, Val(W), Val(WG)) @synchronize end - @inbounds s = nodes[first + g - 1] - _front_trsm!(factor, info, L11, st, s, li, front_ptr, front_nrows, front_ncols, Val(W), Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _front_trsm!(factor, _iview(info, k, bm.nbatch), L11, st, s, li, member_panels(front_ptr, k, bm.nbatch), + front_nrows, front_ncols, Val(W), Val(WG)) @synchronize - @inbounds s = nodes[first + g - 1] - _front_syrk!(factor, stack, st, s, li, front_ptr, front_nrows, front_ncols, cb_ptr, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _front_syrk!(factor, _mview(stack, k, bm.nbatch), st, s, li, member_panels(front_ptr, k, bm.nbatch), front_nrows, + front_ncols, cb_ptr, Val(WG)) end -function _launch_front_cholesky!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, first, count, +function _launch_front_cholesky!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, bm::BatchMap, count, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, asm::Val, ::Val{W}, ::Val{WG}) where {W, WG} _ilog2(WG) kernel! = front_cholesky_kernel!(KernelAbstractions.get_backend(factor), WG) - kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, Int(first), front_ptr, front_nrows, + kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, nodes, bm, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, Int(maxchild), asm, Val(W), - Val(W * (W + 1) ÷ 2), Val(WG); ndrange = WG * count) + Val(W * (W + 1) ÷ 2), Val(WG); ndrange = WG * count * bm.nact) return nothing end @@ -218,16 +230,16 @@ end Regime-B launch: assemble and factor the `count` fronts `symbolic.group_nodes[first:(first + count - 1)]` (widths `≤ W`, -`W ∈ $(REGIME_B_WIDTHS)`, at most `maxchild` children each) with one -[`front_cholesky_kernel!`](@ref) launch, one workgroup per front; `info[s]` -receives each front's status. Asynchronous. +`W ∈ $(REGIME_B_WIDTHS)`, at most `maxchild` children each) of every active +batch member with one [`front_cholesky_kernel!`](@ref) launch, one workgroup +per (front, member); `info[s]` receives each front's status. Asynchronous. """ function factorize_fronts_b!(N::Numeric, S::Symbolic, nzval::AbstractVector, first::Integer, count::Integer, maxchild::Integer, W::Integer) count > 0 || return N _with_width_class(Int(W)) do w _launch_front_cholesky!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, - first, count, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.child_ptr, + batch_map(N; first), count, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.child_ptr, S.child_list, S.relind_ptr, S.relind, maxchild, Val(true), w, Val(FRONT_WORKGROUP)) end @@ -252,7 +264,8 @@ function front_cholesky!(factor::AbstractVector, stack::AbstractVector, info::Ab width::Integer = REGIME_B_MAX_WIDTH, workgroup::Val = Val(FRONT_WORKGROUP)) count > 0 || return info _with_width_class(Int(width)) do w - _launch_front_cholesky!(factor, stack, info, factor, nodes, nodes, nodes, nodes, first, count, front_ptr, + _launch_front_cholesky!(factor, stack, info, factor, nodes, nodes, nodes, nodes, single_batch(; first), count, + front_ptr, front_nrows, front_ncols, cb_ptr, nodes, nodes, nodes, nodes, 0, Val(false), w, workgroup) end diff --git a/src/numeric/ldlt.jl b/src/numeric/ldlt.jl index 32e364c..5640e4e 100644 --- a/src/numeric/ldlt.jl +++ b/src/numeric/ldlt.jl @@ -512,12 +512,13 @@ end """ front_ldlt_kernel!(backend, WG)(factor, stack, info, stats, d, piv, pivot_kind, psign, perm, aux, nzval, amap, - amap_ptr, amap_src, nodes, first, super_ptr, front_ptr, front_nrows, + amap_ptr, amap_src, nodes, bm, super_ptr, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, prm, - Val(WG); ndrange = WG * count) + Val(WG); ndrange = WG * count * bm.nact) -LDLᵀ/LDLᴴ kernel of regimes B and C: workgroup `g` takes front -`s = nodes[first + g - 1]`, zeroes and assembles it (A through the `amap`, then +LDLᵀ/LDLᴴ kernel of regimes B and C: workgroup `G` takes front +`s = nodes[bm.first + g - 1]` of batch member `k` (`g`, `k` from the +[`BatchMap`](@ref) `bm`; the per-member arrays are the member's), zeroes and assembles it (A through the `amap`, then the `maxchild` children's packed contribution blocks in `child_list` order), factors its `w` fully-summed columns in place in the panel with in-block Bunch–Kaufman pivoting, threshold acceptance and perturbation (`prm`, a @@ -527,42 +528,52 @@ statistics, updates its packed contribution block on the update stack (`cb_ptr[s] > 0`) with `F₂₂ − (L₂₁ D) L₂₁ᴴ` and leaves a unit-lower panel. """ @kernel function front_ldlt_kernel!(factor, stack, info, stats, d, piv, pivot_kind, psign, perm, aux, nzval, amap, - amap_ptr, amap_src, nodes, first, super_ptr, front_ptr, front_nrows, front_ncols, + amap_ptr, amap_src, nodes, bm, super_ptr, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, prm::_LDLTDevice{R, HERM}, ::Val{WG}) where {R, HERM, WG} @uniform TT = eltype(factor) @uniform IT = eltype(front_ptr) li = @index(Local, Linear) - g = @index(Group, Linear) + G = @index(Group, Linear) ctl = @localmem IT (_LT_CTL,) pv = @localmem TT (4,) if li == 1 - @inbounds s = nodes[first + g - 1] - _lt_front_setup!(ctl, s, super_ptr, front_ptr, front_nrows, front_ncols) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + _lt_front_setup!(ctl, s, super_ptr, member_panels(front_ptr, _bm_gmember(bm, G), bm.nbatch), front_nrows, + front_ncols) end @synchronize - @inbounds s = nodes[first + g - 1] - _zero_front!(factor, stack, s, li, front_ptr, front_nrows, front_ncols, cb_ptr, Val(WG)) - _lt_init_piv!(piv, ctl, li, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + nb = bm.nbatch + _zero_front!(factor, _mview(stack, k, nb), s, li, member_panels(front_ptr, k, nb), front_nrows, front_ncols, + cb_ptr, Val(WG)) + _lt_init_piv!(_mview(piv, k, nb), ctl, li, Val(WG)) @synchronize - @inbounds s = nodes[first + g - 1] - _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, li, Val(WG)) + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _scatter_front!(factor, _mview(nzval, k, bm.nbatch), amap, amap_ptr, amap_src, s, + _member_shift(front_ptr, s, k, bm.nbatch), li, Val(WG)) @synchronize - for k in 1:maxchild - @inbounds s = nodes[first + g - 1] - _extend_add_child!(factor, stack, s, k, li, front_ptr, front_nrows, front_ncols, cb_ptr, child_ptr, - child_list, relind_ptr, relind, Val(WG)) + for kc in 1:maxchild + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _extend_add_child!(factor, _mview(stack, k, bm.nbatch), s, kc, li, member_panels(front_ptr, k, bm.nbatch), + front_nrows, front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, Val(WG)) @synchronize end for it in 1:ctl[_ST_W] if li == 1 - _lt_choose!(factor, ctl, aux, prm, Val(false), Val(HERM)) + _lt_choose!(factor, ctl, _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, Val(false), Val(HERM)) end @synchronize - _lt_swap!(factor, ctl, piv, li, Val(WG), Val(false), Val(HERM)) + _lt_swap!(factor, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(false), Val(HERM)) @synchronize if li == 1 - _lt_pivot!(factor, ctl, pv, d, pivot_kind, piv, psign, perm, aux, prm, Val(false), Val(HERM)) + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_pivot!(factor, ctl, pv, _mview(d, k, nb), _mview(pivot_kind, k, nb), _mview(piv, k, nb), psign, perm, + _mview(aux, k, nb), prm, Val(false), Val(HERM)) end @synchronize _lt_update!(factor, ctl, pv, li, Val(WG), Val(false), Val(HERM)) @@ -570,10 +581,13 @@ statistics, updates its packed contribution block on the update stack _lt_scale!(factor, ctl, pv, li, Val(WG), Val(false)) @synchronize end - @inbounds s = nodes[first + g - 1] + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] @inbounds coff = Int(cb_ptr[s]) - 1 - _lt_cb_update!(factor, ctl, stack, coff, d, pivot_kind, li, Val(WG), Val(false), Val(HERM)) - _lt_front_finish!(factor, stats, info, ctl, li, Val(WG)) + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_cb_update!(factor, ctl, _mview(stack, k, nb), coff, _mview(d, k, nb), _mview(pivot_kind, k, nb), li, Val(WG), + Val(false), Val(HERM)) + _lt_front_finish!(factor, _mview(stats, k, nb), _iview(info, k, nb), ctl, li, Val(WG)) end # --------------------------------------------------------------------------- @@ -612,12 +626,13 @@ end """ subtree_ldlt_kernel!(backend, WG)(factor, stack, info, stats, d, piv, pivot_kind, psign, perm, aux, nzval, amap, - amap_ptr, amap_src, trees, first, subtree_ptr, subtree_nodes, super_ptr, + amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, super_ptr, front_ptr, front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, - child_list, relind_ptr, relind, prm, Val(NE), Val(WG); ndrange = WG * count) + child_list, relind_ptr, relind, prm, Val(NE), Val(WG); + ndrange = WG * count * bm.nact) -LDLᵀ/LDLᴴ kernel of regime A: workgroup `g` takes subtree -`t = trees[first + g - 1]` and processes its supernodes in order with the +LDLᵀ/LDLᴴ kernel of regime A: workgroup `G` takes subtree +`t = trees[bm.first + g - 1]` of batch member `k` and processes its supernodes in order with the serial stack of packed fronts and contribution blocks in a `@localmem` buffer of `NE` entries (as [`subtree_cholesky_kernel!`](@ref)): zero, scatter A, extend-add the children, factor the front's first `w` columns with the pivoted @@ -627,19 +642,19 @@ statistics, then move the block down or (subtree root) write it to the update stack. """ @kernel function subtree_ldlt_kernel!(factor, stack, info, stats, d, piv, pivot_kind, psign, perm, aux, nzval, amap, - amap_ptr, amap_src, trees, first, subtree_ptr, subtree_nodes, super_ptr, + amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, super_ptr, front_ptr, front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, child_list, relind_ptr, relind, prm::_LDLTDevice{R, HERM}, ::Val{NE}, ::Val{WG}) where {R, HERM, NE, WG} @uniform TT = eltype(factor) @uniform IT = eltype(subtree_nodes) li = @index(Local, Linear) - g = @index(Group, Linear) + G = @index(Group, Linear) buf = @localmem TT (NE,) ctl = @localmem IT (_LT_CTL,) pv = @localmem TT (4,) if li == 1 - @inbounds t = trees[first + g - 1] + @inbounds t = trees[bm.first + _bm_node(bm, G) - 1] @inbounds ctl[_ST_FIRST] = subtree_ptr[t] @inbounds ctl[_ST_COUNT] = subtree_ptr[t + 1] - subtree_ptr[t] end @@ -651,9 +666,10 @@ stack. end @synchronize _subtree_zero!(buf, ctl, li, Val(WG)) - _lt_init_piv!(piv, ctl, li, Val(WG)) + _lt_init_piv!(_mview(piv, _bm_gmember(bm, G), bm.nbatch), ctl, li, Val(WG)) @synchronize - _subtree_scatter!(buf, ctl, nzval, amap, amap_ptr, amap_src, front_ptr, li, Val(WG)) + _subtree_scatter!(buf, ctl, _mview(nzval, _bm_gmember(bm, G), bm.nbatch), amap, amap_ptr, amap_src, front_ptr, + li, Val(WG)) @synchronize for kc in 1:ctl[_ST_NCHILD] _subtree_extend_add!(buf, ctl, kc, child_ptr, child_list, front_nrows, front_ncols, local_cb, relind_ptr, @@ -662,13 +678,16 @@ stack. end for it in 1:ctl[_ST_W] if li == 1 - _lt_choose!(buf, ctl, aux, prm, Val(true), Val(HERM)) + _lt_choose!(buf, ctl, _mview(aux, _bm_gmember(bm, G), bm.nbatch), prm, Val(true), Val(HERM)) end @synchronize - _lt_swap!(buf, ctl, piv, li, Val(WG), Val(true), Val(HERM)) + _lt_swap!(buf, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), li, Val(WG), Val(true), Val(HERM)) @synchronize if li == 1 - _lt_pivot!(buf, ctl, pv, d, pivot_kind, piv, psign, perm, aux, prm, Val(true), Val(HERM)) + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_pivot!(buf, ctl, pv, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), _mview(piv, mb, nb), psign, + perm, _mview(aux, mb, nb), prm, Val(true), Val(HERM)) end @synchronize _lt_update!(buf, ctl, pv, li, Val(WG), Val(true), Val(HERM)) @@ -676,9 +695,14 @@ stack. _lt_scale!(buf, ctl, pv, li, Val(WG), Val(true)) @synchronize end - _lt_cb_update!(buf, ctl, buf, _lt_local_cb_offset(ctl), d, pivot_kind, li, Val(WG), Val(true), Val(HERM)) + mb = _bm_gmember(bm, G) + _lt_cb_update!(buf, ctl, buf, _lt_local_cb_offset(ctl), _mview(d, mb, bm.nbatch), + _mview(pivot_kind, mb, bm.nbatch), li, Val(WG), Val(true), Val(HERM)) @synchronize - _lt_subtree_write!(factor, stack, info, stats, buf, ctl, front_ptr, cb_ptr, li, Val(WG)) + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lt_subtree_write!(factor, _mview(stack, mb, nb), _iview(info, mb, nb), _mview(stats, mb, nb), buf, ctl, + member_panels(front_ptr, mb, nb), cb_ptr, li, Val(WG)) @synchronize for r in 1:ctl[_ST_ROUNDS] _subtree_move!(buf, ctl, r, li, Val(WG)) @@ -690,14 +714,16 @@ end # --------------------------------------------------------------------------- # reductions -@kernel function _abs_max_kernel!(aux, nzval, nnz, ::Val{WG}, ::Val{LOG2WG}) where {WG, LOG2WG} +@kernel function _abs_max_kernel!(aux, nzval, nnz, bm, ::Val{WG}, ::Val{LOG2WG}) where {WG, LOG2WG} @uniform RT = real(eltype(aux)) li = @index(Local, Linear) + G = @index(Group, Linear) best = @localmem RT (WG,) @inbounds begin + nz = _mview(nzval, _bm_gmember(bm, G), bm.nbatch) m = zero(RT) for p in li:WG:nnz - m = max(m, RT(abs(nzval[p]))) + m = max(m, RT(abs(nz[p]))) end best[li] = m end @@ -712,33 +738,40 @@ end @synchronize end if li == 1 - @inbounds aux[1] = eltype(aux)(best[1]) + @inbounds aux[_bm_gmember(bm, G)] = eltype(aux)(best[1]) end end """ - abs_max!(aux, nzval) -> aux + abs_max!(aux, nzval[, bm]) -> aux -`aux[1] = max |nzval[p]|` (one workgroup, `@localmem` tree reduction; the -scale of `pivot_epsilon_alg = "algo1"`). Asynchronous. +`aux[k] = max |nzval of member k|` for the active members `k` of the +[`BatchMap`](@ref) `bm` (default: a single matrix, `aux[1]`), one workgroup per +member, `@localmem` tree reduction; the scale of `pivot_epsilon_alg = "algo1"`. +Asynchronous. """ -function abs_max!(aux::AbstractVector, nzval::AbstractVector) +function abs_max!(aux::AbstractVector, nzval::AbstractVector, bm::BatchMap = single_batch()) WG = STATS_WORKGROUP - _abs_max_kernel!(KernelAbstractions.get_backend(aux), WG)(aux, nzval, length(nzval), Val(WG), Val(_ilog2(WG)); - ndrange = WG) + _abs_max_kernel!(KernelAbstractions.get_backend(aux), WG)(aux, nzval, length(nzval) ÷ bm.nbatch, bm, Val(WG), + Val(_ilog2(WG)); ndrange = WG * bm.nact) return aux end -@kernel function _reduce_stats_kernel!(totals, stats, ns, ::Val{WG}, ::Val{NF}, ::Val{LOG2WG}) where {WG, NF, LOG2WG} +@kernel function _reduce_stats_kernel!(totals, stats, ns, bm, ::Val{WG}, ::Val{NF}, + ::Val{LOG2WG}) where {WG, NF, LOG2WG} li = @index(Local, Linear) + G = @index(Group, Linear) acc = @localmem Int64 (WG * NF,) - @inbounds for q in 1:NF - a = Int64(0) - for s in li:WG:ns - x = stats[(s - 1) * NF + q] - a += q == NF ? Int64(x != 0) : x + @inbounds begin + st = _mview(stats, _bm_gmember(bm, G), bm.nbatch) + for q in 1:NF + a = Int64(0) + for s in li:WG:ns + x = st[(s - 1) * NF + q] + a += q == NF ? Int64(x != 0) : x + end + acc[(q - 1) * WG + li] = a end - acc[(q - 1) * WG + li] = a end @synchronize for lev in 1:LOG2WG @@ -753,7 +786,7 @@ end @synchronize end if li <= NF - @inbounds totals[li] = acc[(li - 1) * WG + 1] + @inbounds _mview(totals, _bm_gmember(bm, G), bm.nbatch)[li] = acc[(li - 1) * WG + 1] end end @@ -762,26 +795,31 @@ end Sum the per-front statistics `numeric.stats` into `numeric.totals` (`npos, nneg, nzero, nperturbed, n2x2`, and the number of fronts with a failed -pivot): one workgroup, `@localmem` tree reduction, no atomics. Asynchronous; -[`pivot_totals`](@ref) reads the result. +pivot) of every active batch member: one workgroup per member, `@localmem` +tree reduction, no atomics. Asynchronous; [`pivot_totals`](@ref) reads the +result. """ function reduce_stats!(N::Numeric, S::Symbolic) WG = STATS_WORKGROUP - _reduce_stats_kernel!(KernelAbstractions.get_backend(N.stats), WG)(N.totals, N.stats, nsupernodes(S), Val(WG), - Val(FRONT_STATS_FIELDS), Val(_ilog2(WG)); - ndrange = WG) + bm = batch_map(N) + _reduce_stats_kernel!(KernelAbstractions.get_backend(N.stats), WG)(N.totals, N.stats, nsupernodes(S), bm, + Val(WG), Val(FRONT_STATS_FIELDS), + Val(_ilog2(WG)); ndrange = WG * bm.nact) return N end """ - pivot_totals(numeric) -> (npos, nneg, nzero, nperturbed, n2x2) + pivot_totals(numeric, k = 1) -> (npos, nneg, nzero, nperturbed, n2x2) -The statistics reduced on the device by [`reduce_stats!`](@ref) (one copy of -`numeric.totals` to the host: a synchronization, done at the phase boundary -by `getparam`). +The statistics of batch member `k` reduced on the device by +[`reduce_stats!`](@ref) (one copy of `numeric.totals` to the host: a +synchronization, done at the phase boundary by `getparam`). """ -function pivot_totals(N::Numeric) - t = Array(N.totals) +pivot_totals(N::Numeric, k::Integer = 1) = pivot_totals(Array(N.totals), N.nbatch, k) + +function pivot_totals(totals::Vector{Int64}, nb::Integer, k::Integer) + 1 <= k <= nb || throw(InvalidValueError("batch member $k outside 1:$nb")) + t = view(totals, ((k - 1) * FRONT_STATS_FIELDS + 1):(k * FRONT_STATS_FIELDS)) return (npos = t[STAT_NPOS], nneg = t[STAT_NNEG], nzero = t[STAT_NZERO], nperturbed = t[STAT_NPERTURBED], n2x2 = t[STAT_N2X2]) end @@ -790,22 +828,24 @@ end # driver function _launch_front_ldlt!(N::Numeric, S::Symbolic, nzval, first, count, maxchild, prm, ::Val{WG}) where {WG} + bm = batch_map(N; first) kernel! = front_ldlt_kernel!(KernelAbstractions.get_backend(N.factor), WG) kernel!(N.factor, N.stack, N.info, N.stats, N.d, N.piv, N.pivot_kind, N.psign, S.perm, N.aux, nzval, S.amap, - S.amap_ptr, S.amap_src, S.group_nodes, Int(first), S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, + S.amap_ptr, S.amap_src, S.group_nodes, bm, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Int(maxchild), prm, Val(WG); - ndrange = WG * count) + ndrange = WG * count * bm.nact) return nothing end function _launch_subtrees_ldlt!(N::Numeric{T}, S::Symbolic, nzval, first, count, prm, ::Val{LB}, ::Val{WG}) where {T, LB, WG} + bm = batch_map(N; first) kernel! = subtree_ldlt_kernel!(KernelAbstractions.get_backend(N.factor), WG) kernel!(N.factor, N.stack, N.info, N.stats, N.d, N.piv, N.pivot_kind, N.psign, S.perm, N.aux, nzval, S.amap, - S.amap_ptr, S.amap_src, S.group_nodes, Int(first), S.subtree_ptr, S.subtree_nodes, S.super_ptr, + S.amap_ptr, S.amap_src, S.group_nodes, bm, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.local_front, S.local_cb, S.child_ptr, S.child_list, S.relind_ptr, S.relind, prm, Val((LB - SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); - ndrange = WG * count) + ndrange = WG * count * bm.nact) return nothing end @@ -857,7 +897,7 @@ function factorize_ldlt!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; opts throw(InvalidValueError("pivot_sign has $(length(ps)) entries, the matrix has $(S.n) rows")) copyto!(N.psign, ps) end - prm.scaled && abs_max!(N.aux, nzval) + prm.scaled && abs_max!(N.aux, nzval, batch_map(N)) if herm _factorize_ldlt_groups!(N, S, nzval, prm, opts.user_host_interrupt) else diff --git a/src/numeric/storage.jl b/src/numeric/storage.jl index 043e6af..bd9e96e 100644 --- a/src/numeric/storage.jl +++ b/src/numeric/storage.jl @@ -20,7 +20,10 @@ group takes the regime-C path (regime C, or a regime-B bin wider than first): group `k` covers the subtree ids `sub_first[k]:sub_last[k]` of `group_nodes` and its kernel has `sub_local[k]` bytes of local memory ([`subtree_local_bytes`](@ref) of its budget). `info_host` is the host staging -buffer of the one `info` read per phase. +buffer of the one `info` read per phase (one entry per batch member). +Batch state (PLAN §3.5): `members_host[1:nact[]]` are the active members of the +last phase (mirrored in `numeric.members`), `runs` the starts of their runs of +consecutive members ([`member_runs!`](@ref), the strided batches of regime C). """ struct NumericPlan step_first::Vector{Int} @@ -34,9 +37,12 @@ struct NumericPlan sub_last::Vector{Int} sub_local::Vector{Int} info_host::Vector{Int32} + members_host::Vector{Int32} + nact::Base.RefValue{Int} + runs::Vector{Int} end -function NumericPlan(S) +function NumericPlan(S, nb::Integer = 1) sp, sc = S.partition, S.schedule ns = nsupernodes(sp) nchild = zeros(Int, ns) @@ -74,7 +80,8 @@ function NumericPlan(S) maxchild[t] = max(maxchild[t], nchild[sc.group_nodes[q]]) end end - return NumericPlan(first, last, maxchild, gfirst, glast, gmaxchild, gwidth, sfirst, slast, slocal, zeros(Int32, 1)) + return NumericPlan(first, last, maxchild, gfirst, glast, gmaxchild, gwidth, sfirst, slast, slocal, zeros(Int32, nb), + Int32.(1:nb), Ref(Int(nb)), [1, nb + 1]) end """ @@ -115,7 +122,16 @@ Numeric storage of a factorization (PLAN §3.2), laid out by the * `info` (`ns + 1` `Int32`): device status of the numeric phase, the `potrf` info of front `s` at `s`, the reduced result (smallest failed factor column, `0` = none) at `ns + 1`; -* `plan`: the host [`NumericPlan`](@ref). +* `plan`: the host [`NumericPlan`](@ref); +* `members` (`nbatch` `Int32`): the active batch members of the last phase + (first `plan.nact[]` entries, see [`BatchMap`](@ref)); +* `nbatch`: the number of matrices of the uniform batch (PLAN §3.5). + +A uniform batch of `nbatch` members multiplies every length above by +`nbatch`, in the layouts of `src/numeric/batch.jl`: interleaved per front for +`factor` ([`panel_offset`](@ref)), interleaved per entry for `info`, member +after member for the others (the regime-C workspace holds the `m×m` blocks of +the members of one front). """ struct Numeric{T, VT <: AbstractVector{T}, VS <: AbstractVector{Int64}, VI <: AbstractVector{Int32}, VK <: AbstractVector{Int8}} @@ -131,6 +147,8 @@ struct Numeric{T, VT <: AbstractVector{T}, VS <: AbstractVector{Int64}, VI <: Ab psign::VK aux::VT plan::NumericPlan + members::VI + nbatch::Int end # `pivot_kind` codes of LDLᵀ/LDLᴴ factor columns @@ -143,31 +161,74 @@ Base.eltype(::Numeric{T}) where {T} = T Base.show(io::IO, N::Numeric{T, VT}) where {T, VT} = print(io, "Numeric{", T, ", ", nameof(VT), "}(factor ", length(N.factor), ", D ", length(N.d), ", stack ", - length(N.stack), " entries)") + length(N.stack), " entries", N.nbatch > 1 ? ", $(N.nbatch) batch members)" : ")") """ - allocate_numeric(symbolic, T, backend = CPU()) -> Numeric{T} + allocate_numeric(symbolic, T, backend = CPU(); nbatch = 1) -> Numeric{T} Allocate (zero-filled) the factor panels, D, update stack, regime-C workspace, per-front statistics and their totals, status vector, pivot order, pivot kinds, pivot sign requests and `aux` of `symbolic`'s [`Layout`](@ref) for element type -`T` on the KernelAbstractions `backend`, and build its [`NumericPlan`](@ref). +`T` on the KernelAbstractions `backend`, for a uniform batch of `nbatch` +matrices (every length times `nbatch`, all members active), and build its [`NumericPlan`](@ref). This is the only allocation of the numeric phase. """ -function allocate_numeric(S::Symbolic, ::Type{T}, backend::KernelAbstractions.Backend = KernelAbstractions.CPU()) where {T} +function allocate_numeric(S::Symbolic{INT}, ::Type{T}, backend::KernelAbstractions.Backend = KernelAbstractions.CPU(); + nbatch::Integer = 1) where {INT, T} L = S.layout + nb = Int(nbatch) + nb >= 1 || throw(InvalidValueError("nbatch must be ≥ 1, got $nb")) S.n < typemax(Int32) || throw(InvalidValueError("n = $(S.n) does not fit the Int32 status vector")) - factor = KernelAbstractions.zeros(backend, T, L.factor_len) - d = KernelAbstractions.zeros(backend, T, L.d_len) - stack = KernelAbstractions.zeros(backend, T, L.stack_len) - work = KernelAbstractions.zeros(backend, T, L.work_len) - stats = KernelAbstractions.zeros(backend, Int64, FRONT_STATS_FIELDS * nsupernodes(S)) - info = KernelAbstractions.zeros(backend, Int32, nsupernodes(S) + 1) - piv = KernelAbstractions.zeros(backend, Int32, S.n) - pivot_kind = KernelAbstractions.zeros(backend, Int8, S.n) - totals = KernelAbstractions.zeros(backend, Int64, FRONT_STATS_FIELDS) + nb * L.factor_len < typemax(INT) && nb * (nsupernodes(S) + 1) < typemax(Int32) || + throw(InvalidValueError("a batch of $nb factors of $(L.factor_len) entries does not fit the index type $INT")) + factor = KernelAbstractions.zeros(backend, T, nb * L.factor_len) + d = KernelAbstractions.zeros(backend, T, nb * L.d_len) + stack = KernelAbstractions.zeros(backend, T, nb * L.stack_len) + work = KernelAbstractions.zeros(backend, T, nb * L.work_len) + stats = KernelAbstractions.zeros(backend, Int64, nb * FRONT_STATS_FIELDS * nsupernodes(S)) + info = KernelAbstractions.zeros(backend, Int32, nb * (nsupernodes(S) + 1)) + piv = KernelAbstractions.zeros(backend, Int32, nb * S.n) + pivot_kind = KernelAbstractions.zeros(backend, Int8, nb * S.n) + totals = KernelAbstractions.zeros(backend, Int64, nb * FRONT_STATS_FIELDS) psign = KernelAbstractions.zeros(backend, Int8, S.n) - aux = KernelAbstractions.zeros(backend, T, 1) + aux = KernelAbstractions.zeros(backend, T, nb) + members = KernelAbstractions.allocate(backend, Int32, nb) + copyto!(members, Int32.(1:nb)) return Numeric{T, typeof(factor), typeof(stats), typeof(info), typeof(pivot_kind)}(factor, d, stack, work, stats, info, piv, pivot_kind, totals, - psign, aux, NumericPlan(S)) + psign, aux, NumericPlan(S, nb), + members, nb) +end + +""" + set_members!(numeric, members) -> numeric + +Make `members` (1-based, increasing, a subset of `1:numeric.nbatch`) the +active batch members of the next phases: copied to `numeric.members` (a +host-to-device copy, only when they change) and split into runs of +consecutive members. Asynchronous. +""" +function set_members!(N::Numeric, members::AbstractVector{<:Integer}) + plan = N.plan + nact = length(members) + 1 <= nact <= N.nbatch || throw(InvalidValueError("$nact active members for a batch of $(N.nbatch)")) + for j in 1:nact + 1 <= members[j] <= N.nbatch && (j == 1 || members[j] > members[j - 1]) || + throw(InvalidValueError("active members must be increasing members of 1:$(N.nbatch), got $members")) + end + nact == plan.nact[] && all(j -> plan.members_host[j] == members[j], 1:nact) && return N + for j in 1:nact + plan.members_host[j] = Int32(members[j]) + end + plan.nact[] = nact + member_runs!(plan.runs, plan.members_host, nact) + copyto!(N.members, 1, plan.members_host, 1, nact) + return N end + +""" + batch_map(numeric; first = 1, nrhs = 1) -> BatchMap + +The [`BatchMap`](@ref) of the active members of `numeric` for a launch. +""" +batch_map(N::Numeric; first::Integer = 1, nrhs::Integer = 1) = + BatchMap(N.members, N.plan.nact[], N.nbatch; first, nrhs) diff --git a/src/numeric/subtree.jl b/src/numeric/subtree.jl index fcae359..61d9c11 100644 --- a/src/numeric/subtree.jl +++ b/src/numeric/subtree.jl @@ -228,12 +228,13 @@ end end """ - subtree_cholesky_kernel!(backend, WG)(factor, stack, info, nzval, amap, amap_ptr, amap_src, trees, first, + subtree_cholesky_kernel!(backend, WG)(factor, stack, info, nzval, amap, amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, front_ptr, front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, child_list, relind_ptr, relind, - Val(NE), Val(WG); ndrange = WG * count) + Val(NE), Val(WG); ndrange = WG * count * bm.nact) -Fused regime-A kernel: workgroup `g` takes subtree `t = trees[first + g - 1]` +Fused regime-A kernel: workgroup `G` takes subtree `t = trees[bm.first + g - 1]` +of batch member `k` (`g`, `k` from the [`BatchMap`](@ref) `bm`) and processes its supernodes `subtree_nodes[subtree_ptr[t]:(subtree_ptr[t+1]-1)]` in order with the serial stack of packed fronts and contribution blocks in a `@localmem` buffer of `NE` entries (offsets `local_front`, `local_cb`): zero, @@ -242,7 +243,7 @@ scatter A, extend-add the children, factor the front's first `w` columns move the contribution block down, or, for the subtree root (`cb_ptr > 0`), write it to the update stack (packed lower triangle). """ -@kernel function subtree_cholesky_kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, trees, first, +@kernel function subtree_cholesky_kernel!(factor, stack, info, nzval, amap, amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, front_ptr, front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, child_list, relind_ptr, relind, ::Val{NE}, ::Val{WG}) where {NE, WG} @@ -250,12 +251,12 @@ write it to the update stack (packed lower triangle). @uniform RT = real(eltype(factor)) @uniform IT = eltype(subtree_nodes) li = @index(Local, Linear) - g = @index(Group, Linear) + G = @index(Group, Linear) buf = @localmem TT (NE,) ctl = @localmem IT (_ST_CTL,) piv = @localmem RT (1,) if li == 1 - @inbounds t = trees[first + g - 1] + @inbounds t = trees[bm.first + _bm_node(bm, G) - 1] @inbounds ctl[_ST_FIRST] = subtree_ptr[t] @inbounds ctl[_ST_COUNT] = subtree_ptr[t + 1] - subtree_ptr[t] end @@ -267,7 +268,8 @@ write it to the update stack (packed lower triangle). @synchronize _subtree_zero!(buf, ctl, li, Val(WG)) @synchronize - _subtree_scatter!(buf, ctl, nzval, amap, amap_ptr, amap_src, front_ptr, li, Val(WG)) + _subtree_scatter!(buf, ctl, _mview(nzval, _bm_gmember(bm, G), bm.nbatch), amap, amap_ptr, amap_src, front_ptr, + li, Val(WG)) @synchronize for kc in 1:ctl[_ST_NCHILD] _subtree_extend_add!(buf, ctl, kc, child_ptr, child_list, front_nrows, front_ncols, local_cb, relind_ptr, @@ -280,7 +282,9 @@ write it to the update stack (packed lower triangle). _subtree_chol_scale!(buf, ctl, piv, j, li, Val(WG)) @synchronize end - _subtree_write!(factor, stack, info, buf, ctl, front_ptr, cb_ptr, li, Val(WG)) + mb = _bm_gmember(bm, G) + _subtree_write!(factor, _mview(stack, mb, bm.nbatch), _iview(info, mb, bm.nbatch), buf, ctl, + member_panels(front_ptr, mb, bm.nbatch), cb_ptr, li, Val(WG)) @synchronize for r in 1:ctl[_ST_ROUNDS] _subtree_move!(buf, ctl, r, li, Val(WG)) @@ -291,11 +295,12 @@ end function _launch_subtrees!(N::Numeric{T}, S::Symbolic, nzval, first, count, ::Val{LB}, ::Val{WG}) where {T, LB, WG} _ilog2(WG) + bm = batch_map(N; first) kernel! = subtree_cholesky_kernel!(KernelAbstractions.get_backend(N.factor), WG) - kernel!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, Int(first), + kernel!(N.factor, N.stack, N.info, nzval, S.amap, S.amap_ptr, S.amap_src, S.group_nodes, bm, S.subtree_ptr, S.subtree_nodes, S.front_ptr, S.front_nrows, S.front_ncols, S.cb_ptr, S.local_front, S.local_cb, S.child_ptr, S.child_list, S.relind_ptr, S.relind, - Val((LB - SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); ndrange = WG * count) + Val((LB - SUBTREE_LOCAL_RESERVE) ÷ sizeof(T)), Val(WG); ndrange = WG * count * bm.nact) return nothing end @@ -312,8 +317,8 @@ end Regime-A launch: assemble and factor the `count` subtrees `symbolic.group_nodes[first:(first + count - 1)]` (subtree ids of one budget -class) with one [`subtree_cholesky_kernel!`](@ref) launch, one workgroup per -subtree, `local_bytes` ([`subtree_local_bytes`](@ref) of the class) of local +class) of every active batch member with one [`subtree_cholesky_kernel!`](@ref) +launch, one workgroup per (subtree, member), `local_bytes` ([`subtree_local_bytes`](@ref) of the class) of local memory per workgroup. Panels and statuses of every supernode of the subtrees are written, and each subtree root's contribution block goes to the update stack. Asynchronous. diff --git a/src/solve/permute.jl b/src/solve/permute.jl index e9f4672..9180748 100644 --- a/src/solve/permute.jl +++ b/src/solve/permute.jl @@ -8,6 +8,12 @@ # an `n` vector, a column-major `n × nrhs` matrix, a strided vector of # `n * nrhs` entries (column-major), or, with `transposed = true`, row-major # data, i.e. a column-major `nrhs × n` matrix (or the strided vector of one). +# +# Uniform batch: the user array holds the `nrhs` right-hand sides of every +# batch member, member after member (`n × nrhs × nbatch`, as a strided vector, +# an `n × (nrhs nbatch)` matrix or a 3-D array reshaped to one); the workspace +# `Y` holds those of the active members only ("compact" columns: member slot +# `j` at `(j - 1) nrhs + 1 : j nrhs`), mapped by a [`BatchMap`](@ref). # call `f(Val(a), Val(b))` with compile-time flags (a runtime `Val(flag)` would dispatch dynamically) @inline function _with_flags(f, a::Bool, b::Bool) @@ -54,68 +60,78 @@ function rhs_count(B::AbstractVecOrMat, n::Integer; transposed::Bool = false) return nrhs end -@kernel function _permute_rhs_kernel!(Y, B, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} +# work item q: entry k of compact column r (user column `_bm_ucol(bm, r)` of the `nrhs` user columns) +@kernel function _permute_rhs_kernel!(Y, B, perm, n, nrhs, ncols, bm, ::Val{TR}, ::Val{CJ}) where {TR, CJ} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 - @inbounds if r <= nrhs - v = _rhs_get(B, perm[k], r, n, nrhs, Val(TR)) + @inbounds if r <= ncols + v = _rhs_get(B, perm[k], _bm_ucol(bm, r), n, nrhs, Val(TR)) Y[k, r] = CJ ? conj(v) : v end end -@kernel function _unpermute_solution_kernel!(X, Y, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} +@kernel function _unpermute_solution_kernel!(X, Y, perm, n, nrhs, ncols, bm, ::Val{TR}, ::Val{CJ}) where {TR, CJ} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 - @inbounds if r <= nrhs + @inbounds if r <= ncols v = Y[k, r] - _rhs_set!(X, CJ ? conj(v) : v, perm[k], r, n, nrhs, Val(TR)) + _rhs_set!(X, CJ ? conj(v) : v, perm[k], _bm_ucol(bm, r), n, nrhs, Val(TR)) end end -function _check_permute(Y, B, perm, transposed) +# (n, user columns, compact columns, batch map): without `bm`, every user column, in order +function _check_permute(Y, B, perm, transposed, bm) n = length(perm) nrhs = rhs_count(B, n; transposed) - size(Y, 1) == n && size(Y, 2) >= nrhs || - throw(DimensionMismatch("the permuted right-hand side is $(size(Y)), needs $n × ≥ $nrhs")) - return n, nrhs + bm === nothing && (bm = single_batch(; nrhs)) + bm.nrhs * bm.nbatch == nrhs || + throw(DimensionMismatch("$nrhs right-hand sides for $(bm.nbatch) batch members with $(bm.nrhs) each")) + ncols = bm.nrhs * bm.nact + size(Y, 1) == n && size(Y, 2) >= ncols || + throw(DimensionMismatch("the permuted right-hand side is $(size(Y)), needs $n × ≥ $ncols")) + return n, nrhs, ncols, bm end """ - permute_rhs!(Y, B, perm; transposed = false, conjugate = false) -> Y + permute_rhs!(Y, B, perm; transposed = false, conjugate = false, bm = nothing) -> Y `Y[k, r] = B[perm[k], r]` (conjugated when `conjugate`, for the solves with `conj(A)` of `solve_mode`) for the `nrhs` right-hand sides of the user array `B` ([`rhs_count`](@ref): vector, matrix, strided vector; row-major when `transposed`) into the first `nrhs` columns of the `n × ≥ nrhs` device matrix -`Y`. One launch; asynchronous. +`Y`. With a [`BatchMap`](@ref) `bm` (uniform batch), `B` holds `bm.nrhs` +right-hand sides for each of the `bm.nbatch` members and only those of the +active members are copied, member slot `j` to the columns +`(j - 1) bm.nrhs + 1 : j bm.nrhs` of `Y`. One launch; asynchronous. """ function permute_rhs!(Y::AbstractMatrix, B::AbstractVecOrMat, perm::AbstractVector; transposed::Bool = false, - conjugate::Bool = false) - n, nrhs = _check_permute(Y, B, perm, transposed) - n * nrhs > 0 || return Y + conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing) + n, nrhs, ncols, bm = _check_permute(Y, B, perm, transposed, bm) + n * ncols > 0 || return Y kernel! = _permute_rhs_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) _with_flags(transposed, conjugate) do tr, cj - kernel!(Y, B, perm, n, nrhs, tr, cj; ndrange = n * nrhs) + kernel!(Y, B, perm, n, nrhs, ncols, bm, tr, cj; ndrange = n * ncols) end return Y end """ - unpermute_solution!(X, Y, perm; transposed = false, conjugate = false) -> X + unpermute_solution!(X, Y, perm; transposed = false, conjugate = false, bm = nothing) -> X `X[perm[k], r] = Y[k, r]` (conjugated when `conjugate`): the inverse of [`permute_rhs!`](@ref), from the -first `nrhs` columns of `Y` into the user array `X` (same layouts). One launch; -asynchronous. +first `nrhs` columns of `Y` into the user array `X` (same layouts; with a +[`BatchMap`](@ref) `bm`, only the columns of the active members of `X` are +written). One launch; asynchronous. """ function unpermute_solution!(X::AbstractVecOrMat, Y::AbstractMatrix, perm::AbstractVector; transposed::Bool = false, - conjugate::Bool = false) - n, nrhs = _check_permute(Y, X, perm, transposed) - n * nrhs > 0 || return X + conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing) + n, nrhs, ncols, bm = _check_permute(Y, X, perm, transposed, bm) + n * ncols > 0 || return X kernel! = _unpermute_solution_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) _with_flags(transposed, conjugate) do tr, cj - kernel!(X, Y, perm, n, nrhs, tr, cj; ndrange = n * nrhs) + kernel!(X, Y, perm, n, nrhs, ncols, bm, tr, cj; ndrange = n * ncols) end return X end diff --git a/src/solve/refinement.jl b/src/solve/refinement.jl index 4ad23f7..ccd57e7 100644 --- a/src/solve/refinement.jl +++ b/src/solve/refinement.jl @@ -14,7 +14,9 @@ # # User arrays come in the layouts of `src/solve/permute.jl` (vector, matrix, # strided vector, row-major when transposed); the residual `R` is a column-major -# `n × nrhs` device matrix in the original numbering. +# `n × nrhs` device matrix in the original numbering. Uniform batch: `R` holds +# the compact columns of the active members (as the solve workspace), each +# residual uses its member's values (`nzval` member after member). "Workgroup size of the residual-norm reduction (one workgroup per right-hand side)." const REFINE_WORKGROUP = 256 @@ -105,35 +107,41 @@ function _grow_refinement(W::RefinementWorkspace{T, R, VI, MT, VR}, nrhs::Intege return RefinementWorkspace{T, R, VI, MT, VR}(W.rowptr, W.colval, W.src, Rm, Bc, norms, zeros(R, 2 * nrhs)) end -@kernel function _residual_kernel!(R, rowptr, colval, src, nzval, X, B, n, nrhs, ::Val{XT}, ::Val{BT}, +# compact column r (user column `_bm_ucol(bm, r)` of the `nrhs` user columns, values of member +# `_bm_cmember(bm, r)`) +@kernel function _residual_kernel!(R, rowptr, colval, src, nzval_all, X, B, n, nrhs, ncols, bm, ::Val{XT}, ::Val{BT}, ::Val{CJ}) where {XT, BT, CJ} q = @index(Global, Linear) i = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 - @inbounds if r <= nrhs + @inbounds if r <= ncols + nzval = _mview(nzval_all, _bm_cmember(bm, r), bm.nbatch) + u = _bm_ucol(bm, r) acc = zero(eltype(R)) for c in Int(rowptr[i]):(Int(rowptr[i + 1]) - 1) s = Int(src[c]) v = nzval[abs(s)] v = xor(s < 0, CJ) ? conj(v) : v - acc += v * _rhs_get(X, Int(colval[c]), r, n, nrhs, Val(XT)) + acc += v * _rhs_get(X, Int(colval[c]), u, n, nrhs, Val(XT)) end - R[i, r] = _rhs_get(B, i, r, n, nrhs, Val(BT)) - acc + R[i, r] = _rhs_get(B, i, u, n, nrhs, Val(BT)) - acc end end -@kernel function _residual_norms_kernel!(norms, R, B, n, nrhs, ::Val{BT}, ::Val{WG}, ::Val{LOG2WG}) where {BT, WG, LOG2WG} +@kernel function _residual_norms_kernel!(norms, R, B, n, nrhs, bm, ::Val{BT}, ::Val{WG}, + ::Val{LOG2WG}) where {BT, WG, LOG2WG} @uniform RT = eltype(norms) li = @index(Local, Linear) r = @index(Group, Linear) nr = @localmem RT (WG,) nb = @localmem RT (WG,) @inbounds begin + u = _bm_ucol(bm, r) sr = zero(RT) sb = zero(RT) for i in li:WG:n sr += abs2(R[i, r]) - sb += abs2(_rhs_get(B, i, r, n, nrhs, Val(BT))) + sb += abs2(_rhs_get(B, i, u, n, nrhs, Val(BT))) end nr[li] = sr nb[li] = sb @@ -155,14 +163,15 @@ end end end -@kernel function _add_correction_kernel!(X, Y, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} +@kernel function _add_correction_kernel!(X, Y, perm, n, nrhs, ncols, bm, ::Val{TR}, ::Val{CJ}) where {TR, CJ} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 - @inbounds if r <= nrhs + @inbounds if r <= ncols i = Int(perm[k]) + u = _bm_ucol(bm, r) y = Y[k, r] - _rhs_set!(X, _rhs_get(X, i, r, n, nrhs, Val(TR)) + (CJ ? conj(y) : y), i, r, n, nrhs, Val(TR)) + _rhs_set!(X, _rhs_get(X, i, u, n, nrhs, Val(TR)) + (CJ ? conj(y) : y), i, u, n, nrhs, Val(TR)) end end @@ -175,48 +184,64 @@ end end end +# the batch map of `nrhs` user columns: given, or every column in order +_user_map(bm::BatchMap, nrhs) = bm +_user_map(::Nothing, nrhs) = single_batch(; nrhs) + """ - residual!(W, nzval, X, B; nrhs, transposed = false, b_transposed = transposed, conjugate = false) -> W.R + residual!(W, nzval, X, B; nrhs, transposed = false, b_transposed = transposed, conjugate = false, + bm = nothing) -> W.R `W.R[:, 1:nrhs] = B - M X` (`conj(M)` when `conjugate`) with the full matrix `M` of the map of `W` and the user values `nzval`; `X`, `B` in the layouts of -[`rhs_count`](@ref) (row-major when `transposed`/`b_transposed`). One launch, +[`rhs_count`](@ref) (row-major when `transposed`/`b_transposed`) with `nrhs` +right-hand sides. Uniform batch ([`BatchMap`](@ref) `bm`): `nzval` holds the +values of the `bm.nbatch` members, and the residuals of the active members go +to the compact columns `1:(bm.nrhs bm.nact)` of `W.R`. One launch, one work item per (row, right-hand side), gather-based (no atomics). Asynchronous. """ function residual!(W::RefinementWorkspace, nzval::AbstractVector, X::AbstractVecOrMat, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false, b_transposed::Bool = transposed, - conjugate::Bool = false) + conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing) n = size(W.R, 1) - n * nrhs > 0 || return W.R + bm = _user_map(bm, nrhs) + ncols = bm.nrhs * bm.nact + n * ncols > 0 || return W.R kernel! = _residual_kernel!(KernelAbstractions.get_backend(W.R), PERMUTE_WORKGROUP) _with_flags(transposed, b_transposed) do xt, bt if conjugate - kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), xt, bt, Val(true); ndrange = n * Int(nrhs)) + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), ncols, bm, xt, bt, Val(true); + ndrange = n * ncols) else - kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), xt, bt, Val(false); ndrange = n * Int(nrhs)) + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), ncols, bm, xt, bt, Val(false); + ndrange = n * ncols) end end return W.R end """ - residual_norms!(W, B; nrhs, transposed = false) -> W.norms_host + residual_norms!(W, B; nrhs, transposed = false, bm = nothing) -> W.norms_host -`‖W.R[:, k]‖₂²` and `‖B[:, k]‖₂²` for `k = 1:nrhs` reduced on the device (one -workgroup per right-hand side, `@localmem` tree reduction) into `W.norms`, then -copied to `W.norms_host` (one host synchronization). +`‖W.R[:, k]‖₂²` and `‖B[:, k]‖₂²` for the compact columns `k` (`1:nrhs`, or +the active members' columns of the [`BatchMap`](@ref) `bm`) reduced on the +device (one workgroup per right-hand side, `@localmem` tree reduction) into +`W.norms`, then copied to `W.norms_host` (one host synchronization). """ -function residual_norms!(W::RefinementWorkspace, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false) +function residual_norms!(W::RefinementWorkspace, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false, + bm::Union{Nothing, BatchMap} = nothing) n = size(W.R, 1) + bm = _user_map(bm, nrhs) + ncols = bm.nrhs * bm.nact WG = REFINE_WORKGROUP kernel! = _residual_norms_kernel!(KernelAbstractions.get_backend(W.R), WG) if transposed - kernel!(W.norms, W.R, B, n, Int(nrhs), Val(true), Val(WG), Val(_ilog2(WG)); ndrange = WG * Int(nrhs)) + kernel!(W.norms, W.R, B, n, Int(nrhs), bm, Val(true), Val(WG), Val(_ilog2(WG)); ndrange = WG * ncols) else - kernel!(W.norms, W.R, B, n, Int(nrhs), Val(false), Val(WG), Val(_ilog2(WG)); ndrange = WG * Int(nrhs)) + kernel!(W.norms, W.R, B, n, Int(nrhs), bm, Val(false), Val(WG), Val(_ilog2(WG)); ndrange = WG * ncols) end - copyto!(W.norms_host, 1, W.norms, 1, 2 * Int(nrhs)) + copyto!(W.norms_host, 1, W.norms, 1, 2 * ncols) return W.norms_host end @@ -249,19 +274,23 @@ function copy_rhs!(C::AbstractMatrix, B::AbstractVecOrMat; nrhs::Integer, transp end """ - add_correction!(X, Y, perm; nrhs, transposed = false, conjugate = false) -> X + add_correction!(X, Y, perm; nrhs, transposed = false, conjugate = false, bm = nothing) -> X `X[perm[k], r] += Y[k, r]` (`conj(Y[k, r])` when `conjugate`): the solution of the correction system, in factor order in the workspace `Y`, added to the user -array `X`. One launch, conflict-free (`perm` is a permutation). +array `X` of `nrhs` right-hand sides (uniform batch: the compact columns of `Y` +to the active members' columns of `X`, [`BatchMap`](@ref) `bm`). One launch, +conflict-free (`perm` is a permutation). """ function add_correction!(X::AbstractVecOrMat, Y::AbstractMatrix, perm::AbstractVector; nrhs::Integer, - transposed::Bool = false, conjugate::Bool = false) + transposed::Bool = false, conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing) n = length(perm) - n * nrhs > 0 || return X + bm = _user_map(bm, nrhs) + ncols = bm.nrhs * bm.nact + n * ncols > 0 || return X kernel! = _add_correction_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) _with_flags(transposed, conjugate) do tr, cj - kernel!(X, Y, perm, n, Int(nrhs), tr, cj; ndrange = n * Int(nrhs)) + kernel!(X, Y, perm, n, Int(nrhs), ncols, bm, tr, cj; ndrange = n * ncols) end return X end @@ -290,17 +319,21 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac b_transposed::Bool = transposed, conjugate::Bool = false, deterministic::Bool = false, interrupt::Union{Nothing, Threads.Atomic{Bool}} = nothing, progress::Base.RefValue{Int} = Ref(0)) - nrhs = rhs_count(X, S.n; transposed) - max_rhs(W) >= nrhs && max_rhs(ws) >= nrhs || - throw(DimensionMismatch("the refinement workspace holds $(max_rhs(W)) right-hand sides, need $nrhs")) + nu = rhs_count(X, S.n; transposed) # user columns: nrhs per member × nbatch + nu % N.nbatch == 0 || throw(DimensionMismatch("$nu right-hand sides for a batch of $(N.nbatch) members")) + nrhs = nu ÷ N.nbatch + ncols = _ncols(N, nrhs) + max_rhs(W) >= ncols && max_rhs(ws) >= ncols || + throw(DimensionMismatch("the refinement workspace holds $(max_rhs(W)) right-hand sides, need $ncols")) + bm = batch_map(N; nrhs) steps = 0 progress[] = 0 for _ in 1:nsteps _poll_interrupt(interrupt) - residual!(W, nzval, X, B; nrhs, transposed, b_transposed, conjugate) + residual!(W, nzval, X, B; nrhs = nu, transposed, b_transposed, conjugate, bm) if tol > 0 - norms = residual_norms!(W, B; nrhs, transposed = b_transposed) - rel = _max_relative_residual(norms, nrhs) + norms = residual_norms!(W, B; nrhs = nu, transposed = b_transposed, bm) + rel = _max_relative_residual(norms, ncols) _log(LOG_DEBUG, () -> "refinement: step $steps, relative residual $rel") rel <= tol && break end @@ -308,7 +341,7 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac forward_sweep!(ws, S, N; nrhs, deterministic) diagonal_sweep!(ws, S, N; nrhs) backward_sweep!(ws, S, N; nrhs) - add_correction!(X, ws.Y, S.perm; nrhs, transposed, conjugate) + add_correction!(X, ws.Y, S.perm; nrhs = nu, transposed, conjugate, bm) steps += 1 progress[] = steps end diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index 2acf5cb..297ff6a 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -13,7 +13,11 @@ # interface (vendor by default), plus a scatter / gather kernel. # # The right-hand sides are the second grid dimension of every kernel (flattened -# into the 1-D ndrange, see `_sv_front`). +# into the 1-D ndrange, see `_sv_front`). Uniform batch: the columns of `Y` are +# those of the active members (`nrhs` per member, member slot after member slot, +# see `src/solve/permute.jl`), so the batch index is part of that dimension; a +# workgroup reads the factor, D and pivot order of its column's member +# (`_bm_cmember`, `src/numeric/batch.jl`). # # The forward sweep has write conflicts: the fronts of one launch update shared # ancestor rows. Two variants: @@ -403,20 +407,21 @@ end # kernels """ - solve_fwd_kernel!(backend, WG)(Y, U, factor, piv, list, first, count, subtree_ptr, subtree_nodes, super_ptr, + solve_fwd_kernel!(backend, WG)(Y, U, factor, piv, list, bm, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, rowval, front_ptr, front_nrows, front_ncols, child_ptr, child_list, relind_ptr, relind, Val(SUB), Val(DET), Val(LDL), Val(WG); - ndrange = WG * count * nrhs) + ndrange = WG * count * ncols) -Forward sweep `L z = y` of right-hand side `r` by workgroup `G = g + count (r - 1)` -over its fronts: the subtree `list[first + g - 1]`, its +Forward sweep `L z = y` of right-hand side `r` (a column of `Y`, of batch +member `_bm_cmember(bm, r)`) by workgroup `G = g + count (r - 1)` +over its fronts: the subtree `list[bm.first + g - 1]`, its supernodes in processing order (`SUB = true`, regime A), or the single front -`list[first + g - 1]` (regime B). Per front: (deterministic, `DET`) zero its +`list[bm.first + g - 1]` (regime B). Per front: (deterministic, `DET`) zero its update buffer and pull the children's; (`LDL`) reorder its columns of `Y` by its local pivot order `piv`; TRSV with `L11` on its columns of `Y`; GEMV with `L21` into its update buffer or (atomic variant) into the rows of `Y` below. """ -@kernel function solve_fwd_kernel!(Y, U, factor, piv, list, first, count, subtree_ptr, subtree_nodes, super_ptr, +@kernel function solve_fwd_kernel!(Y, U, factor, piv, list, bm, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, rowval, front_ptr, front_nrows, front_ncols, child_ptr, child_list, relind_ptr, relind, ::Val{SUB}, ::Val{DET}, ::Val{LDL}, ::Val{WG}) where {SUB, DET, LDL, WG} @@ -425,13 +430,14 @@ its local pivot order `piv`; TRSV with `L11` on its columns of `Y`; GEMV with G = @index(Group, Linear) ctl = @localmem IT (_SV_CTL,) if li == 1 - _sv_start!(ctl, _sv_front(G, count), list, first, subtree_ptr, subtree_nodes, super_ptr, Val(SUB)) + _sv_start!(ctl, _sv_front(G, count), list, bm.first, subtree_ptr, subtree_nodes, super_ptr, Val(SUB)) end @synchronize for k in 1:ctl[_SV_COUNT] if li == 1 v = _sv_node(ctl, k, list, subtree_nodes, Val(SUB), Val(false)) - _sv_setup!(ctl, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr) + _sv_setup!(ctl, v, super_ptr, rowptr, member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), + front_nrows, front_ncols, child_ptr) end @synchronize if DET @@ -446,7 +452,8 @@ its local pivot order `piv`; TRSV with `L11` on its columns of `Y`; GEMV with if LDL _sv_piv_save!(Y, U, ctl, _sv_rhs(G, count), li, Val(WG)) @synchronize - _sv_piv_apply!(Y, U, piv, ctl, _sv_rhs(G, count), li, Val(WG), Val(true)) + _sv_piv_apply!(Y, U, _mview(piv, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), ctl, _sv_rhs(G, count), li, + Val(WG), Val(true)) @synchronize end for j in 1:ctl[_SV_W] @@ -461,16 +468,17 @@ its local pivot order `piv`; TRSV with `L11` on its columns of `Y`; GEMV with end """ - solve_bwd_kernel!(backend, WG)(Y, U, factor, piv, list, first, count, subtree_ptr, subtree_nodes, super_ptr, + solve_bwd_kernel!(backend, WG)(Y, U, factor, piv, list, bm, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, rowval, front_ptr, front_nrows, front_ncols, child_ptr, Val(SUB), - Val(LDL), Val(CONJ), Val(WG); ndrange = WG * count * nrhs) + Val(LDL), Val(CONJ), Val(WG); ndrange = WG * count * ncols) -Backward sweep `Lᴴ x = z` (`Lᵀ` when `CONJ = false`) of right-hand side `r` by workgroup `G = g + count (r - 1)` -over its fronts (the subtree `list[first + g - 1]` in reverse processing order, `SUB = true`, or one +Backward sweep `Lᴴ x = z` (`Lᵀ` when `CONJ = false`) of right-hand side `r` (of batch member +`_bm_cmember(bm, r)`) by workgroup `G = g + count (r - 1)` +over its fronts (the subtree `list[bm.first + g - 1]` in reverse processing order, `SUB = true`, or one front): gather the finished rows below, GEMV with `L21ᴴ`, TRSV with `L11ᴴ`, then (`LDL`) restore the order before the front's local pivoting. """ -@kernel function solve_bwd_kernel!(Y, U, factor, piv, list, first, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, +@kernel function solve_bwd_kernel!(Y, U, factor, piv, list, bm, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, rowval, front_ptr, front_nrows, front_ncols, child_ptr, ::Val{SUB}, ::Val{LDL}, ::Val{CONJ}, ::Val{WG}) where {SUB, LDL, CONJ, WG} @uniform IT = eltype(front_ptr) @@ -478,13 +486,14 @@ order before the front's local pivoting. G = @index(Group, Linear) ctl = @localmem IT (_SV_CTL,) if li == 1 - _sv_start!(ctl, _sv_front(G, count), list, first, subtree_ptr, subtree_nodes, super_ptr, Val(SUB)) + _sv_start!(ctl, _sv_front(G, count), list, bm.first, subtree_ptr, subtree_nodes, super_ptr, Val(SUB)) end @synchronize for k in 1:ctl[_SV_COUNT] if li == 1 v = _sv_node(ctl, k, list, subtree_nodes, Val(SUB), Val(true)) - _sv_setup!(ctl, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr) + _sv_setup!(ctl, v, super_ptr, rowptr, member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), + front_nrows, front_ncols, child_ptr) end @synchronize _sv_bwd_gather!(Y, factor, rowval, ctl, _sv_rhs(G, count), li, Val(WG), Val(CONJ)) @@ -498,15 +507,16 @@ order before the front's local pivoting. if LDL _sv_piv_save!(Y, U, ctl, _sv_rhs(G, count), li, Val(WG)) @synchronize - _sv_piv_apply!(Y, U, piv, ctl, _sv_rhs(G, count), li, Val(WG), Val(false)) + _sv_piv_apply!(Y, U, _mview(piv, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), ctl, _sv_rhs(G, count), li, + Val(WG), Val(false)) @synchronize end end end # LDLᵀ regime-C path: reorder the front's columns of Y by its local pivot order (`FWD`) or restore the -# order (backward), one workgroup per right-hand side -@kernel function _sv_piv_kernel!(Y, U, piv, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr, +# order (backward), one workgroup per right-hand side (the pivot order of its column's batch member) +@kernel function _sv_piv_kernel!(Y, U, piv, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr, bm, ::Val{FWD}, ::Val{WG}) where {FWD, WG} @uniform IT = eltype(front_ptr) li = @index(Local, Linear) @@ -518,15 +528,18 @@ end @synchronize _sv_piv_save!(Y, U, ctl, r, li, Val(WG)) @synchronize - _sv_piv_apply!(Y, U, piv, ctl, r, li, Val(WG), Val(FWD)) + _sv_piv_apply!(Y, U, _mview(piv, _bm_cmember(bm, r), bm.nbatch), ctl, r, li, Val(WG), Val(FWD)) end # diagonal sweep of LDLᵀ/LDLᴴ: Y ← D⁻¹ Y (1×1 and 2×2 blocks), one work item per (column, right-hand side) -@kernel function _sv_diag_kernel!(Y, d, pivot_kind, n, nrhs, ::Val{HERM}) where {HERM} +@kernel function _sv_diag_kernel!(Y, d_all, pivot_kind_all, n, nrhs, bm, ::Val{HERM}) where {HERM} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 @inbounds if r <= nrhs + mb = _bm_cmember(bm, r) + d = _mview(d_all, mb, bm.nbatch) + pivot_kind = _mview(pivot_kind_all, mb, bm.nbatch) kind = pivot_kind[k] if kind == PIVOT_KIND_2X2_FIRST a, b, c = d[k], d[n + k], d[k + 1] @@ -563,6 +576,7 @@ end end # regime-C path, forward: subtract tmp (= L21 y) from the update buffer (DET) or from the rows of Y below +# (the columns `1:nrhs` of Y, U and tmp) @kernel function _sv_scatter_kernel!(Y, U, tmp, rowval, rp, m, nrhs, ::Val{DET}) where {DET} q = @index(Global, Linear) k = (q - 1) % m + 1 @@ -588,15 +602,21 @@ end # --------------------------------------------------------------------------- # drivers +# +# `nrhs` is the number of right-hand sides per batch member; the sweeps work on +# the `nrhs × nact` columns of the active members (`_ncols`). + +_ncols(N::Numeric, nrhs::Integer) = Int(nrhs) * N.plan.nact[] function _check_solve(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric{T}, nrhs::Integer) where {T} _is_ldlt_structure(S.structure) ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) - length(N.factor) == S.layout.factor_len || + length(N.factor) == N.nbatch * S.layout.factor_len || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) size(ws.Y, 1) == S.n && size(ws.U, 1) == length(S.partition.rowval) && size(ws.tmp, 1) == ws.plan.maxm || throw(InvalidValueError("the solve workspace was not allocated for this analysis")) - 0 <= nrhs <= max_rhs(ws) || - throw(InvalidValueError("$nrhs right-hand sides, the solve workspace holds $(max_rhs(ws))")) + 0 <= _ncols(N, nrhs) <= max_rhs(ws) || + throw(InvalidValueError("$nrhs right-hand sides for $(N.plan.nact[]) batch members, the solve workspace " * + "holds $(max_rhs(ws))")) backend = typeof(KernelAbstractions.get_backend(ws.Y)) typeof(KernelAbstractions.get_backend(N.factor)) == backend && typeof(KernelAbstractions.get_backend(S.rowval)) == backend || @@ -605,19 +625,23 @@ function _check_solve(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric{T}, nrhs::I return nothing end -# dense implementations of the regime-C-path fronts, resolved once per sweep (as `_front_impls`) +# dense implementations of the regime-C-path fronts, resolved once per sweep (as `_front_impls`; the strided +# batched `b*` for a uniform batch, `:none` with `impl = :generic`) function _solve_impls(N::Numeric, S::Symbolic, impl::Symbol) impl === :auto && !S.schedule.vendor_c && (impl = :ka) - return (trsm = select_impl(:trsm, N.factor, impl), gemm = select_impl(:gemm, N.factor, impl)) + batched = N.nbatch > 1 && impl !== :generic + bimpl(op) = batched ? (impl === :ka ? :ka : select_impl(op, N.factor, :auto)) : :none + return (trsm = select_impl(:trsm, N.factor, impl), gemm = select_impl(:gemm, N.factor, impl), + btrsm = bimpl(:trsm_strided_batched), bgemm = bimpl(:gemm_strided_batched)) end function _launch_fwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, det::Val, ldl::Val) WG = SOLVE_WORKGROUP solve_fwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, ws.U, N.factor, N.piv, S.group_nodes, first, count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, - S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, - S.relind, sub, det, ldl, Val(WG); ndrange = WG * count * nrhs) + ws.Y, ws.U, N.factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, + S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, + S.relind_ptr, S.relind, sub, det, ldl, Val(WG); ndrange = WG * count * _ncols(N, nrhs)) return nothing end @@ -625,9 +649,9 @@ function _launch_bwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, c ldl::Val, cj::Val) WG = SOLVE_WORKGROUP solve_bwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, ws.U, N.factor, N.piv, S.group_nodes, first, count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, - S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, sub, ldl, cj, Val(WG); - ndrange = WG * count * nrhs) + ws.Y, ws.U, N.factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, + S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, sub, ldl, cj, Val(WG); + ndrange = WG * count * _ncols(N, nrhs)) return nothing end @@ -668,60 +692,119 @@ function _launch_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, return nothing end -# LDLᵀ regime-C path: local pivot order of front s on the first nrhs columns of Y +# LDLᵀ regime-C path: local pivot order of front s on the columns of Y (each with its member's order) function _piv_dense!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, forward::Bool) WG = SOLVE_WORKGROUP kernel! = _sv_piv_kernel!(KernelAbstractions.get_backend(ws.Y), WG) - args = (ws.Y, ws.U, N.piv, s, S.super_ptr, S.rowptr, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr) + args = (ws.Y, ws.U, N.piv, s, S.super_ptr, S.rowptr, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, + batch_map(N; nrhs)) if forward - kernel!(args..., Val(true), Val(WG); ndrange = WG * nrhs) + kernel!(args..., Val(true), Val(WG); ndrange = WG * _ncols(N, nrhs)) else - kernel!(args..., Val(false), Val(WG); ndrange = WG * nrhs) + kernel!(args..., Val(false), Val(WG); ndrange = WG * _ncols(N, nrhs)) end return nothing end -# the f×w panel of s and its blocks L11, L21, and s's columns of Y -function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int) +# the blocks L11, L21 of the panels of front `s` of the `count` consecutive batch members starting at member slot +# `j0`, the front's rows of their columns of Y and the m rows of their columns of tmp: matrices for one member +# (the single-matrix layout, `count = 1`), `×count` strided batches otherwise +function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int) sc = S.schedule f, w = sc.rows[s], sc.width[s] - p0 = S.layout.panel_ptr[s] + m = f - w + nb = N.nbatch + p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j0], nb) c0 = S.partition.super_ptr[s] - P = reshape(view(N.factor, p0:(p0 + f * w - 1)), f, w) - return f, w, view(P, 1:w, 1:w), view(P, (w + 1):f, 1:w), view(ws.Y, c0:(c0 + w - 1), 1:nrhs) + cols = ((j0 - 1) * nrhs + 1):((j0 + count - 1) * nrhs) + if count == 1 + P = reshape(view(N.factor, p0:(p0 + f * w - 1)), f, w) + return f, w, view(P, 1:w, 1:w), view(P, (w + 1):f, 1:w), view(ws.Y, c0:(c0 + w - 1), cols), + view(ws.tmp, 1:m, cols) + end + P = reshape(view(N.factor, p0:(p0 + f * w * count - 1)), f, w, count) + Y3 = reshape(view(ws.Y, :, cols), size(ws.Y, 1), nrhs, count) + T3 = reshape(view(ws.tmp, :, cols), size(ws.tmp, 1), nrhs, count) + return f, w, view(P, 1:w, 1:w, :), view(P, (w + 1):f, 1:w, :), view(Y3, c0:(c0 + w - 1), :, :), + view(T3, 1:m, :, :) +end + +# the dense calls of front `s` for each run of consecutive active members: one strided-batched call per run, or +# (single matrix, `impl = :generic`) one call per member; `forward` selects `_dense_fwd!` or `_dense_bwd!` +function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple, forward::Bool, + tr::Char) + if N.nbatch == 1 + _dense_call!(ws, S, N, s, nrhs, 1, 1, p, forward, tr, false) + return nothing + end + runs = N.plan.runs + for r in 1:(length(runs) - 1) + j0, j1 = runs[r], runs[r + 1] - 1 + if p.btrsm === :none + for j in j0:j1 + _dense_call!(ws, S, N, s, nrhs, j, 1, p, forward, tr, false) + end + else + _dense_call!(ws, S, N, s, nrhs, j0, j1 - j0 + 1, p, forward, tr, true) + end + end + return nothing +end + +function _dense_call!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int, + p::NamedTuple, forward::Bool, tr::Char, batched::Bool) where {T} + f, w, L11, L21, Yc, Tm = _dense_front(ws, S, N, s, nrhs, j0, batched ? count : 1) + m = f - w + if forward + if batched + trsm_strided_batched!('L', 'L', 'N', 'N', one(T), L11, Yc; impl = p.btrsm) + m > 0 && gemm_strided_batched!(Tm, L21, Yc, one(T), zero(T); impl = p.bgemm) + else + _trsm_impl!(p.trsm, 'L', 'L', 'N', 'N', one(T), L11, Yc) + m > 0 && _gemm_impl!(p.gemm, 'N', 'N', one(T), L21, Yc, zero(T), Tm) + end + elseif batched + m > 0 && gemm_strided_batched!(Yc, L21, Tm, -one(T), one(T); transA = tr, impl = p.bgemm) + trsm_strided_batched!('L', 'L', tr, 'N', one(T), L11, Yc; impl = p.btrsm) + else + m > 0 && _gemm_impl!(p.gemm, tr, 'N', -one(T), L21, Tm, one(T), Yc) + _trsm_impl!(p.trsm, 'L', 'L', tr, 'N', one(T), L11, Yc) + end + return nothing end function _fwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, det::Bool, p::NamedTuple, kind::Int) where {T} - f, w, L11, L21, Yc = _dense_front(ws, S, N, s, nrhs) + sc = S.schedule + w = sc.width[s] + m = sc.rows[s] - w kind == _SV_CHOLESKY || _piv_dense!(ws, S, N, s, nrhs, true) - _trsm_impl!(p.trsm, 'L', 'L', 'N', 'N', one(T), L11, Yc) - m = f - w + _dense_runs!(ws, S, N, s, nrhs, p, true, 'N') m > 0 || return nothing - Tm = view(ws.tmp, 1:m, 1:nrhs) - _gemm_impl!(p.gemm, 'N', 'N', one(T), L21, Yc, zero(T), Tm) + ncols = _ncols(N, nrhs) rp = S.partition.rowptr[s] + w kernel! = _sv_scatter_kernel!(KernelAbstractions.get_backend(ws.Y), SOLVE_WORKGROUP) if det - kernel!(ws.Y, ws.U, ws.tmp, S.rowval, rp, m, nrhs, Val(true); ndrange = m * nrhs) + kernel!(ws.Y, ws.U, ws.tmp, S.rowval, rp, m, ncols, Val(true); ndrange = m * ncols) else - kernel!(ws.Y, ws.U, ws.tmp, S.rowval, rp, m, nrhs, Val(false); ndrange = m * nrhs) + kernel!(ws.Y, ws.U, ws.tmp, S.rowval, rp, m, ncols, Val(false); ndrange = m * ncols) end return nothing end function _bwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple, kind::Int) where {T} - f, w, L11, L21, Yc = _dense_front(ws, S, N, s, nrhs) - m = f - w + sc = S.schedule + w = sc.width[s] + m = sc.rows[s] - w tr = kind == _SV_LDL_SYM ? 'T' : 'C' if m > 0 + ncols = _ncols(N, nrhs) rp = S.partition.rowptr[s] + w kernel! = _sv_gather_kernel!(KernelAbstractions.get_backend(ws.Y), SOLVE_WORKGROUP) - kernel!(ws.tmp, ws.Y, S.rowval, rp, m, nrhs; ndrange = m * nrhs) - _gemm_impl!(p.gemm, tr, 'N', -one(T), L21, view(ws.tmp, 1:m, 1:nrhs), one(T), Yc) + kernel!(ws.tmp, ws.Y, S.rowval, rp, m, ncols; ndrange = m * ncols) end - _trsm_impl!(p.trsm, 'L', 'L', tr, 'N', one(T), L11, Yc) + _dense_runs!(ws, S, N, s, nrhs, p, false, tr) kind == _SV_CHOLESKY || _piv_dense!(ws, S, N, s, nrhs, false) return nothing end @@ -738,9 +821,12 @@ interface (`impl` as in [`factorize!`](@ref)) and a scatter kernel. The atomic variant runs unless `deterministic` is set or the backend lacks an atomic add for `T` (complex `T`, issue #36); the deterministic variant goes through the per-front buffers `ws.U` and an owner-pull of the children (one extra launch -per regime-C step). Asynchronous; allocates nothing on the device. +per regime-C step). Uniform batch: `nrhs` right-hand sides per active member +of `numeric` (`nrhs × nact` columns of `ws.Y`, each solved with its member's +factor; regime-C-path fronts use strided-batched `trsm`/`gemm` over runs of +consecutive members). Asynchronous; allocates nothing on the device. """ -function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws), +function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[], deterministic::Bool = false, impl::Symbol = :auto) _check_solve(ws, S, N, nrhs) nrhs > 0 || return ws @@ -757,7 +843,7 @@ function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integ _sv_pull_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( ws.Y, ws.U, S.group_nodes, a, b - a + 1, S.rowptr, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Val(WG); - ndrange = WG * (b - a + 1) * Int(nrhs)) + ndrange = WG * (b - a + 1) * _ncols(N, nrhs)) end for q in a:b _fwd_dense!(ws, S, N, nodes[q], Int(nrhs), det, p, kind) @@ -777,17 +863,20 @@ columns of `ws.Y`: the identity for Cholesky (`"SPD"`/`"HPD"`); for LDLᵀ/LDLᴴ one launch, one work item per (factor column, right-hand side), dividing by the 1×1 pivots of `numeric.d` and solving the 2×2 blocks (`d[k]`, `d[n + k]`, `d[k + 1]`; Hermitian or complex symmetric as the -structure). Asynchronous. +structure). Uniform batch: `nrhs` per active member, each with its member's D. +Asynchronous. """ -function diagonal_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws)) +function diagonal_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[]) _check_solve(ws, S, N, nrhs) kind = _solve_kind(S, eltype(ws)) (kind == _SV_CHOLESKY || nrhs == 0 || S.n == 0) && return ws + ncols = _ncols(N, nrhs) + bm = batch_map(N; nrhs) kernel! = _sv_diag_kernel!(KernelAbstractions.get_backend(ws.Y), SOLVE_WORKGROUP) if kind == _SV_LDL_HERM - kernel!(ws.Y, N.d, N.pivot_kind, S.n, Int(nrhs), Val(true); ndrange = S.n * Int(nrhs)) + kernel!(ws.Y, N.d, N.pivot_kind, S.n, ncols, bm, Val(true); ndrange = S.n * ncols) else - kernel!(ws.Y, N.d, N.pivot_kind, S.n, Int(nrhs), Val(false); ndrange = S.n * Int(nrhs)) + kernel!(ws.Y, N.d, N.pivot_kind, S.n, ncols, bm, Val(false); ndrange = S.n * ncols) end return ws end @@ -800,9 +889,9 @@ first `nrhs` columns of `ws.Y`, the forward plan in reverse: per step the regime `gemm` with `L21ᴴ`, `trsm` with `L11ᴴ`) and one launch over the regime-B fronts, then one launch over all regime-A subtrees (supernodes in reverse processing order). Gather-based and conflict-free: no atomics, deterministic. -Asynchronous. +Uniform batch as in [`forward_sweep!`](@ref). Asynchronous. """ -function backward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws), +function backward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[], impl::Symbol = :auto) _check_solve(ws, S, N, nrhs) nrhs > 0 || return ws @@ -833,18 +922,24 @@ its local pivot orders) of [`factorize!`](@ref): [`unpermute_solution!`](@ref) into `X`. `B` and `X` are device arrays in the layouts of [`rhs_count`](@ref) (vector, `n × nrhs` matrix, strided vector; row-major with `transposed = true`), with at most `max_rhs(ws)` right-hand -sides; `X === B` is allowed. Asynchronous; allocates nothing on the device. +sides; `X === B` is allowed. Uniform batch: `B` and `X` hold the right-hand +sides of all `numeric.nbatch` members (member after member), only the active +members are solved. Asynchronous; allocates nothing on the device. """ function sweep_solve!(X::AbstractVecOrMat, ws::SolveWorkspace, S::Symbolic, N::Numeric, B::AbstractVecOrMat; transposed::Bool = false, deterministic::Bool = false, impl::Symbol = :auto) - nrhs = rhs_count(B, S.n; transposed) - rhs_count(X, S.n; transposed) == nrhs || - throw(DimensionMismatch("X has $(rhs_count(X, S.n; transposed)) right-hand sides, B has $nrhs")) + ncols = rhs_count(B, S.n; transposed) + rhs_count(X, S.n; transposed) == ncols || + throw(DimensionMismatch("X has $(rhs_count(X, S.n; transposed)) right-hand sides, B has $ncols")) + ncols % N.nbatch == 0 || + throw(DimensionMismatch("$ncols right-hand sides for a batch of $(N.nbatch) members")) + nrhs = ncols ÷ N.nbatch _check_solve(ws, S, N, nrhs) - permute_rhs!(ws.Y, B, S.perm; transposed) + bm = batch_map(N; nrhs) + permute_rhs!(ws.Y, B, S.perm; transposed, bm) forward_sweep!(ws, S, N; nrhs, deterministic, impl) diagonal_sweep!(ws, S, N; nrhs) backward_sweep!(ws, S, N; nrhs, impl) - unpermute_solution!(X, ws.Y, S.perm; transposed) + unpermute_solution!(X, ws.Y, S.perm; transposed, bm) return X end diff --git a/src/solver.jl b/src/solver.jl index a983bd6..b00ee21 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -60,14 +60,28 @@ right-hand side and the solution (as for `solve_mode`, see [`solve_conjugated`](@ref)). Implemented at this point: structures `"SPD"` (real `T`), `"HPD"`, `"S"` -(LDLᵀ; complex symmetric for complex `T`) and `"H"` (LDLᴴ), a single matrix -(no uniform batch), the phases `"reordering"`, `"symbolic_factorization"`, +(LDLᵀ; complex symmetric for complex `T`) and `"H"` (LDLᴴ), single matrices +and uniform batches (see below), the phases `"reordering"`, `"symbolic_factorization"`, `"analysis"`, `"factorization"`, `"refactorization"`, `"solve"`, `"solve_fwd_perm"`, `"solve_fwd"`, `"solve_diag"`, `"solve_bwd"`, `"solve_bwd_perm"` and `"solve_refinement"` ([`execute!`](@ref)). Structure `"G"` and the Schur phases raise [`NotSupportedError`](@ref) when they are executed. +Uniform batch (PLAN §1.6, §3.5, ≅ CUDSS.jl's uniform batch): `nbatch` +matrices with the pattern of `rowptr`/`colval` and values `nzval`, either a +vector of `nbatch · nnz` entries (member after member) or an `nnz × nbatch` +matrix; `nbatch` is deduced from the values (`length(nzval) ÷ length(colval)`), +and `"ubatch_size"` must be `0` (deduced) or equal to it. One analysis serves +every member; each member has its own factor, `"info"`, `"inertia"`, +`"npivots"` and `"pivot_stats"` (vectors over the members). Right-hand sides +and solutions hold `nrhs` columns per member, member after member: a strided +vector of `n · nrhs · nbatch` entries, an `n × (nrhs · nbatch)` matrix, an +`n × nrhs × nbatch` array or a [`MatrixDescriptor`](@ref) with `nbatch`. +`"ubatch_index"` (0-based member, `-1` = all) and `"ubatch_mask"` (one 0/1 +flag per member) restrict the factorization and solve phases to a subset of +the members: the other members' factors and solutions are left untouched. + Fields: `A` (the current [`CSR`](@ref), re-pointed by [`update!`](@ref)), `structure`, `view` (as given), `options` ([`Options`](@ref), set through [`setparam!`](@ref)), `backend`, `nbatch`, `fresh_factorization` (`true` until @@ -88,7 +102,7 @@ mutable struct DirectSolver{T, INT, M <: CSR{T, INT}, B <: KernelAbstractions.Ba backend::B nbatch::Int fresh_factorization::Bool - info::Int + info::Vector{Int} stage::Int host_rowptr::Vector{INT} host_colval::Vector{INT} @@ -126,13 +140,12 @@ function DirectSolver(A::CSR{T, INT}, structure, view; index = A.index) where {T A = CSR(A.rowptr, A.colval, A.nzval, A.nrows, A.ncols; index = base, transposed = A.transposed) end nb = nbatch(A) - nb == 1 || throw(NotSupportedError("uniform batches (nbatch = $nb) are not implemented yet (T17)")) backend = KernelAbstractions.get_backend(A) SY, NU, WS = _symbolic_type(backend, INT), _numeric_type(backend, T), _workspace_type(backend, T) RF = _refinement_type(backend, T, INT) return DirectSolver{T, INT, typeof(A), typeof(backend), SY, NU, WS, RF}( - A, s, v, Options(), backend, nb, true, 0, STAGE_NONE, INT[], INT[], nothing, nothing, nothing, nothing, - nothing, nothing, -1) + A, s, v, Options(), backend, nb, true, zeros(Int, nb), STAGE_NONE, INT[], INT[], nothing, nothing, nothing, + nothing, nothing, nothing, -1) end function DirectSolver(rowptr::AbstractVector{<:Integer}, colval::AbstractVector{<:Integer}, nzval::AbstractVecOrMat, @@ -152,11 +165,12 @@ function Base.show(io::IO, solver::DirectSolver{T, INT}) where {T, INT} print(io, "DirectSolver{", T, ", ", INT, "}(n = ", size(solver, 1), ", nnz = ", nnz(solver.A), ", structure = \"", convert(String, solver.structure), "\", view = '", convert(Char, solver.view), "', backend = ", nameof(typeof(solver.backend)), ", ", _STAGE_NAMES[solver.stage + 1]) + solver.nbatch > 1 && print(io, ", nbatch = ", solver.nbatch) if solver.host_symbolic !== nothing print(io, ", ", nsupernodes(solver.host_symbolic), " supernodes, nnz(L) = ", solver.host_symbolic.partition.nnz_L) end - solver.stage == STAGE_FACTORIZED && print(io, ", info = ", solver.info) + solver.stage == STAGE_FACTORIZED && print(io, ", info = ", _info_value(solver)) print(io, ")") return nothing end @@ -305,7 +319,9 @@ function _check_analysis_supported(solver::DirectSolver{T}) where {T} T <: Complex && s == STRUCTURE_SPD && throw(InvalidValueError("a complex positive definite matrix needs structure \"HPD\"")) opts = solver.options - opts.ubatch_size > 1 && throw(NotSupportedError("uniform batches (ubatch_size > 1) are not implemented yet (T17)")) + opts.ubatch_size == 0 || opts.ubatch_size == solver.nbatch || + throw(InvalidValueError("ubatch_size = $(opts.ubatch_size), but the matrix values hold $(solver.nbatch) " * + "batch member(s) (length(nzval) ÷ nnz)")) opts.matching_alg == MATCHING_NONE || throw(NotSupportedError("matching_alg = \"$(convert(String, opts.matching_alg))\" is not implemented yet (T21)")) opts.schur_mode == 0 || throw(NotSupportedError("schur_mode = 1 is not implemented yet (T20)")) @@ -326,8 +342,8 @@ function _reorder!(solver::DirectSolver{T}) where {T} solver.host_colval = Array(A.colval) P = SymmetricPattern(solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) - # 2×2 pivot pairs ("S"/"H"): the only host copy of the values, at analysis - pp = analysis_pairs(P, solver.host_rowptr, solver.host_colval, A.nzval, A.nrows, solver.structure, + # 2×2 pivot pairs ("S"/"H"): the only host copy of the values, at analysis (of the first batch member) + pp = analysis_pairs(P, solver.host_rowptr, solver.host_colval, _first_member(A), A.nrows, solver.structure, solver.options; view = _stored_view(solver), index = A.index) solver.ordering = compute_ordering(P, solver.options; T, pp.pairs, pp.candidates) solver.host_symbolic = solver.symbolic = solver.numeric = solver.workspace = solver.refinement = nothing @@ -336,6 +352,10 @@ function _reorder!(solver::DirectSolver{T}) where {T} return solver end +# the values of the first batch member (all of them for a single matrix) +_first_member(A::CSR) = nbatch(A) == 1 ? A.nzval : A.nzval isa AbstractMatrix ? view(A.nzval, :, 1) : + view(A.nzval, 1:nnz(A)) + function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} solver.stage >= STAGE_REORDERED || throw(_phase_error("symbolic_factorization", "needs \"reordering\" (or run \"analysis\")")) @@ -352,16 +372,16 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} layout = build_layout(sp, sc) Sh = Symbolic(sp, sc, layout, solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) - nrhs = solver.workspace === nothing ? 1 : max_rhs(solver.workspace) + nrhs = solver.workspace === nothing ? solver.nbatch : max_rhs(solver.workspace) Sd = adapt(solver.backend, Sh, INT) solver.host_symbolic = Sh solver.symbolic = Sd - solver.numeric = allocate_numeric(Sd, T, solver.backend) + solver.numeric = allocate_numeric(Sd, T, solver.backend; nbatch = solver.nbatch) solver.workspace = allocate_solve(Sd, T, solver.backend, nrhs) solver.refinement = nothing solver.stage = STAGE_ANALYZED solver.fresh_factorization = true - solver.info = 0 + fill!(solver.info, 0) _log(LOG_INFO, () -> "symbolic_factorization: $(nsupernodes(Sh)) supernodes, nnz(L) = $(sp.nnz_L), " * "flops = $(sp.flops), $(_elapsed(tic))") return solver @@ -372,42 +392,72 @@ function _factorize!(solver::DirectSolver, p::Phase) solver.stage >= STAGE_ANALYZED || throw(_phase_error(name, "needs \"analysis\"")) p == PHASE_REFACTORIZATION && solver.stage < STAGE_FACTORIZED && throw(_phase_error(name, "needs a previous \"factorization\"")) - p == PHASE_REFACTORIZATION && (solver.info = 0) + N, S = solver.numeric, solver.symbolic + _set_members!(solver) + if p == PHASE_REFACTORIZATION + for j in 1:N.plan.nact[] + solver.info[N.plan.members_host[j]] = 0 + end + end tic = time_ns() try - solver.info = _numeric_phase!(solver.numeric, solver.symbolic, solver.A.nzval, solver.options) + _numeric_phase!(N, S, solver.A.nzval, solver.options) catch err if err isa InterruptedError # the panels are partly overwritten: back to "analyzed", a "factorization" must follow solver.stage = STAGE_ANALYZED solver.fresh_factorization = true - solver.info = 0 + fill!(solver.info, 0) _log(LOG_INFO, () -> "$name: interrupted") end rethrow() end + member_info!(solver.info, N, S) solver.stage = STAGE_FACTORIZED p == PHASE_FACTORIZATION && (solver.fresh_factorization = false) - _log(LOG_INFO, () -> "$name: info = $(solver.info), $(_elapsed(tic)) (launches, asynchronous)") + _log(LOG_INFO, () -> "$name: info = $(_info_value(solver)), $(_elapsed(tic)) (launches, asynchronous)") return solver end +# the members the next numeric or solve phase processes (`ubatch_index`, `ubatch_mask`) +function _set_members!(solver::DirectSolver) + opts = solver.options + nb = solver.nbatch + if opts.ubatch_index == -1 && opts.ubatch_mask === nothing + N = solver.numeric + N.plan.nact[] == nb || set_members!(N, 1:nb) + return N + end + members = active_members(nb, opts.ubatch_index, opts.ubatch_mask) + isempty(members) && throw(InvalidValueError("ubatch_index and ubatch_mask select no batch member")) + return set_members!(solver.numeric, members) +end + +# the "info" data parameter: an `Int` for a single matrix, a vector over the members of a batch +_info_value(solver::DirectSolver) = solver.nbatch == 1 ? solver.info[1] : copy(solver.info) + _elapsed(tic) = string(round((time_ns() - tic) / 1.0e6; digits = 3), " ms") # function barrier: concrete storage types for the numeric phase _numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractVector, opts::Options) = factorize!(N, S, nzval; opts) _numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractMatrix, opts::Options) = factorize!(N, S, vec(nzval); opts) -# user data of a right-hand side / solution: (array, transposed) +# user data of a right-hand side / solution: (array, transposed); an `n × nrhs × nbatch` array is the +# `n × (nrhs nbatch)` matrix with the same memory (uniform batch) _rhs_data(X::AbstractVecOrMat) = (X, false) function _rhs_data(X::MatrixDescriptor) X.data === nothing && throw(InvalidValueError("MatrixDescriptor has no data; call update! first")) - X.nbatch == 1 || throw(NotSupportedError("batched right-hand sides are not implemented yet (T17)")) - X.data isa AbstractVecOrMat || throw(NotSupportedError("3-D right-hand sides are not implemented yet (T17)")) - return (X.data, X.transposed) + X.nbatch == 1 || !X.transposed || + throw(NotSupportedError("row-major (transposed) right-hand sides of a uniform batch are not supported")) + return (_rhs_matrix(X.data), X.transposed) end -_rhs_data(X::AbstractArray) = throw(NotSupportedError("$(ndims(X))-D right-hand sides are not implemented yet (T17)")) -_rhs_data(X) = throw(InvalidValueError("the solve phases need a vector, a matrix or a MatrixDescriptor, got $(typeof(X))")) +_rhs_data(X::AbstractArray{<:Any, 3}) = (_rhs_matrix(X), false) +_rhs_data(X::AbstractArray) = throw(InvalidValueError("$(ndims(X))-D right-hand sides are not supported")) +_rhs_data(X) = throw(InvalidValueError("the solve phases need a vector, a matrix, an n × nrhs × nbatch array or a " * + "MatrixDescriptor, got $(typeof(X))")) + +_rhs_matrix(X::AbstractVecOrMat) = X +_rhs_matrix(X::AbstractArray{<:Any, 3}) = reshape(X, size(X, 1), size(X, 2) * size(X, 3)) function _check_rhs_backend(solver::DirectSolver, X, name) typeof(KernelAbstractions.get_backend(X)) == typeof(solver.backend) || @@ -463,31 +513,39 @@ function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::A S, N = solver.symbolic, solver.numeric opts = solver.options n = S.n - nrhs = rhs_count(B, n; transposed) - rhs_count(X, n; transposed) == nrhs || - throw(DimensionMismatch("X has $(rhs_count(X, n; transposed)) right-hand sides, B has $nrhs")) - ws = _workspace!(solver, nrhs) + nu = rhs_count(B, n; transposed) # user columns: nrhs per member × nbatch + rhs_count(X, n; transposed) == nu || + throw(DimensionMismatch("X has $(rhs_count(X, n; transposed)) right-hand sides, B has $nu")) + nb = solver.nbatch + nb == 1 || !transposed || + throw(NotSupportedError("row-major (transposed) right-hand sides of a uniform batch are not supported")) + nu % nb == 0 || + throw(DimensionMismatch("$nu right-hand side columns do not split over the $nb members of the batch")) + nrhs = nu ÷ nb + _set_members!(solver) + ws = _workspace!(solver, nu) + bm = batch_map(N; nrhs) det = opts.deterministic_mode == 1 cj = solve_conjugated(solver.structure, eltype(ws), solver.A.transposed, opts.solve_mode) refine = opts.ir_n_steps > 0 && (p == PHASE_SOLVE || p == PHASE_SOLVE_REFINEMENT) if p == PHASE_SOLVE_REFINEMENT && refine && Base.mightalias(X, B) throw(InvalidValueError("\"solve_refinement\" needs the original right-hand side: X and B must not alias")) end - W = refine ? _refinement!(solver, nrhs) : nothing + W = refine ? _refinement!(solver, nu) : nothing if p == PHASE_SOLVE Bs, bt = B, transposed if refine && Base.mightalias(X, B) # the solve overwrites B: keep a copy for the residual - copy_rhs!(W.Bc, B; nrhs, transposed) + copy_rhs!(W.Bc, B; nrhs = nu, transposed) Bs, bt = W.Bc, false end - permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj) + permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj, bm) forward_sweep!(ws, S, N; nrhs, deterministic = det) diagonal_sweep!(ws, S, N; nrhs) backward_sweep!(ws, S, N; nrhs) - unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj) + unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj, bm) refine ? _refine_phase!(solver, W, ws, X, Bs, transposed, bt, cj, det) : (solver.ir_steps = 0) elseif p == PHASE_SOLVE_FWD_PERM - permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj) + permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj, bm) elseif p == PHASE_SOLVE_FWD forward_sweep!(ws, S, N; nrhs, deterministic = det) elseif p == PHASE_SOLVE_DIAG @@ -495,7 +553,7 @@ function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::A elseif p == PHASE_SOLVE_BWD backward_sweep!(ws, S, N; nrhs) elseif p == PHASE_SOLVE_BWD_PERM - unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj) + unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj, bm) else # PHASE_SOLVE_REFINEMENT refine ? _refine_phase!(solver, W, ws, X, B, transposed, transposed, cj, det) : (solver.ir_steps = 0) end @@ -570,7 +628,6 @@ function _output_task(name) name in ("perm_matching", "scale_row", "scale_col") && return "T21" name in ("schur_shape", "schur_matrix") && return "T20" name == "nd_partition_tree" && return "T24" - name == "ubatch_mask" && return "T17" return "M12" # hybrid_device_memory_min end @@ -580,7 +637,12 @@ end Set a configuration or user-input data parameter of `solver` (≅ `cudss_set`): everything [`setparam!`](@ref)`(::Options, …)` accepts, stored in `solver.options` and used by the next phase that reads it, plus the data -parameter `"info"` (an integer; CUDSS.jl resets it before a refactorization). +parameter `"info"` (an integer, or for a uniform batch an integer for every +member or a vector of `nbatch` integers; CUDSS.jl resets it before a +refactorization). Uniform batch: `"ubatch_size"` (`0` or the batch size), +`"ubatch_index"` (0-based member, `-1` = all) and `"ubatch_mask"` (`nbatch` +flags 0/1, or `nothing`) select the members of the next phases; the index and +the mask are checked against the batch size when a phase runs. `"pivot_sign"` (PLAN §1.7: a vector of `n` entries in `(-1, 0, 1)`, host or device, or `nothing`) is checked against the size of the matrix; the next `"factorization"`/`"refactorization"` copies it to the device. Setting @@ -592,9 +654,16 @@ which replaces cuDSS's set-buffer-then-get protocol. """ function setparam!(solver::DirectSolver, name::AbstractString, value) if name == "info" - (value isa Integer && !(value isa Bool)) || - throw(InvalidValueError("invalid value $(repr(value)) for parameter \"info\"; expected an integer")) - solver.info = Int(value) + if value isa AbstractVector && solver.nbatch > 1 + length(value) == solver.nbatch && all(x -> x isa Integer && !(x isa Bool), value) || + throw(InvalidValueError("invalid value $(repr(value)) for parameter \"info\"; expected " * + "$(solver.nbatch) integers")) + solver.info .= value + else + (value isa Integer && !(value isa Bool)) || + throw(InvalidValueError("invalid value $(repr(value)) for parameter \"info\"; expected an integer")) + fill!(solver.info, Int(value)) + end elseif name == "schur_matrix" throw(NotSupportedError("the data parameter \"schur_matrix\" is not implemented yet (T20)")) elseif name == "pivot_sign" && value !== nothing && length(value) != size(solver, 1) @@ -623,6 +692,7 @@ The data parameters computed by the solver: | name | value | available after | | --- | --- | --- | | `"info"` | `Int`: `0`, or the original column (1-based) of the first non-positive pivot | always | +| `"ubatch_mask"` | the `"ubatch_mask"` set (or `nothing`) | always | | `"ir_n_steps"` | `Int`: refinement steps performed by the last `"solve"`/`"solve_refinement"` (≤ the configured `ir_n_steps`, fewer when `ir_tol` stopped early); the configured value before the first solve and after `setparam!(solver, "ir_n_steps", k)` | always | | `"lu_nnz"` | `Int64`: nonzeros of `L` (diagonal included, amalgamation zeros excluded) | analysis | | `"flops"` | `Float64`: factorization flops of the stored panels | analysis | @@ -635,6 +705,11 @@ The data parameters computed by the solver: | `"inertia"` | `Tuple{INT, INT}`: `(npos, nneg)` of D (after perturbation, so the inertia of `A + E`; read it with `"npivots"`); `(0, 0)` for complex symmetric `"S"`; Cholesky: `(number of positive pivots, 0)` | factorization | | `"pivot_stats"` | `NamedTuple` `(npos, nneg, nzero, nperturbed, n2x2)` of `Int64` (PLAN §1.7) | factorization | +Uniform batch (`nbatch > 1`): `"info"`, `"npivots"`, `"inertia"` and +`"pivot_stats"` are vectors with one entry per batch member, `"diag"` is the +`n · nbatch` vector of the members' diagonals one after the other; the +analysis outputs are shared by the members. + The pivot statistics are reduced on the device ([`reduce_stats!`](@ref)) and copied to the host when read (one synchronization). The reordering permutation after `"reordering"` alone is the ordering algorithm's; `"symbolic_factorization"` composes it with the supernodal @@ -645,7 +720,7 @@ one before the phase that computes it raises [`FactorizationError`](@ref). """ function getparam(solver::DirectSolver, name::AbstractString) spec = parameter_spec(name) - name == "info" && return solver.info + name == "info" && return _info_value(solver) name == "ir_n_steps" && solver.ir_steps >= 0 && return solver.ir_steps spec.status === :output || spec.status === :solver || return getparam(solver.options, name) name in SOLVER_OUTPUTS || @@ -698,40 +773,52 @@ function getparam!(buffer::AbstractVector, solver::DirectSolver, name::AbstractS return buffer end -@kernel function _factor_diag_kernel!(d, factor, super_ptr, front_ptr, front_nrows, ns) - s = @index(Global, Linear) - @inbounds if s <= ns +# work item (s, member k of nb): the diagonal of the panel of supernode s of member k into d[(k - 1) n + …] +@kernel function _factor_diag_kernel!(d, factor, super_ptr, front_ptr, front_nrows, ns, nb) + q = @index(Global, Linear) + s = (q - 1) % ns + 1 + k = (q - 1) ÷ ns + 1 + @inbounds if k <= nb + dk = _mview(d, k, nb) c0 = Int(super_ptr[s]) w = Int(super_ptr[s + 1]) - c0 f = Int(front_nrows[s]) - p0 = Int(front_ptr[s]) + p0 = Int(member_panels(front_ptr, k, nb)[s]) for j in 0:(w - 1) - d[c0 + j] = factor[p0 + j * f + j] + dk[c0 + j] = factor[p0 + j * f + j] end end end -# npivots / inertia / pivot_stats from the statistics reduced on the device (Cholesky: reduced on demand) +# npivots / inertia / pivot_stats from the statistics reduced on the device (Cholesky: reduced on demand), +# per member for a batch function _pivot_output(solver::DirectSolver{T, INT}, name) where {T, INT} S, N = solver.symbolic, solver.numeric - _is_ldlt_structure(S.structure) || reduce_stats!(N, S) - st = pivot_totals(N) - name == "npivots" && return INT(st.nperturbed) - name == "inertia" && return (INT(st.npos), INT(st.nneg)) - return st + if !_is_ldlt_structure(S.structure) + _set_members!(solver) + reduce_stats!(N, S) + end + totals = Array(N.totals) + out(st) = name == "npivots" ? INT(st.nperturbed) : name == "inertia" ? (INT(st.npos), INT(st.nneg)) : st + solver.nbatch == 1 && return out(pivot_totals(totals, 1, 1)) + return [out(pivot_totals(totals, solver.nbatch, k)) for k in 1:solver.nbatch] end -# diagonal of L (Cholesky) or D (LDLᵀ/LDLᴴ) in factor order, on the solver's backend -# (Cholesky: one launch, one work item per supernode) +# diagonal of L (Cholesky) or D (LDLᵀ/LDLᴴ) in factor order, on the solver's backend, the members of a batch one +# after the other (Cholesky: one launch, one work item per (supernode, member)) function _factor_diag(solver::DirectSolver{T}) where {T} S, N = solver.symbolic, solver.numeric - d = KernelAbstractions.zeros(solver.backend, T, S.n) + n, nb = S.n, solver.nbatch + d = KernelAbstractions.zeros(solver.backend, T, n * nb) if _is_ldlt_structure(S.structure) - copyto!(d, 1, N.d, 1, S.n) + for k in 1:nb + copyto!(d, (k - 1) * n + 1, N.d, (k - 1) * 2n + 1, n) + end return d end ns = nsupernodes(S) ns > 0 || return d - _factor_diag_kernel!(solver.backend, 64)(d, N.factor, S.super_ptr, S.front_ptr, S.front_nrows, ns; ndrange = ns) + _factor_diag_kernel!(solver.backend, 64)(d, N.factor, S.super_ptr, S.front_ptr, S.front_nrows, ns, nb; + ndrange = ns * nb) return d end From 88a453845b29bfe0ab073b5ce8e59164e3147b8b Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:39:54 +0000 Subject: [PATCH 2/4] T17: batched vendor pointers, solve dense path fix, uniform-batch tests Co-Authored-By: Claude Opus 5.5 --- ext/SparseDirectSolverCUDAExt.jl | 37 +++- src/dense/vendor.jl | 10 +- src/matrix.jl | 8 +- src/numeric/factorize.jl | 25 ++- src/numeric/storage.jl | 35 +++- src/solve/sweeps.jl | 89 +++++---- test/backends.jl | 29 +++ test/matrices.jl | 18 ++ test/ported/cudss_uniform_batch.jl | 134 ++++++++++++++ test/test_api.jl | 17 +- test/test_ubatch.jl | 288 +++++++++++++++++++++++++++++ test/utils.jl | 29 +++ 12 files changed, 663 insertions(+), 56 deletions(-) create mode 100644 test/ported/cudss_uniform_batch.jl create mode 100644 test/test_ubatch.jl diff --git a/ext/SparseDirectSolverCUDAExt.jl b/ext/SparseDirectSolverCUDAExt.jl index 159ef77..89ec4c9 100644 --- a/ext/SparseDirectSolverCUDAExt.jl +++ b/ext/SparseDirectSolverCUDAExt.jl @@ -3,7 +3,8 @@ module SparseDirectSolverCUDAExt # CUDA support: CSR adapters (T02) and the vendor dense bindings of the dense # layer (T03, PLAN §2.6): cuBLAS gemm/syrk/herk/trsm/strided-batched gemm/ # batched trsm, cuSOLVER potrf (also with a device info)/getrf/sytrf/potrfBatched, -# cuBLAS getrfBatched. +# cuBLAS getrfBatched; the batched trsm/potrf also on member pointers built once +# (uniform batch, T17). using SparseDirectSolver using SparseDirectSolver: CSR, DirectSolver, INDEX_ONE, INDEX_ZERO, InvalidValueError @@ -225,24 +226,36 @@ function _batch_pointers(A::StridedCuArray{T, 3}) where {T} return CuArray([base + (i - 1) * s for i in 1:size(A, 3)]) end +# built once per regime-C front of a uniform batch (T17): the numeric phase then allocates nothing +SDS.vendor_batch_pointers(A::StridedCuArray{T, 3}) where {T <: CuBlasT} = _batch_pointers(A) + for (fname, fname_64, elty) in ((:cublasStrsmBatched, :cublasStrsmBatched_64, :Float32), (:cublasDtrsmBatched, :cublasDtrsmBatched_64, :Float64), (:cublasCtrsmBatched, :cublasCtrsmBatched_64, :ComplexF32), (:cublasZtrsmBatched, :cublasZtrsmBatched_64, :ComplexF64)) + @eval function SDS.vendor_trsm_batched_ptrs!(side::Char, uplo::Char, trans::Char, diag::Char, α, m::Integer, + n::Integer, Aptrs::StridedCuArray{CuPtr{$elty}, 1}, lda::Integer, + Bptrs::StridedCuArray{CuPtr{$elty}, 1}, ldb::Integer, count::Integer) + count == 0 && return Bptrs + a = _blas_scalar($elty, α) + if cuBLAS.version() >= v"12.0" + cuBLAS.$fname_64(cuBLAS.handle(), side, uplo, trans, diag, m, n, a, Aptrs, max(1, lda), Bptrs, max(1, ldb), + count) + else + cuBLAS.$fname(cuBLAS.handle(), side, uplo, trans, diag, m, n, a, Aptrs, max(1, lda), Bptrs, max(1, ldb), + count) + end + return Bptrs + end @eval function SDS.vendor_trsm_batched!(side::Char, uplo::Char, trans::Char, diag::Char, α, A::StridedCuArray{$elty, 3}, B::StridedCuArray{$elty, 3}) m, n, nb = size(B) nb == 0 && return B - lda, ldb = max(1, stride(A, 2)), max(1, stride(B, 2)) GC.@preserve A B begin Aptrs = _batch_pointers(A) Bptrs = _batch_pointers(B) - a = _blas_scalar($elty, α) - if cuBLAS.version() >= v"12.0" - cuBLAS.$fname_64(cuBLAS.handle(), side, uplo, trans, diag, m, n, a, Aptrs, lda, Bptrs, ldb, nb) - else - cuBLAS.$fname(cuBLAS.handle(), side, uplo, trans, diag, m, n, a, Aptrs, lda, Bptrs, ldb, nb) - end + SDS.vendor_trsm_batched_ptrs!(side, uplo, trans, diag, α, m, n, Aptrs, stride(A, 2), Bptrs, stride(B, 2), + nb) CUDACore.unsafe_free!(Aptrs) CUDACore.unsafe_free!(Bptrs) end @@ -252,13 +265,19 @@ end for (fname, elty) in ((:cusolverDnSpotrfBatched, :Float32), (:cusolverDnDpotrfBatched, :Float64), (:cusolverDnCpotrfBatched, :ComplexF32), (:cusolverDnZpotrfBatched, :ComplexF64)) + @eval function SDS.vendor_potrf_batched_ptrs!(uplo::Char, n::Integer, Aptrs::StridedCuArray{CuPtr{$elty}, 1}, + lda::Integer, info::StridedCuArray{Cint, 1}, count::Integer) + count == 0 && return info + cuSOLVER.$fname(cuSOLVER.dense_handle(), uplo, n, Aptrs, max(1, lda), info, count) + return info + end @eval function SDS.vendor_potrf_batched!(uplo::Char, A::StridedCuArray{$elty, 3}, info::CuVector{Cint}) n = _check_square_batch(A) nb = size(A, 3) nb == 0 && return info GC.@preserve A begin Aptrs = _batch_pointers(A) - cuSOLVER.$fname(cuSOLVER.dense_handle(), uplo, n, Aptrs, max(1, stride(A, 2)), info, nb) + SDS.vendor_potrf_batched_ptrs!(uplo, n, Aptrs, stride(A, 2), info, nb) CUDACore.unsafe_free!(Aptrs) end return info diff --git a/src/dense/vendor.jl b/src/dense/vendor.jl index 612b551..e081737 100644 --- a/src/dense/vendor.jl +++ b/src/dense/vendor.jl @@ -9,7 +9,8 @@ const VENDOR_FUNCTIONS = (:vendor_gemm!, :vendor_syrk!, :vendor_herk!, :vendor_trsm!, :vendor_potrf!, :vendor_potrf_info!, :vendor_getrf!, :vendor_sytrf!, :vendor_gemm_strided_batched!, - :vendor_trsm_batched!, :vendor_potrf_batched!, :vendor_getrf_batched!) + :vendor_trsm_batched!, :vendor_potrf_batched!, :vendor_getrf_batched!, + :vendor_batch_pointers, :vendor_trsm_batched_ptrs!, :vendor_potrf_batched_ptrs!) for f in VENDOR_FUNCTIONS @eval $f(args...) = throw(NotSupportedError(string($(string(f)), " has no vendor binding for argument types ", @@ -29,6 +30,13 @@ end vendor_trsm_batched!(side, uplo, trans, diag, α, A, B) -> B # 3-D arrays vendor_potrf_batched!(uplo, A, info) -> info # 3-D A, Int32 info vendor_getrf_batched!(A, ipiv, info) -> info # 3-D A, Int32 ipiv matrix + vendor_batch_pointers(A) -> ptrs # device vector of the member pointers of 3-D A + vendor_trsm_batched_ptrs!(side, uplo, trans, diag, α, m, n, Aptrs, lda, Bptrs, ldb, count) + vendor_potrf_batched_ptrs!(uplo, n, Aptrs, lda, info, count) -> info # Int32 info, count entries + +The `_ptrs` variants take device vectors of member pointers built once by +`vendor_batch_pointers` (the uniform-batch regime C, T17: no allocation in the +numeric phase); `vendor_trsm_batched!`/`vendor_potrf_batched!` build them per call. Vendor dense routines behind the `impl = :vendor` path of the dense interface. The generic methods throw `NotSupportedError`; backends add methods for their diff --git a/src/matrix.jl b/src/matrix.jl index db45d96..4983444 100644 --- a/src/matrix.jl +++ b/src/matrix.jl @@ -263,7 +263,9 @@ end Point `desc` at `x` without copying (≅ `cudss_update(matrix, x)`). `x` must hold `nrows * ncols * nbatch` entries of type `T`: either as a strided vector, or with the descriptor's shape (`(nrows, ncols)`, `(ncols, nrows)` when transposed, plus -a trailing `nbatch` dimension for 3-D arrays). Anything else throws +a trailing `nbatch` dimension for 3-D arrays; a uniform batch also accepts the +`nrows × (ncols nbatch)` matrix with the same memory, e.g. `n × nbatch` for +`MatrixDescriptor(T, n; nbatch)`, as CUDSS.jl). Anything else throws `InvalidValueError`. """ function update!(desc::MatrixDescriptor{T, A}, x::AbstractArray) where {T, A} @@ -276,7 +278,9 @@ function update!(desc::MatrixDescriptor{T, A}, x::AbstractArray) where {T, A} if ndims(x) > 1 shape = desc.transposed ? (desc.ncols, desc.nrows) : (desc.nrows, desc.ncols) expected = ndims(x) == 2 ? shape : (shape..., desc.nbatch) - ndims(x) <= 3 && size(x) == expected || + batch_matrix = ndims(x) == 2 && desc.nbatch > 1 && !desc.transposed && + size(x) == (desc.nrows, desc.ncols * desc.nbatch) + ndims(x) <= 3 && (size(x) == expected || batch_matrix) || throw(InvalidValueError("array of size $(size(x)) does not match descriptor shape $expected")) end desc.data = x diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index becce03..7c58d8a 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -143,10 +143,22 @@ function _factor_panels_c_batched!(N::Numeric{T}, S::Symbolic, s::Int, j0::Int, P3 = reshape(view(N.factor, (off + 1):(off + f * w * count)), f, w, count) F11 = view(P3, 1:w, 1:w, :) ioff = (s - 1) * nb + k0 - 1 # info of front s, member k0 (interleaved) - _potrf_batched_impl!(p.bpotrf, F11, N.info, ioff) + ptrs = N.plan.vendor_ptrs[s] # (F11, F21) member pointers, or nothing + if p.bpotrf === :vendor && ptrs !== nothing + vendor_potrf_batched_ptrs!('L', w, view(ptrs[1], k0:(k0 + count - 1)), f, + view(N.info, (ioff + 1):(ioff + count)), count) + ka_chol_check_info!(N.info, ioff, F11) + else + _potrf_batched_impl!(p.bpotrf, F11, N.info, ioff) + end m > 0 || return nothing F21 = view(P3, (w + 1):f, 1:w, :) - trsm_strided_batched!('R', 'L', 'C', 'N', one(T), F11, F21; impl = p.btrsm) + if p.btrsm === :vendor && ptrs !== nothing + vendor_trsm_batched_ptrs!('R', 'L', 'C', 'N', one(T), m, w, view(ptrs[1], k0:(k0 + count - 1)), f, + view(ptrs[2], k0:(k0 + count - 1)), f, count) + else + trsm_strided_batched!('R', 'L', 'C', 'N', one(T), F11, F21; impl = p.btrsm) + end c0 = L.cb_ptr[s] c0 > 0 || return nothing # no parent: nothing to update C = reshape(view(N.work, 1:(m * m * count)), m, m, count) @@ -203,7 +215,14 @@ the analysis used `factorization_alg = "algo1"`, no vendor calls). Then Returns `info = 0` on success, else the original column of the smallest non-positive pivot of the factor (as [`ref_factorize!`](@ref)); fronts above a -failed one hold garbage. `opts.user_host_interrupt` is polled before every +failed one hold garbage. Uniform batch (`numeric.nbatch > 1`, PLAN §3.5): the +active members of `numeric` ([`set_members!`](@ref)) are factorized by the same +launches with the batch index as one more grid dimension (`nzval` holds the +values of every member, member after member), a regime-C front by +strided-batched `potrf`/`trsm`/`gemm` over each run of consecutive members +(vendor `potrf`/`trsm` on the member pointers built at allocation; per member +with `impl = :generic`); the return value is the first nonzero member `info` +(see [`member_info!`](@ref) for all of them). `opts.user_host_interrupt` is polled before every launch group (a host read): when set, [`InterruptedError`](@ref) is raised and the factor is incomplete. Structure `"SPD"` (real) or `"HPD"`. Nothing is allocated on the device; the panels are bitwise reproducible for a fixed diff --git a/src/numeric/storage.jl b/src/numeric/storage.jl index bd9e96e..a5d426a 100644 --- a/src/numeric/storage.jl +++ b/src/numeric/storage.jl @@ -24,6 +24,10 @@ buffer of the one `info` read per phase (one entry per batch member). Batch state (PLAN §3.5): `members_host[1:nact[]]` are the active members of the last phase (mirrored in `numeric.members`), `runs` the starts of their runs of consecutive members ([`member_runs!`](@ref), the strided batches of regime C). +`vendor_ptrs[s]` (regime-C front `s` of a batch on a backend with batched vendor +`potrf`/`trsm`, else `nothing`): the device vectors of member pointers +`(F11, F21)` of its `nbatch` panels, built once at allocation +([`vendor_batch_pointers`](@ref)). """ struct NumericPlan step_first::Vector{Int} @@ -40,6 +44,7 @@ struct NumericPlan members_host::Vector{Int32} nact::Base.RefValue{Int} runs::Vector{Int} + vendor_ptrs::Vector{Any} end function NumericPlan(S, nb::Integer = 1) @@ -81,7 +86,7 @@ function NumericPlan(S, nb::Integer = 1) end end return NumericPlan(first, last, maxchild, gfirst, glast, gmaxchild, gwidth, sfirst, slast, slocal, zeros(Int32, nb), - Int32.(1:nb), Ref(Int(nb)), [1, nb + 1]) + Int32.(1:nb), Ref(Int(nb)), [1, nb + 1], Any[nothing for _ in 1:ns]) end """ @@ -193,10 +198,34 @@ function allocate_numeric(S::Symbolic{INT}, ::Type{T}, backend::KernelAbstractio aux = KernelAbstractions.zeros(backend, T, nb) members = KernelAbstractions.allocate(backend, Int32, nb) copyto!(members, Int32.(1:nb)) + plan = NumericPlan(S, nb) + nb > 1 && _vendor_batch_pointers!(plan.vendor_ptrs, factor, S, nb) return Numeric{T, typeof(factor), typeof(stats), typeof(info), typeof(pivot_kind)}(factor, d, stack, work, stats, info, piv, pivot_kind, totals, - psign, aux, NumericPlan(S, nb), - members, nb) + psign, aux, plan, members, nb) +end + +# member pointers of the F11 and F21 blocks of the `nb` panels of every regime-C front, when the backend has batched +# vendor `potrf`/`trsm` (and the analysis uses vendor calls in regime C) +function _vendor_batch_pointers!(ptrs::Vector{Any}, factor::AbstractVector{T}, S::Symbolic, nb::Int) where {T} + sc, L = S.schedule, S.layout + _is_ldlt_structure(S.structure) && return ptrs # LDLᵀ/LDLᴴ: KA kernels in every regime + sc.vendor_c || return ptrs + caps = capabilities(KernelAbstractions.get_backend(factor), T) + caps.vendor_potrf_batched || caps.vendor_trsm_batched || return ptrs + for s in 1:nsupernodes(S) + takes_c_path(sc, s) || continue + f, w = sc.rows[s], sc.width[s] + off = panel_offset(L.panel_ptr, s, 1, nb) - 1 + P = reshape(view(factor, (off + 1):(off + f * w * nb)), f, w, nb) + try + ptrs[s] = (vendor_batch_pointers(view(P, 1:w, 1:w, :)), vendor_batch_pointers(view(P, (w + 1):f, 1:w, :))) + catch err + err isa NotSupportedError || rethrow() + return ptrs # no pointer binding: the per-call batched calls + end + end + return ptrs end """ diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index 297ff6a..0cc5787 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -436,8 +436,8 @@ its local pivot order `piv`; TRSV with `L11` on its columns of `Y`; GEMV with for k in 1:ctl[_SV_COUNT] if li == 1 v = _sv_node(ctl, k, list, subtree_nodes, Val(SUB), Val(false)) - _sv_setup!(ctl, v, super_ptr, rowptr, member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), - front_nrows, front_ncols, child_ptr) + P = member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch) + _sv_setup!(ctl, v, super_ptr, rowptr, P, front_nrows, front_ncols, child_ptr) end @synchronize if DET @@ -492,8 +492,8 @@ order before the front's local pivoting. for k in 1:ctl[_SV_COUNT] if li == 1 v = _sv_node(ctl, k, list, subtree_nodes, Val(SUB), Val(true)) - _sv_setup!(ctl, v, super_ptr, rowptr, member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch), - front_nrows, front_ncols, child_ptr) + P = member_panels(front_ptr, _bm_cmember(bm, _sv_rhs(G, count)), bm.nbatch) + _sv_setup!(ctl, v, super_ptr, rowptr, P, front_nrows, front_ncols, child_ptr) end @synchronize _sv_bwd_gather!(Y, factor, rowval, ctl, _sv_rhs(G, count), li, Val(WG), Val(CONJ)) @@ -706,35 +706,40 @@ function _piv_dense!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs:: return nothing end -# the blocks L11, L21 of the panels of front `s` of the `count` consecutive batch members starting at member slot -# `j0`, the front's rows of their columns of Y and the m rows of their columns of tmp: matrices for one member -# (the single-matrix layout, `count = 1`), `×count` strided batches otherwise -function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int) +# the blocks L11, L21 of the panel of front `s` of the batch member in slot `j`, the front's rows of the member's +# columns of Y and the m rows of its columns of tmp (the single-matrix path) +function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j::Int) sc = S.schedule f, w = sc.rows[s], sc.width[s] m = f - w - nb = N.nbatch - p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j0], nb) + p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j], N.nbatch) + c0 = S.partition.super_ptr[s] + cols = ((j - 1) * nrhs + 1):(j * nrhs) + P = reshape(view(N.factor, p0:(p0 + f * w - 1)), f, w) + return view(P, 1:w, 1:w), view(P, (w + 1):f, 1:w), view(ws.Y, c0:(c0 + w - 1), cols), view(ws.tmp, 1:m, cols) +end + +# the same for the `count` consecutive members in the slots `j0:(j0 + count - 1)`, as `×count` strided batches +function _dense_front_batch(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int) + sc = S.schedule + f, w = sc.rows[s], sc.width[s] + m = f - w + p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j0], N.nbatch) c0 = S.partition.super_ptr[s] cols = ((j0 - 1) * nrhs + 1):((j0 + count - 1) * nrhs) - if count == 1 - P = reshape(view(N.factor, p0:(p0 + f * w - 1)), f, w) - return f, w, view(P, 1:w, 1:w), view(P, (w + 1):f, 1:w), view(ws.Y, c0:(c0 + w - 1), cols), - view(ws.tmp, 1:m, cols) - end P = reshape(view(N.factor, p0:(p0 + f * w * count - 1)), f, w, count) Y3 = reshape(view(ws.Y, :, cols), size(ws.Y, 1), nrhs, count) T3 = reshape(view(ws.tmp, :, cols), size(ws.tmp, 1), nrhs, count) - return f, w, view(P, 1:w, 1:w, :), view(P, (w + 1):f, 1:w, :), view(Y3, c0:(c0 + w - 1), :, :), - view(T3, 1:m, :, :) + return view(P, 1:w, 1:w, :), view(P, (w + 1):f, 1:w, :), view(Y3, c0:(c0 + w - 1), :, :), view(T3, 1:m, :, :) end # the dense calls of front `s` for each run of consecutive active members: one strided-batched call per run, or -# (single matrix, `impl = :generic`) one call per member; `forward` selects `_dense_fwd!` or `_dense_bwd!` +# (single matrix, `impl = :generic`) one call per member function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple, forward::Bool, tr::Char) + m = S.schedule.rows[s] - S.schedule.width[s] if N.nbatch == 1 - _dense_call!(ws, S, N, s, nrhs, 1, 1, p, forward, tr, false) + _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, 1)...) return nothing end runs = N.plan.runs @@ -742,30 +747,19 @@ function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs: j0, j1 = runs[r], runs[r + 1] - 1 if p.btrsm === :none for j in j0:j1 - _dense_call!(ws, S, N, s, nrhs, j, 1, p, forward, tr, false) + _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, j)...) end else - _dense_call!(ws, S, N, s, nrhs, j0, j1 - j0 + 1, p, forward, tr, true) + _dense_call_batch!(p, forward, tr, m, _dense_front_batch(ws, S, N, s, nrhs, j0, j1 - j0 + 1)...) end end return nothing end -function _dense_call!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int, - p::NamedTuple, forward::Bool, tr::Char, batched::Bool) where {T} - f, w, L11, L21, Yc, Tm = _dense_front(ws, S, N, s, nrhs, j0, batched ? count : 1) - m = f - w +function _dense_call!(p::NamedTuple, forward::Bool, tr::Char, m::Int, L11::AbstractMatrix{T}, L21, Yc, Tm) where {T} if forward - if batched - trsm_strided_batched!('L', 'L', 'N', 'N', one(T), L11, Yc; impl = p.btrsm) - m > 0 && gemm_strided_batched!(Tm, L21, Yc, one(T), zero(T); impl = p.bgemm) - else - _trsm_impl!(p.trsm, 'L', 'L', 'N', 'N', one(T), L11, Yc) - m > 0 && _gemm_impl!(p.gemm, 'N', 'N', one(T), L21, Yc, zero(T), Tm) - end - elseif batched - m > 0 && gemm_strided_batched!(Yc, L21, Tm, -one(T), one(T); transA = tr, impl = p.bgemm) - trsm_strided_batched!('L', 'L', tr, 'N', one(T), L11, Yc; impl = p.btrsm) + _trsm_impl!(p.trsm, 'L', 'L', 'N', 'N', one(T), L11, Yc) + m > 0 && _gemm_impl!(p.gemm, 'N', 'N', one(T), L21, Yc, zero(T), Tm) else m > 0 && _gemm_impl!(p.gemm, tr, 'N', -one(T), L21, Tm, one(T), Yc) _trsm_impl!(p.trsm, 'L', 'L', tr, 'N', one(T), L11, Yc) @@ -773,6 +767,31 @@ function _dense_call!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nr return nothing end +function _dense_call_batch!(p::NamedTuple, forward::Bool, tr::Char, m::Int, L11::AbstractArray{T, 3}, L21, Yc, + Tm) where {T} + if forward + _trsm_batched!(p, 'N', L11, Yc) + m > 0 && gemm_strided_batched!(Tm, L21, Yc, one(T), zero(T); impl = p.bgemm) + else + m > 0 && gemm_strided_batched!(Yc, L21, Tm, -one(T), one(T); transA = tr, impl = p.bgemm) + _trsm_batched!(p, tr, L11, Yc) + end + return nothing +end + +# the strided-batched trsm of the solve: the vendor's pointer-array trsm would need member pointers of `Y` per call +# (a device allocation), so the vendor path solves member by member +function _trsm_batched!(p::NamedTuple, tr::Char, L11::AbstractArray{T, 3}, Yc::AbstractArray{T, 3}) where {T} + if p.btrsm === :vendor + for b in axes(Yc, 3) + _trsm_impl!(p.trsm, 'L', 'L', tr, 'N', one(T), view(L11, :, :, b), view(Yc, :, :, b)) + end + else + trsm_strided_batched!('L', 'L', tr, 'N', one(T), L11, Yc; impl = p.btrsm) + end + return nothing +end + function _fwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, det::Bool, p::NamedTuple, kind::Int) where {T} sc = S.schedule diff --git a/test/backends.jl b/test/backends.jl index 4a1a30b..b568867 100644 --- a/test/backends.jl +++ b/test/backends.jl @@ -106,6 +106,29 @@ api_matrix(backend, A::SparseMatrixCSC{<:Any, INT}) where {INT} = api_matrix(bac api_matrix(backend, A::SparseMatrixCSC, ::Type{INT}) where {INT} = to_device(backend, A, INT) api_matrix(::CPU, A::SparseMatrixCSC{T}, ::Type{INT}) where {T, INT} = CSR(SparseMatrixCSC{T, INT}(A)) +""" + api_batch_matrix(backend, members, view, INT = Int32) + +The uniform batch `members` (one sparsity pattern, see `batch_members`) as the +public API takes it on `backend` (T17): the `view` triangle's CSR pattern with +the values of every member one after the other, as a [`CSR`](@ref) on the CPU +backend and a vendor CSR matrix with a long `nzVal` on a GPU (as in CUDSS.jl). +""" +function api_batch_matrix(backend, members, view, ::Type{INT} = Int32) where {INT} + rowptr, colval, nzval = batch_csr(members, view, INT) + return api_csr(backend, rowptr, colval, nzval, size(members[1], 1)) +end + +""" + api_csr(backend, rowptr, colval, nzval, n) + +The square `n × n` CSR matrix of the host arrays `rowptr`, `colval` (one-based, +their element type is the index type) and `nzval` (a uniform batch when longer +than `colval`) as the public API takes it on `backend`: a [`CSR`](@ref) on the +CPU backend, a vendor CSR matrix on a GPU (T17). +""" +api_csr(::CPU, rowptr, colval, nzval, n) = CSR(copy(rowptr), copy(colval), copy(vec(nzval)), n, n) + """ device_allocated(backend, f) -> Union{Int, Missing} @@ -129,6 +152,9 @@ if CUDA_LOADED end to_host(A::CuSparseMatrixCSR) = host_csc(A.rowPtr, A.colVal, A.nzVal, size(A)...) index_eltype(A::CuSparseMatrixCSR) = eltype(A.rowPtr) + api_csr(::CUDABackend, rowptr::Vector{INT}, colval::Vector{INT}, nzval, n) where {INT} = + CuSparseMatrixCSR{eltype(nzval), INT}(CuVector{INT}(rowptr), CuVector{INT}(colval), CuVector(vec(nzval)), + (n, n)) end if AMDGPU_LOADED @@ -141,6 +167,9 @@ if AMDGPU_LOADED end to_host(A::ROCSparseMatrixCSR) = host_csc(A.rowPtr, A.colVal, A.nzVal, size(A)...) index_eltype(A::ROCSparseMatrixCSR) = eltype(A.rowPtr) + api_csr(::ROCBackend, rowptr::Vector{INT}, colval::Vector{INT}, nzval, n) where {INT} = + ROCSparseMatrixCSR{eltype(nzval), INT}(ROCVector{INT}(rowptr), ROCVector{INT}(colval), ROCVector(vec(nzval)), + (n, n)) end let gpus = String[] diff --git a/test/matrices.jl b/test/matrices.jl index 2e7b7a2..564ca79 100644 --- a/test/matrices.jl +++ b/test/matrices.jl @@ -301,6 +301,24 @@ function ubatch_example(::Type{T}) where {T} return (A = A, rowptr = rowptr, colval = colval, nzval = nzval, b = b, Λ = Λ, n = 3, nbatch = 3) end +""" + batch_members(A, nb; rng) -> Vector{SparseMatrixCSC} + +`nb` matrices with the sparsity pattern of `A` (a uniform batch, T17): member +1 is `A`, member `k > 1` is the congruence `Dₖ A Dₖ` with a random positive +diagonal `Dₖ` (entries in `[1/2, 2]`), so every member keeps the structure of +`A` (SPD/HPD, symmetric or Hermitian indefinite with the inertia of `A`, +complex symmetric). +""" +function batch_members(A::SparseMatrixCSC{T}, nb::Integer; rng::AbstractRNG = Random.default_rng()) where {T} + members = [copy(A)] + for _ in 2:nb + D = Diagonal(real(T)(0.5) .+ real(T)(1.5) .* rand(rng, real(T), size(A, 1))) + push!(members, sparse(D * A * D)) + end + return members +end + # --------------------------------------------------------------------------- # dense generators (T03) diff --git a/test/ported/cudss_uniform_batch.jl b/test/ported/cudss_uniform_batch.jl new file mode 100644 index 0000000..dbecc81 --- /dev/null +++ b/test/ported/cudss_uniform_batch.jl @@ -0,0 +1,134 @@ +# Port of CUDSS.jl test/test_uniform_batch_cudss.jl (T17): `uniform_batch_ldlt()` +# and `uniform_batch_cholesky()`, each with the cuDSS API (`DirectSolver` + +# `"ubatch_size"` + `MatrixDescriptor(T, n[, nrhs]; nbatch)`) and the generic +# API (`ldlt`/`cholesky` on a matrix whose values are longer than its pattern), +# strided and non-strided storage, views 'L', 'U', 'F'. +# +# Changes: `uniform_batch_lu()` is skipped (structure "G" is T19). The data are +# CUDSS.jl's, written once as `nnz × nbatch` matrices (the strided storage is +# their `vec`, exactly CUDSS.jl's strided vectors). Matrices come from +# `api_csr` (a `CSR` on the CPU backend); CUDSS.jl's `As_gpu.nzVal = …` is a new +# matrix on the same pattern. Residuals are checked per member with +# `relres ≤ tol(T)` (CUDSS.jl: `norm` of the absolute residuals `≤ √eps(R)`). + +# pattern and values (nnz × nbatch, then the refactorization values) of the 5×5 examples +function ported_ubatch_ldlt_data(::Type{T}, uplo) where {T} + real_t = T <: Real + if uplo == 'L' + rowptr, colval = [1, 2, 3, 6, 7, 9], [1, 2, 1, 2, 3, 4, 3, 5] + nz = real_t ? [4 2; 3 3; 1 1; 2 1; 5 6; 1 4; 1 2; 2 8] : + [4 2; 3 3; 1+im 1-im; 2-im 1+im; 5 6; 1 4; 1+im 2-im; 2 8] + new = real_t ? [-4 -2; -3 -3; -1 -1; -2 -1; -5 -6; -1 -4; -1 -2; -2 -8] : + [-4 -2; -3 -3; -1-im -1+im; -2+im -1-im; -5 -6; -1 -4; -1-im -2+im; -2 -8] + elseif uplo == 'U' + rowptr, colval = [1, 3, 5, 7, 8, 9], [1, 3, 2, 3, 3, 5, 4, 5] + nz = real_t ? [4 2; 1 1; 3 3; 2 1; 5 6; 1 2; 1 4; 2 8] : + [4 2; 1-im 1+im; 3 3; 2+im 1-im; 5 6; 1-im 2+im; 1 4; 2 8] + new = real_t ? [-4 -2; -1 -1; -3 -3; -2 -1; -5 -6; -1 -2; -1 -4; -2 -8] : + [-4 -2; -1+im -1-im; -3 -3; -2-im -1+im; -5 -6; -1+im -2-im; -1 -4; -2 -8] + else + rowptr, colval = [1, 3, 5, 9, 10, 12], [1, 3, 2, 3, 1, 2, 3, 5, 4, 3, 5] + nz = real_t ? [4 2; 1 1; 3 3; 2 1; 1 1; 2 1; 5 6; 1 2; 1 4; 1 2; 2 8] : + [4 2; 1-im 1+im; 3 3; 2+im 1-im; 1+im 1-im; 2-im 1+im; 5 6; 1-im 2+im; 1 4; 1+im 2-im; 2 8] + new = real_t ? [-4 -2; -1 -1; -3 -3; -2 -1; -1 -1; -2 -1; -5 -6; -1 -2; -1 -4; -1 -2; -2 -8] : + [-4 -2; -1+im -1-im; -3 -3; -2-im -1+im; -1-im -1+im; -2+im -1-im; -5 -6; -1+im -2-im; -1 -4; + -1-im -2+im; -2 -8] + end + return rowptr, colval, T.(nz), T.(new) +end + +function ported_ubatch_cholesky_data(::Type{T}, uplo) where {T} + if uplo == 'L' + rowptr, colval = [1, 2, 3, 6, 7, 9], [1, 2, 1, 2, 3, 4, 3, 5] + nz = [4 2; 3 3; 1 1; 2 1; 5 6; 1 2; 1 4; 2 8] + new = [8 6; 6 9; 2 3; 4 3; 10 18; 2 12; 2 6; 4 24] + elseif uplo == 'U' + rowptr, colval = [1, 3, 5, 7, 8, 9], [1, 3, 2, 3, 3, 5, 4, 5] + nz = [4 2; 1 1; 3 3; 2 1; 5 6; 1 2; 1 4; 2 8] + new = [8 6; 2 3; 6 9; 4 3; 10 18; 2 6; 2 12; 4 24] + else + rowptr, colval = [1, 3, 5, 9, 10, 12], [1, 3, 2, 3, 1, 2, 3, 5, 4, 3, 5] + nz = [4 2; 1 1; 3 3; 2 1; 1 1; 2 1; 5 6; 1 2; 1 4; 1 2; 2 8] + new = [8 6; 2 3; 6 9; 4 3; 2 3; 4 3; 10 18; 2 6; 2 12; 2 6; 4 24] + end + return rowptr, colval, T.(nz), T.(new) +end + +# the full Hermitian members of a batch on the host (CUDSS.jl: A + Aᴴ - Diagonal(A) for a triangle) +function ported_ubatch_members(rowptr, colval, nz::AbstractMatrix, uplo) + n = length(rowptr) - 1 + return map(1:size(nz, 2)) do k + A = host_csc(rowptr, colval, nz[:, k], n, n) + uplo == 'F' ? A : A + A' - Diagonal(A) + end +end + +function ported_ubatch(backend, ::Type{T}, ::Type{INT}, uplo, generic::Bool, strided::Bool, + cholesky_case::Bool) where {T, INT} + n, nbatch = 5, 2 + nrhs = cholesky_case ? 1 : 2 + rowptr, colval, nz, new_nz = cholesky_case ? ported_ubatch_cholesky_data(T, uplo) : ported_ubatch_ldlt_data(T, uplo) + rowptr, colval = Vector{INT}(rowptr), Vector{INT}(colval) + structure = cholesky_case ? "HPD" : "H" + if cholesky_case + B = T.(reshape([7, 12, 25, 4, 13, 13, 15, 29, 8, 14], n, nrhs, nbatch)) + new_B = B + elseif T <: Real + B = T.(cat([7 -7; 12 -12; 25 -25; 4 -4; 13 -13], [13 -13; 15 -15; 29 -29; 8 -8; 14 -14]; dims = 3)) + new_B = B[:, :, [2, 1]] + else + B = T.(cat([7+im -7+im; 12+im -12+im; 25+im -25+im; 4+im -4+im; 13+im -13+im], + [13-im -13-im; 15-im -15-im; 29-im -29-im; 8-im -8-im; 14-im -14-im]; dims = 3)) + new_B = B[:, :, [2, 1]] + end + # strided storage: vectors; otherwise an nnz × nbatch matrix of values and n × nbatch / n × nrhs × nbatch arrays + shape_x = strided ? (n * nrhs * nbatch,) : cholesky_case ? (n, nbatch) : (n, nrhs, nbatch) + nz_dev(x) = to_device(backend, strided ? vec(x) : x) + Bs_gpu = to_device(backend, reshape(B, shape_x)) + Xs_gpu = to_device(backend, zeros(T, shape_x)) + + if generic + As_gpu = api_csr(backend, rowptr, colval, vec(nz), n) + solver = cholesky_case ? cholesky(As_gpu; view = uplo) : ldlt(As_gpu; view = uplo) + @test solver.nbatch == nbatch + ldiv!(Xs_gpu, solver, Bs_gpu) + else + solver = DirectSolver(to_device(backend, rowptr), to_device(backend, colval), nz_dev(nz), structure, uplo) + setparam!(solver, "ubatch_size", nbatch) + Bs_desc = cholesky_case ? MatrixDescriptor(T, n; nbatch) : MatrixDescriptor(T, n, nrhs; nbatch) + Xs_desc = cholesky_case ? MatrixDescriptor(T, n; nbatch) : MatrixDescriptor(T, n, nrhs; nbatch) + update!(Bs_desc, Bs_gpu) + update!(Xs_desc, Xs_gpu) + execute!("analysis", solver, Xs_desc, Bs_desc) + execute!("factorization", solver, Xs_desc, Bs_desc; asynchronous = false) + execute!("solve", solver, Xs_desc, Bs_desc; asynchronous = false) + end + @test maximum(batch_relres(ported_ubatch_members(rowptr, colval, nz, uplo), to_host(Xs_gpu), B)) <= tol(T) + + new_Bs_gpu = to_device(backend, reshape(new_B, shape_x)) + if generic + As_new = api_csr(backend, rowptr, colval, vec(new_nz), n) + cholesky_case ? cholesky!(solver, As_new) : ldlt!(solver, As_new) + Xs_gpu .= new_Bs_gpu + ldiv!(solver, Xs_gpu) + else + update!(solver, to_device(backend, rowptr), to_device(backend, colval), nz_dev(new_nz)) + execute!("refactorization", solver, Xs_desc, Bs_desc; asynchronous = false) + update!(Bs_desc, new_Bs_gpu) + execute!("solve", solver, Xs_desc, Bs_desc; asynchronous = false) + end + @test maximum(batch_relres(ported_ubatch_members(rowptr, colval, new_nz, uplo), to_host(Xs_gpu), new_B)) <= + tol(T) + return nothing +end + +for (name, cholesky_case) in (("LDLᵀ and LDLᴴ", false), ("Cholesky", true)) + @testset "uniform batch $name ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + for generic in (false, true), strided in (false, true), uplo in ('L', 'U', 'F') + @testset "$(generic ? "Generic" : "cuDSS") API, strided = $strided, view = $uplo" begin + ported_ubatch(backend, T, INT, uplo, generic, strided, cholesky_case) + end + end + end +end diff --git a/test/test_api.jl b/test/test_api.jl index 897be39..2fb2ec3 100644 --- a/test/test_api.jl +++ b/test/test_api.jl @@ -53,8 +53,14 @@ end @test thrown(() -> DirectSolver(C, s, 'X')) isa InvalidValueError @test thrown(() -> DirectSolver(C, s, 'L'; index = 'Q')) isa InvalidValueError @test thrown(() -> DirectSolver(CSR(C.rowptr, C.colval, C.nzval, 60, 61), s, 'L')) isa InvalidValueError + # values of two members: a uniform batch (T17, test_ubatch.jl) batch = CSR(C.rowptr, C.colval, to_device(backend, repeat(Array(C.nzval), 2)), 60, 60) - @test thrown(() -> DirectSolver(batch, s, 'L')) isa NotSupportedError + sb = DirectSolver(batch, s, 'L') + @test sb.nbatch == 2 + execute!("analysis", sb, nothing, nothing) + execute!("factorization", sb, nothing, nothing) + bb = rand(T, 60, 2) + @test maximum(batch_relres([A, A], api_solve(backend, sb, bb), bb)) <= tol(T) end @testset "phases ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES @@ -141,18 +147,23 @@ end @test thrown(() -> execute!("solve", solver, nothing, b)) isa InvalidValueError @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 41)))) isa DimensionMismatch @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 40, 2)))) isa DimensionMismatch - @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 40, 2, 2)))) isa NotSupportedError + # an n × nrhs × nbatch array is a uniform-batch layout (T17): 4 columns for X's 1 + @test thrown(() -> execute!("solve", solver, x, to_device(backend, rand(T, 40, 2, 2)))) isa DimensionMismatch other = T <: Complex ? ComplexF64 === T ? ComplexF32 : ComplexF64 : Float64 === T ? Float32 : Float64 bo = to_device(backend, rand(other, 40)) @test thrown(() -> execute!("solve", solver, bo, bo)) isa InvalidValueError backend isa CPU || @test thrown(() -> execute!("solve", solver, rand(T, 40), rand(T, 40))) isa InvalidValueError # options and structures that are not implemented yet - for (name, value) in (("matching_alg", "algo1"), ("schur_mode", 1), ("ubatch_size", 2), ("schedule", "syncfree")) + for (name, value) in (("matching_alg", "algo1"), ("schur_mode", 1), ("schedule", "syncfree")) s2 = DirectSolver(api_matrix(backend, tril(A)), spd_structure(T), 'L') setparam!(s2, name, value) @test thrown(() -> execute!("analysis", s2, x, b)) isa NotSupportedError end + # uniform batches since T17: ubatch_size must match the batch size of the values (1 here) + s2 = DirectSolver(api_matrix(backend, tril(A)), spd_structure(T), 'L') + setparam!(s2, "ubatch_size", 2) + @test thrown(() -> execute!("analysis", s2, x, b)) isa InvalidValueError setparam!(solver, "solve_alg", "algo1") @test thrown(() -> execute!("solve", solver, x, b)) isa NotSupportedError setparam!(solver, "solve_alg", "default") diff --git a/test/test_ubatch.jl b/test/test_ubatch.jl new file mode 100644 index 0000000..c3f1971 --- /dev/null +++ b/test/test_ubatch.jl @@ -0,0 +1,288 @@ +# T17: uniform batch (PLAN §1.6, §3.5): one analysis for `nbatch` matrices with +# one sparsity pattern; factors, `info`, `inertia` and pivot statistics per +# member; the right-hand side layouts of CUDSS.jl (strided vector, +# `n × (nrhs nbatch)` matrix, `n × nrhs × nbatch` array, +# `MatrixDescriptor(T, n, nrhs; nbatch)`); `ubatch_index` and `ubatch_mask`. +# Members come from `batch_members` (test/matrices.jl), device matrices from +# `api_batch_matrix` (test/backends.jl). + +ubatch_spd(::Type{T}) where {T} = random_spd(T, 120, 0.03) +# "S": real symmetric or complex symmetric (not Hermitian) indefinite, diagonally dominant +ubatch_sym(::Type{T}) where {T} = random_symindef(T, 120, 0.03; hermitian = T <: Real) + +# the three paths of the numeric phase: regimes A + B (default), B only, C only (strided-batched dense calls) +ubatch_regimes() = (("default", Options()), ("regime B", Options(subtree_budgets = Int[])), + ("regime C", Options(subtree_budgets = Int[], factorization_alg = "algo2"))) + +function ubatch_solver(backend, members, structure, ::Type{INT} = Int32; view = 'L', opts = Options(), + params = ()) where {INT} + solver = DirectSolver(api_batch_matrix(backend, members, view, INT), structure, view) + solver.options = opts + for (name, value) in params + setparam!(solver, name, value) + end + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + return solver +end + +# solve on the device with the host right-hand side `B` (any batch layout) into a copy of `X0` +function ubatch_solve(backend, solver, B; X0 = zero(B)) + Bd = to_device(backend, B) + Xd = to_device(backend, X0) + execute!("solve", solver, Xd, Bd; asynchronous = false) + return to_host(Xd) +end + +@testset "batches of SPD and symmetric systems ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + for (structure, A) in ((spd_structure(T), ubatch_spd(T)), ("S", ubatch_sym(T))), nb in (1, 2, 3, 16, 64) + @testset "$structure nb = $nb" begin + members = batch_members(A, nb) + solver = ubatch_solver(backend, members, structure) + @test solver.nbatch == nb == nbatch(solver.A) + @test getparam(solver, "info") == (nb == 1 ? 0 : zeros(Int, nb)) + # the cuDSS multi-RHS batch bug (CUDSS.jl "Known issues"): every column of every member is right + for nrhs in (1, 2, 4) + B = rand(T, size(A, 1), nrhs, nb) + @test maximum(batch_relres(members, ubatch_solve(backend, solver, B), B)) <= tol(T) + end + end + end +end + +@testset "regimes A, B, C and refinement ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + nb, nrhs = 5, 3 + for (structure, A) in ((spd_structure(T), ubatch_spd(T)), ("S", ubatch_sym(T))), (rname, opts) in ubatch_regimes() + @testset "$structure $rname" begin + members = batch_members(A, nb) + solver = ubatch_solver(backend, members, structure; opts) + B = rand(T, size(A, 1), nrhs, nb) + @test maximum(batch_relres(members, ubatch_solve(backend, solver, B), B)) <= tol(T) + # refinement uses each member's values; the steps are counted once for the batch + setparam!(solver, "ir_n_steps", 2) + @test maximum(batch_relres(members, ubatch_solve(backend, solver, B), B)) <= tol(T) + @test getparam(solver, "ir_n_steps") == 2 + # the diagonal of every member, member after member + d = to_host(getparam(solver, "diag")) + @test length(d) == nb * size(A, 1) + setparam!(solver, "ubatch_index", 3) + execute!("refactorization", solver, nothing, nothing) # member 4 again: same factor + @test to_host(getparam(solver, "diag")) == d + end + end +end + +@testset "per-member inertia and pivot statistics ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + nb = 3 + A = random_symindef(T, 60, 0.05) # "S" (real) / "H" (complex) + members = batch_members(A, nb) + members[2] = -members[2] # the opposite inertia + solver = ubatch_solver(backend, members, sym_structure(T)) + inertia = getparam(solver, "inertia") + stats = getparam(solver, "pivot_stats") + npiv = getparam(solver, "npivots") + @test length(inertia) == length(stats) == length(npiv) == nb + for k in 1:nb + @test npiv[k] == 0 + @test inertia[k] == eigen_npos_nneg(members[k]) + @test (stats[k].npos, stats[k].nneg) == inertia[k] + single = DirectSolver(api_matrix(backend, tril(members[k]), Int32), sym_structure(T), 'L') + execute!("analysis", single, nothing, nothing) + execute!("factorization", single, nothing, nothing) + @test getparam(single, "pivot_stats") == stats[k] + end + @test inertia[1] == reverse(inertia[2]) + # Cholesky: npos = n for every member + spd = batch_members(ubatch_spd(T), nb) + s2 = ubatch_solver(backend, spd, spd_structure(T)) + @test getparam(s2, "inertia") == fill((Int32(120), Int32(0)), nb) + @test all(st -> st.nperturbed == 0 && st.npos == 120, getparam(s2, "pivot_stats")) +end + +@testset "ubatch_index and ubatch_mask ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + nb, nrhs = 4, 2 + for (structure, A) in ((spd_structure(T), ubatch_spd(T)), ("S", ubatch_sym(T))), (rname, opts) in ubatch_regimes() + @testset "$structure $rname" begin + n = size(A, 1) + members = batch_members(A, nb) + solver = ubatch_solver(backend, members, structure; opts) + S = solver.symbolic + snapshot() = [SDS.member_numeric(solver.numeric, S, k) for k in 1:nb] + same(a, b) = a.factor == b.factor && a.d == b.d && a.piv == b.piv && a.pivot_kind == b.pivot_kind && + a.stats == b.stats && a.info == b.info + before = snapshot() + fresh = batch_members(A, nb) # new values for every member + update!(solver, api_batch_matrix(backend, fresh, 'L')) + # ubatch_index = 2 (0-based): only member 3 is refactorized and solved + setparam!(solver, "ubatch_index", 2) + execute!("refactorization", solver, nothing, nothing) + after = snapshot() + @test all(k -> same(after[k], before[k]), (1, 2, 4)) + @test after[3].factor != before[3].factor + B = rand(T, n, nrhs, nb) + X0 = fill(T(7), n, nrhs, nb) + X = ubatch_solve(backend, solver, B; X0) + @test relres(fresh[3], X[:, :, 3], B[:, :, 3]) <= tol(T) + @test X[:, :, [1, 2, 4]] == X0[:, :, [1, 2, 4]] + # ubatch_mask: members 1 and 4 + setparam!(solver, "ubatch_index", -1) + setparam!(solver, "ubatch_mask", [1, 0, 0, 1]) + @test getparam(solver, "ubatch_mask") == [1, 0, 0, 1] + execute!("refactorization", solver, nothing, nothing) + masked = snapshot() + @test same(masked[2], after[2]) && same(masked[3], after[3]) + @test masked[1].factor != after[1].factor && masked[4].factor != after[4].factor + X = ubatch_solve(backend, solver, B; X0) + @test relres(fresh[1], X[:, :, 1], B[:, :, 1]) <= tol(T) + @test relres(fresh[4], X[:, :, 4], B[:, :, 4]) <= tol(T) + @test X[:, :, [2, 3]] == X0[:, :, [2, 3]] + # both: their intersection (member 4) + setparam!(solver, "ubatch_index", 3) + X = ubatch_solve(backend, solver, B; X0) + @test relres(fresh[4], X[:, :, 4], B[:, :, 4]) <= tol(T) + @test X[:, :, 1:3] == X0[:, :, 1:3] + # all members again: member 2 still has its old factor, the others the new one + setparam!(solver, "ubatch_index", -1) + setparam!(solver, "ubatch_mask", nothing) + X = ubatch_solve(backend, solver, B; X0) + @test relres(members[2], X[:, :, 2], B[:, :, 2]) <= tol(T) + @test all(k -> relres(fresh[k], X[:, :, k], B[:, :, k]) <= tol(T), (1, 3, 4)) + # an index or a mask outside the batch + setparam!(solver, "ubatch_index", nb) + @test thrown(() -> execute!("refactorization", solver, nothing, nothing)) isa InvalidValueError + setparam!(solver, "ubatch_index", -1) + setparam!(solver, "ubatch_mask", [1, 0]) + @test thrown(() -> execute!("solve", solver, to_device(backend, X), to_device(backend, B))) isa + InvalidValueError + setparam!(solver, "ubatch_mask", zeros(Int, nb)) + @test thrown(() -> execute!("refactorization", solver, nothing, nothing)) isa InvalidValueError + end + end +end + +@testset "right-hand side layouts ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + nb = 3 + for (structure, A) in ((spd_structure(T), ubatch_spd(T)), ("S", ubatch_sym(T))), nrhs in (1, 2) + n = size(A, 1) + members = batch_members(A, nb) + # the atomic forward sweep sums in a run-dependent order on GPUs (T16): compare bitwise deterministically + solver = ubatch_solver(backend, members, structure; params = (("deterministic_mode", 1),)) + B3 = rand(T, n, nrhs, nb) + X3 = ubatch_solve(backend, solver, B3) # n × nrhs × nbatch + @test maximum(batch_relres(members, X3, B3)) <= tol(T) + @test ubatch_solve(backend, solver, vec(B3)) == vec(X3) # strided (n nrhs nbatch,) + @test ubatch_solve(backend, solver, reshape(B3, n, nrhs * nb)) == reshape(X3, n, nrhs * nb) + # MatrixDescriptor(T, n, nrhs; nbatch) (≅ CudssMatrix(T, n, nrhs; nbatch)) on 3-D, matrix and strided data + for shape in ((n, nrhs, nb), (n, nrhs * nb), (n * nrhs * nb,)) + Bd = to_device(backend, reshape(B3, shape)) + Xd = to_device(backend, zeros(T, shape)) + bdesc = nrhs == 1 ? MatrixDescriptor(T, n; nbatch = nb) : MatrixDescriptor(T, n, nrhs; nbatch = nb) + xdesc = nrhs == 1 ? MatrixDescriptor(T, n; nbatch = nb) : MatrixDescriptor(T, n, nrhs; nbatch = nb) + update!(bdesc, Bd) + update!(xdesc, Xd) + execute!("solve", solver, xdesc, bdesc; asynchronous = false) + @test vec(to_host(Xd)) == vec(X3) + end + wrong = to_device(backend, zeros(T, n * nb, nrhs)) + @test thrown(() -> update!(MatrixDescriptor(T, n, nrhs; nbatch = nb), wrong)) isa InvalidValueError + # in place (X === B), the LinearAlgebra layer + Bd = to_device(backend, B3) + execute!("solve", solver, Bd, Bd) + @test to_host(Bd) == X3 + # the number of columns must split over the members; row-major batches are not supported + @test thrown(() -> ubatch_solve(backend, solver, rand(T, n, nb + 1))) isa DimensionMismatch + Bt = to_device(backend, rand(T, nrhs * nb, n)) + bt = MatrixDescriptor(Bt; transposed = true) + @test thrown(() -> execute!("solve", solver, bt, bt)) isa NotSupportedError + end +end + +@testset "a failed member ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + nb, j = 4, 37 + A = ubatch_spd(T) + n = size(A, 1) + for (r, (rname, opts)) in enumerate(ubatch_regimes()) + @testset "$rname" begin + members = batch_members(A, nb) + M = copy(members[3]) + M[j, j] = -M[j, j] # not positive definite: fails at column j + members[3] = M + solver = ubatch_solver(backend, members, spd_structure(T); opts) + @test getparam(solver, "info") == [0, 0, j, 0] + single = DirectSolver(api_matrix(backend, tril(M), Int32), spd_structure(T), 'L') + single.options = ubatch_regimes()[r][2] + execute!("analysis", single, nothing, nothing) + execute!("factorization", single, nothing, nothing) + @test getparam(single, "info") == j + # the other members solve, and only member 3 is refactorized after a fix + B = rand(T, n, 2, nb) + X = ubatch_solve(backend, solver, B) + @test all(k -> relres(members[k], X[:, :, k], B[:, :, k]) <= tol(T), (1, 2, 4)) + members[3] = batch_members(A, 2)[2] + update!(solver, api_batch_matrix(backend, members, 'L')) + setparam!(solver, "ubatch_mask", [0, 0, 1, 0]) + execute!("refactorization", solver, nothing, nothing) + @test getparam(solver, "info") == zeros(Int, nb) + setparam!(solver, "ubatch_mask", nothing) + X = ubatch_solve(backend, solver, B) + @test maximum(batch_relres(members, X, B)) <= tol(T) + end + end + # the LinearAlgebra layer reports every member's info + members = batch_members(A, 2) + M = copy(members[2]) + M[j, j] = -M[j, j] + members[2] = M + e = thrown(() -> cholesky(api_batch_matrix(backend, members, 'L'); view = 'L', check = true)) + @test e isa FactorizationError && e.info == [0, j] +end + +@testset "LinearAlgebra layer and errors ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + nb = 3 + A = ubatch_spd(T) + n = size(A, 1) + members = batch_members(A, nb) + Ad = api_batch_matrix(backend, members, 'L') + # auto-detection from length(nzval) ÷ length(colval) + F = cholesky(Ad; view = 'L') + @test F.nbatch == nb && getparam(F, "info") == zeros(Int, nb) + B = rand(T, n, nb) + Xd = to_device(backend, zeros(T, n, nb)) + ldiv!(Xd, F, to_device(backend, B)) + @test maximum(batch_relres(members, to_host(Xd), B)) <= tol(T) + @test maximum(batch_relres(members, to_host(F \ to_device(backend, B)), B)) <= tol(T) + B3 = rand(T, n, 2, nb) + @test maximum(batch_relres(members, to_host(F \ to_device(backend, B3)), B3)) <= tol(T) + @test thrown(() -> logabsdet(F)) isa NotSupportedError + fresh = batch_members(A, nb) + cholesky!(F, api_batch_matrix(backend, fresh, 'L')) + @test maximum(batch_relres(fresh, to_host(F \ to_device(backend, B)), B)) <= tol(T) + G = ldlt(api_batch_matrix(backend, batch_members(ubatch_sym(T), nb), 'L'); view = 'L') + @test G.nbatch == nb + # values given as an nnz × nbatch matrix + rowptr, colval, nzval = batch_csr(members, 'L', Int32) + Cm = CSR(to_device(backend, rowptr), to_device(backend, colval), to_device(backend, reshape(nzval, :, nb))) + @test nbatch(Cm) == nb + sm = DirectSolver(Cm, spd_structure(T), 'L') + execute!("analysis", sm, nothing, nothing) + execute!("factorization", sm, nothing, nothing) + @test maximum(batch_relres(members, ubatch_solve(backend, sm, B), B)) <= tol(T) + # ubatch_size: 0 (deduced) or the batch size + s = DirectSolver(Ad, spd_structure(T), 'L') + setparam!(s, "ubatch_size", nb + 1) + @test thrown(() -> execute!("analysis", s, nothing, nothing)) isa InvalidValueError + setparam!(s, "ubatch_size", nb) + execute!("analysis", s, nothing, nothing) + execute!("factorization", s, nothing, nothing) + @test getparam(s, "info") == zeros(Int, nb) + setparam!(s, "info", [1, 2, 3]) + @test getparam(s, "info") == [1, 2, 3] + setparam!(s, "info", 0) + @test getparam(s, "info") == zeros(Int, nb) + @test thrown(() -> setparam!(s, "info", [1, 2])) isa InvalidValueError + # update! keeps the batch size + @test thrown(() -> update!(s, api_batch_matrix(backend, members[1:2], 'L'))) isa InvalidValueError + @test occursin("nbatch = $nb", sprint(show, s)) +end diff --git a/test/utils.jl b/test/utils.jl index b3e50a5..68874ff 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -436,3 +436,32 @@ function madnlp_inertia_loop_device(backend, A::SparseMatrixCSC{T}, nh, nj, opts end return (false, iterations, solver, δw) end + +""" + batch_csr(members, view, INT = Int) -> (rowptr, colval, nzval) + +Host CSR arrays of the uniform batch `members` (same pattern, see +`batch_members`): the shared one-based pattern of the `view` triangle and the +values of every member one after the other (`length(nzval) = nnz · nbatch`), as +CUDSS.jl's strided uniform-batch storage (T17). +""" +function batch_csr(members::AbstractVector{<:SparseMatrixCSC{T}}, view, ::Type{INT} = Int) where {T, INT} + Cs = [SparseDirectSolver.CSR(SparseMatrixCSC{T, INT}(triangle_view(A, view))) for A in members] + all(C -> C.rowptr == Cs[1].rowptr && C.colval == Cs[1].colval, Cs) || + error("batch_csr: the members do not share a sparsity pattern") + return Cs[1].rowptr, Cs[1].colval, reduce(vcat, (C.nzval for C in Cs)) +end + +""" + batch_relres(members, X, B) -> Vector + +`relres` of every member `k` of a uniform batch: `X`, `B` are host arrays of +`n × nrhs × nbatch` entries in any shape (strided vector, `n × (nrhs nbatch)` +matrix, 3-D array), member `k` owning the `k`-th block of `n · nrhs` entries. +""" +function batch_relres(members, X::AbstractArray, B::AbstractArray) + n, nb = size(members[1], 1), length(members) + Xr = reshape(Array(X), n, :, nb) + Br = reshape(Array(B), n, :, nb) + return [relres(members[k], Xr[:, :, k], Br[:, :, k]) for k in 1:nb] +end From ae51ff6ac4240220cb162f728d78eb5534268a26 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:07:41 +0000 Subject: [PATCH 3/4] T17: contiguous batched solve views, test fix, docs Co-Authored-By: Claude Opus 5.5 --- src/solve/sweeps.jl | 12 +++++++++--- src/solver.jl | 2 ++ test/test_ubatch.jl | 2 +- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index 0cc5787..d5655ae 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -726,13 +726,19 @@ function _dense_front_batch(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, m = f - w p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j0], N.nbatch) c0 = S.partition.super_ptr[s] - cols = ((j0 - 1) * nrhs + 1):((j0 + count - 1) * nrhs) P = reshape(view(N.factor, p0:(p0 + f * w * count - 1)), f, w, count) - Y3 = reshape(view(ws.Y, :, cols), size(ws.Y, 1), nrhs, count) - T3 = reshape(view(ws.tmp, :, cols), size(ws.tmp, 1), nrhs, count) + Y3 = _column_batch(ws.Y, (j0 - 1) * nrhs, nrhs, count) + T3 = _column_batch(ws.tmp, (j0 - 1) * nrhs, nrhs, count) return view(P, 1:w, 1:w, :), view(P, (w + 1):f, 1:w, :), view(Y3, c0:(c0 + w - 1), :, :), view(T3, 1:m, :, :) end +# the columns `c0 + 1 : c0 + nrhs count` of the matrix `M` as an `size(M, 1) × nrhs × count` array (a contiguous +# range of its memory, so a strided device array for the vendor calls) +function _column_batch(M::AbstractMatrix, c0::Int, nrhs::Int, count::Int) + m = size(M, 1) + return reshape(view(vec(M), (c0 * m + 1):((c0 + nrhs * count) * m)), m, nrhs, count) +end + # the dense calls of front `s` for each run of consecutive active members: one strided-batched call per run, or # (single matrix, `impl = :generic`) one call per member function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple, forward::Bool, diff --git a/src/solver.jl b/src/solver.jl index b00ee21..e0ba86f 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -81,6 +81,8 @@ vector of `n · nrhs · nbatch` entries, an `n × (nrhs · nbatch)` matrix, an `"ubatch_index"` (0-based member, `-1` = all) and `"ubatch_mask"` (one 0/1 flag per member) restrict the factorization and solve phases to a subset of the members: the other members' factors and solutions are left untouched. +`"pivot_sign"` and the other options apply to every member; the 2×2 pivot +pairs of `"S"`/`"H"` (`"pivot_pairs"`) come from the first member's values. Fields: `A` (the current [`CSR`](@ref), re-pointed by [`update!`](@ref)), `structure`, `view` (as given), `options` ([`Options`](@ref), set through diff --git a/test/test_ubatch.jl b/test/test_ubatch.jl index c3f1971..dec7b8a 100644 --- a/test/test_ubatch.jl +++ b/test/test_ubatch.jl @@ -113,7 +113,7 @@ end same(a, b) = a.factor == b.factor && a.d == b.d && a.piv == b.piv && a.pivot_kind == b.pivot_kind && a.stats == b.stats && a.info == b.info before = snapshot() - fresh = batch_members(A, nb) # new values for every member + fresh = batch_members(A, nb + 1)[2:end] # new values for every member (member 1 of a draw is A) update!(solver, api_batch_matrix(backend, fresh, 'L')) # ubatch_index = 2 (0-based): only member 3 is refactorized and solved setparam!(solver, "ubatch_index", 2) From a7d210767a8e46bbe9ec55b553f2fdd5db9ace35 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:09:18 +0000 Subject: [PATCH 4/4] T17: report Co-Authored-By: Claude Opus 5.5 --- TASKS.md | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 74 insertions(+), 1 deletion(-) diff --git a/TASKS.md b/TASKS.md index 91253e6..8e18ba4 100644 --- a/TASKS.md +++ b/TASKS.md @@ -2272,7 +2272,7 @@ repeat it once T21 has landed. The bar is cuDSS with `"matching_alg" = --- -## T17 — Uniform batch (v1 cut line) `[ ]` +## T17 — Uniform batch (v1 cut line) `[!]` **Reads**: PLAN §1.6, §3.5, §1.3 (`ubatch_*`), §1.4 (`ubatch_mask`); `../CUDSS.jl/docs/src/uniform_batch.md`, `test/test_uniform_batch_cudss.jl`. @@ -2303,6 +2303,79 @@ repeat it once T21 has landed. The bar is cuDSS with `"matching_alg" = **This completes v1 (PLAN M0–M6).** The owner re-evaluates the plan before the tasks below are refined. +### Report + +- Status: [!] (done on the KA CPU backend; CUDA from CI; deviations below) +- What was built: + - `src/numeric/batch.jl` (new): `BatchMap` (active members, `nact`, `nbatch`, the launch's `first`, `nrhs` per + member; Adapt-able, `members = nothing` for single-matrix helper launches), member decoding of the flattened + group index (`_bm_node`/`_bm_slot`/`_bm_cmember`/`_bm_ucol`), member views (`_mview` member-major, + `_iview` interleaved), `MemberPanels` (the panel pointers of member `k`, passed to the kernels in place of + `front_ptr`), `panel_offset`, `active_members` (`ubatch_index` 0-based ∩ `ubatch_mask`), `member_runs!`. + - Layout (PLAN §3.5): factor panels interleaved per front (the `nb` panels of supernode `s` contiguous, so a + regime-C front of a run of consecutive members is one `f×w×count` strided batch); `info` interleaved per + entry (batched `potrf` statuses contiguous); `d`, `piv`, `pivot_kind`, update stack, regime-C workspace, + `stats`, `totals`, `aux`, `nzval` and the right-hand sides member after member. `nb = 1` is exactly the old + layout and the old code paths (regime C, solve dense path). + - `Numeric` gains `members` (device `Int32`) and `nbatch`; `NumericPlan` gains `members_host`, `nact`, `runs`, + `vendor_ptrs`. `allocate_numeric(S, T, backend; nbatch)`, `set_members!`, `batch_map`, `member_info!`, + `member_numeric` (host single-member copy), `panel(S, N, s, k)`, `pivot_totals(N, k)`. + - Every numeric kernel (zero/scatter/extend-add, `pack_add`, fused regime-B Cholesky, regime-A Cholesky, both + LDLᵀ kernels, `cholesky_stats`, `reduce_stats`, `abs_max`) and every solve kernel (permute/unpermute, + forward/backward sweeps, pivot order, diagonal, residual, residual norms, add-correction, `"diag"`) takes the + batch index as one more grid dimension, flattened into the 1-D ndrange (as the right-hand sides already were; + 2-D ndranges allocate on the KA CPU backend). The `BatchMap` replaces the `first` argument, so no kernel goes + over the 32-argument specialization limit; the launches stay allocation-free (CPU budgets unchanged). + - Regime C with `nb > 1`: per run of consecutive active members, strided-batched `potrf` (`potrf_batched_info!`, + new in the dense interface: vendor or KA `ka_potrf!`), `trsm_strided_batched!`, `gemm_strided_batched!` + (`F21 F21ᴴ` into `count` workspace blocks) and a batched `pack_add!` (forces the herk real diagonal); + `impl = :generic` falls back to the single-matrix calls per member. Vendor `potrf`/`trsm` run on member pointer + arrays built once at allocation (`vendor_batch_pointers`, `vendor_potrf_batched_ptrs!`, + `vendor_trsm_batched_ptrs!`, new vendor functions; CUDA extension implements them, the existing allocating + `vendor_*_batched!` now delegate to them). Solve regime C: strided-batched `gemm`; `trsm` batched on the KA + path, member by member on the vendor path (see deviations). + - `DirectSolver`: `nbatch` deduced from `length(nzval) ÷ nnz` or an `nnz × nbatch` matrix; `ubatch_size` must be + 0 or equal (`InvalidValueError`); `info` is a per-member vector (`getparam` returns `Int` for `nb = 1`, a vector + otherwise; `setparam!` accepts either); `npivots`, `inertia`, `pivot_stats` per member; `"diag"` = members' + diagonals one after the other; `ubatch_index`/`ubatch_mask` restrict factorization and solve phases (inactive + members' factors, `info` and solution columns untouched). Right-hand sides: strided vectors, `n × (nrhs nb)` + matrices, `n × nrhs × nb` arrays, `MatrixDescriptor(T, n[, nrhs]; nbatch)` (`update!` now also accepts the + `n × (ncols nbatch)` matrix of a batched descriptor, as `CudssMatrix(T, n; nbatch)` + `CuMatrix` in CUDSS.jl). + Refinement per member (residual with each member's values). LinearAlgebra layer: auto-detection through + `DirectSolver`, `ldiv!`/`\` on 3-D arrays, `check` reports the info vector; `logabsdet` of a batch raises + `NotSupportedError`. + - Tests: `test/test_ubatch.jl` (new), `test/ported/cudss_uniform_batch.jl` (new: LDLᴴ and Cholesky parts of + CUDSS.jl's `test_uniform_batch_cudss.jl`, both APIs, strided/non-strided, views L/U/F; LU skipped, T19); + shared helpers `batch_members` (test/matrices.jl), `batch_csr`, `batch_relres` (test/utils.jl), `api_csr`, + `api_batch_matrix` (test/backends.jl). +- Tests: `SDS_TEST_GPU=0 julia --project=. -e 'using Pkg; Pkg.test()'` (CPU): 62074 pass, 0 fail, 1 broken (the + T16 `@test_broken`), ≈ 14 min; the counts combine the full run (62050 pass, 24 failures of one wrong test + assertion in `test_ubatch`, member 1 of a fresh `batch_members` draw is `A` itself) and the rerun of + `test_ubatch` after fixing that assertion (1092 pass) — `test_ported`, `test_solve`, `test_refinement` were + rerun on the final code as well (9655 pass, 1 broken). `test_ubatch` alone: 1092 pass. CUDA/AMDGPU: pending + CI on the PR (the CUDA pointer-array bindings are untested locally). +- Measurements: KA CPU backend, 1 thread, `laplacian2d(40, 40)` (`n = 1600`), 64 members: batched + refactorization 209 ms / solve 36 ms vs 64 single solvers 212 ms / 39 ms — the same arithmetic, as expected on + the CPU; the batch's gain is the launch count on GPUs (one launch per group for all members instead of 64). +- Deviations from PLAN.md / this task: + - Solve phase, regime C: cuBLAS has no strided `trsm`, and pointer arrays of `Y` depend on `nrhs` (would be a + device allocation per solve), so the vendor path solves `trsm` member by member (the `gemm` stays strided + batched; the KA path is batched). The numeric phase uses the prebuilt pointer arrays as asked. + - Row-major (`transposed`) right-hand sides of a batch raise `NotSupportedError`; `logabsdet` of a batch too. + - The 2×2 pivot pairs of `"S"`/`"H"` (T14 analysis, value based) are computed from the first member's values. + - `test/test_api.jl`: three assertions that encoded "not implemented until T17" (`DirectSolver` on batched + values, a 3-D right-hand side, `ubatch_size = 2`) now assert the implemented behaviour (a solved batch, + `DimensionMismatch` for 4 columns against 1, `InvalidValueError` for a size mismatch). +- Open issues / follow-ups: + - The CUDA bindings `vendor_batch_pointers`/`vendor_*_batched_ptrs!` and the batched regime C on CUDA are + exercised only by CI. T23 (AMDGPU/oneAPI/Metal) needs the same three bindings, or falls back to the + allocating batched calls / KA (handled: a missing binding leaves `vendor_ptrs` empty). + - The GitHub App token expired before the last push: the final commits are local and published by the + workflow's "Publish the branch" step. +- Suggested plan changes: + - PLAN §3.5: record the layout choice (panels interleaved per front, `info` interleaved, the rest member-major) + and the solve-phase vendor `trsm` per member. + --- ## Post-v1 tasks (to be refined after the v1 evaluation)