diff --git a/src/SparseDirectSolver.jl b/src/SparseDirectSolver.jl index cd9e3d8..d358b05 100644 --- a/src/SparseDirectSolver.jl +++ b/src/SparseDirectSolver.jl @@ -49,6 +49,7 @@ include("symbolic/pairs.jl") include("symbolic/supernodes.jl") include("symbolic/schedule.jl") include("symbolic/layout.jl") +include("symbolic/device.jl") include("symbolic/maps.jl") # matching and scaling (PLAN §1.3, milestone M9), host diff --git a/src/solver.jl b/src/solver.jl index 3d70c81..9bfbf98 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -454,6 +454,10 @@ function _check_analysis_supported(solver::DirectSolver{T}) where {T} return nothing end +# the backend when the symbolic maps may be built on it (device.jl), else nothing +_device_backend(solver::DirectSolver) = + solver.backend isa KernelAbstractions.GPU ? solver.backend : nothing + function _reorder!(solver::DirectSolver{T}) where {T} _check_analysis_supported(solver) _poll_interrupt(solver.options.user_host_interrupt) @@ -468,7 +472,7 @@ function _reorder!(solver::DirectSolver{T}) where {T} solver.analysis_colval = m === nothing || m.symmetric ? solver.host_colval : matched_colval(m, solver.host_colval, A.index) P = SymmetricPattern(solver.host_rowptr, solver.analysis_colval, A.nrows, solver.structure; - view = _stored_view(solver), index = A.index) + view = _stored_view(solver), index = A.index, device = _device_backend(solver)) if solver.options.schur_mode == 1 # Schur complement mode: the Schur rows and columns last, no 2×2 pivot pairs flags = schur_flags(solver.options.user_schur_indices, A.nrows) @@ -535,7 +539,7 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} A = solver.A ord = solver.ordering P = SymmetricPattern(solver.host_rowptr, solver.analysis_colval, A.nrows, solver.structure; - view = _stored_view(solver), index = A.index) + view = _stored_view(solver), index = A.index, device = _device_backend(solver)) if solver.schur === nothing sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) schur = 0 @@ -548,7 +552,7 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} max_local = max_local_bytes(solver.backend)) layout = build_layout(sp, sc; ldlt = _is_ldlt_structure(solver.structure)) Sh = Symbolic(sp, sc, layout, solver.host_rowptr, solver.analysis_colval, A.nrows, solver.structure; - view = _stored_view(solver), index = A.index) + view = _stored_view(solver), index = A.index, device = _device_backend(solver)) nrhs = solver.workspace === nothing ? solver.nbatch : max_rhs(solver.workspace) Sd = adapt(solver.backend, Sh, INT) solver.host_symbolic = Sh diff --git a/src/symbolic/device.jl b/src/symbolic/device.jl new file mode 100644 index 0000000..5aa77a6 --- /dev/null +++ b/src/symbolic/device.jl @@ -0,0 +1,239 @@ +# Device construction of the symbolic maps (PLAN §2.3 step 6 hot spots). +# The host analysis stays the reference; a GPU solver routes three +# data-parallel pieces through KernelAbstractions kernels with element-exact +# results (the tests assert equality against the host originals): +# +# * `device_symmetric_pattern` — the both-triangle adjacency of +# `SymmetricPattern` by packed-key sort and deduplication; +# * `device_assembly_map` — `assembly_map` for the symmetric +# structures, one thread per CSR row, binary search per entry; +# * `device_group_amap` — `_group_amap` (unsigned) as a packed-key +# sort reproducing the host `(abs, neg, p)` order exactly. +# +# The packed 64-bit keys need `n`, `nnz` < 2^32 and factor offsets < 2^31; +# [`device_maps_supported`](@ref) checks the bounds and the callers fall back +# to the host route when it returns `false` (a larger-than-2^31-entry factor +# is a > 16 GiB panel buffer, beyond a single-GPU analysis anyway). + +""" + device_maps_supported(backend, structure, view, n, nnz, factor_len) -> Bool + +Whether the device symbolic-map routes apply: a GPU backend, a symmetric +structure (`"SPD"`, `"HPD"`, `"S"`, `"H"`), a triangular stored view, and +sizes that fit the packed 64-bit sort keys. +""" +function device_maps_supported(backend, structure, view, n::Integer, nnz::Integer, + factor_len::Integer) + backend isa KernelAbstractions.GPU || return false + s = _structure(structure) + (s == STRUCTURE_SPD || s == STRUCTURE_HPD || s == STRUCTURE_SYMMETRIC || + s == STRUCTURE_HERMITIAN) || return false + _matrix_view(view) == VIEW_FULL && return false + return n < 2^32 && nnz < 2^32 && factor_len < 2^31 +end + +# device vector of `x` as `I` +function _device_vec(backend, ::Type{I}, x) where {I <: Integer} + y = KernelAbstractions.allocate(backend, I, length(x)) + isempty(x) || copyto!(y, I.(x)) + return y +end + +# --------------------------------------------------------------------------- +# symmetric pattern + +@kernel function _pattern_keys_kernel!(keys, @Const(rowptr), @Const(colval), base::Int32, + upper::Bool) + r = @index(Global) + @inbounds for p in (rowptr[r] + (one(eltype(rowptr)) - base)):(rowptr[r + 1] - base) + c = Int64(colval[p]) + (1 - Int64(base)) + q = Int64(p) + # the host `_entry_filter`: diagonal dropped by the pattern, entries + # outside the declared triangle dropped by the filter + if c == r || (upper ? c < r : c > r) + keys[2q - 1] = typemax(UInt64) + keys[2q] = typemax(UInt64) + else + keys[2q - 1] = (UInt64(r) << 32) | UInt64(c) + keys[2q] = (UInt64(c) << 32) | UInt64(r) + end + end +end + +""" + device_symmetric_pattern(backend, rowptr, colval, n, view, index) -> (colptr, rowval) + +The `SymmetricPattern` arrays (1-based, sorted, deduplicated, no diagonal) +of the stored entries, built on `backend`. `rowptr`/`colval` are host or +device vectors with base `index`. +""" +function device_symmetric_pattern(backend, rowptr, colval, n::Integer, view, index) + m = length(colval) + base = _index_base(index) == INDEX_ONE ? Int32(1) : Int32(0) + upper = _matrix_view(view) == VIEW_UPPER + I = m + 1 <= typemax(Int32) && n + 1 <= typemax(Int32) ? Int32 : Int64 + drp = _device_vec(backend, I, Array(rowptr)) + dcv = _device_vec(backend, I, Array(colval)) + keys = KernelAbstractions.allocate(backend, UInt64, 2m) + n > 0 && _pattern_keys_kernel!(backend)(keys, drp, dcv, base, upper; ndrange = Int(n)) + KernelAbstractions.synchronize(backend) + isempty(keys) || sort!(keys) # empty: device sorts reject length 0 + hk = Array(keys) + nvalid = searchsortedfirst(hk, typemax(UInt64)) - 1 + colptr = zeros(Int, n + 1) + rowval = Vector{Int}(undef, nvalid) + colptr[1] = 1 + nv = 0 + prev = typemax(UInt64) + @inbounds for i in 1:nvalid + k = hk[i] + k == prev && continue # duplicate of a 'F' view + prev = k + nv += 1 + rowval[nv] = Int(k & 0xffffffff) + colptr[Int(k >> 32) + 1] += 1 + end + resize!(rowval, nv) + for j in 1:n + colptr[j + 1] += colptr[j] + end + return colptr, rowval +end + +# --------------------------------------------------------------------------- +# assembly map + +@kernel function _amap_kernel!(amap, @Const(rowptr), @Const(colval), @Const(iperm), + @Const(col2sn), @Const(snrowptr), @Const(snrowval), + @Const(superptr), @Const(panelptr), base::Int32, + upper::Bool, neg::Bool) + r = @index(Global) + T = eltype(amap) + @inbounds for p in (rowptr[r] + (one(eltype(rowptr)) - base)):(rowptr[r + 1] - base) + c = Int(colval[p]) + (1 - Int(base)) + if upper ? r <= c : r >= c + i = Int(iperm[r]) + j = Int(iperm[c]) + flip = i < j + if flip + i, j = j, i + end + s = Int(col2sn[j]) + lo = Int(snrowptr[s]) + hi = Int(snrowptr[s + 1]) - 1 + f = hi - lo + 1 + a = lo # binary search for i + b = hi + while a < b + mid = (a + b) >> 1 + if Int(snrowval[mid]) < i + a = mid + 1 + else + b = mid + end + end + off = Int(panelptr[s]) + (j - Int(superptr[s])) * f + (a - lo + 1) - 1 + amap[p] = T(flip && neg ? -off : off) + else + amap[p] = zero(T) + end + end +end + +""" + device_assembly_map(backend, sp, layout, rowptr, colval, n, structure; view, index) + -> amap::Vector{Int} + +[`assembly_map`](@ref) on `backend` for the symmetric structures with a +triangular stored view; identical output. The caller checks +[`device_maps_supported`](@ref). +""" +function device_assembly_map(backend, sp::SupernodePartition, layout::Layout, + rowptr, colval, n::Integer, structure; + view = VIEW_FULL, index = INDEX_ONE) + s = _structure(structure) + v = _matrix_view(view) + base = _index_base(index) == INDEX_ONE ? Int32(1) : Int32(0) + I = Int32 # bounds checked by the caller + drp = _device_vec(backend, I, Array(rowptr)) + dcv = _device_vec(backend, I, Array(colval)) + amap = KernelAbstractions.allocate(backend, I, length(colval)) + n > 0 && _amap_kernel!(backend)(amap, drp, dcv, _device_vec(backend, I, sp.iperm), + _device_vec(backend, I, sp.col2sn), + _device_vec(backend, I, sp.rowptr), + _device_vec(backend, I, sp.rowval), + _device_vec(backend, I, sp.super_ptr), + _device_vec(backend, I, layout.panel_ptr), + base, v == VIEW_UPPER, + _is_hermitian(s); ndrange = Int(n)) + KernelAbstractions.synchronize(backend) + return Vector{Int}(Array(amap)) +end + +# --------------------------------------------------------------------------- +# amap grouping + +@kernel function _group_keys_kernel!(keys, @Const(amap)) + p = @index(Global) + @inbounds begin + a = Int(amap[p]) + if a == 0 + keys[p] = typemax(UInt64) # dropped entries sort last + else + off = UInt64(a < 0 ? -a : a) + # the host order: (abs(amap[p]), amap[p] < 0, p) + keys[p] = (off << 33) | (UInt64(a < 0) << 32) | UInt64(p) + end + end +end + +@kernel function _group_count_kernel!(cnt, @Const(keys), @Const(panelptr), nvalid::Int64) + gi = @index(Global) + @inbounds if gi <= nvalid + off = Int((keys[gi] >> 33) & 0x7fffffff) + a = 1 # searchsortedlast(panel_ptr, off) + b = length(panelptr) + while a < b + mid = (a + b + 1) >> 1 + if Int(panelptr[mid]) <= off + a = mid + else + b = mid - 1 + end + end + KernelAbstractions.@atomic cnt[a + 1] += one(eltype(cnt)) + end +end + +""" + device_group_amap(backend, amap, layout, ns) -> (amap_ptr, order) + +`_group_amap` (unsigned) on `backend`: identical order and per-owner +counts. The signed (`"G"`) grouping stays on the host. +""" +function device_group_amap(backend, amap::Vector{Int}, layout::Layout, ns::Int) + m = length(amap) + damap = _device_vec(backend, Int32, amap) + keys = KernelAbstractions.allocate(backend, UInt64, m) + m > 0 && _group_keys_kernel!(backend)(keys, damap; ndrange = m) + KernelAbstractions.synchronize(backend) + isempty(keys) || sort!(keys) # empty: device sorts reject length 0 + hk = Array(keys) + nvalid = searchsortedfirst(hk, typemax(UInt64)) - 1 + order = Vector{Int}(undef, nvalid) + @inbounds for i in 1:nvalid + order[i] = Int(hk[i] & 0xffffffff) + end + cnt = KernelAbstractions.allocate(backend, Int32, ns + 1) + fill!(cnt, Int32(0)) + m > 0 && _group_count_kernel!(backend)(cnt, keys, _device_vec(backend, Int32, layout.panel_ptr), + Int64(nvalid); ndrange = m) + KernelAbstractions.synchronize(backend) + hc = Array(cnt) + amap_ptr = Vector{Int}(undef, ns + 1) + amap_ptr[1] = 1 + for s in 1:ns + amap_ptr[s + 1] = amap_ptr[s] + Int(hc[s + 1]) + end + return amap_ptr, order +end diff --git a/src/symbolic/maps.jl b/src/symbolic/maps.jl index 19f3013..7821b91 100644 --- a/src/symbolic/maps.jl +++ b/src/symbolic/maps.jl @@ -193,10 +193,19 @@ Assemble the host [`Symbolic`](@ref) from the supernodes, schedule and layout of a pattern and its user CSR arrays (only the pattern is read). """ function Symbolic(sp::SupernodePartition, sc::Schedule, layout::Layout, rowptr::AbstractVector{<:Integer}, - colval::AbstractVector{<:Integer}, n::Integer, structure; view = VIEW_FULL, index = INDEX_ONE) + colval::AbstractVector{<:Integer}, n::Integer, structure; view = VIEW_FULL, index = INDEX_ONE, + device = nothing) ns = nsupernodes(sp) - amap = assembly_map(sp, layout, rowptr, colval, n, structure; view, index) - amap_ptr, amap_src = _group_amap(amap, layout, ns, _structure(structure) == STRUCTURE_GENERAL) + signed = _structure(structure) == STRUCTURE_GENERAL + use_device = device !== nothing && !signed && + device_maps_supported(device, structure, view, n, length(colval), layout.factor_len) + if use_device + amap = device_assembly_map(device, sp, layout, rowptr, colval, n, structure; view, index) + amap_ptr, amap_src = device_group_amap(device, amap, layout, ns) + else + amap = assembly_map(sp, layout, rowptr, colval, n, structure; view, index) + amap_ptr, amap_src = _group_amap(amap, layout, ns, signed) + end child_ptr = zeros(Int, ns + 1) child_ptr[1] = 1 for s in 1:ns @@ -229,7 +238,7 @@ function Symbolic(sp::SupernodePartition, sc::Schedule, layout::Layout, rowptr:: end """ - symbolic_analysis(A::CSR, structure, view = 'F'; opts = Options(), T = eltype(A), max_local = typemax(Int)) + symbolic_analysis(A::CSR, structure, view = 'F'; opts = Options(), T = eltype(A), device = nothing, max_local = typemax(Int)) -> Symbolic{Int, Vector{Int}} The whole host analysis of PLAN §2.3: [`SymmetricPattern`](@ref), @@ -244,18 +253,18 @@ which bound the regime-A budgets ([`resolve_subtree_budgets`](@ref)); the default leaves them unbounded (host analyses, the CPU backend). """ function symbolic_analysis(A::CSR, structure, view = VIEW_FULL; opts::Options = Options(), T::Type = eltype(A), - max_local::Integer = typemax(Int)) + max_local::Integer = typemax(Int), device = nothing) A.nrows == A.ncols || throw(InvalidValueError("the matrix must be square, got $(A.nrows) × $(A.ncols)")) rowptr = Array(A.rowptr) colval = Array(A.colval) - P = SymmetricPattern(rowptr, colval, A.nrows, structure; view, index = A.index) + P = SymmetricPattern(rowptr, colval, A.nrows, structure; view, index = A.index, device) pp = analysis_pairs(P, rowptr, colval, A.nzval, A.nrows, structure, opts; view, index = A.index) ord = compute_ordering(P, opts; T, pp.pairs, pp.candidates) sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(structure), elsize = schedule_elsize(structure, T), max_local) layout = build_layout(sp, sc; ldlt = _is_ldlt_structure(_structure(structure))) - return Symbolic(sp, sc, layout, rowptr, colval, A.nrows, structure; view, index = A.index) + return Symbolic(sp, sc, layout, rowptr, colval, A.nrows, structure; view, index = A.index, device) end function _index_vector(backend::KernelAbstractions.Backend, x::AbstractVector{<:Integer}, ::Type{INT}, diff --git a/src/symbolic/pattern.jl b/src/symbolic/pattern.jl index 59d97c1..b88f895 100644 --- a/src/symbolic/pattern.jl +++ b/src/symbolic/pattern.jl @@ -24,7 +24,7 @@ present). SymmetricPattern(n, colptr, rowval) SymmetricPattern(A::CSR, structure; view = 'F') - SymmetricPattern(rowptr, colval, n, structure; view = 'F', index = 'O') + SymmetricPattern(rowptr, colval, n, structure; view = 'F', index = 'O', device = nothing) The first form wraps arrays that already satisfy the invariants (checked). The others build the pattern from a CSR pattern, following PLAN §2.3 step 1: @@ -148,9 +148,13 @@ function _selected_entries(rp, cv, n, structure, view) end function SymmetricPattern(rowptr::AbstractVector{<:Integer}, colval::AbstractVector{<:Integer}, n::Integer, - structure; view = VIEW_FULL, index = INDEX_ONE) + structure; view = VIEW_FULL, index = INDEX_ONE, device = nothing) s = _structure(structure) v = _matrix_view(view) + if device !== nothing && device_maps_supported(device, s, v, n, length(colval), 0) + colptr_d, rowval_d = device_symmetric_pattern(device, rowptr, colval, n, v, index) + return SymmetricPattern(Int(n), colptr_d, rowval_d) + end rp, cv = _host_pattern(rowptr, colval, n, index) rows, cols, _ = _selected_entries(rp, cv, n, s, v) # off-diagonal entries in both orientations, then column-sort and deduplicate diff --git a/test/test_symbolic_device.jl b/test/test_symbolic_device.jl new file mode 100644 index 0000000..b34a6fc --- /dev/null +++ b/test/test_symbolic_device.jl @@ -0,0 +1,68 @@ +# Device construction of the symbolic maps (src/symbolic/device.jl): the +# kernels must produce ELEMENT-IDENTICAL results to the host originals, on +# every backend (the CPU backend runs the same kernels through KA). + +random_pattern_csr(n, density, view) = begin + A = random_spd(Float64, n, density) + T = view == 'U' ? triu(A) : tril(A) + C = SDS.CSR(SparseMatrixCSC{Float64, Int}(T)) + (Vector{Int}(C.rowptr), Vector{Int}(C.colval)) +end + +@testset "device symbolic maps ($(backend_name(backend)))" for backend in BACKENDS + for (n, density) in ((0, 0.0), (1, 0.0), (40, 0.15), (300, 0.02), (1500, 0.004)), + view in ('L', 'U'), structure in ("SPD", "S") + + rowptr, colval = random_pattern_csr(n, density, view) + # pattern: host route vs device route + Ph = SDS.SymmetricPattern(rowptr, colval, n, structure; view, index = 'O') + Pd = SDS.SymmetricPattern(rowptr, colval, n, structure; view, index = 'O', device = backend) + @test Ph == Pd # GPU: device route; CPU: gated to host + n == 0 && continue + # maps: build the partition once, then both map routes on it + opts = Options() + ord = SDS.compute_ordering(Ph, opts; alg = :amd) + sp = SDS.supernode_partition(SDS.factor_pattern(Ph, ord), ord.perm, opts) + sc = SDS.build_schedule(sp, opts, Float64; + reserve = SDS.subtree_local_reserve(structure), + elsize = SDS.schedule_elsize(structure, Float64)) + layout = SDS.build_layout(sp, sc; ldlt = structure == "S") + amap_h = SDS.assembly_map(sp, layout, rowptr, colval, n, structure; view, index = 'O') + ptr_h, src_h = SDS._group_amap(amap_h, layout, SDS.nsupernodes(sp)) + # the kernels run on every backend here (the solver gate excludes the + # CPU backend only as a performance choice; correctness is universal) + amap_d = SDS.device_assembly_map(backend, sp, layout, rowptr, colval, n, structure; + view, index = 'O') + ptr_d, src_d = SDS.device_group_amap(backend, amap_h, layout, SDS.nsupernodes(sp)) + cp_d, rv_d = SDS.device_symmetric_pattern(backend, rowptr, colval, n, view, 'O') + @test length(amap_h) > 0 # the comparisons are not vacuous + @test amap_d == amap_h + @test ptr_d == ptr_h + @test src_d == src_h + @test cp_d == Ph.colptr + @test rv_d == Ph.rowval + end + # end-to-end: an analyzed solver's stored maps equal the host originals + # recomputed from its own state (on GPU backends the stored maps came + # through the device route; on the CPU backend this is a host identity) + for T in eltypes_among((Float64,)) + A = random_spd(T, 200, 0.03) + solver = DirectSolver(api_matrix(backend, tril(A), Int32), spd_structure(T), 'L') + execute!("analysis", solver, nothing, nothing) + Sh = solver.host_symbolic + amap_ref = SDS.assembly_map(Sh.partition, Sh.layout, solver.host_rowptr, solver.analysis_colval, + 200, spd_structure(T); view = 'L', index = solver.A.index) + @test length(amap_ref) > 0 + @test Sh.amap == amap_ref + ptr_ref, src_ref = SDS._group_amap(amap_ref, Sh.layout, SDS.nsupernodes(Sh.partition)) + @test Sh.amap_ptr == ptr_ref + @test Sh.amap_src == src_ref + # the factorization still works through the device-built maps + execute!("factorization", solver, nothing, nothing) + b = rand(T, 200) + bd = to_device(backend, b) + xd = similar(bd) + execute!("solve", solver, xd, bd; asynchronous = false) + @test relres(A, to_host(xd), b) <= tol(T) + end +end