T24: ND partition-tree export/import and canonical supernode postorder (WIP)

README.md

@@ -53,13 +53,16 @@ CUDA and on AMDGPU (ROCm):
   entries on the diagonal, the symmetric structures a symmetrically scaled
   matrix (inertia preserved) with 2×2 pivot pairs from the matching;
   `perm_matching`, `scale_row`, `scale_col`;
+* the ND partition tree in the cuDSS encoding (`nd_partition_tree`,
+  `user_nd_partition_tree` with `user_perm`): an analysis under a stored
+  permutation skips the ordering and reproduces the stored analysis exactly;
 * the public API: `DirectSolver`, `execute!` with cuDSS phase strings, named
   phase wrappers, `update!`, `setparam!`/`getparam`, and the `LinearAlgebra`
   layer (`cholesky`, `cholesky!`, `ldlt`, `ldlt!`, `lu`, `lu!`, `ldiv!`, `\`, `logabsdet`),
   checked by the test suite of CUDSS.jl ported to this package; phase logging
   through `SDS_LOG_LEVEL`.
 
-Not there yet: non-uniform batches, ND partition-tree export, mixed precision,
+Not there yet: non-uniform batches, mixed precision,
 hybrid host memory, delayed pivots, and the oneAPI and Metal
 extensions. Unsupported structures, phases and parameters raise
 `NotSupportedError` rather than falling back silently. The remaining gap to

docs/src/man/parameters.md

@@ -50,6 +50,31 @@ The outputs and when they become available are tabulated in the docstring of
 [`getparam`](@ref). For a uniform batch the per-factorization outputs are
 vectors with one entry per member.
 
+## Caching an ordering
+
+The ordering is usually the most expensive part of the analysis. Store the
+permutation and the nested-dissection partition tree of one analysis and pass
+them to a later one (another process, the same sparsity pattern), as with
+cuDSS:
+
+```julia
+analyze!(solver)
+perm = getparam(solver, "perm_reorder_row")
+tree = getparam(solver, "nd_partition_tree")   # 2^nd_nlevels - 1 sizes, cuDSS encoding
+
+later = DirectSolver(A, "SPD", 'L')
+setparam!(later, "user_perm", perm)
+setparam!(later, "user_nd_partition_tree", tree)  # checked against perm
+analyze!(later)                                   # no ordering; same supernodes and lu_nnz
+```
+
+The supernode partition, `lu_nnz` and the schedule depend only on the
+elimination tree of the permutation, so the second analysis reproduces the
+first exactly. The tree is validated (sizes, and every column's dependencies
+inside its node's ancestors) and is optional: `user_perm` alone gives the same
+analysis. The 2×2 pivot pairs that `"S"`/`"H"` choose at analysis are not part
+of the encoding and are not applied under a `user_perm`.
+
 ## Differences from cuDSS
 
 * `"info"` reports the original column of the first failed pivot.
@@ -58,3 +83,6 @@ vectors with one entry per member.
 * `pivot_sign` chooses the sign of the perturbation of each row, so that a
   perturbed KKT system keeps the inertia an interior-point method expects.
 * `pivot_stats` returns `(npos, nneg, nzero, nperturbed, n2x2)` in one read.
+* `nd_partition_tree` is available after the symbolic factorization (cuDSS:
+  after the reordering) and exists for every ordering, not only nested
+  dissection.

src/SparseDirectSolver.jl

@@ -47,6 +47,7 @@ include("symbolic/etree.jl")
 include("symbolic/ordering.jl")
 include("symbolic/pairs.jl")
 include("symbolic/supernodes.jl")
+include("symbolic/ndtree.jl")
 include("symbolic/schedule.jl")
 include("symbolic/layout.jl")
 include("symbolic/maps.jl")

src/solver.jl

@@ -446,8 +446,6 @@ function _check_analysis_supported(solver::DirectSolver{T}) where {T}
                                 "schur_mode = 1"))
     opts.schur_mode == 0 || solver.nbatch == 1 ||
         throw(NotSupportedError("schur_mode = 1 is not supported for uniform batches"))
-    opts.user_nd_partition_tree === nothing ||
-        throw(NotSupportedError("user_nd_partition_tree is not implemented yet (T24)"))
     opts.schedule == SCHEDULE_SYNCFREE && throw(NotSupportedError("schedule = \"syncfree\" is not implemented yet"))
     opts.factor_precision === nothing || opts.factor_precision === real(T) ||
         throw(NotSupportedError("factor_precision = $(opts.factor_precision) for $T input is not implemented yet"))
@@ -490,13 +488,26 @@ function _reorder!(solver::DirectSolver{T}) where {T}
         solver.ordering = compute_ordering(P, solver.options; T, pp.pairs, pp.candidates)
         solver.schur = nothing
     end
+    _check_user_tree(solver.options, P, solver.ordering)
     solver.host_symbolic = solver.symbolic = solver.numeric = solver.workspace = solver.refinement = nothing
     solver.matching = nothing
     solver.stage = STAGE_REORDERED
     _log(LOG_INFO, () -> "reordering: n = $(A.nrows), nnz = $(nnz(A)), $(_elapsed(tic))")
     return solver
 end
 
+# "user_nd_partition_tree" (with "user_perm"): validated against the etree of the ordering it describes; the
+# analysis depends only on that etree (`supernode_partition`), so the tree is checked, not otherwise used
+function _check_user_tree(opts::Options, P::SymmetricPattern, ord::Ordering)
+    tree = opts.user_nd_partition_tree
+    tree === nothing && return nothing
+    opts.user_perm === nothing &&
+        throw(InvalidValueError("user_nd_partition_tree needs \"user_perm\" (the permutation the tree describes, " *
+                                "the \"perm_reorder_row\" of the analysis that exported it)"))
+    check_nd_partition_tree(tree, opts.nd_nlevels, etree(factor_pattern(P, ord), ord.perm))
+    return nothing
+end
+
 # the device matching state of the analysis `Sh` (host) / `Sd` (device maps)
 function _matching_state(solver::DirectSolver{T, INT}, Sh::Symbolic, Sd::Symbolic) where {T, INT}
     m = solver.host_matching
@@ -863,15 +874,12 @@ end
 # parameters
 
 # data parameters the solver computes (PLAN §1.4, §1.7)
-const SOLVER_OUTPUTS = ("lu_nnz", "flops", "nsuperpanels", "memory_estimates", "perm_reorder_row",
+const SOLVER_OUTPUTS = ("lu_nnz", "flops", "nsuperpanels", "memory_estimates", "nd_partition_tree", "perm_reorder_row",
                         "perm_reorder_col", "perm_row", "perm_col", "diag", "npivots", "inertia", "pivot_stats",
                         "schur_shape", "schur_matrix", "perm_matching", "scale_row", "scale_col")
 
 # task that provides the other computed data parameters
-function _output_task(name)
-    name == "nd_partition_tree" && return "T24"
-    return "M12"   # hybrid_device_memory_min
-end
+_output_task(name) = "M12"   # hybrid_device_memory_min
 
 """
     setparam!(solver::DirectSolver, name::String, value)
@@ -950,6 +958,7 @@ The data parameters computed by the solver:
 | `"flops"` | `Float64`: factorization flops of the stored panels | analysis |
 | `"nsuperpanels"` | `Int`: supernodes after amalgamation | analysis |
 | `"memory_estimates"` | `Vector{Int64}` (16 entries, see [`memory_estimates`](@ref)) | analysis |
+| `"nd_partition_tree"` | `Vector{Int}`: the partition tree of the analysis in the cuDSS encoding, `2^k - 1` column counts with `k = nd_nlevels` (at the time of reading), leaves first and root last, for the permutation `"perm_reorder_row"` ([`nd_partition_tree`](@ref)); with that permutation as `"user_perm"` and the tree as `"user_nd_partition_tree"` a later analysis of the same matrix skips the ordering and reproduces the supernodes, `"lu_nnz"` and the schedule exactly (the 2×2 pivot pairs of `"S"`/`"H"` are not part of the encoding and are not applied under a `"user_perm"`) | analysis |
 | `"perm_reorder_row"`, `"perm_reorder_col"` | `Vector{Int}`: the fill-reducing permutation, 1-based (`perm[k]` = original index of the `k`-th pivot) | reordering |
 | `"perm_row"`, `"perm_col"` | `Vector{Int}`: the final permutation of the factor (= the reordering for Cholesky and LDLᵀ/LDLᴴ); LU (`"G"`) after a factorization: `perm_row[k]` is the original row of factor row `k` (the reordering composed with the local row interchanges of batch member 1), `perm_col` the reordering (with matching: composed with the matching, `perm_matching[reordering]`), so `A[perm_row, perm_col] = L D U` (with matching `Dr[perm_row] A[perm_row, perm_col] Dc[perm_col] = L D U`, the scalings `Dr`, `Dc` of `"scale_row"`/`"scale_col"` as diagonal matrices in the original numbering). For a CSC input (`CuSparseMatrixCSC`, [`csr_of_transpose`](@ref)) `A` is the stored CSR matrix, the transpose of the matrix given | analysis |
 | `"perm_matching"` | `Vector{Int}`: the matching permutation ([`Matching`](@ref)): row `i` is matched to column `perm_matching[i]` (`"G"`: `A[:, perm_matching]` has the matched entries on its diagonal; symmetric structures: only its cycles are used, for the 2×2 pivot pairs) | reordering (`matching_alg ≠ "default"`) |
@@ -977,8 +986,7 @@ 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
 renumbering, and from then on both permutations are the one the factor uses.
-Data parameters of later tasks (`"nd_partition_tree"`, `"hybrid_device_memory_min"`) raise
-[`NotSupportedError`](@ref); reading one before the phase that computes it raises [`FactorizationError`](@ref), and
+The data parameter `"hybrid_device_memory_min"` (M12) raises [`NotSupportedError`](@ref); reading one before the phase that computes it raises [`FactorizationError`](@ref), and
 the matching outputs without matching [`InvalidValueError`](@ref).
 """
 function getparam(solver::DirectSolver, name::AbstractString)
@@ -1025,6 +1033,7 @@ function getparam(solver::DirectSolver, name::AbstractString)
     name == "flops" && return sp.flops
     name == "nsuperpanels" && return nsuperpanels(sp)
     name == "memory_estimates" && return _memory_estimates(solver)
+    name == "nd_partition_tree" && return nd_partition_tree(sp, solver.options.nd_nlevels)
     if name == "perm_row" && solver.structure == STRUCTURE_GENERAL && solver.stage >= STAGE_FACTORIZED
         piv = Array(view(solver.numeric.piv, 1:Sh.n))    # local row order of member 1
         return sp.perm[piv]

src/symbolic/etree.jl

@@ -36,18 +36,28 @@ function etree(P::SymmetricPattern, perm::AbstractVector{<:Integer})
 end
 
 """
-    postorder(parent) -> post::Vector{Int}
+    postorder(parent; key = nothing) -> post::Vector{Int}
 
 A postorder of the forest `parent` (`0` = root): `post[k]` is the node visited
-`k`-th; children are visited in increasing index order and every node comes
-after all its descendants. Non-recursive.
+`k`-th; the roots and the children of every node are visited in increasing
+index order, or in increasing `key` order when a vector `key` of distinct values
+is given, and every node comes after all its descendants. Non-recursive.
+
+With `key = perm` (the original column of every node of the etree of
+`A[perm, perm]`) the visit sequence, in original columns, depends only on the
+tree and not on which topological order `perm` lists it in: the supernode step
+([`supernode_partition`](@ref)) uses it so that an analysis under its own
+output permutation reproduces itself (T24).
 """
-function postorder(parent::AbstractVector{<:Integer})
+function postorder(parent::AbstractVector{<:Integer}; key::Union{Nothing, AbstractVector{<:Integer}} = nothing)
     n = length(parent)
+    key === nothing || length(key) == n ||
+        throw(InvalidValueError("postorder: key has length $(length(key)), expected $n"))
+    seq = key === nothing ? (1:n) : sortperm(key)     # nodes in visiting order among siblings
     head = zeros(Int, n)     # first child
     next = zeros(Int, n)     # next sibling
     # insert in reverse so the child lists come out in increasing order
-    for j in n:-1:1
+    for j in Iterators.reverse(seq)
         p = parent[j]
         p == 0 && continue
         next[j] = head[p]
@@ -56,7 +66,7 @@ function postorder(parent::AbstractVector{<:Integer})
     post = Vector{Int}(undef, n)
     stack = Int[]
     k = 0
-    for root in 1:n
+    for root in seq
         parent[root] == 0 || continue
         push!(stack, root)
         while !isempty(stack)

src/symbolic/ndtree.jl

@@ -0,0 +1,169 @@
+# Nested-dissection partition tree in the cuDSS encoding (PLAN §1.4, T24):
+# `"nd_partition_tree"` (export) and `"user_nd_partition_tree"` (import, with
+# `"user_perm"`), so an ordering can be cached between runs.
+#
+# Encoding (cuDSS `CUDSS_DATA_ND_PARTITION_TREE`): a complete binary tree with
+# `k = nd_nlevels` levels stored as a flat array of `2^k - 1` sizes in level
+# order from the bottom to the top: the `2^(k-1)` leaves first, the root last.
+# Every node owns a contiguous range of columns of the permuted matrix; the
+# columns of a node's subtree are its left subtree, its right subtree, then the
+# node itself (the separator), so a parent is eliminated after both children.
+#
+# Host only, plain `Int` arrays.
+
+"""
+    ND_TREE_MAX_LEVELS
+
+Largest `nd_nlevels` for which [`nd_partition_tree`](@ref) builds the tree
+(`2^24 - 1` entries); larger values raise [`InvalidValueError`](@ref).
+"""
+const ND_TREE_MAX_LEVELS = 24
+
+function _check_nd_levels(k::Integer)
+    1 <= k <= ND_TREE_MAX_LEVELS ||
+        throw(InvalidValueError("nd_nlevels = $k: the ND partition tree needs 1 ≤ nd_nlevels ≤ $ND_TREE_MAX_LEVELS"))
+    return Int(k)
+end
+
+# flat (1-based) position of the node with heap index `h` (root 1, children 2h and 2h + 1) in a tree of `k` levels
+@inline function _nd_flat(k::Int, h::Int)
+    d = 8 * sizeof(Int) - 1 - leading_zeros(h)          # depth, root 0
+    return (1 << k) - (1 << (d + 1)) + (h - (1 << d)) + 1
+end
+
+"""
+    nd_partition_tree(sp::SupernodePartition, nlevels) -> Vector{Int}
+
+The partition tree of the analysis `sp` in the cuDSS encoding with
+`k = nlevels` levels (`2^k - 1` sizes, leaves first, root last), for the
+permutation `sp.perm` (the `"perm_reorder_row"` of the analysis). It is read
+off the supernodal elimination tree, whose numbering is a postorder: a node's
+separator is the chain of supernodes at the top of its subtree forest down to
+the first supernode with two or more children (empty when the forest already
+has several roots), and the remaining subtrees are split into a left and a
+right child at the root boundary that best balances their columns. A node at
+the last level takes all the columns of its subtree forest. Nodes past the
+bottom of the elimination tree have size `0`. The tree satisfies the
+dependency rule that [`check_nd_partition_tree`](@ref) tests: every column's
+etree parent lies in the column's node or in one of its ancestors. It exists for
+every ordering, not only nested dissection.
+"""
+function nd_partition_tree(sp::SupernodePartition, nlevels::Integer)
+    k = _check_nd_levels(nlevels)
+    ns = nsupernodes(sp)
+    first = collect(1:ns)                 # first supernode of the subtree of s (postorder: first[s]:s)
+    sizes = ones(Int, ns)
+    for s in 1:ns
+        p = sp.snparent[s]
+        p == 0 && continue
+        p > s || throw(InvalidValueError("nd_partition_tree: the supernodal numbering is not topological"))
+        first[p] = min(first[p], first[s])
+        sizes[p] += sizes[s]
+    end
+    all(s -> s - first[s] + 1 == sizes[s], 1:ns) ||
+        throw(InvalidValueError("nd_partition_tree: the supernodal numbering is not a postorder"))
+    tree = zeros(Int, (1 << k) - 1)
+    cols(lo, hi) = lo > hi ? 0 : sp.super_ptr[hi + 1] - sp.super_ptr[lo]
+    # supernodes lo:hi (whole subtrees) -> node with heap index h
+    function build!(lo, hi, h)
+        idx = _nd_flat(k, h)
+        if h >= (1 << (k - 1))                       # last level: the whole forest
+            tree[idx] = cols(lo, hi)
+            return
+        end
+        top = hi
+        while hi >= lo && first[hi] == lo            # a single tree: its root joins the separator
+            hi -= 1
+        end
+        tree[idx] = cols(hi + 1, top)
+        m = lo                                       # left lo:(m - 1), right m:hi
+        if hi >= lo                                  # two or more roots: split at a root boundary
+            total = cols(lo, hi)
+            best = typemax(Int)
+            r = hi
+            while r >= lo
+                f = first[r]
+                if f > lo
+                    e = abs(2 * cols(lo, f - 1) - total)
+                    e < best && (best = e; m = f)
+                end
+                r = f - 1
+            end
+        end
+        build!(lo, m - 1, 2h)
+        build!(m, hi, 2h + 1)
+        return
+    end
+    build!(1, ns, 1)
+    return tree
+end
+
+"""
+    nd_tree_nodes(tree, n) -> node::Vector{Int}
+
+The node of every column `1:n` of the permuted matrix under the partition tree
+`tree` (cuDSS encoding, `length(tree) = 2^k - 1`), as a heap index (root `1`,
+children of `h` at `2h` and `2h + 1`). Raises [`InvalidValueError`](@ref) when
+the length is not `2^k - 1`, an entry is negative or the sizes do not add up to
+`n`.
+"""
+function nd_tree_nodes(tree::AbstractVector{<:Integer}, n::Integer)
+    len = length(tree)
+    (len >= 1 && ispow2(len + 1)) ||
+        throw(InvalidValueError("user_nd_partition_tree has $len entries, expected 2^k - 1 (k = nd_nlevels)"))
+    k = trailing_zeros(len + 1)
+    t = Vector{Int}(tree)
+    all(>=(0), t) || throw(InvalidValueError("user_nd_partition_tree has a negative entry"))
+    sum(t) == n ||
+        throw(InvalidValueError("user_nd_partition_tree sizes add up to $(sum(t)), the matrix has $n columns"))
+    nh = (1 << k) - 1
+    total = zeros(Int, nh)                           # columns of the subtree of heap node h
+    for h in nh:-1:1
+        total[h] = t[_nd_flat(k, h)] + (2h <= nh ? total[2h] + total[2h + 1] : 0)
+    end
+    node = Vector{Int}(undef, n)
+    start = zeros(Int, nh)                           # first column of the subtree of h
+    start[1] = 1
+    for h in 1:nh
+        if 2h <= nh
+            start[2h] = start[h]
+            start[2h + 1] = start[h] + total[2h]
+        end
+        s0 = start[h] + total[h] - t[_nd_flat(k, h)]   # the node's own columns come last
+        for j in s0:(start[h] + total[h] - 1)
+            node[j] = h
+        end
+    end
+    return node
+end
+
+"""
+    check_nd_partition_tree(tree, nlevels, parent) -> nothing
+
+Validate an imported partition tree (`"user_nd_partition_tree"`) against
+`nd_nlevels = nlevels` and the elimination tree `parent` of the matrix permuted
+by `"user_perm"` ([`etree`](@ref)): `2^nlevels - 1` non-negative sizes adding up
+to `n` ([`nd_tree_nodes`](@ref)), and every column's etree parent in the
+column's node or in an ancestor node, so that the dependencies follow the tree
+(cuDSS: "dependencies between the subsets should correspond to the tree
+structure"). Raises [`InvalidValueError`](@ref) otherwise.
+"""
+function check_nd_partition_tree(tree::AbstractVector{<:Integer}, nlevels::Integer, parent::AbstractVector{<:Integer})
+    k = _check_nd_levels(nlevels)
+    length(tree) == (1 << k) - 1 ||
+        throw(InvalidValueError("user_nd_partition_tree has $(length(tree)) entries, expected 2^nd_nlevels - 1 = " *
+                                "$((1 << k) - 1) (nd_nlevels = $k)"))
+    node = nd_tree_nodes(tree, length(parent))
+    for j in eachindex(parent)
+        p = parent[j]
+        p == 0 && continue
+        a, b = node[p], node[j]
+        while b > a
+            b >>= 1
+        end
+        b == a || throw(InvalidValueError("user_nd_partition_tree does not match user_perm: column $j of the " *
+                                          "permuted matrix updates column $p, which is not in its node or an " *
+                                          "ancestor node"))
+    end
+    return nothing
+end