symbolic: device construction of the pattern and assembly maps

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

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

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

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},

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