diff --git a/TASKS.md b/TASKS.md index acc4a30..11e2607 100644 --- a/TASKS.md +++ b/TASKS.md @@ -1983,7 +1983,7 @@ whenever `atomic_add` is `false`. Close #36 in this task's PR. --- -## T15 — GPU LDLᵀ/LDLᴴ in regimes A/B/C, diagonal solve, statistics `[ ]` +## T15 — GPU LDLᵀ/LDLᴴ in regimes A/B/C, diagonal solve, statistics `[!]` **Reads**: PLAN §2.4, §3.3, §2.6 (sytrf), §1.4. @@ -2010,6 +2010,135 @@ whenever `atomic_add` is `false`. Close #36 in this task's PR. * MadNLP-style loop from T14 on the GPU; the ported CUDSS.jl symmetric tests pass for `T ∈ ELTYPES`. +### Report + +- Status: [!] (done on the CPU backend; regime C runs the KA kernel for `"S"` too, no vendor `sytrf`, see deviations; + CUDA from CI) +- What was built: + - `src/numeric/ldlt.jl` (new): the T14 in-front factorization as KA kernels, pivot for pivot. + `front_ldlt_kernel!` (regimes B and C): one workgroup (`LDLT_WORKGROUP = 128`) per front of a launch group, + fused zero/scatter/owner-pull extend-add (the T09 device functions), then the pivoted factorization of the + `w` fully-summed columns in place in the panel, then `F₂₂ ← F₂₂ − (L₂₁D)L₂₁ᴴ` into the front's packed block on + the update stack (`W = L₂₁D` formed on the fly from `d`/`pivot_kind`, 2×2 blocks included), unit diagonal. + `subtree_ldlt_kernel!` (regime A): the T11 subtree walk with packed fronts and blocks in `@localmem`, the same + pivot steps on the packed front, the block update in place in the trailing triangle. Per pivot step (a loop of + `w` iterations with a uniform trip count; a 2×2 pivot leaves the last ones idle): work item 1 chooses the pivot + (`_lt_choose` = the reference's `_choose_pivot`, `_accept_1x1`, `_accept_2x2`, `_best_1x1` on the lower + triangle), the workgroup swaps rows/columns (LAPACK-style symmetric interchange of the lower triangle), work + item 1 stores D (perturbation `±ε` with the `pivot_sign` of the original row, `psign[perm[piv[g]]]`), the pivot + kinds and the counts, the workgroup applies the rank-1/rank-2 update to the remaining fully-summed columns, then + scales the pivot columns into L. Pivot choice, control words and per-front counts live in `@localmem` and go + to `numeric.stats`. No atomics, deterministic. `_LDLTDevice{R, H}` carries `pivot_type`, `pivot_threshold`, + `pivot_epsilon` and the scaled flag (`H`: Hermitian/real vs complex symmetric as a type parameter). + `factorize_ldlt!(numeric, symbolic, nzval; opts)` (copies `pivot_sign` to `numeric.psign`, `abs_max!` for + `pivot_epsilon_alg = "algo1"`, one launch per regime-A class and per B/C launch group, `reduce_stats!`; returns 0, + no host synchronization), `reduce_stats!` (one workgroup, `@localmem` tree reduction of the per-front statistics + into `numeric.totals`), `pivot_totals` (reads them), `abs_max!`. `factorize!(N, S, nzval; impl, opts)` + dispatches to it for `"S"`/`"H"`. + - `Numeric` gains `totals` (6 `Int64`), `psign` (`n` `Int8`) and `aux` (one `T`: `max|aᵢⱼ|`); `allocate_numeric`, + `host_numeric`, `memory_estimates` slot 11 updated. `ref_ldlt!` fills `totals` too. + - `src/solve/sweeps.jl`: the forward kernel reorders a front's slice of `Y` by its local pivot order before its + TRSV, the backward kernel restores it after its TRSV (`Val(LDL)`; scratch: rows `1:w` of the front's + gather-list buffer in `U`, unused before), as `ref_solve_ldlt!` does; `Val(CONJ) = false` solves with `Lᵀ` for + complex symmetric `"S"`; the regime-C path gets a pivot-order kernel each way and `'T'` dense calls for complex + `"S"`. The unit diagonal is stored, so the TRSVs divide by 1 (exact) and need no change. `diagonal_sweep!` is + one launch (1×1 and 2×2 blocks of D); identity for Cholesky. + - API (`src/solver.jl`, `src/generic.jl`, CUDA extension): structures `"S"` (complex: complex symmetric) and + `"H"` accepted (`"G"` refused, T19); `execute!("solve_diag", …)`; `getparam` for `"npivots"` (`INT`), `"inertia"` + (`Tuple{INT, INT}`, as CUDSS.jl), `"pivot_stats"` (NamedTuple) from the device totals (Cholesky: reduced on + demand, `(npos, 0)`), `"diag"` = `d[1:n]` for LDLᵀ; `setparam!(solver, "pivot_sign", v)` checks the length + (host or device vector), the next factorization copies it to the device; `ldlt(A::CSR; view, check)`, + `ldlt!(solver, A)`, `logabsdet` for LDLᵀ/LDLᴴ (`det D`, 2×2 blocks, sign ±1 or complex); CUDA extension: + `ldlt(::CuSparseMatrixCSR)` and its `Symmetric`/`Hermitian` wrappers. + - Tests: `test/test_numeric_ldlt.jl` (new); `test/ported/cudss_execution.jl` and `cudss_generic.jl` gain the + "Symmetric -- Hermitian" parts; `cudss_solver.jl` runs `"S"`/`"H"` too and checks `"npivots"`/`"inertia"`; + helpers `ldlt_setup`, `d_error`, `ldlt_alloc_budget`, `madnlp_inertia_loop_device` in `test/utils.jl`, + `solve_alloc_budget(…; ldlt)`; `MADNLP_ORDERINGS` moved from `test_reference_ldlt.jl` to `test/matrices.jl`. +- Tests: `SDS_TEST_GPU=0 julia --project=. -e 'using Pkg; Pkg.test()'` (Julia 1.13.1, ubuntu-latest, CPU): + 59877 pass / 0 fail / 0 broken (11 min 54 s); test_numeric_ldlt 1822 pass (31 s in the suite, ~6 min + standalone, mostly compilation), test_ported 8104 pass. Under CI's flags (`coverage = true`, `--check-bounds=yes`) + test_numeric_ldlt + test_api: 3768 pass / 0 fail. Julia 1.10 not run locally (CI). Covered, per backend and `T ∈ ELTYPES` + (`"S"` real, `"H"` complex): + - Device factor vs `ref_ldlt!` on the T14 matrices (`random_symindef(400,0.01)`, `kkt(300,100,1e-2)`, + `kkt(300,100,1e-8)` default and interleaved ordering, `random_symindef(300,0.02)`, indefinite-H KKT δ = 0, + the 2×2 case `hessian_scale = 1e-3` interleaved, complex symmetric for complex `T`), each under three analyses + (default, regimes A+B+C with small budgets and C thresholds, C only): `piv == `, `pivot_kind ==`, `stats ==`, + totals `==`, D and panels within `panel_tol(T)` (100 eps) of the reference; inertia equals the eigenvalue + signs on the `n ≤ 300` cases; `relres ≤ tol(T)` for `nrhs ∈ {1, 5}`, atomic and deterministic sweeps + (`kkt(300,100,1e-8)` ComplexF32 after one refinement step, as in T14 / #66). + - Regime-A groups with Int64 maps; views `'U'`/`'F'` give `==` factors. + - `singular_block_matrix` (structurally zero and stored zero, both `pivot_sign` signs) under the default analysis + and with the zero pivot in a regime-C root front (`"algo2"`; asserted `snparent == 0`, `takes_c_path`): + `piv`, kinds `==` the reference, D within `panel_tol(T)`, perturbed D entries `==` the reference, perturbed entry `±ε` exactly at row `j`, inertia, consistent-RHS relres ≤ 1e-6; + `pivot_epsilon`, scaled `"algo1"` (device `max|a|`), `'N'`, `'D'`; wrong `pivot_sign` length. + - `'D'`/`'N'` on quasi-definite KKT (no 2×2, no perturbation, inertia `(nh, nj)`, `'N'` no interchange). + - Determinism: refactorization with new values and back gives `==` factor, D, `piv`; deterministic solves `==`; + no allocation beyond the budgets in `factorize!` (64 B measured) and the solve (0 B) on the CPU backend. + - Sweeps one at a time (permute, forward, diagonal, backward) `==` the solve, against `extract_ldlt`'s `L` and + `D`; every dense implementation on the regime-C path. + - MadNLP-style inertia loop through the handle layer (`DirectSolver`, `nzval` rewritten, `"refactorization"`, + `getparam("inertia")`, `"npivots"`) for the three T14 orderings: same iterations and `δw` as the reference loop, + final solve `relres ≤ tol(T)`. + - Public API for `INT ∈ INTTYPES`: phases, `npivots`/`inertia`/`pivot_stats` (types, values against the + reference, phase-order errors), `diag`, `getparam!`, the five sub-phases `==` `"solve"`, `logabsdet` against the + dense one (also complex symmetric), `pivot_sign` set through the solver (length error, device vector, perturbed + sign and inertia), `ldlt(::CSR)`. + - Ported: `cudss_execution` "Symmetric -- Hermitian" (views F/L/U × pivoting C/R/N, refactorization after a + diagonal shift), `cudss_generic` `ldlt`/`ldlt!`/`ldiv!`/`\`, `cudss_solver` for `"S"`/`"H"`; `T ∈ ELTYPES`, + `INT ∈ INTTYPES`. + - CUDA/AMDGPU: pending CI on the PR. CI fix round 1: the first CUDA run failed 12 checks in "perturbation and + pivot_sign (CUDA, Float32)", which compared the whole of D bitwise with the host reference; nvcc contracts the + Float32 Schur updates to FMA, so the non-perturbed pivots differ by rounding. These checks now compare D within + `panel_tol(T)` like every other device-vs-reference test, and the perturbed entries (the subject of the test) + still bitwise. CPU after the fix: 59973 pass / 0 fail / 0 broken. +- Measurements (ubuntu-latest KA CPU backend, Float64, best of 3; device = `factorize!` on the CPU backend, so only + launch counts and equality are meaningful, the GPU numbers are owed): + + ```text + matrix fronts A/B/C launches ref_ldlt! device ref_solve! sweep_solve! same piv + random_symindef(2000,0.002) 194/33/23 34 159 ms 292 ms 0.72 ms 1.45 ms yes + kkt(3000,1000,1e-8) 151/80/35 31 12.9 s 13.5 s 6.0 ms 6.8 ms yes + kkt(3000,1000,1e-8), C only 0/0/302 37 2.06 s 7.7 s 5.4 ms 5.8 ms yes + laplacian2d(100,100) - 3I 3305/27/1 20 13.3 ms 31.8 ms 2.1 ms 5.3 ms yes + ``` + + The 13 s of the KKT case are the reference's `_best_1x1` fallback scan (`O(w²f)` per front) on fronts with + unpaired dual rows, inherited by the device kernel (#75). `@allocated factorize!` (LDLᵀ, second call): 64 B. +- Deviations from PLAN.md / this task: + - **No vendor `sytrf` for regime C.** Regime-B and regime-C groups both run the fused KA `front_ldlt_kernel!` + (one workgroup per front), for `"S"` as for `"H"`. LAPACK/cuSOLVER `sytrf` factors `F₁₁` alone with its own + Bunch–Kaufman sequence (no `pivot_threshold` against the rows of `F₂₁`, its own tie-breaking, blocked updates), + so `D`/`pivot_kind` would not equal the reference, which this task's first test requires, and the "post-check + and KA redo" needs a copy of the assembled block (`sytrf` overwrites it). The large-front performance (vendor + BLAS-3 for `L₂₁` and the contribution block, or `sytrf` with a documented different sequence) is follow-up #75 + (T25). + - Regime B keeps the panel in global memory (pivot decisions, control words and counts in `@localmem`), the + pivot search is serial on work item 1 (the reference's logic verbatim); both are T25 items (#75). + - Bunch–Kaufman `α` is now `real(T)(α)` in the reference too (was a Float64 constant, so Float32 comparisons ran + in Float64): the device kernels need the same arithmetic, and Metal has no Float64. + - `getparam(solver, "inertia")` for SPD/HPD returns `(positive pivots, 0)` and `"npivots"` 0 (not specified). + - `Numeric` gains three more fields (`totals`, `psign`, `aux`; PLAN §3.2). + - The first regime-A kernel had 34 arguments; its KA CPU launch was not specialized (dynamic dispatch, ~400 B per + workgroup). The Hermitian flag is now a type parameter of the pivot parameters and `n` comes from + `length(perm)`, so the kernel has 32 arguments and allocates nothing (32 is what was verified, not the exact + limit). + - Test adaptations required by the new features: `test_api.jl` no longer expects `"S"`/`"H"` (now `"G"`), + `"solve_diag"`, `"npivots"`, `"inertia"`, `"pivot_stats"` to raise `NotSupportedError`, and asserts their values + instead; `cudss_solver.jl` likewise. No assertion was loosened. +- Open issues / follow-ups: + - #75 (found-by-agent): serial pivot search with the `O(w²f)` fallback scan, one workgroup per regime-C front, no + vendor `sytrf`, `F₁₁` not in local memory (T25). + - On a GPU, the pivot sequence can differ from the CPU reference if FMA contraction changes the rounding of a + near-tie in the pivot tests; the equality tests use well-separated generators, CI will tell. + - T16: `"solve_diag"` exists now; the complex `solve_mode` and refinement remain. `ir_n_steps` still warns. +- Suggested plan changes: + - PLAN §2.4 / §3.3: regime C of `"S"` uses the KA in-front kernel (same pivot sequence as the oracle); vendor + calls only for the BLAS-3 parts (`L₂₁`, contribution block), or `sytrf` when equality with the reference is + not required. + - PLAN §3.2: `Numeric` also holds `totals`, `psign`, `aux`. + - PLAN §2.7: KA CPU launches with many arguments (34 here) are not specialized and allocate; keep kernels at + ≤ 32 arguments. + --- ## T16 — Iterative refinement, solve sub-phases, `solve_mode`, interrupt, logging `[ ]` diff --git a/ext/SparseDirectSolverCUDAExt.jl b/ext/SparseDirectSolverCUDAExt.jl index 0fe417e..159ef77 100644 --- a/ext/SparseDirectSolverCUDAExt.jl +++ b/ext/SparseDirectSolverCUDAExt.jl @@ -109,6 +109,19 @@ LinearAlgebra.cholesky(A::Symmetric{T, <:CuSparseMatrixCSR{T}}, p::NoPivot = NoP LinearAlgebra.cholesky(A::Hermitian{T, <:CuSparseMatrixCSR{T}}, p::NoPivot = NoPivot(); check::Bool = false) where {T} = cholesky(CSR(A.data), p; view = A.uplo, check) +""" + ldlt(A::CuSparseMatrixCSR; view = 'F', check = false) -> DirectSolver + ldlt(Symmetric(A::CuSparseMatrixCSR)) / ldlt(Hermitian(A::CuSparseMatrixCSR)) + +LDLᵀ/LDLᴴ factorization of `A` on the GPU (≅ CUDSS.jl's `ldlt`); the wrappers +pass their `uplo` as the view. See `ldlt(::CSR)`. +""" +LinearAlgebra.ldlt(A::CuSparseMatrixCSR; view::Char = 'F', check::Bool = false) = ldlt(CSR(A); view, check) +LinearAlgebra.ldlt(A::Symmetric{T, <:CuSparseMatrixCSR{T}}; check::Bool = false) where {T <: Real} = + ldlt(CSR(A.data); view = A.uplo, check) +LinearAlgebra.ldlt(A::Hermitian{T, <:CuSparseMatrixCSR{T}}; check::Bool = false) where {T} = + ldlt(CSR(A.data); view = A.uplo, check) + # --------------------------------------------------------------------------- # vendor dense bindings (see `src/dense/vendor.jl` for the contracts) diff --git a/src/SparseDirectSolver.jl b/src/SparseDirectSolver.jl index f49cff9..5c77ab7 100644 --- a/src/SparseDirectSolver.jl +++ b/src/SparseDirectSolver.jl @@ -60,6 +60,7 @@ include("numeric/assembly.jl") include("numeric/front.jl") include("numeric/subtree.jl") include("numeric/factorize.jl") +include("numeric/ldlt.jl") include("numeric/extract.jl") # solve phase on the device (PLAN §2.5) diff --git a/src/generic.jl b/src/generic.jl index fea12da..5c9cfbc 100644 --- a/src/generic.jl +++ b/src/generic.jl @@ -1,5 +1,5 @@ # Julia-native layer (PLAN §1.5, §3.1) on top of the handle layer: `cholesky`, -# `cholesky!`, `ldiv!`, `\`, `logabsdet`, `diag`, `nnz`. Methods take the +# `cholesky!`, `ldlt`, `ldlt!`, `ldiv!`, `\`, `logabsdet`, `diag`, `nnz`. Methods take the # in-package `CSR`; the CUDA extension adds `CuSparseMatrixCSR` and its # `Symmetric`/`Hermitian` wrappers, as CUDSS.jl does. A `SparseMatrixCSC` is not # accepted by `cholesky` (that method belongs to CHOLMOD): wrap it with `CSR(A)`. @@ -33,6 +33,40 @@ function _check_info(solver::DirectSolver, check::Bool) return nothing end +""" + ldlt(A::CSR; view = 'F', check = false) -> DirectSolver + +LDLᵀ (real `T`, structure `"S"`) or LDLᴴ (complex `T`, `"H"`) factorization of +the sparse symmetric/Hermitian matrix `A` on its backend (≅ CUDSS.jl's `ldlt`): +a [`DirectSolver`](@ref) after `"analysis"` and `"factorization"` +(synchronized), with in-front Bunch–Kaufman pivoting and static perturbation of +tiny pivots (read `getparam(solver, "npivots")` and `"inertia"`). `view` as in +[`cholesky`](@ref). The factorization does not fail (`"info"` stays 0), so +`check` has no effect; it is accepted for symmetry with `cholesky`. +""" +function LinearAlgebra.ldlt(A::CSR{T}; view::Char = 'F', check::Bool = false) where {T} + solver = DirectSolver(A, T <: Real ? "S" : "H", view) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing; asynchronous = false) + _check_info(solver, check) + return solver +end + +""" + ldlt!(solver::DirectSolver, A; check = false) -> solver + +Factorize the new values `A` reusing the analysis of `solver`, as +[`cholesky!`](@ref) (`"factorization"` the first time after an analysis, +`"refactorization"` afterwards, synchronized). +""" +function LinearAlgebra.ldlt!(solver::DirectSolver, A; check::Bool = false) + update!(solver, A) + phase = solver.fresh_factorization ? "factorization" : "refactorization" + execute!(phase, solver, nothing, nothing; asynchronous = false) + _check_info(solver, check) + return solver +end + """ cholesky!(solver::DirectSolver, A; check = false) -> solver @@ -86,8 +120,8 @@ Base.:\(solver::DirectSolver{T}, B::Union{Array{Complex{T}, 1}, Array{Complex{T} """ diag(solver::DirectSolver) -Diagonal of the factor `L` in factor order (the `"diag"` data parameter), a -vector on the solver's backend. +Diagonal of the factor `L` (Cholesky) or of `D` (LDLᵀ/LDLᴴ) in factor order +(the `"diag"` data parameter), a vector on the solver's backend. """ LinearAlgebra.diag(solver::DirectSolver) = getparam(solver, "diag") @@ -101,17 +135,45 @@ SparseArrays.nnz(solver::DirectSolver) = Int(getparam(solver, "lu_nnz")) """ logabsdet(solver::DirectSolver) -> (log|det A|, sign) -From the diagonal of the Cholesky factor: `log|det A| = 2 Σ log Lₖₖ`, `sign = -one(T)` (`A` is positive definite). Raises [`FactorizationError`](@ref) when the +Cholesky: from the diagonal of the factor, `log|det A| = 2 Σ log Lₖₖ`, `sign = +one(T)` (`A` is positive definite). LDLᵀ/LDLᴴ: `det A = det D` (the +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`). """ function LinearAlgebra.logabsdet(solver::DirectSolver{T}) where {T} solver.stage >= STAGE_FACTORIZED && solver.info == 0 || throw(FactorizationError(solver.info, "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) end +function _logabsdet_ldlt(solver::DirectSolver{T}) where {T} + N = solver.numeric + n = size(solver, 1) + d, kind = Array(N.d), Array(N.pivot_kind) + herm = !(T <: Complex) || solver.structure == STRUCTURE_HERMITIAN + la, sgn = zero(real(T)), one(T) + k = 1 + while k <= n + if kind[k] == PIVOT_KIND_2X2_FIRST + b = d[n + k] + x = d[k] * d[k + 1] - (herm ? conj(b) : b) * b + herm && (x = T(real(x))) + k += 2 + else + x = d[k] + k += 1 + end + la += log(abs(x)) + sgn *= herm ? T(sign(real(x))) : x / abs(x) + end + return la, sgn +end + """ logdet(solver::DirectSolver) diff --git a/src/numeric/extract.jl b/src/numeric/extract.jl index e5608d0..e3c1db6 100644 --- a/src/numeric/extract.jl +++ b/src/numeric/extract.jl @@ -4,7 +4,8 @@ """ host_numeric(numeric) -> Numeric{T, Vector{T}} -Host copy of `numeric` (panels, D, update stack, workspace, statistics, status, pivot order and kinds; the +Host copy of `numeric` (panels, D, update stack, workspace, statistics and totals, status, pivot order, kinds and sign +requests, `aux`; the plan is shared), usable by [`ref_solve!`](@ref) and [`extract_L`](@ref). A host `numeric` is returned as is. """ @@ -13,7 +14,8 @@ function host_numeric(N::Numeric{T}) where {T} factor, d, stack, work = Array(N.factor), Array(N.d), Array(N.stack), Array(N.work) 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, N.plan) + pivot_kind, Array(N.totals), + Array(N.psign), Array(N.aux), N.plan) end """ diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index d11605d..ce1bbf1 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -61,7 +61,7 @@ function cholesky_stats!(N::Numeric, S::Symbolic) end function _check_numeric(N::Numeric{T}, S::Symbolic, nzval::AbstractVector) where {T} - _check_reference_cholesky(S, T) + _is_ldlt_structure(S.structure) ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) sizeof(T) <= S.elsize || throw(InvalidValueError("the analysis was built for $(S.elsize)-byte elements, the numeric storage has " * "$(sizeof(T))-byte $T")) @@ -70,7 +70,7 @@ function _check_numeric(N::Numeric{T}, S::Symbolic, nzval::AbstractVector) where length(N.factor) == S.layout.factor_len && length(N.stack) == S.layout.stack_len && length(N.info) == nsupernodes(S) + 1 || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) - length(N.work) == S.layout.work_len || + length(N.work) == S.layout.work_len && length(N.piv) == 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) @@ -117,10 +117,14 @@ function _factor_panel_c!(factor::AbstractVector{T}, stack::AbstractVector{T}, w return nothing end +_check_ldlt_eltype(::Type{T}) where {T} = T <: LinearAlgebra.BlasFloat || + throw(NotSupportedError("LDLᵀ/LDLᴴ supports BLAS element types only, got $T")) + """ - factorize!(numeric, symbolic, nzval; impl = :auto) -> info::Int + factorize!(numeric, symbolic, nzval; impl = :auto, opts = Options()) -> info::Int -Multifrontal Cholesky `P A Pᵀ = L Lᴴ` on the device: `nzval` are the stored +Structures `"S"`/`"H"`: [`factorize_ldlt!`](@ref) with the pivoting options +`opts` (`impl` is not used). Structures `"SPD"`/`"HPD"`: multifrontal Cholesky `P A Pᵀ = L Lᴴ` on the device: `nzval` are the stored values of A (same pattern, view and index base as the analysis) on the backend of `numeric`, and `symbolic` has its maps on that backend (`adapt(backend, symbolic, INT)`). For every launch group of the schedule, @@ -145,7 +149,9 @@ failed one hold garbage. Structure `"SPD"` (real) or `"HPD"`. Nothing is allocated on the device; the panels are bitwise reproducible for a fixed `impl` and backend (regimes A and B are deterministic by construction). """ -function factorize!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; impl::Symbol = :auto) where {T} +function factorize!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; impl::Symbol = :auto, + opts::Options = Options()) where {T} + _is_ldlt_structure(S.structure) && return factorize_ldlt!(N, S, nzval; opts) _check_numeric(N, S, nzval) p = _front_impls(N, S, impl) plan = N.plan @@ -174,4 +180,5 @@ function factorize!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; impl::Sym return k == 0 ? 0 : S.partition.perm[k] end -factorize!(N::Numeric, S::Symbolic, A::CSR; impl::Symbol = :auto) = factorize!(N, S, vec(A.nzval); impl) +factorize!(N::Numeric, S::Symbolic, A::CSR; impl::Symbol = :auto, opts::Options = Options()) = + factorize!(N, S, vec(A.nzval); impl, opts) diff --git a/src/numeric/ldlt.jl b/src/numeric/ldlt.jl new file mode 100644 index 0000000..bdd6e3e --- /dev/null +++ b/src/numeric/ldlt.jl @@ -0,0 +1,865 @@ +# Device LDLᵀ/LDLᴴ (PLAN §2.4, §3.3): the in-front factorization of the CPU +# reference (`src/reference/ldlt.jl`, T14) as KA kernels, with the same pivot +# sequence: Bunch–Kaufman 1×1/2×2 pivots inside the fully-summed block, +# `pivot_threshold` acceptance against the whole remaining front column, static +# perturbation `±ε` with the `pivot_sign` policy, and the contribution block +# updated once per front with `F₂₂ ← F₂₂ − (L₂₁ D) L₂₁ᴴ`. +# +# * Regime A: `subtree_ldlt_kernel!`, one workgroup per subtree, the serial +# stack of packed fronts and contribution blocks in `@localmem` (as the +# Cholesky kernel of `src/numeric/subtree.jl`). +# * Regimes B and C: `front_ldlt_kernel!`, one workgroup per front of a launch +# group, fused assembly (zero, scatter A, owner-pull extend-add), then the +# pivoted factorization of the panel in global memory and the update of the +# front's packed contribution block on the update stack. +# +# Per pivot step (a loop over the `w` columns of the front with a uniform trip +# count; a 2×2 pivot leaves the last iterations idle): work item 1 chooses the +# pivot (serial search, the reference's `_choose_pivot`), the workgroup swaps +# rows/columns, work item 1 stores D (perturbing a tiny 1×1 pivot) and counts +# the statistics, the workgroup applies the rank-1/rank-2 update to the +# remaining fully-summed columns and then scales the pivot columns into L. +# Only the lower triangle of the front is stored; the upper entries the +# algorithm reads are `conj` (Hermitian) or plain (complex symmetric) mirrors. +# Pivot decisions, control words and the per-front counts live in `@localmem`; +# the counts go to `numeric.stats` and `reduce_stats!` sums them. No atomics: +# every destination has one writer per segment, so the result is deterministic. + +"Workgroup size of the LDLᵀ/LDLᴴ front kernel (regimes B and C)." +const LDLT_WORKGROUP = 128 + +# pivot search of the device kernels (`_ldlt_pivot_type` of the reference) +const _LT_PIVOT_NONE = 0 +const _LT_PIVOT_DIAGONAL = 1 +const _LT_PIVOT_BK = 2 + +# control slots of the LDLᵀ kernels, after the regime-A slots of `subtree.jl` +# (the front kernel uses `_ST_NODE`, `_ST_F`, `_ST_W` and `_ST_LF` = panel offset) +const _LT_C0 = _ST_CTL + 1 # first factor column super_ptr[v] +const _LT_K = _ST_CTL + 2 # current local column +const _LT_STEP = _ST_CTL + 3 # size of the pivot taken at _LT_K (1 or 2; 0 before the first) +const _LT_C = _ST_CTL + 4 # chosen 1×1 pivot column (2×2: k) +const _LT_R = _ST_CTL + 5 # 2×2 partner column (0: 1×1 pivot) +const _LT_STAT = _ST_CTL + 5 # _LT_STAT + q: statistic q ∈ 1:5 (STAT_NPOS … STAT_N2X2) of the front +const _LT_CTL = _ST_CTL + 10 + +# resolved pivoting parameters passed to the kernels (isbits); `H`: Hermitian (or real symmetric) vs complex +# symmetric (a type parameter, so the kernels need no extra `Val` argument: launches with more than 32 +# arguments are not specialized and allocate on the KA CPU backend) +struct _LDLTDevice{R, H} + ptype::Int # _LT_PIVOT_NONE, _LT_PIVOT_DIAGONAL or _LT_PIVOT_BK + u::R # pivot_threshold + eps::R # pivot_epsilon (multiplied by max |aᵢⱼ| = real(aux[1]) when `scaled`) + scaled::Bool # pivot_epsilon_alg = "algo1" +end + +# Hermitian (or real symmetric) vs complex symmetric +_ldlt_herm(S::Symbolic, ::Type{T}) where {T} = !(T <: Complex) || S.structure == STRUCTURE_HERMITIAN + +# `H` must be a compile-time constant (`Val`): with a runtime flag the parameter type, and every call that +# takes it, would be resolved at run time +function _ldlt_device_params(S::Symbolic, ::Type{T}, opts::Options, ::Val{H} = Val(_ldlt_herm(S, T))) where {T, H} + _is_ldlt_structure(S.structure) || + throw(InvalidValueError("LDLᵀ/LDLᴴ needs structure \"S\" or \"H\", got \"$(convert(String, S.structure))\"")) + R = real(T) + p = _ldlt_pivot_type(opts.pivot_type) + ptype = p == PIVOT_NONE ? _LT_PIVOT_NONE : p == PIVOT_DIAGONAL ? _LT_PIVOT_DIAGONAL : _LT_PIVOT_BK + return _LDLTDevice{R, H}(ptype, R(opts.pivot_threshold), R(resolved_pivot_epsilon(opts, R)), + opts.pivot_epsilon_alg == PIVOT_EPSILON_SCALED) +end + +# effective perturbation ε (as `_ldlt_params`: ε · max |aᵢⱼ| when scaled and the matrix is not zero) +@inline function _lt_eps(p::_LDLTDevice, aux) + p.scaled || return p.eps + @inbounds a = real(aux[1]) + return a > 0 ? p.eps * a : p.eps +end + +# --------------------------------------------------------------------------- +# lower triangle of a front: packed `f×f` in local memory (regime A) or the +# `f×w` column-major panel in global memory (regimes B/C, columns ≤ w only) + +struct _PackedFront{A} + a::A + off::Int + f::Int +end + +struct _PanelFront{A} + a::A + off::Int + f::Int +end + +@inline _fpos(F::_PackedFront, i, j) = F.off + _packed(i, j, F.f) +@inline _fpos(F::_PanelFront, i, j) = F.off + (j - 1) * F.f + i +@inline _fget(F, i, j) = @inbounds F.a[_fpos(F, i, j)] +@inline function _fset!(F, i, j, v) + @inbounds F.a[_fpos(F, i, j)] = v + return nothing +end +@inline _cj(x, ::Val{true}) = conj(x) +@inline _cj(x, ::Val{false}) = x +# |F[i, j]| of the symmetric/Hermitian front, from its lower triangle +@inline _fabs(F, i, j) = abs(i >= j ? _fget(F, i, j) : _fget(F, j, i)) +# F[i, j] from the lower triangle (upper entries: conj for Hermitian, plain for complex symmetric) +@inline _fsym(F, i, j, h::Val) = i >= j ? _fget(F, i, j) : _cj(_fget(F, j, i), h) + +# the front of the current node from the control words: offset `_ST_LF` (1-based) into `fa` +@inline _lt_front(fa, ctl, ::Val{true}) = @inbounds _PackedFront(fa, Int(ctl[_ST_LF]) - 1, Int(ctl[_ST_F])) +@inline _lt_front(fa, ctl, ::Val{false}) = @inbounds _PanelFront(fa, Int(ctl[_ST_LF]) - 1, Int(ctl[_ST_F])) + +# --------------------------------------------------------------------------- +# pivot choice (work item 1; the reference's `_colmax`, `_accept_1x1`, `_accept_2x2`, +# `_best_1x1`, `_choose_pivot` on the lower triangle) + +@inline function _lt_colmax(F, c, k, f, skip) + m = zero(real(eltype(F.a))) + for i in k:f + (i == c || i == skip) && continue + m = max(m, _fabs(F, i, c)) + end + return m +end + +@inline function _lt_accept_1x1(F, c, k, f, u, ε) + d = _fabs(F, c, c) + return d >= ε && d >= u * _lt_colmax(F, c, k, f, 0) +end + +@inline _lt_det2(F, k, r, h::Val) = _fget(F, k, k) * _fget(F, r, r) - _fsym(F, k, r, h) * _fget(F, r, k) + +@inline function _lt_nonsingular_2x2(F, k, r, ε, h::Val) + det = _lt_det2(F, k, r, h) + return !iszero(det) && isfinite(det) && max(_fabs(F, k, k), _fabs(F, r, k), _fabs(F, r, r)) >= ε +end + +@inline function _lt_accept_2x2(F, k, r, f, u, ε, h::Val) + _lt_nonsingular_2x2(F, k, r, ε, h) || return false + det = _lt_det2(F, k, r, h) + e11 = abs(_fget(F, r, r) / det) + e12 = abs(_fsym(F, k, r, h) / det) + e21 = abs(_fget(F, r, k) / det) + e22 = abs(_fget(F, k, k) / det) + m1 = _lt_colmax(F, k, k, f, r) + m2 = _lt_colmax(F, r, k, f, k) + return u * (e11 * m1 + e21 * m2) <= 1 && u * (e12 * m1 + e22 * m2) <= 1 +end + +@inline function _lt_best_1x1(F, k, w, f, u, ε) + best = 0 + bv = zero(real(eltype(F.a))) + for j in k:w + _lt_accept_1x1(F, j, k, f, u, ε) || continue + a = _fabs(F, j, j) + if best == 0 || a > bv + best = j + bv = a + end + end + return best +end + +# `(c, 0)`: 1×1 pivot at column c; `(k, r)`: 2×2 pivot on columns k and r (as `_choose_pivot`) +@inline function _lt_choose(F, k, w, f, p::_LDLTDevice{R}, ε, h::Val) where {R} + p.ptype == _LT_PIVOT_NONE && return (k, 0) + u = p.u + if p.ptype == _LT_PIVOT_DIAGONAL + _lt_accept_1x1(F, k, k, f, u, ε) && return (k, 0) + best = _lt_best_1x1(F, k, w, f, u, ε) + best != 0 && return (best, 0) + big = k + for j in (k + 1):w + _fabs(F, j, j) > _fabs(F, big, big) && (big = j) + end + return (_fabs(F, big, big) >= ε ? big : k, 0) + end + α = R(BUNCH_KAUFMAN_ALPHA) + λ = zero(R) + r = 0 + for i in (k + 1):w + a = _fabs(F, i, k) + if a > λ + λ = a + r = i + end + end + akk = _fabs(F, k, k) + c, r2 = k, 0 + if !(r == 0 || akk >= α * λ) + σ = zero(R) # largest off-diagonal of column r in the block + for i in k:w + i != r && (σ = max(σ, _fabs(F, i, r))) + end + if akk * σ >= α * λ^2 + # 1×1 at k + elseif _fabs(F, r, r) >= α * σ + c = r + else + r2 = r + end + end + (r2 == 0 ? _lt_accept_1x1(F, c, k, f, u, ε) : _lt_accept_2x2(F, k, r2, f, u, ε, h)) && return (c, r2) + best = _lt_best_1x1(F, k, w, f, u, ε) + best != 0 && return (best, 0) + r2 != 0 && _lt_nonsingular_2x2(F, k, r2, ε, h) && return (c, r2) + return (c, 0) +end + +# --------------------------------------------------------------------------- +# phases of one pivot step (work item `li` of `WG`; `PK`: packed local front) + +# reset the per-front control words (work item 1) +@inline function _lt_reset!(ctl, c0) + @inbounds begin + IT = eltype(ctl) + ctl[_LT_C0] = c0 % IT + ctl[_LT_K] = one(IT) + ctl[_LT_STEP] = zero(IT) + ctl[_LT_C] = zero(IT) + ctl[_LT_R] = zero(IT) + for q in 1:5 + ctl[_LT_STAT + q] = zero(IT) + end + end + return nothing +end + +# identity local pivot order of the front's columns +@inline function _lt_init_piv!(piv, ctl, li, ::Val{WG}) where {WG} + @inbounds begin + c0 = Int(ctl[_LT_C0]) + for j in li:WG:Int(ctl[_ST_W]) + piv[c0 + j - 1] = Int32(c0 + j - 1) + end + end + return nothing +end + +# work item 1: advance to the next column and choose its pivot +@inline function _lt_choose!(fa, ctl, aux, p::_LDLTDevice, pk::Val, h::Val) + @inbounds begin + IT = eltype(ctl) + k = Int(ctl[_LT_K]) + Int(ctl[_LT_STEP]) + ctl[_LT_K] = k % IT + ctl[_LT_STEP] = zero(IT) + w = Int(ctl[_ST_W]) + if k <= w + F = _lt_front(fa, ctl, pk) + c, r = _lt_choose(F, k, w, Int(ctl[_ST_F]), p, _lt_eps(p, aux), h) + ctl[_LT_C] = c % IT + ctl[_LT_R] = r % IT + end + end + return nothing +end + +# symmetric interchange of the columns/rows p < q of the front (lower triangle), as `_swap_front!` +@inline function _lt_swap!(fa, ctl, piv, li, ::Val{WG}, pk::Val, h::Val) where {WG} + @inbounds begin + k = Int(ctl[_LT_K]) + if k <= Int(ctl[_ST_W]) + r = Int(ctl[_LT_R]) + p = r == 0 ? k : k + 1 + q = r == 0 ? Int(ctl[_LT_C]) : r + if p != q + F = _lt_front(fa, ctl, pk) + for i in li:WG:Int(ctl[_ST_F]) + if i < p + x = _fget(F, p, i) + _fset!(F, p, i, _fget(F, q, i)) + _fset!(F, q, i, x) + elseif i == p + x = _fget(F, p, p) + _fset!(F, p, p, _fget(F, q, q)) + _fset!(F, q, q, x) + elseif i < q + x = _fget(F, i, p) + _fset!(F, i, p, _cj(_fget(F, q, i), h)) + _fset!(F, q, i, _cj(x, h)) + elseif i == q + _fset!(F, q, p, _cj(_fget(F, q, p), h)) + else + x = _fget(F, i, p) + _fset!(F, i, p, _fget(F, i, q)) + _fset!(F, i, q, x) + end + end + if li == 1 + c0 = Int(ctl[_LT_C0]) + x = piv[c0 + p - 1] + piv[c0 + p - 1] = piv[c0 + q - 1] + piv[c0 + q - 1] = x + end + end + end + end + return nothing +end + +# work item 1: store the pivot block in D (perturbing a tiny 1×1 pivot), the pivot kinds and the +# statistics; `pv` receives d (1×1) or the entries (e11, e12, e21, e22) of the inverse 2×2 block +@inline function _lt_pivot!(fa, ctl, pv, d, pivot_kind, piv, psign, perm, aux, p::_LDLTDevice, pk::Val, + h::Val{H}) where {H} + @inbounds begin + n = length(perm) + IT = eltype(ctl) + T = eltype(fa) + k = Int(ctl[_LT_K]) + if k <= Int(ctl[_ST_W]) + F = _lt_front(fa, ctl, pk) + g = Int(ctl[_LT_C0]) + k - 1 + if ctl[_LT_R] == 0 + ε = _lt_eps(p, aux) + x = _fget(F, k, k) + dk = H ? T(real(x)) : x + kind = PIVOT_KIND_1X1 + if !(abs(dk) >= ε) # tiny (or NaN): perturb + iszero(dk) && (ctl[_LT_STAT + STAT_NZERO] += one(IT)) + dk = _perturbation_sign(dk, Int(psign[perm[piv[g]]]), H) * ε + kind = PIVOT_KIND_PERTURBED + ctl[_LT_STAT + STAT_NPERTURBED] += one(IT) + end + _fset!(F, k, k, dk) + d[g] = dk + d[n + g] = zero(T) + pivot_kind[g] = kind + if H + if real(dk) > 0 + ctl[_LT_STAT + STAT_NPOS] += one(IT) + elseif real(dk) < 0 + ctl[_LT_STAT + STAT_NNEG] += one(IT) + end + end + pv[1] = dk + ctl[_LT_STEP] = one(IT) + else + a = _fget(F, k, k) + b = _fget(F, k + 1, k) + c2 = _fget(F, k + 1, k + 1) + if H + a = T(real(a)) + c2 = T(real(c2)) + end + up = _cj(b, h) + _fset!(F, k, k, a) + _fset!(F, k + 1, k + 1, c2) + _fset!(F, k + 1, k, zero(T)) # the 2×2 block lives in D, L's block is the identity + det = a * c2 - up * b + pv[1] = c2 / det + pv[2] = -up / det + pv[3] = -b / det + pv[4] = a / det + d[g] = a + d[g + 1] = c2 + d[n + g] = b + d[n + g + 1] = zero(T) + pivot_kind[g] = PIVOT_KIND_2X2_FIRST + pivot_kind[g + 1] = PIVOT_KIND_2X2_SECOND + ctl[_LT_STAT + STAT_N2X2] += one(IT) + if H + if real(det) < 0 + ctl[_LT_STAT + STAT_NPOS] += one(IT) + ctl[_LT_STAT + STAT_NNEG] += one(IT) + elseif real(a) > 0 + ctl[_LT_STAT + STAT_NPOS] += IT(2) + else + ctl[_LT_STAT + STAT_NNEG] += IT(2) + end + end + ctl[_LT_STEP] = IT(2) + end + end + end + return nothing +end + +# update of the remaining fully-summed columns (lower part, rows to f) with the unscaled pivot columns: +# F[i, j] -= l_i F[k, j] (1×1) or l1_i F[k, j] + l2_i F[k+1, j] (2×2), as the reference +@inline function _lt_update!(fa, ctl, pv, li, ::Val{WG}, pk::Val, h::Val) where {WG} + @inbounds begin + k = Int(ctl[_LT_K]) + w = Int(ctl[_ST_W]) + if k <= w + F = _lt_front(fa, ctl, pk) + step = Int(ctl[_LT_STEP]) + kk = k + step - 1 + nr = Int(ctl[_ST_F]) - kk + nc = w - kk + for q in (li - 1):WG:(nr * nc - 1) + i = kk + 1 + q % nr + j = kk + 1 + q ÷ nr + i >= j || continue + if step == 1 + l = _fget(F, i, k) / pv[1] + _fset!(F, i, j, _fget(F, i, j) - l * _cj(_fget(F, j, k), h)) + else + x1 = _fget(F, i, k) + x2 = _fget(F, i, k + 1) + l1 = x1 * pv[1] + x2 * pv[3] + l2 = x1 * pv[2] + x2 * pv[4] + _fset!(F, i, j, _fget(F, i, j) - (l1 * _cj(_fget(F, j, k), h) + l2 * _cj(_fget(F, j, k + 1), h))) + end + end + end + end + return nothing +end + +# scale the pivot columns below the pivot block into L +@inline function _lt_scale!(fa, ctl, pv, li, ::Val{WG}, pk::Val) where {WG} + @inbounds begin + k = Int(ctl[_LT_K]) + if k <= Int(ctl[_ST_W]) + F = _lt_front(fa, ctl, pk) + step = Int(ctl[_LT_STEP]) + for i in (k + step + li - 1):WG:Int(ctl[_ST_F]) + if step == 1 + _fset!(F, i, k, _fget(F, i, k) / pv[1]) + else + x1 = _fget(F, i, k) + x2 = _fget(F, i, k + 1) + _fset!(F, i, k, x1 * pv[1] + x2 * pv[3]) + _fset!(F, i, k + 1, x1 * pv[2] + x2 * pv[4]) + end + end + end + end + return nothing +end + +# F₂₂ ← F₂₂ − (L₂₁ D) L₂₁ᴴ on the packed m×m block `ca[coff .+ (1:m(m+1)/2)]` (`coff < 0`: no block) +@inline function _lt_cb_update!(fa, ctl, ca, coff, d, pivot_kind, li, ::Val{WG}, pk::Val, + h::Val{H}) where {WG, H} + @inbounds begin + n = length(pivot_kind) + T = eltype(fa) + f = Int(ctl[_ST_F]) + w = Int(ctl[_ST_W]) + m = f - w + if coff >= 0 && m > 0 + F = _lt_front(fa, ctl, pk) + c0 = Int(ctl[_LT_C0]) + for q in (li - 1):WG:(m * m - 1) + jj = q ÷ m + 1 + ii = q - (jj - 1) * m + 1 + ii >= jj || continue + i = w + ii + j = w + jj + acc = zero(T) + k = 1 + while k <= w + g = c0 + k - 1 + if pivot_kind[g] == PIVOT_KIND_2X2_FIRST + a, b, c = d[g], d[n + g], d[g + 1] + up = _cj(b, h) + x1 = _fget(F, i, k) + x2 = _fget(F, i, k + 1) + acc += (x1 * a + x2 * b) * _cj(_fget(F, j, k), h) + (x1 * up + x2 * c) * _cj(_fget(F, j, k + 1), h) + k += 2 + else + acc += _fget(F, i, k) * d[g] * _cj(_fget(F, j, k), h) + k += 1 + end + end + pos = coff + _packed(ii, jj, m) + ca[pos] = (H && ii == jj) ? T(real(ca[pos]) - real(acc)) : ca[pos] - acc + end + end + end + return nothing +end + +# work item 1: per-front statistics and status +@inline function _lt_stats!(stats, info, ctl, s) + @inbounds begin + base = (Int(s) - 1) * FRONT_STATS_FIELDS + for q in 1:5 + stats[base + q] = Int64(ctl[_LT_STAT + q]) + end + stats[base + STAT_INFO] = 0 + info[s] = Int32(0) + end + return nothing +end + +# --------------------------------------------------------------------------- +# regimes B/C: one workgroup per front, panel in global memory + +@inline function _lt_front_setup!(ctl, s, super_ptr, front_ptr, front_nrows, front_ncols) + @inbounds begin + IT = eltype(ctl) + ctl[_ST_NODE] = s % IT + ctl[_ST_F] = front_nrows[s] % IT + ctl[_ST_W] = front_ncols[s] % IT + ctl[_ST_LF] = front_ptr[s] % IT + _lt_reset!(ctl, super_ptr[s]) + end + return nothing +end + +# unit diagonal of the panel, statistics and status +@inline function _lt_front_finish!(factor, stats, info, ctl, li, ::Val{WG}) where {WG} + @inbounds begin + F = _lt_front(factor, ctl, Val(false)) + for j in li:WG:Int(ctl[_ST_W]) + _fset!(F, j, j, one(eltype(factor))) + end + li == 1 && _lt_stats!(stats, info, ctl, ctl[_ST_NODE]) + end + return nothing +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, + front_ncols, cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, prm, + Val(WG); ndrange = WG * count) + +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 +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 +`_LDLTDevice{R, HERM}`; `HERM`: Hermitian or real symmetric, else complex symmetric), +writes D, the local pivot order `piv`, the pivot kinds and the front's +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, + 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) + 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) + 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)) + @synchronize + @inbounds s = nodes[first + g - 1] + _scatter_front!(factor, nzval, amap, amap_ptr, amap_src, s, 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)) + @synchronize + end + for it in 1:ctl[_ST_W] + if li == 1 + _lt_choose!(factor, ctl, aux, prm, Val(false), Val(HERM)) + end + @synchronize + _lt_swap!(factor, ctl, piv, 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)) + end + @synchronize + _lt_update!(factor, ctl, pv, li, Val(WG), Val(false), Val(HERM)) + @synchronize + _lt_scale!(factor, ctl, pv, li, Val(WG), Val(false)) + @synchronize + end + @inbounds s = nodes[first + 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)) +end + +# --------------------------------------------------------------------------- +# regime A: one workgroup per subtree, fronts in local memory + +# offset (0-based) of the trailing packed m×m triangle of the local packed front +@inline _lt_local_cb_offset(ctl) = + @inbounds Int(ctl[_ST_LF]) - 1 + Int(ctl[_ST_W]) * (2 * Int(ctl[_ST_F]) - Int(ctl[_ST_W]) + 1) ÷ 2 + +# write the unit-lower panel, statistics and status; the subtree root writes its contribution block +@inline function _lt_subtree_write!(factor, stack, info, stats, buf, ctl, front_ptr, cb_ptr, li, ::Val{WG}) where {WG} + @inbounds begin + v = ctl[_ST_NODE] + f = Int(ctl[_ST_F]) + w = Int(ctl[_ST_W]) + lf = Int(ctl[_ST_LF]) - 1 + p0 = Int(front_ptr[v]) + T = eltype(factor) + for q in (li - 1):WG:(f * w - 1) + j = q ÷ f + 1 + i = q - (j - 1) * f + 1 + factor[p0 + q] = i > j ? buf[lf + _packed(i, j, f)] : i == j ? one(T) : zero(T) + end + li == 1 && _lt_stats!(stats, info, ctl, v) + c0 = cb_ptr[v] + if c0 > 0 + m = f - w + src = _lt_local_cb_offset(ctl) + 1 + for q in (li - 1):WG:(m * (m + 1) ÷ 2 - 1) + stack[c0 + q] = buf[src + q] + end + end + end + return nothing +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, + 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) + +LDLᵀ/LDLᴴ kernel of regime A: workgroup `g` takes subtree +`t = trees[first + g - 1]` 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 +LDLᵀ/LDLᴴ of [`front_ldlt_kernel!`](@ref), update the trailing contribution +block, write the unit-lower panel, D, `piv`, the pivot kinds and the +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, + 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) + buf = @localmem TT (NE,) + ctl = @localmem IT (_LT_CTL,) + pv = @localmem TT (4,) + if li == 1 + @inbounds t = trees[first + g - 1] + @inbounds ctl[_ST_FIRST] = subtree_ptr[t] + @inbounds ctl[_ST_COUNT] = subtree_ptr[t + 1] - subtree_ptr[t] + end + @synchronize + for k in 1:ctl[_ST_COUNT] + if li == 1 + _subtree_setup!(ctl, k, subtree_nodes, front_nrows, front_ncols, local_front, local_cb, child_ptr) + @inbounds _lt_reset!(ctl, super_ptr[ctl[_ST_NODE]]) + end + @synchronize + _subtree_zero!(buf, ctl, li, Val(WG)) + _lt_init_piv!(piv, ctl, li, Val(WG)) + @synchronize + _subtree_scatter!(buf, ctl, nzval, 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, + relind, li, Val(WG)) + @synchronize + end + for it in 1:ctl[_ST_W] + if li == 1 + _lt_choose!(buf, ctl, aux, prm, Val(true), Val(HERM)) + end + @synchronize + _lt_swap!(buf, ctl, piv, 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)) + end + @synchronize + _lt_update!(buf, ctl, pv, li, Val(WG), Val(true), Val(HERM)) + @synchronize + _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)) + @synchronize + _lt_subtree_write!(factor, stack, info, stats, buf, ctl, front_ptr, cb_ptr, li, Val(WG)) + @synchronize + for r in 1:ctl[_ST_ROUNDS] + _subtree_move!(buf, ctl, r, li, Val(WG)) + @synchronize + end + end +end + +# --------------------------------------------------------------------------- +# reductions + +@kernel function _abs_max_kernel!(aux, nzval, nnz, ::Val{WG}, ::Val{LOG2WG}) where {WG, LOG2WG} + @uniform RT = real(eltype(aux)) + li = @index(Local, Linear) + best = @localmem RT (WG,) + @inbounds begin + m = zero(RT) + for p in li:WG:nnz + m = max(m, RT(abs(nzval[p]))) + end + best[li] = m + end + @synchronize + for lev in 1:LOG2WG + @inbounds begin + h = WG >> lev + if li <= h + best[li] = max(best[li], best[li + h]) + end + end + @synchronize + end + if li == 1 + @inbounds aux[1] = eltype(aux)(best[1]) + end +end + +""" + abs_max!(aux, nzval) -> aux + +`aux[1] = max |nzval[p]|` (one workgroup, `@localmem` tree reduction; the +scale of `pivot_epsilon_alg = "algo1"`). Asynchronous. +""" +function abs_max!(aux::AbstractVector, nzval::AbstractVector) + WG = STATS_WORKGROUP + _abs_max_kernel!(KernelAbstractions.get_backend(aux), WG)(aux, nzval, length(nzval), Val(WG), Val(_ilog2(WG)); + ndrange = WG) + return aux +end + +@kernel function _reduce_stats_kernel!(totals, stats, ns, ::Val{WG}, ::Val{NF}, ::Val{LOG2WG}) where {WG, NF, LOG2WG} + li = @index(Local, 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 + end + acc[(q - 1) * WG + li] = a + end + @synchronize + for lev in 1:LOG2WG + @inbounds begin + h = WG >> lev + if li <= h + for q in 1:NF + acc[(q - 1) * WG + li] += acc[(q - 1) * WG + li + h] + end + end + end + @synchronize + end + if li <= NF + @inbounds totals[li] = acc[(li - 1) * WG + 1] + end +end + +""" + reduce_stats!(numeric, symbolic) -> numeric + +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. +""" +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) + return N +end + +""" + pivot_totals(numeric) -> (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`). +""" +function pivot_totals(N::Numeric) + t = Array(N.totals) + return (npos = t[STAT_NPOS], nneg = t[STAT_NNEG], nzero = t[STAT_NZERO], nperturbed = t[STAT_NPERTURBED], + n2x2 = t[STAT_N2X2]) +end + +# --------------------------------------------------------------------------- +# driver + +function _launch_front_ldlt!(N::Numeric, S::Symbolic, nzval, first, count, maxchild, prm, ::Val{WG}) where {WG} + 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.cb_ptr, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Int(maxchild), prm, Val(WG); + ndrange = WG * count) + return nothing +end + +function _launch_subtrees_ldlt!(N::Numeric{T}, S::Symbolic, nzval, first, count, prm, ::Val{LB}, + ::Val{WG}) where {T, LB, WG} + 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.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) + return nothing +end + +function _factorize_ldlt_groups!(N::Numeric, S::Symbolic, nzval::AbstractVector, prm) + plan = N.plan + for k in eachindex(plan.sub_first) + a, b = plan.sub_first[k], plan.sub_last[k] + _with_local_bytes(plan.sub_local[k]) do lb + _launch_subtrees_ldlt!(N, S, nzval, a, b - a + 1, prm, lb, Val(SUBTREE_WORKGROUP)) + end + end + for k in eachindex(plan.group_first) + a, b = plan.group_first[k], plan.group_last[k] + _launch_front_ldlt!(N, S, nzval, a, b - a + 1, plan.group_maxchild[k], prm, Val(LDLT_WORKGROUP)) + end + return nothing +end + +""" + factorize_ldlt!(numeric, symbolic, nzval; opts = Options()) -> 0 + +Multifrontal `P A Pᵀ = L D Lᴴ` (structure `"H"`, or `"S"` with real `T`) or +`L D Lᵀ` (complex symmetric `"S"`) on the device, with the in-front pivoting +and the static perturbation of the reference [`ref_ldlt!`](@ref) (same pivot +choice, same storage): `opts.pivot_type`, `pivot_threshold`, `pivot_epsilon`, +`pivot_epsilon_alg` and `pivot_sign` (copied to `numeric.psign`; a length +other than `n` raises [`InvalidValueError`](@ref)). Regime-A groups are one +[`subtree_ldlt_kernel!`](@ref) launch per budget class; every regime-B and +regime-C launch group is one [`front_ldlt_kernel!`](@ref) launch (fused +assembly and factorization, one workgroup per front; no vendor calls, see the +T15 report); then [`reduce_stats!`](@ref). With `pivot_epsilon_alg = "algo1"` +[`abs_max!`](@ref) computes the scale first. Fills `numeric.factor` +(unit-lower panels), `d`, `piv`, `pivot_kind`, `stats` and `totals`. The +factorization always completes (`info = 0`); the phase allocates nothing on +the device, never synchronizes with the host, and is deterministic. +""" +function factorize_ldlt!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; opts::Options = Options()) where {T} + _check_numeric(N, S, nzval) + herm = _ldlt_herm(S, T) + prm = _ldlt_device_params(S, T, opts, Val(true)) # validates the options + ps = opts.pivot_sign + if ps === nothing + fill!(N.psign, Int8(0)) + else + length(ps) == S.n || + 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) + if herm + _factorize_ldlt_groups!(N, S, nzval, prm) + else + _factorize_ldlt_groups!(N, S, nzval, _ldlt_device_params(S, T, opts, Val(false))) + end + reduce_stats!(N, S) + return 0 +end diff --git a/src/numeric/storage.jl b/src/numeric/storage.jl index 7f09897..043e6af 100644 --- a/src/numeric/storage.jl +++ b/src/numeric/storage.jl @@ -106,6 +106,12 @@ Numeric storage of a factorization (PLAN §3.2), laid out by the * `stats` (`FRONT_STATS_FIELDS × ns` `Int64`, column `s` = front `s`): `(npos, nneg, nzero, nperturbed, n2x2, info)`, `info` = the local column of the first failed pivot of the front (`0` = none); +* `totals` (`FRONT_STATS_FIELDS` `Int64`): the sums of `stats` over the fronts + ([`reduce_stats!`](@ref); `info` field: number of failed fronts); +* `psign` (`n` `Int8`, LDLᵀ/LDLᴴ): the `pivot_sign` request per original row + (`0` = none), copied from the options at every factorization; +* `aux` (1 entry): `max |aᵢⱼ|` of the values of the last LDLᵀ/LDLᴴ + factorization with `pivot_epsilon_alg = "algo1"` (scaled perturbation); * `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`; @@ -121,6 +127,9 @@ struct Numeric{T, VT <: AbstractVector{T}, VS <: AbstractVector{Int64}, VI <: Ab info::VI piv::VI pivot_kind::VK + totals::VS + psign::VK + aux::VT plan::NumericPlan end @@ -140,7 +149,7 @@ Base.show(io::IO, N::Numeric{T, VT}) where {T, VT} = allocate_numeric(symbolic, T, backend = CPU()) -> Numeric{T} Allocate (zero-filled) the factor panels, D, update stack, regime-C workspace, per-front -statistics, status vector, pivot order and pivot kinds of `symbolic`'s [`Layout`](@ref) for element type +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). This is the only allocation of the numeric phase. """ @@ -155,7 +164,10 @@ function allocate_numeric(S::Symbolic, ::Type{T}, backend::KernelAbstractions.Ba 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) + psign = KernelAbstractions.zeros(backend, Int8, S.n) + aux = KernelAbstractions.zeros(backend, T, 1) return Numeric{T, typeof(factor), typeof(stats), typeof(info), typeof(pivot_kind)}(factor, d, stack, work, stats, - info, piv, pivot_kind, - NumericPlan(S)) + info, piv, pivot_kind, totals, + psign, aux, NumericPlan(S)) end diff --git a/src/reference/ldlt.jl b/src/reference/ldlt.jl index 1747f9e..c82b792 100644 --- a/src/reference/ldlt.jl +++ b/src/reference/ldlt.jl @@ -191,9 +191,21 @@ function ref_ldlt!(N::Numeric{T, Vector{T}}, S::Symbolic, nzval::AbstractVector; end end end + _host_totals!(N) return 0 end +# the totals of the per-front statistics (host storage; `reduce_stats!` on the device) +function _host_totals!(N::Numeric) + for q in 1:FRONT_STATS_FIELDS + N.totals[q] = sum((N.stats[(s - 1) * FRONT_STATS_FIELDS + q] for s in 1:(length(N.stats) ÷ FRONT_STATS_FIELDS)); + init = Int64(0)) + end + N.totals[STAT_INFO] = count(s -> N.stats[(s - 1) * FRONT_STATS_FIELDS + STAT_INFO] != 0, + 1:(length(N.stats) ÷ FRONT_STATS_FIELDS)) + return N +end + ref_ldlt!(N::Numeric, S::Symbolic, A::CSR; kwargs...) = ref_ldlt!(N, S, vec(A.nzval); kwargs...) # X D for X with the columns of the factor and D block diagonal (1×1 and 2×2 blocks) @@ -276,8 +288,8 @@ function _choose_pivot(F::AbstractMatrix{T}, k::Int, w::Int, prm::_LDLTParams) w end return (abs(F[big, big]) >= prm.eps ? big : k, 0) end - # Bunch–Kaufman inside the block - α = BUNCH_KAUFMAN_ALPHA + # Bunch–Kaufman inside the block (α in the working precision, as the device kernels: no Float64 on Metal) + α = real(T)(BUNCH_KAUFMAN_ALPHA) λ, r = zero(real(T)), 0 for i in (k + 1):w a = abs(F[i, k]) diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index e61ee65..2acf5cb 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -30,7 +30,16 @@ # phase. No atomics; bitwise reproducible. # # The backward sweep only gathers from finished ancestor rows: conflict-free. -# The diagonal sweep is the identity for Cholesky (D is used by LDLᵀ, T15). +# The diagonal sweep is the identity for Cholesky; for LDLᵀ/LDLᴴ it applies +# D⁻¹ (1×1 and 2×2 blocks). +# +# LDLᵀ/LDLᴴ (`LDL = true`): the panels are unit lower triangular (the TRSVs +# divide by the stored 1), and pivoting permuted the columns of each +# supernode (`numeric.piv`). As the reference `ref_solve_ldlt!` does, a front +# reorders its slice of the right-hand side by its local pivot order before +# its forward TRSV and restores the order after its backward TRSV, through the +# rows `1:w` of its gather-list buffer in `U` (unused otherwise). Complex +# symmetric `"S"` (`CONJ = false`) solves with `Lᵀ` instead of `Lᴴ`. # # As in the numeric kernels, per-node values sit in a `@localmem` control array # written by work item 1 before a barrier, so every loop containing a barrier @@ -313,8 +322,9 @@ end return nothing end -# backward gather + GEMV: y_j -= Σ_i conj(L21[i, j]) x[row_i] (the rows below belong to finished ancestors) -@inline function _sv_bwd_gather!(Y, factor, rowval, ctl, r, li, ::Val{WG}) where {WG} +# backward gather + GEMV: y_j -= Σ_i conj(L21[i, j]) x[row_i] (the rows below belong to finished ancestors; +# no conj for complex symmetric LDLᵀ, `cj = Val(false)`) +@inline function _sv_bwd_gather!(Y, factor, rowval, ctl, r, li, ::Val{WG}, cj::Val = Val(true)) where {WG} @inbounds begin f = ctl[_SV_F] w = ctl[_SV_W] @@ -325,7 +335,7 @@ end col = p0 + (j - 1) * f t = zero(eltype(Y)) for i in (w + 1):f - t += conj(factor[col + i]) * Y[rowval[rp + i], r] + t += _cj(factor[col + i], cj) * Y[rowval[rp + i], r] end Y[c0 + j, r] -= t end @@ -334,27 +344,56 @@ end end # backward TRSV with L11ᴴ, step jj (column j = w - jj + 1): Y[c0 + k] -= conj(L[j, k]) x_j for k < j -@inline function _sv_bwd_col!(Y, factor, ctl, jj, r, li, ::Val{WG}) where {WG} +@inline function _sv_bwd_col!(Y, factor, ctl, jj, r, li, ::Val{WG}, cj::Val = Val(true)) where {WG} @inbounds begin f = ctl[_SV_F] j = ctl[_SV_W] - jj + 1 c0 = ctl[_SV_C0] - 1 p0 = ctl[_SV_P0] - 1 - xj = Y[c0 + j, r] / conj(factor[p0 + (j - 1) * f + j]) + xj = Y[c0 + j, r] / _cj(factor[p0 + (j - 1) * f + j], cj) for k in li:WG:(j - 1) - Y[c0 + k, r] -= conj(factor[p0 + (k - 1) * f + j]) * xj + Y[c0 + k, r] -= _cj(factor[p0 + (k - 1) * f + j], cj) * xj end end return nothing end -@inline function _sv_bwd_diag!(Y, factor, ctl, r, li, ::Val{WG}) where {WG} +@inline function _sv_bwd_diag!(Y, factor, ctl, r, li, ::Val{WG}, cj::Val = Val(true)) where {WG} @inbounds begin f = ctl[_SV_F] c0 = ctl[_SV_C0] - 1 p0 = ctl[_SV_P0] - 1 for j in li:WG:ctl[_SV_W] - Y[c0 + j, r] /= conj(factor[p0 + (j - 1) * f + j]) + Y[c0 + j, r] /= _cj(factor[p0 + (j - 1) * f + j], cj) + end + end + return nothing +end + +# LDLᵀ local pivot order, first half: copy the front's slice of Y to the rows 1:w of its buffer in U +@inline function _sv_piv_save!(Y, U, ctl, r, li, ::Val{WG}) where {WG} + @inbounds begin + c0 = ctl[_SV_C0] - 1 + rp = ctl[_SV_RP] - 1 + for j in li:WG:ctl[_SV_W] + U[rp + j, r] = Y[c0 + j, r] + end + end + return nothing +end + +# second half: forward Y[c0 + k] = (saved)[piv[c0 + k] - c0]; backward Y[c0 + piv[c0 + k] - c0] = (saved)[k] +@inline function _sv_piv_apply!(Y, U, piv, ctl, r, li, ::Val{WG}, ::Val{FWD}) where {WG, FWD} + @inbounds begin + c0 = ctl[_SV_C0] - 1 + rp = ctl[_SV_RP] - 1 + for k in li:WG:ctl[_SV_W] + p = Int(piv[c0 + k]) - c0 + if FWD + Y[c0 + k, r] = U[rp + p, r] + else + Y[c0 + p, r] = U[rp + k, r] + end end end return nothing @@ -364,21 +403,23 @@ end # kernels """ - solve_fwd_kernel!(backend, WG)(Y, U, factor, list, first, count, subtree_ptr, subtree_nodes, super_ptr, + solve_fwd_kernel!(backend, WG)(Y, U, factor, piv, list, first, 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(WG); ndrange = WG * count * nrhs) + relind_ptr, relind, Val(SUB), Val(DET), Val(LDL), Val(WG); + ndrange = WG * count * nrhs) 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 supernodes in processing order (`SUB = true`, regime A), or the single front `list[first + g - 1]` (regime B). Per front: (deterministic, `DET`) zero its -update buffer and pull the children's; TRSV with `L11` on its columns of `Y`; -GEMV with `L21` into its update buffer or (atomic variant) into the rows of -`Y` below. +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, list, first, 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{WG}) where {SUB, DET, WG} +@kernel function solve_fwd_kernel!(Y, U, factor, piv, list, first, 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} @uniform IT = eltype(front_ptr) li = @index(Local, Linear) G = @index(Group, Linear) @@ -402,6 +443,12 @@ GEMV with `L21` into its update buffer or (atomic variant) into the rows of @synchronize end end + 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)) + @synchronize + end for j in 1:ctl[_SV_W] _sv_fwd_col!(Y, factor, ctl, j, _sv_rhs(G, count), li, Val(WG)) @synchronize @@ -414,17 +461,18 @@ GEMV with `L21` into its update buffer or (atomic variant) into the rows of end """ - solve_bwd_kernel!(backend, WG)(Y, factor, list, first, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, - rowval, front_ptr, front_nrows, front_ncols, child_ptr, Val(SUB), Val(WG); - ndrange = WG * count * nrhs) + solve_bwd_kernel!(backend, WG)(Y, U, factor, piv, list, first, 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) -Backward sweep `Lᴴ x = z` of right-hand side `r` by workgroup `G = g + count (r - 1)` +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 -front): gather the finished rows below, GEMV with `L21ᴴ`, TRSV with `L11ᴴ`. +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, factor, list, first, count, subtree_ptr, subtree_nodes, super_ptr, rowptr, rowval, - front_ptr, front_nrows, front_ncols, child_ptr, ::Val{SUB}, - ::Val{WG}) where {SUB, WG} +@kernel function solve_bwd_kernel!(Y, U, factor, piv, list, first, 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) li = @index(Local, Linear) G = @index(Group, Linear) @@ -439,14 +487,57 @@ front): gather the finished rows below, GEMV with `L21ᴴ`, TRSV with `L11ᴴ`. _sv_setup!(ctl, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr) end @synchronize - _sv_bwd_gather!(Y, factor, rowval, ctl, _sv_rhs(G, count), li, Val(WG)) + _sv_bwd_gather!(Y, factor, rowval, ctl, _sv_rhs(G, count), li, Val(WG), Val(CONJ)) @synchronize for jj in 1:ctl[_SV_W] - _sv_bwd_col!(Y, factor, ctl, jj, _sv_rhs(G, count), li, Val(WG)) + _sv_bwd_col!(Y, factor, ctl, jj, _sv_rhs(G, count), li, Val(WG), Val(CONJ)) @synchronize end - _sv_bwd_diag!(Y, factor, ctl, _sv_rhs(G, count), li, Val(WG)) + _sv_bwd_diag!(Y, factor, ctl, _sv_rhs(G, count), li, Val(WG), Val(CONJ)) @synchronize + 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)) + @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, + ::Val{FWD}, ::Val{WG}) where {FWD, WG} + @uniform IT = eltype(front_ptr) + li = @index(Local, Linear) + r = @index(Group, Linear) + ctl = @localmem IT (_SV_CTL,) + if li == 1 + _sv_setup!(ctl, v, super_ptr, rowptr, front_ptr, front_nrows, front_ncols, child_ptr) + 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)) +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} + q = @index(Global, Linear) + k = (q - 1) % n + 1 + r = (q - 1) ÷ n + 1 + @inbounds if r <= nrhs + kind = pivot_kind[k] + if kind == PIVOT_KIND_2X2_FIRST + a, b, c = d[k], d[n + k], d[k + 1] + up = HERM ? conj(b) : b + det = a * c - up * b + y1, y2 = Y[k, r], Y[k + 1, r] + Y[k, r] = (c * y1 - up * y2) / det + Y[k + 1, r] = (a * y2 - b * y1) / det + elseif kind != PIVOT_KIND_2X2_SECOND + Y[k, r] /= d[k] + end end end @@ -499,7 +590,7 @@ end # drivers function _check_solve(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric{T}, nrhs::Integer) where {T} - _check_reference_cholesky(S, T) + _is_ldlt_structure(S.structure) ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) length(N.factor) == 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 || @@ -521,35 +612,71 @@ function _solve_impls(N::Numeric, S::Symbolic, impl::Symbol) end function _launch_fwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, - det::Val) + det::Val, ldl::Val) WG = SOLVE_WORKGROUP solve_fwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, ws.U, N.factor, 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, Val(WG); ndrange = WG * count * nrhs) + 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) return nothing end -function _launch_bwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val) +function _launch_bwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, + ldl::Val, cj::Val) WG = SOLVE_WORKGROUP solve_bwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, N.factor, 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, Val(WG); ndrange = WG * count * nrhs) + 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) return nothing end +# solve kind of the analysis: Cholesky, LDLᵀ/LDLᴴ with Lᴴ (Hermitian or real), LDLᵀ complex symmetric +const _SV_CHOLESKY = 0 +const _SV_LDL_HERM = 1 +const _SV_LDL_SYM = 2 + +_solve_kind(S::Symbolic, ::Type{T}) where {T} = + !_is_ldlt_structure(S.structure) ? _SV_CHOLESKY : + (T <: Complex && S.structure == STRUCTURE_SYMMETRIC) ? _SV_LDL_SYM : _SV_LDL_HERM + # the flags as compile-time constants through explicit branches (closures over runtime flags allocate) function _launch_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, - sub::Bool, det::Bool, forward::Bool) + sub::Bool, det::Bool, forward::Bool, kind::Int) + if kind == _SV_CHOLESKY + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(false), Val(true)) + elseif kind == _SV_LDL_HERM + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(true)) + else + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(false)) + end + return nothing +end + +function _launch_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, + sub::Bool, det::Bool, forward::Bool, ldl::Val, cj::Val) if !forward - sub ? _launch_bwd!(ws, S, N, first, count, nrhs, Val(true)) : - _launch_bwd!(ws, S, N, first, count, nrhs, Val(false)) + sub ? _launch_bwd!(ws, S, N, first, count, nrhs, Val(true), ldl, cj) : + _launch_bwd!(ws, S, N, first, count, nrhs, Val(false), ldl, cj) elseif sub - det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(true)) : - _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(false)) + det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(true), ldl) : + _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(false), ldl) else - det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(true)) : - _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(false)) + det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(true), ldl) : + _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(false), ldl) + end + return nothing +end + +# LDLᵀ regime-C path: local pivot order of front s on the first nrhs columns of Y +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) + if forward + kernel!(args..., Val(true), Val(WG); ndrange = WG * nrhs) + else + kernel!(args..., Val(false), Val(WG); ndrange = WG * nrhs) end return nothing end @@ -565,8 +692,9 @@ function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs: end function _fwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, det::Bool, - p::NamedTuple) where {T} + p::NamedTuple, kind::Int) where {T} f, w, L11, L21, Yc = _dense_front(ws, S, N, s, nrhs) + 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 m > 0 || return nothing @@ -582,16 +710,19 @@ function _fwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrh return nothing end -function _bwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple) where {T} +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 + tr = kind == _SV_LDL_SYM ? 'T' : 'C' if m > 0 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, 'C', 'N', -one(T), L21, view(ws.tmp, 1:m, 1:nrhs), one(T), Yc) + _gemm_impl!(p.gemm, tr, 'N', -one(T), L21, view(ws.tmp, 1:m, 1:nrhs), one(T), Yc) end - _trsm_impl!(p.trsm, 'L', 'L', 'C', 'N', one(T), L11, Yc) + _trsm_impl!(p.trsm, 'L', 'L', tr, 'N', one(T), L11, Yc) + kind == _SV_CHOLESKY || _piv_dense!(ws, S, N, s, nrhs, false) return nothing end @@ -615,6 +746,7 @@ function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integ nrhs > 0 || return ws det = deterministic || !ws.atomic p = _solve_impls(N, S, impl) + kind = _solve_kind(S, eltype(ws)) plan = ws.plan nodes = S.schedule.group_nodes for k in eachindex(plan.kind) @@ -628,10 +760,10 @@ function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integ ndrange = WG * (b - a + 1) * Int(nrhs)) end for q in a:b - _fwd_dense!(ws, S, N, nodes[q], Int(nrhs), det, p) + _fwd_dense!(ws, S, N, nodes[q], Int(nrhs), det, p, kind) end else - _launch_sweep!(ws, S, N, a, b - a + 1, Int(nrhs), plan.kind[k] == SOLVE_SUBTREES, det, true) + _launch_sweep!(ws, S, N, a, b - a + 1, Int(nrhs), plan.kind[k] == SOLVE_SUBTREES, det, true, kind) end end return ws @@ -640,19 +772,31 @@ end """ diagonal_sweep!(ws, symbolic, numeric; nrhs = max_rhs(ws)) -> ws -Diagonal step `D W = Z` between the sweeps: the identity for Cholesky -(`"SPD"`/`"HPD"`); LDLᵀ/LDLᴴ fill it in T15. +Diagonal step `D W = Z` between the sweeps, in place on the first `nrhs` +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. """ function diagonal_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws)) _check_solve(ws, S, N, nrhs) + kind = _solve_kind(S, eltype(ws)) + (kind == _SV_CHOLESKY || nrhs == 0 || S.n == 0) && return ws + 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)) + else + kernel!(ws.Y, N.d, N.pivot_kind, S.n, Int(nrhs), Val(false); ndrange = S.n * Int(nrhs)) + end return ws end """ backward_sweep!(ws, symbolic, numeric; nrhs = max_rhs(ws), impl = :auto) -> ws -Backward sweep `Lᴴ X = W` in place on the first `nrhs` columns of `ws.Y`, the -forward plan in reverse: per step the regime-C-path fronts (gather kernel, +Backward sweep `Lᴴ X = W` (`Lᵀ` for complex symmetric LDLᵀ) in place on the +first `nrhs` columns of `ws.Y`, the forward plan in reverse: per step the regime-C-path fronts (gather kernel, `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. @@ -663,16 +807,17 @@ function backward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Inte _check_solve(ws, S, N, nrhs) nrhs > 0 || return ws p = _solve_impls(N, S, impl) + kind = _solve_kind(S, eltype(ws)) plan = ws.plan nodes = S.schedule.group_nodes for k in reverse(eachindex(plan.kind)) a, b = plan.first[k], plan.last[k] if plan.kind[k] == SOLVE_DENSE for q in b:-1:a - _bwd_dense!(ws, S, N, nodes[q], Int(nrhs), p) + _bwd_dense!(ws, S, N, nodes[q], Int(nrhs), p, kind) end else - _launch_sweep!(ws, S, N, a, b - a + 1, Int(nrhs), plan.kind[k] == SOLVE_SUBTREES, false, false) + _launch_sweep!(ws, S, N, a, b - a + 1, Int(nrhs), plan.kind[k] == SOLVE_SUBTREES, false, false, kind) end end return ws @@ -681,7 +826,8 @@ end """ sweep_solve!(X, ws, symbolic, numeric, B; transposed = false, deterministic = false, impl = :auto) -> X -Solve `A X = B` with the factor `P A Pᵀ = L Lᴴ` of [`factorize!`](@ref): +Solve `A X = B` with the factor `P A Pᵀ = L Lᴴ` (or the LDLᵀ/LDLᴴ factor with +its local pivot orders) of [`factorize!`](@ref): [`permute_rhs!`](@ref) into `ws.Y`, [`forward_sweep!`](@ref), [`diagonal_sweep!`](@ref), [`backward_sweep!`](@ref), [`unpermute_solution!`](@ref) into `X`. `B` and `X` are device arrays in the diff --git a/src/solver.jl b/src/solver.jl index 83ab09b..d71ff2f 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -54,12 +54,13 @@ of its transpose with the view flipped (`'L'` ↔ `'U'`); for complex Hermitian matrices this needs the conjugated solve of `solve_mode` (T16) and raises [`NotSupportedError`](@ref) for now. -Implemented at this point (v0.1): structures `"SPD"` (real `T`) and `"HPD"`, a -single matrix (no uniform batch), the phases `"reordering"`, -`"symbolic_factorization"`, `"analysis"`, `"factorization"`, -`"refactorization"`, `"solve"`, `"solve_fwd_perm"`, `"solve_fwd"`, -`"solve_bwd"` and `"solve_bwd_perm"` ([`execute!`](@ref)). Other structures and -phases raise [`NotSupportedError`](@ref) when they are executed. +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"`, +`"analysis"`, `"factorization"`, `"refactorization"`, `"solve"`, +`"solve_fwd_perm"`, `"solve_fwd"`, `"solve_diag"`, `"solve_bwd"` and +`"solve_bwd_perm"` ([`execute!`](@ref)). Structure `"G"` and the other phases +raise [`NotSupportedError`](@ref) when they are executed. Fields: `A` (the current [`CSR`](@ref), re-pointed by [`update!`](@ref)), `structure`, `view` (as given), `options` ([`Options`](@ref), set through @@ -221,13 +222,16 @@ Execute `phase` on `solver` (≅ `cudss(phase, solver, x, b)`), PLAN §1.2: allocates the numeric storage and the solve workspace (needs `"reordering"`); * `"analysis"`: both; * `"factorization"`: numeric factorization with the current values of the - matrix (needs the analysis); sets `"info"` (`0` or the original column of the - first non-positive pivot) and clears `fresh_factorization`; + matrix (needs the analysis); sets `"info"` (`0` or, for `"SPD"`/`"HPD"`, the + original column of the first non-positive pivot; LDLᵀ/LDLᴴ perturbs tiny + pivots instead and always reports `0`, see `"npivots"`) and clears + `fresh_factorization`; * `"refactorization"`: the same with the analysis and storage reused, `"info"` reset first (needs a previous `"factorization"`); * `"solve"`: `X = A⁻¹ B`; the sub-phases `"solve_fwd_perm"` (permute `B` into the - workspace), `"solve_fwd"` (forward sweep), `"solve_bwd"` (backward sweep) and - `"solve_bwd_perm"` (inverse permutation into `X`) do the same in four calls + workspace), `"solve_fwd"` (forward sweep), `"solve_diag"` (`D⁻¹` of LDLᵀ/LDLᴴ, + the identity for Cholesky), `"solve_bwd"` (backward sweep) and + `"solve_bwd_perm"` (inverse permutation into `X`) do the same in five calls with the same `X`, `B`. `X` and `B` are only read by the solve phases (pass anything, e.g. `nothing`, @@ -236,10 +240,10 @@ to the others): vectors, `n × nrhs` matrices, strided vectors or backend; `X === B` is allowed. A failed factorization does not throw (as in cuDSS): check `getparam(solver, "info")`; solving with it gives garbage. -`"solve_diag"`, `"solve_refinement"`, `"solve_fwd_schur"` and -`"solve_bwd_schur"` raise [`NotSupportedError`](@ref) until their task (T16, -T20); unknown phase strings raise `ArgumentError`. Executing a phase before -the phases it depends on raises [`FactorizationError`](@ref). +`"solve_refinement"`, `"solve_fwd_schur"` and `"solve_bwd_schur"` raise +[`NotSupportedError`](@ref) until their task (T16, T20); unknown phase +strings raise `ArgumentError`. Executing a phase before the phases it depends +on raises [`FactorizationError`](@ref). With `asynchronous = false` the backend is synchronized before returning (`KernelAbstractions.synchronize`). """ @@ -254,12 +258,12 @@ function execute!(phase::AbstractString, solver::DirectSolver, X, B; asynchronou _symbolic!(solver) elseif p == PHASE_FACTORIZATION || p == PHASE_REFACTORIZATION _factorize!(solver, p) - elseif p == PHASE_SOLVE || p == PHASE_SOLVE_FWD_PERM || p == PHASE_SOLVE_FWD || p == PHASE_SOLVE_BWD || - p == PHASE_SOLVE_BWD_PERM + elseif p == PHASE_SOLVE || p == PHASE_SOLVE_FWD_PERM || p == PHASE_SOLVE_FWD || p == PHASE_SOLVE_DIAG || + p == PHASE_SOLVE_BWD || p == PHASE_SOLVE_BWD_PERM _solve!(solver, p, X, B) else throw(NotSupportedError("phase \"$phase\" is not implemented yet " * - "($(p == PHASE_SOLVE_DIAG || p == PHASE_SOLVE_REFINEMENT ? "T16" : "T20"))")) + "($(p == PHASE_SOLVE_REFINEMENT ? "T16" : "T20"))")) end asynchronous || KernelAbstractions.synchronize(solver.backend) return nothing @@ -268,10 +272,10 @@ end # options that no phase implements yet: refuse them instead of silently ignoring them function _check_analysis_supported(solver::DirectSolver{T}) where {T} s = solver.structure - (s == STRUCTURE_SPD || s == STRUCTURE_HPD) || - throw(NotSupportedError("structure \"$(convert(String, s))\" is not implemented yet (LDLᵀ/LDLᴴ: T14/T15, " * - "LU: T19); use \"SPD\" or \"HPD\"")) - T <: Complex && s != STRUCTURE_HPD && + s == STRUCTURE_GENERAL && + throw(NotSupportedError("structure \"G\" (LU) is not implemented yet (T19); use \"S\", \"H\", \"SPD\" or " * + "\"HPD\"")) + 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)")) @@ -334,15 +338,15 @@ function _factorize!(solver::DirectSolver, p::Phase) p == PHASE_REFACTORIZATION && solver.stage < STAGE_FACTORIZED && throw(_phase_error(name, "needs a previous \"factorization\"")) p == PHASE_REFACTORIZATION && (solver.info = 0) - solver.info = _numeric_phase!(solver.numeric, solver.symbolic, solver.A.nzval) + solver.info = _numeric_phase!(solver.numeric, solver.symbolic, solver.A.nzval, solver.options) solver.stage = STAGE_FACTORIZED p == PHASE_FACTORIZATION && (solver.fresh_factorization = false) return solver end # function barrier: concrete storage types for the numeric phase -_numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractVector) = factorize!(N, S, nzval) -_numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractMatrix) = factorize!(N, S, vec(nzval)) +_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) _rhs_data(X::AbstractVecOrMat) = (X, false) @@ -411,6 +415,8 @@ function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::A permute_rhs!(ws.Y, B, S.perm; transposed) elseif p == PHASE_SOLVE_FWD forward_sweep!(ws, S, N; nrhs, deterministic = det) + elseif p == PHASE_SOLVE_DIAG + diagonal_sweep!(ws, S, N; nrhs) elseif p == PHASE_SOLVE_BWD backward_sweep!(ws, S, N; nrhs) else # PHASE_SOLVE_BWD_PERM @@ -462,13 +468,12 @@ end # --------------------------------------------------------------------------- # parameters -# data parameters the v0.1 solver computes (PLAN §1.4) +# data parameters the solver computes (PLAN §1.4, §1.7) const SOLVER_OUTPUTS = ("lu_nnz", "flops", "nsuperpanels", "memory_estimates", "perm_reorder_row", - "perm_reorder_col", "perm_row", "perm_col", "diag") + "perm_reorder_col", "perm_row", "perm_col", "diag", "npivots", "inertia", "pivot_stats") # task that provides the other computed data parameters function _output_task(name) - name in ("npivots", "inertia", "pivot_stats") && return "T14/T15" name in ("perm_matching", "scale_row", "scale_col") && return "T21" name in ("schur_shape", "schur_matrix") && return "T20" name == "nd_partition_tree" && return "T24" @@ -483,6 +488,9 @@ 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). +`"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. Computed data parameters (`"lu_nnz"`, `"diag"`, `"perm_row"`, …) cannot be set (`ArgumentError`): read them with [`getparam`](@ref) or [`getparam!`](@ref), which replaces cuDSS's set-buffer-then-get protocol. @@ -494,6 +502,8 @@ function setparam!(solver::DirectSolver, name::AbstractString, value) solver.info = Int(value) 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) + throw(InvalidValueError("pivot_sign has $(length(value)) entries, the matrix has $(size(solver, 1)) rows")) else setparam!(solver.options, name, value) end @@ -523,13 +533,17 @@ The data parameters computed by the solver: | `"memory_estimates"` | `Vector{Int64}` (16 entries, see [`memory_estimates`](@ref)) | 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) | analysis | -| `"diag"` | vector of `T` on the solver's backend: the diagonal of `L` in factor order | factorization | - -The reordering permutation after `"reordering"` alone is the ordering -algorithm's; `"symbolic_factorization"` composes it with the supernodal +| `"diag"` | vector of `T` on the solver's backend: the diagonal of `L` (Cholesky) or of `D` (LDLᵀ/LDLᴴ; for a 2×2 block its two diagonal entries) in factor order | factorization | +| `"npivots"` | `INT`: perturbed pivots (LDLᵀ/LDLᴴ; `0` for Cholesky) | factorization | +| `"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 | + +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 (`"npivots"`, `"inertia"`, `"perm_matching"`, -`"scale_row"`, `"schur_shape"`, …) raise [`NotSupportedError`](@ref); reading +Data parameters of later tasks (`"perm_matching"`, `"scale_row"`, +`"schur_shape"`, …) raise [`NotSupportedError`](@ref); reading one before the phase that computes it raises [`FactorizationError`](@ref). """ function getparam(solver::DirectSolver, name::AbstractString) @@ -547,6 +561,10 @@ function getparam(solver::DirectSolver, name::AbstractString) _need_stage(solver, STAGE_FACTORIZED, name) return _factor_diag(solver) end + if name in ("npivots", "inertia", "pivot_stats") + _need_stage(solver, STAGE_FACTORIZED, name) + return _pivot_output(solver, name) + end _need_stage(solver, STAGE_ANALYZED, name) Sh = solver.host_symbolic sp = Sh.partition @@ -595,10 +613,25 @@ end end end -# diagonal of L in factor order, on the solver's backend (one launch, one work item per supernode) +# npivots / inertia / pivot_stats from the statistics reduced on the device (Cholesky: reduced on demand) +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 +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) function _factor_diag(solver::DirectSolver{T}) where {T} S, N = solver.symbolic, solver.numeric d = KernelAbstractions.zeros(solver.backend, T, S.n) + if _is_ldlt_structure(S.structure) + copyto!(d, 1, N.d, 1, S.n) + 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) diff --git a/src/symbolic/maps.jl b/src/symbolic/maps.jl index 7c93a36..8c5cdf0 100644 --- a/src/symbolic/maps.jl +++ b/src/symbolic/maps.jl @@ -305,7 +305,7 @@ element type `T` and device indices `INT` (default: the index type of | 8 | D (`2n` entries of `T`) | | 9 | update stack (`Layout.stack_len` entries of `T`) | | 10 | device maps (`INT`) | -| 11 | per-front statistics (6 `Int64` per supernode), status (`ns + 1` `Int32`), pivot order (`n` `Int32`) and pivot kinds (`n` `Int8`) | +| 11 | per-front statistics (6 `Int64` per supernode) and their totals (6 `Int64`), status (`ns + 1` `Int32`), pivot order (`n` `Int32`), pivot kinds and sign requests (`n` `Int8` each), `aux` (one `T`) | | 12 | largest regime-A local memory in use (per workgroup, [`subtree_local_bytes`](@ref) of its class) | | 13–16 | 0 (reserved) | @@ -319,8 +319,8 @@ function memory_estimates(S::Symbolic{INT0}, ::Type{T}, ::Type{INT} = INT0) wher est[8] = Int64(L.d_len) * sizeof(T) est[9] = Int64(L.stack_len) * sizeof(T) est[10] = device_map_bytes(S, INT) - est[11] = Int64(nsupernodes(S)) * FRONT_STATS_FIELDS * sizeof(Int64) + - Int64(nsupernodes(S) + 1) * sizeof(Int32) + Int64(S.n) * (sizeof(Int32) + sizeof(Int8)) + est[11] = Int64(nsupernodes(S) + 1) * FRONT_STATS_FIELDS * sizeof(Int64) + + Int64(nsupernodes(S) + 1) * sizeof(Int32) + Int64(S.n) * (sizeof(Int32) + 2 * sizeof(Int8)) + sizeof(T) est[12] = Int64(maximum((subtree_local_bytes(sc.budgets[c]) for c in sc.subtree_class); init = 0)) est[1] = est[7] + est[8] + est[10] + est[11] + Int64(L.work_len) * sizeof(T) est[2] = est[1] + est[9] diff --git a/test/matrices.jl b/test/matrices.jl index e4147b1..72a843e 100644 --- a/test/matrices.jl +++ b/test/matrices.jl @@ -163,6 +163,11 @@ block instead of eliminating a (near-)zero dual pivot on its own. """ kkt_interleaved_perm(nh::Integer, nj::Integer) = vcat([[i, nh + i] for i in 1:nj]..., (nj + 1):nh) +# orderings of the MadNLP-style loops on kkt_matrix(T, 200, 100, 0; hessian = :indefinite) (T14, T15) +const MADNLP_ORDERINGS = (("default ordering", Options()), + ("default ordering, pivot_pairs = all", Options(pivot_pairs = "all")), + ("interleaved", Options(user_perm = kkt_interleaved_perm(200, 100)))) + """ random_general(T, n, density; rng) diff --git a/test/ported/cudss_execution.jl b/test/ported/cudss_execution.jl index cb2278c..8735b08 100644 --- a/test/ported/cudss_execution.jl +++ b/test/ported/cudss_execution.jl @@ -1,8 +1,12 @@ -# Port of `cudss_execution()`, testset "SPD -- HPD" (CUDSS.jl test/test_cudss.jl). +# Port of `cudss_execution()`, testsets "SPD -- HPD" (T13) and "Symmetric -- Hermitian" +# (T15) (CUDSS.jl test/test_cudss.jl). # # Changes: global pivoting `pivot_type = 'C'/'R'` (with `reordering_alg = # "algo2"` in CUDSS.jl) is not planned (PLAN §3.3): setting it raises # `NotSupportedError`, and the case then runs with the default pivot type. +# "Symmetric -- Hermitian": the matrix is `random_symindef(T, n, 0.01)` +# (symmetric/Hermitian indefinite) instead of `sprand + I` symmetrized, and +# CUDSS.jl's shift `A + Diagonal(d)` with `d = rand(R, n)` is kept. function ported_execution_spd(backend, ::Type{T}, ::Type{INT}, view, pivot) where {T, INT} n, p = 100, 5 @@ -45,6 +49,56 @@ function ported_execution_spd(backend, ::Type{T}, ::Type{INT}, view, pivot) wher return nothing end +function ported_execution_sym(backend, ::Type{T}, ::Type{INT}, view, pivot) where {T, INT} + n, p = 100, 5 + R = real(T) + A_cpu = random_symindef(T, n, 0.01) + X_cpu = zeros(T, n, p) + B_cpu = rand(T, n, p) + + A_gpu = api_matrix(backend, triangle_view(A_cpu, view), INT) + X_gpu = to_device(backend, X_cpu) + B_gpu = to_device(backend, B_cpu) + + structure = T <: Real ? "S" : "H" + solver = DirectSolver(A_gpu, structure, view) + if pivot in ('C', 'R') + @test_throws NotSupportedError setparam!(solver, "pivot_type", pivot) + else + setparam!(solver, "pivot_type", pivot) + @test getparam(solver, "pivot_type") == pivot + end + + execute!("analysis", solver, X_gpu, B_gpu) + execute!("factorization", solver, X_gpu, B_gpu) + execute!("solve", solver, X_gpu, B_gpu) + + @test relres(A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + + # In-place LDLᵀ / LDLᴴ + d_cpu = rand(R, n) + A_cpu2 = A_cpu + Diagonal(d_cpu) + update!(solver, api_matrix(backend, triangle_view(A_cpu2, view), INT)) + + C_cpu = rand(T, n, p) + C_gpu = to_device(backend, C_cpu) + + execute!("refactorization", solver, X_gpu, C_gpu) + execute!("solve", solver, X_gpu, C_gpu) + + @test relres(A_cpu2, to_host(X_gpu), C_cpu) <= tol(T) + return nothing +end + +@testset "Symmetric -- Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + @testset "view = $view" for view in ('F', 'L', 'U') + @testset "Pivoting = $pivot" for pivot in ('C', 'R', 'N') + ported_execution_sym(backend, T, INT, view, pivot) + end + end +end + @testset "SPD -- HPD ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES @testset "view = $view" for view in ('F', 'L', 'U') @testset "Pivoting = $pivot" for pivot in ('C', 'R', 'N') diff --git a/test/ported/cudss_generic.jl b/test/ported/cudss_generic.jl index e710fc9..e723139 100644 --- a/test/ported/cudss_generic.jl +++ b/test/ported/cudss_generic.jl @@ -1,4 +1,5 @@ -# Port of `cudss_generic()`, testset "SPD -- HPD" (CUDSS.jl test/test_cudss.jl). +# Port of `cudss_generic()`, testsets "SPD -- HPD" (T13) and "Symmetric -- +# Hermitian" (`ldlt`, `ldlt!`, T15) (CUDSS.jl test/test_cudss.jl). # # Changes: on the CPU backend the matrix is a `CSR` (`cholesky(::SparseMatrixCSC)` # is CHOLMOD's); `CudssMatrix(x)` is `MatrixDescriptor(x)`. `rand(R)` scalings @@ -64,6 +65,73 @@ function ported_generic_spd(backend, ::Type{T}, ::Type{INT}, view) where {T, INT return nothing end +function ported_generic_sym(backend, ::Type{T}, ::Type{INT}, view) where {T, INT} + n, p = 100, 5 + R = real(T) + A_cpu = random_symindef(T, n, 0.01) + B_cpu = rand(T, n, p) + X_cpu = rand(T, n, p) + + A_gpu = api_matrix(backend, triangle_view(A_cpu, view), INT) + B_gpu = to_device(backend, B_cpu) + scaled(c) = api_matrix(backend, triangle_view(c * A_cpu, view), INT) + + @testset "ldlt!" begin + structure = T <: Real ? "S" : "H" + solver = DirectSolver(A_gpu, structure, view) + X_gpu = to_device(backend, X_cpu) + execute!("analysis", solver, X_gpu, B_gpu) + solver = ldlt!(solver, A_gpu) + @test !solver.fresh_factorization + end + + @testset "ldiv!" begin + solver = ldlt(A_gpu; view) + X_gpu = to_device(backend, X_cpu) + ldiv!(X_gpu, solver, B_gpu) + @test relres(A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + + c = rand(R) + R(0.5) + ldlt!(solver, scaled(c)) + X_gpu .= B_gpu + ldiv!(solver, X_gpu) + @test relres(c * A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + + c = rand(R) + R(0.5) + ldlt!(solver, scaled(c)) + X_gpu .= B_gpu + X_desc = MatrixDescriptor(X_gpu) + ldiv!(solver, X_desc) + @test relres(c * A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + + c = rand(R) + R(0.5) + ldlt!(solver, scaled(c)) + X_desc = MatrixDescriptor(X_gpu) + B_desc = MatrixDescriptor(B_gpu) + ldiv!(X_desc, solver, B_desc) + @test relres(c * A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + end + + @testset "\\" begin + solver = ldlt(A_gpu; view) + X_gpu = solver \ B_gpu + @test relres(A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + + c = rand(R) + R(0.5) + ldlt!(solver, scaled(c)) + X_gpu = solver \ B_gpu + @test relres(c * A_cpu, to_host(X_gpu), B_cpu) <= tol(T) + end + return nothing +end + +@testset "Symmetric -- Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + @testset "view = $view" for view in ('F', 'L', 'U') + ported_generic_sym(backend, T, INT, view) + end +end + @testset "SPD -- HPD ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES @testset "view = $view" for view in ('F', 'L', 'U') ported_generic_spd(backend, T, INT, view) diff --git a/test/ported/cudss_solver.jl b/test/ported/cudss_solver.jl index 508a762..bf045f2 100644 --- a/test/ported/cudss_solver.jl +++ b/test/ported/cudss_solver.jl @@ -1,17 +1,18 @@ # Port of `cudss_solver()` with `cudss_solver_data_parameters` and # `cudss_solver_config_parameters` (CUDSS.jl test/test_cudss.jl): the get/set -# loop over every data and configuration parameter, restricted to the SPD/HPD -# structures. +# loop over every data and configuration parameter. # # Changes: -# * structures "SPD" and "HPD" only (real `T`: both; complex `T`: "HPD", and -# "SPD" is checked to be rejected at analysis); +# * structures "S", "H", "SPD" and "HPD" ("G" is T19; complex `T`: "SPD" is +# checked to be rejected at analysis); the matrix is SPD/HPD, so `"inertia"` +# is `(n, 0)` for real "S" and "H" (complex symmetric "S": `(0, 0)`, no +# inertia) and `"diag"` is real positive except for complex "S"; # * `matching_alg = "algo6"` before the analysis is skipped: matching is T21 and # the analysis raises `NotSupportedError` for it (checked); # * cuDSS's set-buffer-then-get protocol for vector data is # `getparam!(buffer, solver, name)`; the matching outputs (`"perm_matching"`, -# `"scale_row"`, `"scale_col"`), `"npivots"`, `"inertia"` and -# `"hybrid_device_memory_min"` are not implemented yet (`NotSupportedError`); +# `"scale_row"`, `"scale_col"`) and `"hybrid_device_memory_min"` are not +# implemented yet (`NotSupportedError`); # * configuration values that the package does not accept raise errors instead # of being passed to the library: `factorization_alg`/`pivot_epsilon_alg` # `"algo3"`–`"algo5"` and `solve_alg` `"algo2"`–`"algo5"` (`InvalidValueError`, @@ -26,8 +27,7 @@ const PORTED_ACCEPTED_ALGOS = Dict( "pivot_epsilon_alg" => ("default", "algo1", "algo2"), ) -const PORTED_NOT_IMPLEMENTED_DATA = ("npivots", "inertia", "perm_matching", "scale_row", "scale_col", - "hybrid_device_memory_min") +const PORTED_NOT_IMPLEMENTED_DATA = ("perm_matching", "scale_row", "scale_col", "hybrid_device_memory_min") function ported_solver_data_parameters(backend, solver, structure, n, memory_estimates, buffer_int::Vector{INT}, buffer_R, buffer_T) where {INT} @@ -54,7 +54,8 @@ function ported_solver_data_parameters(backend, solver, structure, n, memory_est end if parameter == "diag" getparam!(buffer_T, solver, parameter) - @test all(x -> real(x) > 0 && imag(x) == 0, buffer_T) + csym = structure == "S" && eltype(buffer_T) <: Complex # D of a complex symmetric matrix + csym || @test all(x -> real(x) > 0 && imag(x) == 0, buffer_T) end if parameter == "memory_estimates" getparam!(memory_estimates, solver, parameter) @@ -80,6 +81,9 @@ function ported_solver_data_parameters(backend, solver, structure, n, memory_est parameter ∈ ("perm_reorder_row", "perm_reorder_col") && @test isperm(val) parameter == "diag" && @test length(val) == n parameter == "memory_estimates" && @test length(val) == 16 + parameter == "npivots" && @test val == 0 + csym = structure == "S" && eltype(buffer_T) <: Complex + parameter == "inertia" && @test val == (csym ? (0, 0) : (n, 0)) end end end @@ -165,7 +169,7 @@ end INT in INTTYPES n = 20 A_cpu = random_spd(T, n, 1.0) - @testset "structure = $structure" for structure in (T <: Real ? ("SPD", "HPD") : ("HPD",)) + @testset "structure = $structure" for structure in (T <: Real ? ("S", "H", "SPD", "HPD") : ("S", "H", "HPD")) @testset "view = $view" for view in ('L', 'U', 'F') A_gpu = api_matrix(backend, triangle_view(A_cpu, view), INT) ported_solver(backend, T, INT, A_gpu, structure, view, n) diff --git a/test/test_api.jl b/test/test_api.jl index 4422a5e..34a634c 100644 --- a/test/test_api.jl +++ b/test/test_api.jl @@ -122,9 +122,11 @@ end # unknown phase strings: ArgumentError; phases of later tasks: NotSupportedError @test thrown(() -> execute!("solving", solver, x, b)) isa ArgumentError @test thrown(() -> execute!("Solve", solver, x, b)) isa ArgumentError - for phase in ("solve_diag", "solve_refinement", "solve_fwd_schur", "solve_bwd_schur") + for phase in ("solve_refinement", "solve_fwd_schur", "solve_bwd_schur") @test thrown(() -> execute!(phase, solver, x, b)) isa NotSupportedError end + # "solve_diag" (T15) is the identity for Cholesky + @test execute!("solve_diag", solver, x, b) === nothing # a new analysis resets the factorization execute!("analysis", solver, x, b) @test solver.fresh_factorization @@ -149,7 +151,7 @@ end setparam!(solver, "solve_alg", "algo1") @test thrown(() -> execute!("solve", solver, x, b)) isa NotSupportedError setparam!(solver, "solve_alg", "default") - s3 = DirectSolver(api_matrix(backend, tril(A)), sym_structure(T), 'L') + s3 = DirectSolver(api_matrix(backend, A), "G", 'F') # LU: T19 @test thrown(() -> execute!("analysis", s3, x, b)) isa NotSupportedError if T <: Complex s4 = DirectSolver(api_matrix(backend, tril(A)), "SPD", 'L') @@ -210,10 +212,14 @@ end for name in ("lu_nnz", "perm_row", "diag", "nsuperpanels", "memory_estimates") @test thrown(() -> setparam!(solver, name, 1)) isa ArgumentError end - for name in ("npivots", "inertia", "perm_matching", "scale_row", "scale_col", "schur_shape", "schur_matrix", - "nd_partition_tree", "hybrid_device_memory_min", "pivot_stats") + for name in ("perm_matching", "scale_row", "scale_col", "schur_shape", "schur_matrix", "nd_partition_tree", + "hybrid_device_memory_min") @test thrown(() -> getparam(solver, name)) isa NotSupportedError end + # pivot statistics of a successful Cholesky factorization (T15) + @test getparam(solver, "inertia") == (n, 0) && getparam(solver, "inertia") isa Tuple{INT, INT} + @test getparam(solver, "npivots") == 0 && getparam(solver, "npivots") isa INT + @test getparam(solver, "pivot_stats") == (npos = n, nneg = 0, nzero = 0, nperturbed = 0, n2x2 = 0) @test thrown(() -> getparam(solver, "no_such_parameter")) isa ArgumentError @test thrown(() -> setparam!(solver, "no_such_parameter", 1)) isa ArgumentError # configuration parameters go to the solver's options diff --git a/test/test_numeric_ldlt.jl b/test/test_numeric_ldlt.jl new file mode 100644 index 0000000..02cfd1a --- /dev/null +++ b/test/test_numeric_ldlt.jl @@ -0,0 +1,323 @@ +# T15: device LDLᵀ/LDLᴴ in regimes A/B/C (fused kernels, in-front Bunch–Kaufman, perturbation, +# pivot_sign), the diagonal solve and the pivot statistics, against the T14 reference. +# Structure "S" for real T, "H" for complex T (sym_structure), plus complex symmetric "S". + +# regime mixes of the analysis (merged into each case's options) +const LDLT_REGIMES = (("default", (;)), + ("A+B+C, small budgets", (subtree_budgets = [8192, 16384], regime_c_width = 16, + regime_c_rows = 128)), + ("C only", (factorization_alg = "algo2", subtree_budgets = Int[]))) + +# the T14 matrices: (name, matrix, option keywords, structure, compare the inertia with the eigenvalues) +function ldlt_cases(::Type{T}) where {T} + nh, nj = 300, 100 + cases = Any[("random_symindef(400,0.01)", random_symindef(T, 400, 0.01), (;), sym_structure(T), false), + ("kkt(300,100,1e-2)", kkt_matrix(T, nh, nj, 1.0e-2), (;), sym_structure(T), false), + ("kkt(300,100,1e-8)", kkt_matrix(T, nh, nj, 1.0e-8), (;), sym_structure(T), false), + ("kkt(300,100,1e-8) interleaved", kkt_matrix(T, nh, nj, 1.0e-8), + (user_perm = kkt_interleaved_perm(nh, nj),), sym_structure(T), false), + ("random_symindef(300,0.02)", random_symindef(T, 300, 0.02), (;), sym_structure(T), true), + ("kkt indefinite H, δ = 0", kkt_matrix(T, 200, 100, 0.0; hessian = :indefinite), (;), + sym_structure(T), true), + ("kkt 1e-3 H, δ = 0, interleaved (2×2)", + kkt_matrix(T, 200, 100, 0.0; hessian = :indefinite, hessian_scale = 1.0e-3), + (user_perm = kkt_interleaved_perm(200, 100),), sym_structure(T), true)] + if T <: Complex + push!(cases, ("complex symmetric", random_symindef(T, 300, 0.02; hermitian = false), (;), "S", false)) + end + return cases +end + +@testset "device factor equals the reference ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + for (name, A, kw, structure, eig) in ldlt_cases(T), (rname, rkw) in LDLT_REGIMES + @testset "$name, $rname" begin + opts = Options(; kw..., rkw...) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts, structure) + @test SDS.factorize!(Nd, Sd, nz; opts) == 0 + Nh = SDS.host_numeric(Nd) + # the same pivot sequence: local orders and pivot kinds are equal, D and L to rounding + @test Nh.piv == Nr.piv + @test Nh.pivot_kind == Nr.pivot_kind + @test d_error(Nh, Nr) <= panel_tol(T) + @test panel_error(Nh, Nr) <= panel_tol(T) + @test Nh.stats == Nr.stats + st = SDS.pivot_totals(Nd) + @test st == SDS.pivot_stats(Nr) && Nh.totals == Nr.totals + herm = structure == "H" || T <: Real + herm && eig && @test (st.npos, st.nneg) == eigen_npos_nneg(A) + herm && @test st.npos + st.nneg == size(A, 1) + # device solve: forward (local pivot orders), diagonal, backward sweeps + ws = SDS.allocate_solve(Sd, T, backend, 5) + for nrhs in (1, 5), det in (false, true) + b = nrhs == 1 ? rand(T, size(A, 1)) : rand(T, size(A, 1), nrhs) + x = device_solve(backend, ws, Sd, Nd, b; deterministic = det) + if name == "kkt(300,100,1e-8)" && T == ComplexF32 + # T14 / issue #66: this draw needs one refinement step in ComplexF32 (the reference too) + x = x + device_solve(backend, ws, Sd, Nd, b - A * x; deterministic = det) + end + @test relres(A, x, b) <= tol(T) + end + end + end +end + +@testset "regime-A groups and Int64 maps ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + A = kkt_matrix(T, 200, 100, 0.0; hessian = :indefinite, hessian_scale = 1.0e-3) + opts = Options(user_perm = kkt_interleaved_perm(200, 100), subtree_budgets = [8192, 16384, 32768, 49152]) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A, Int64; opts) + @test !isempty(Nd.plan.sub_first) && SDS.pivot_stats(Nr).n2x2 > 0 + SDS.factorize!(Nd, Sd, nz; opts) + Nh = SDS.host_numeric(Nd) + @test Nh.piv == Nr.piv && Nh.pivot_kind == Nr.pivot_kind + @test d_error(Nh, Nr) <= panel_tol(T) && panel_error(Nh, Nr) <= panel_tol(T) + # views of the matrix give the same factor + for view in ('U', 'F') + _, _, Sv, Nv, nzv = ldlt_setup(backend, A, Int64; view, opts) + SDS.factorize!(Nv, Sv, nzv; opts) + @test to_host(Nv.factor) == Nh.factor && to_host(Nv.d) == Nh.d && to_host(Nv.piv) == Nh.piv + end +end + +@testset "perturbation and pivot_sign ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + n, j = 60, 23 + R = real(T) + # the zero pivot in its own root front, regime C ("algo2"), or as the analysis schedules it + for (rname, rkw) in (("default", (;)), ("regime-C root", (factorization_alg = "algo2", subtree_budgets = Int[]))), + stored_zero in (false, true), sgn in (1, -1) + A = singular_block_matrix(T, n, j; stored_zero) + psign = zeros(Int8, n) + psign[j] = sgn + opts = Options(; pivot_sign = psign, rkw...) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts) + @test SDS.factorize!(Nd, Sd, nz; opts) == 0 + Nh = SDS.host_numeric(Nd) + # the same pivot sequence; D to rounding (the device may contract to FMA), perturbed pivots exactly + @test Nh.pivot_kind == Nr.pivot_kind && Nh.piv == Nr.piv + @test d_error(Nh, Nr) <= panel_tol(T) + pert = Nr.pivot_kind .== SDS.PIVOT_KIND_PERTURBED + @test Nh.d[1:n][pert] == Nr.d[1:n][pert] + st = SDS.pivot_totals(Nd) + @test st.nperturbed >= 1 && st.nzero >= 1 + k = findfirst(==(SDS.PIVOT_KIND_PERTURBED), Nh.pivot_kind) + _, _, p = SDS.extract_ldlt(S, Nh) + @test k !== nothing && p[k] == j + @test real(Nh.d[k]) == sgn * R(SDS.default_pivot_epsilon(R)) && iszero(imag(Nh.d[k])) + @test (st.npos, st.nneg) == (sgn > 0 ? (n, 0) : (n - 1, 1)) + if rname == "regime-C root" + s = findfirst(s -> k in SDS.sncols(S.partition, s), 1:SDS.nsupernodes(S)) + @test S.partition.snparent[s] == 0 && SDS.takes_c_path(S.schedule, s) + end + x0 = rand(T, n) + x0[j] = 0 + b = A * x0 + ws = SDS.allocate_solve(Sd, T, backend, 1) + @test relres(A, device_solve(backend, ws, Sd, Nd, b), b) <= 1.0e-6 + end + # pivot_epsilon, the scaled variant (ε max|aᵢⱼ| on the device), 'N' and 'D' + A = singular_block_matrix(T, n, j) + for opts in (Options(pivot_epsilon = 1.0e-3, pivot_sign = Int8.(-(1:n .== j))), + Options(pivot_epsilon = 1.0e-3, pivot_epsilon_alg = "algo1"), + Options(pivot_type = 'N'), Options(pivot_type = 'D')) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + Nh = SDS.host_numeric(Nd) + @test Nh.pivot_kind == Nr.pivot_kind && Nh.piv == Nr.piv + @test d_error(Nh, Nr) <= panel_tol(T) + pert = Nr.pivot_kind .== SDS.PIVOT_KIND_PERTURBED + @test Nh.d[1:n][pert] == Nr.d[1:n][pert] + @test SDS.pivot_totals(Nd).nperturbed == 1 + end + # pivot_sign of the wrong length + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A) + @test thrown(() -> SDS.factorize!(Nd, Sd, nz; opts = Options(pivot_sign = ones(Int8, n - 1)))) isa + InvalidValueError +end + +@testset "pivot_type 'D' and 'N' on quasi-definite KKT ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + nh, nj = 300, 100 + A = kkt_matrix(T, nh, nj, 1.0e-2) + for pt in ('D', 'N') + opts = Options(pivot_type = pt) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + Nh = SDS.host_numeric(Nd) + @test Nh.piv == Nr.piv && Nh.pivot_kind == Nr.pivot_kind && d_error(Nh, Nr) <= panel_tol(T) + st = SDS.pivot_totals(Nd) + @test st.n2x2 == 0 && st.nperturbed == 0 && (st.npos, st.nneg) == (nh, nj) + pt == 'N' && @test Nh.piv == 1:(nh + nj) + end +end + +@testset "determinism and refactorization ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + A = kkt_matrix(T, 200, 100, 0.0; hessian = :indefinite, hessian_scale = 1.0e-3) + opts = Options(user_perm = kkt_interleaved_perm(200, 100), subtree_budgets = [8192, 16384], regime_c_width = 16, + regime_c_rows = 128) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + factor, d, piv = to_host(Nd.factor), to_host(Nd.d), to_host(Nd.piv) + # new values, then the old ones again: bitwise the same factor + nz2 = to_device(backend, 2 .* to_host(nz)) + SDS.factorize!(Nd, Sd, nz2; opts) + @test to_host(Nd.d) ≈ 2 .* d + SDS.factorize!(Nd, Sd, nz; opts) + @test to_host(Nd.factor) == factor && to_host(Nd.d) == d && to_host(Nd.piv) == piv + # solves: the deterministic sweep is bitwise reproducible + ws = SDS.allocate_solve(Sd, T, backend, 2) + b = rand(T, 300, 2) + @test device_solve(backend, ws, Sd, Nd, b; deterministic = true) == + device_solve(backend, ws, Sd, Nd, b; deterministic = true) + # no allocation in the numeric phase and the solve (CPU backend) + if backend isa CPU + SDS.factorize!(Nd, Sd, nz; opts) + @test (@allocated SDS.factorize!(Nd, Sd, nz; opts)) <= ldlt_alloc_budget(S) + x = similar(b) + SDS.sweep_solve!(x, ws, Sd, Nd, b) + @test (@allocated SDS.sweep_solve!(x, ws, Sd, Nd, b)) <= solve_alloc_budget(S, ws; ldlt = true) + end +end + +@testset "solve sweeps: diagonal step and dense path ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + A = kkt_matrix(T, 200, 100, 0.0; hessian = :indefinite, hessian_scale = 1.0e-3) + for (rname, rkw) in LDLT_REGIMES + opts = Options(; user_perm = kkt_interleaved_perm(200, 100), rkw...) + S, Nr, Sd, Nd, nz = ldlt_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + ws = SDS.allocate_solve(Sd, T, backend, 3) + b = rand(T, 300, 3) + bd = to_device(backend, b) + # the sweeps one at a time: Y = P b, L Z = Y (local pivot orders), D W = Z, Lᴴ V = W + SDS.permute_rhs!(ws.Y, bd, Sd.perm) + Y = to_host(ws.Y) + SDS.forward_sweep!(ws, Sd, Nd; deterministic = true) + Z = to_host(ws.Y) + SDS.diagonal_sweep!(ws, Sd, Nd) + W = to_host(ws.Y) + SDS.backward_sweep!(ws, Sd, Nd) + xd = similar(bd) + SDS.unpermute_solution!(xd, ws.Y, Sd.perm) + @test to_host(xd) == device_solve(backend, ws, Sd, Nd, b; deterministic = true) + @test relres(A, to_host(xd), b) <= tol(T) + # against the extracted factors (A[q, q] = L D Lᴴ, q = perm ∘ piv): Y[piv] = L Z, Z = D W + L, D, _ = SDS.extract_ldlt(S, SDS.host_numeric(Nd)) + piv = Int.(to_host(Nd.piv)) + @test norm(L * Z - Y[piv, :]) <= tol(T) * norm(Y) + @test norm(D * W - Z) <= tol(T) * norm(Z) + # every dense implementation of the regime-C path + if rname == "C only" + for impl in SDS.dense_impls(:trsm, backend, T) + impl in SDS.dense_impls(:gemm, backend, T) || continue + @test relres(A, device_solve(backend, ws, Sd, Nd, b; impl), b) <= tol(T) + end + end + end +end + +@testset "MadNLP-style inertia correction on the device ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + nh, nj = 200, 100 + A = kkt_matrix(T, nh, nj, 0.0; hessian = :indefinite) + for (label, opts) in MADNLP_ORDERINGS + @testset "$label" begin + done, iterations, solver, δw = madnlp_inertia_loop_device(backend, A, nh, nj, opts) + @test done && iterations > 1 + @test getparam(solver, "npivots") == 0 + # the same answer as the reference loop + done_r, iterations_r, _, _, δw_r = madnlp_inertia_loop(A, nh, nj, opts) + @test done_r && iterations == iterations_r && δw == δw_r + Aw = A + spdiagm(0 => [fill(T(δw), nh); zeros(T, nj)]) + b = rand(T, nh + nj) + x = to_device(backend, zeros(T, nh + nj)) + execute!("solve", solver, x, to_device(backend, b)) + @test relres(Aw, to_host(x), b) <= tol(T) + end + end +end + +@testset "public API ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES + nh, nj = 200, 100 + n = nh + nj + A = kkt_matrix(T, nh, nj, 0.0; hessian = :indefinite, hessian_scale = 1.0e-3) + solver = DirectSolver(api_matrix(backend, tril(A), INT), sym_structure(T), 'L') + setparam!(solver, "user_perm", kkt_interleaved_perm(nh, nj)) + setparam!(solver, "deterministic_mode", 1) + for name in ("npivots", "inertia", "pivot_stats") + @test thrown(() -> getparam(solver, name)) isa FactorizationError + end + b = to_device(backend, rand(T, n, 2)) + x, y = similar(b), similar(b) + execute!("analysis", solver, x, b) + execute!("factorization", solver, x, b) + @test getparam(solver, "info") == 0 + # the pivot statistics against the reference factorization of the same analysis + S = solver.host_symbolic + Nr = SDS.allocate_numeric(S, T) + SDS.ref_ldlt!(Nr, S, SDS.CSR(tril(A)).nzval; opts = solver.options) + st = SDS.pivot_stats(Nr) + @test st.n2x2 > 0 + @test getparam(solver, "pivot_stats") == st + @test getparam(solver, "inertia") == (st.npos, st.nneg) == eigen_npos_nneg(A) + @test getparam(solver, "inertia") isa Tuple{INT, INT} + @test getparam(solver, "npivots") == 0 && getparam(solver, "npivots") isa INT + d = getparam(solver, "diag") + @test typeof(KernelAbstractions.get_backend(d)) == typeof(backend) + @test norm(to_host(d) - Nr.d[1:n]) <= panel_tol(T) * norm(Nr.d) + hd = zeros(T, n) + getparam!(hd, solver, "diag") + @test hd == to_host(d) + # the sub-phases, "solve_diag" included, compose to "solve" + execute!("solve", solver, x, b) + @test relres(A, to_host(x), to_host(b)) <= tol(T) + for phase in ("solve_fwd_perm", "solve_fwd", "solve_diag", "solve_bwd", "solve_bwd_perm") + execute!(phase, solver, y, b) + end + @test to_host(y) == to_host(x) + # logabsdet = log|det A| with the sign of det A + la, sg = logabsdet(solver) + la_ref, sg_ref = logabsdet(Matrix(A)) + @test la ≈ la_ref rtol = sqrt(eps(real(T))) + @test sg ≈ sg_ref + # pivot_sign: wrong length rejected; host or device vector used by the next factorization + @test thrown(() -> setparam!(solver, "pivot_sign", ones(Int8, n - 1))) isa InvalidValueError + setparam!(solver, "pivot_sign", to_device(backend, [ones(Int8, nh); -ones(Int8, nj)])) + @test getparam(solver, "pivot_sign") == [ones(Int8, nh); -ones(Int8, nj)] + execute!("refactorization", solver, x, b) + @test getparam(solver, "inertia") == (st.npos, st.nneg) + # a zero pivot is perturbed with the requested sign + j = 23 + Z = singular_block_matrix(T, 60, j) + for sgn in (1, -1) + zs = DirectSolver(api_matrix(backend, tril(Z), INT), sym_structure(T), 'L') + setparam!(zs, "pivot_sign", Int8.(sgn .* (1:60 .== j))) + execute!("analysis", zs, nothing, nothing) + execute!("factorization", zs, nothing, nothing) + @test getparam(zs, "npivots") == 1 + @test getparam(zs, "inertia") == (sgn > 0 ? (60, 0) : (59, 1)) + ps = getparam(zs, "pivot_stats") + @test ps.nperturbed == 1 && ps.nzero == 1 + end + # the LinearAlgebra layer + F = ldlt(SDS.CSR(SparseMatrixCSC{T, INT}(A)); view = 'F') + backend isa CPU && @test relres(A, F \ to_host(b), to_host(b)) <= tol(T) + @test F.structure == SDS.convert(SDS.Structure, sym_structure(T)) +end + +@testset "complex symmetric through the API ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in COMPLEX_ELTYPES + A = random_symindef(T, 300, 0.02; hermitian = false) + solver = DirectSolver(api_matrix(backend, tril(A), Int32), "S", 'L') + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + @test getparam(solver, "inertia") == (0, 0) + b = rand(T, 300, 3) + x = to_device(backend, zeros(T, 300, 3)) + execute!("solve", solver, x, to_device(backend, b)) + @test relres(A, to_host(x), b) <= tol(T) + la, sg = logabsdet(solver) + la_ref, sg_ref = logabsdet(Matrix(A)) + @test la ≈ la_ref rtol = sqrt(eps(real(T))) + @test sg ≈ sg_ref rtol = sqrt(eps(real(T))) +end diff --git a/test/test_reference_ldlt.jl b/test/test_reference_ldlt.jl index 5fdbfb8..e8ec8b5 100644 --- a/test/test_reference_ldlt.jl +++ b/test/test_reference_ldlt.jl @@ -238,10 +238,6 @@ end end end -const MADNLP_ORDERINGS = (("default ordering", Options()), - ("default ordering, pivot_pairs = all", Options(pivot_pairs = "all")), - ("interleaved", Options(user_perm = kkt_interleaved_perm(200, 100)))) - @testset "MadNLP-style inertia correction ($label): $T" for T in ELTYPES, (label, opts) in MADNLP_ORDERINGS nh, nj = 200, 100 # primal regularization δw on an indefinite H, no dual regularization (δ = 0); the diff --git a/test/utils.jl b/test/utils.jl index 237bd6e..b3e50a5 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -284,19 +284,22 @@ function device_solve(backend, ws, Sd, Nd, b::AbstractArray; kwargs...) end """ - solve_alloc_budget(S, ws) -> Int + solve_alloc_budget(S, ws; ldlt = false) -> Int [`ka_cpu_alloc_budget`](@ref) for one `sweep_solve!` with the plan of `ws` on the CPU backend: two permutation launches, a forward and a backward launch per kernel launch of the plan, and per regime-C-path front up to three dense calls or kernels each way plus one pull launch per dense step; `@localmem` of -the control array (`_SV_CTL` indices) per launch. +the control array (`_SV_CTL` indices) per launch. `ldlt = true` (T15) adds +the diagonal launch and, per regime-C-path front, the local pivot order kernel +each way. """ -function solve_alloc_budget(S, ws) +function solve_alloc_budget(S, ws; ldlt::Bool = false) plan = ws.plan - launches = 2 + launches = ldlt ? 3 : 2 + per_front = ldlt ? 8 : 6 for k in eachindex(plan.kind) - launches += plan.kind[k] == SparseDirectSolver.SOLVE_DENSE ? 6 * (plan.last[k] - plan.first[k] + 1) + 1 : 2 + launches += plan.kind[k] == SparseDirectSolver.SOLVE_DENSE ? per_front * (plan.last[k] - plan.first[k] + 1) + 1 : 2 end return ka_cpu_alloc_budget(launches, launches * SparseDirectSolver._SV_CTL * sizeof(Int64)) end @@ -358,3 +361,78 @@ function madnlp_inertia_loop(A::SparseMatrixCSC{T}, nh, nj, opts) where {T} end return (false, iterations, S, N, δw) end + +# --- device LDLᵀ/LDLᴴ (T15) --- + +""" + ldlt_setup(backend, A, INT = Int32; view = 'L', opts = Options(), structure = sym_structure(T)) + -> (S, Nr, Sd, Nd, nz) + +Host analysis `S` of the `view` triangle of the symmetric/Hermitian matrix `A` +with its reference LDLᵀ/LDLᴴ factor `Nr` ([`reference_ldlt`](@ref)), the +analysis adapted to `backend` with `INT` maps, device storage `Nd` and the +values `nz` on `backend` (T15). +""" +function ldlt_setup(backend, A::SparseMatrixCSC{T}, ::Type{INT} = Int32; view = 'L', opts = Options(), + structure = sym_structure(T)) where {T, INT} + S, Nr, info, C = reference_ldlt(A; view, opts, structure) + info == 0 || error("ldlt_setup: the reference factorization failed") + Sd = SparseDirectSolver.adapt(backend, S, INT) + Nd = SparseDirectSolver.allocate_numeric(Sd, T, backend) + return S, Nr, Sd, Nd, to_device(backend, C.nzval) +end + +""" + d_error(Nh, Nr) + +`max |D_device - D_ref| / max |D_ref|` (both entries `d[k]` and the 2×2 +subdiagonals `d[n + k]`) of a host copy `Nh` of a device LDLᵀ factor and the +reference factor `Nr`; compare with [`panel_tol`](@ref) (T15). +""" +d_error(Nh, Nr) = maximum(abs, Nh.d - Nr.d) / maximum(abs, Nr.d) + +""" + ldlt_alloc_budget(S) -> Int + +[`ka_cpu_alloc_budget`](@ref) for one LDLᵀ/LDLᴴ `factorize!` on the CPU +backend: one launch per launch group, the `max |aᵢⱼ|` and statistics +reductions; `@localmem` of the largest regime-A kernel (48 KiB) per group plus +the two reductions (T15). +""" +function ldlt_alloc_budget(S) + launches = length(S.schedule.groups) + 2 + localmem = launches * maximum(SparseDirectSolver.SUBTREE_LOCAL_SIZES) + + SparseDirectSolver.STATS_WORKGROUP * SparseDirectSolver.FRONT_STATS_FIELDS * sizeof(Int64) + return ka_cpu_alloc_budget(launches, localmem) +end + +""" + madnlp_inertia_loop_device(backend, A, nh, nj, opts, INT = Int32) -> (done, iterations, solver, δw) + +The MadNLP-style inertia correction of [`madnlp_inertia_loop`](@ref) through +the handle layer on `backend` (T15): one `DirectSolver` (`"S"`/`"H"`, lower +triangle, options `opts`), one `"analysis"`; per iteration the regularized +values are written into the solver's `nzval`, then `"factorization"` / +`"refactorization"`, and `getparam(solver, "inertia")`, `"npivots"` decide. +""" +function madnlp_inertia_loop_device(backend, A::SparseMatrixCSC{T}, nh, nj, opts, ::Type{INT} = Int32) where {T, INT} + C = SparseDirectSolver.CSR(tril(A)) + dpos = [C.rowptr[i] - 1 + findfirst(==(i), C.colval[C.rowptr[i]:(C.rowptr[i + 1] - 1)]) for i in 1:nh] + solver = DirectSolver(api_matrix(backend, tril(A), INT), sym_structure(T), 'L') + solver.options = opts + execute!("analysis", solver, nothing, nothing) + δw, iterations = 0.0, 0 + nz = similar(C.nzval) + while iterations < 50 + iterations += 1 + nz .= C.nzval + nz[dpos] .+= T(δw) + copyto!(solver.A.nzval, nz) + execute!(iterations == 1 ? "factorization" : "refactorization", solver, nothing, nothing) + if getparam(solver, "inertia") == (nh, nj) && getparam(solver, "npivots") == 0 + return (true, iterations, solver, δw) + end + δw = δw == 0 ? 1.0e-4 : 8 * δw + end + return (false, iterations, solver, δw) +end