diff --git a/TASKS.md b/TASKS.md index 02fff95..91253e6 100644 --- a/TASKS.md +++ b/TASKS.md @@ -2141,7 +2141,7 @@ whenever `atomic_add` is `false`. Close #36 in this task's PR. --- -## T16 — Iterative refinement, solve sub-phases, `solve_mode`, interrupt, logging `[ ]` +## T16 — Iterative refinement, solve sub-phases, `solve_mode`, interrupt, logging `[!]` **Reads**: PLAN §1.2, §1.3 (`ir_*`, `solve_mode`), §1.4 (`user_host_interrupt`, `ir_n_steps` data), §2.5. @@ -2182,6 +2182,94 @@ repeat it once T21 has landed. The bar is cuDSS with `"matching_alg" = * Setting the interrupt flag before `factorize!` raises `InterruptedError`; clearing it and calling `factorize!` again succeeds. +### Report + +- Status: [!] (done on the CPU backend; the "100× on a badly scaled SPD matrix" assertion is `@test_broken`, see + deviations; CUDA from CI) +- What was built: + - `src/solve/refinement.jl` (new): `RefinementWorkspace` (device CSR of the full matrix over the contributions + of the user's `nzval`, signed source index = conjugated mirror entry for `"H"`/`"HPD"`, built from the existing + `full_pattern_map` by `refinement_map`; residual `R`, a copy `Bc` of `B`, column norms on device + host copy), + `allocate_refinement`, kernels `residual!` (gather SpMV, one work item per (row, rhs), no atomics), + `residual_norms!` (one workgroup per rhs, `@localmem` tree reduction, one host read), `add_correction!`, + `copy_rhs!`, the loop `refine!` (poll interrupt → residual → early-exit test only when `ir_tol > 0` → permute, + forward/diagonal/backward sweeps → `X += correction`), and `solve_conjugated` (when `op(A)` is `conj(M)`). + - `src/solve/permute.jl`: `permute_rhs!`/`unpermute_solution!` take `conjugate`; `_with_flags` turns runtime + flags into compile-time `Val`s without dynamic dispatch. + - `src/solver.jl`: `DirectSolver` gains the type parameter `RF` and the fields `refinement` (allocated lazily + by the first refining solve, reset by a new analysis, grown with `nrhs`) and `ir_steps`. + `"solve_refinement"` is a phase; `"solve"` = sweeps + refinement; the six sub-phases compose bitwise to + `"solve"`. `"solve"` with `X === B` (aliasing via `Base.mightalias`) copies `B` first; + `"solve_refinement"` with aliasing raises `InvalidValueError`. `solve_mode` 0/1/2 for all four structures: + `op(A)` is always `M` or `conj(M)` of the stored matrix `M`, solved as `conj(M⁻¹ conj(b))` (conjugating + permutations), and the residual uses the conjugated values. This also lifts the T13 restriction on complex + Hermitian CSC input (`A = Mᵀ = conj(M)`). `getparam(solver, "ir_n_steps")` = steps performed by the last + solve; the configured value before a solve and after `setparam!(solver, "ir_n_steps", k)`. + `ir_mode = "fgmres"` raises `NotSupportedError` (T18). + - `user_host_interrupt`: polled (host read, no sync) before every launch group of `factorize!`/`factorize_ldlt!`, + at the start of `"reordering"`/`"symbolic_factorization"` and before every refinement step. An interrupted + factorization/refactorization sets the solver back to "analyzed" and resets `fresh_factorization` (solve and + `"refactorization"` then raise `FactorizationError`; `"factorization"`, and `cholesky!`/`ldlt!`, work once the + flag is cleared). An interrupted refinement leaves the last completed iterate in `X` and reports the completed + steps in `"ir_n_steps"` (`refine!(…; progress)`, review round 1). + - `src/logging.jl` (new): `SDS_LOG_LEVEL` (`0`/`none`, `1`/`info`, `2`/`debug`, read in `__init__`), + `SparseDirectSolver.set_log_level!` (unexported). Phase summaries (reordering, symbolic factorization, + factorization with host wall time of the launches, refinement steps) as `@info` at level ≥ 1, per-step + relative residuals at level 2; below the level they are `@debug` records (visible with `JULIA_DEBUG`). + - `src/generic.jl`: `cholesky`/`ldlt` set `ir_n_steps = 2` on the solvers they create (handle layer stays 0). + - `test/test_refinement.jl` (new), `test/matrices.jl`: `badly_scaled_spd(T, n, density; exponent = 4)`. + `bench/refinement.jl` (new): the relres table below. + - `test/test_api.jl`: the assertions that `"solve_refinement"`, complex `solve_mode = 1`, complex Hermitian CSC + input raise `NotSupportedError` and that `ir_n_steps` warns are replaced by assertions of the results. +- Tests (after review round 1): `SDS_TEST_GPU=0 SDS_TEST_ONLY=test_refinement`: 517 pass, 1 broken (CPU). Full + `SDS_TEST_GPU=0 julia --project=. -e 'using Pkg; Pkg.test()'`: 60490 pass, 0 fail, 1 broken (the one above), 13.8 min. CUDA/AMDGPU: pending CI on the PR. + CI round 2 (CUDA H200): the five bitwise `==` checks between solves (X === B, row-major, strided, sub-phases vs + `"solve"`) failed for Float64 by last-bit differences: the default atomic forward sweep sums in a run-dependent + order on GPUs. Those two testsets now set `deterministic_mode = 1` (the documented reproducible path); test counts + unchanged. +- Measurements: relres `‖b − Ax‖/‖b‖`, `b = A·1`, handle layer `"S"`, view `'L'`, default pivoting, KA CPU + backend, after `ir_n_steps` steps (`bench/refinement.jl`; dumps generated in this session with + `bench/dump_madnlp_kkt.jl`, MadNLP/ExaModelsPower current releases): + + | K2 dump | n | npivots | ir 0 | ir 2 | ir 5 | + | --- | ---: | ---: | ---: | ---: | ---: | + | case118 k2_1 | 3150 | 5 | 0.042 | 3.0e-10 | 1.1e-15 | + | case118 k2_10 | 3150 | 3 | 2.1e-5 | 3.8e-11 | 3.9e-17 | + | case118 k2_20 | 3150 | 3 | 7.1e-7 | 1.5e-13 | 3.8e-18 | + | case1354 k2_1 | 33811 | 34 | 0.14 | 6.5e-4 | 7.9e-7 | + | case1354 k2_10 | 33811 | 48 | 0.042 | 5.0e-3 | 5.7e-5 | + | case1354 k2_20 | 33811 | 32 | 7.3e-3 | 8.5e-5 | 1.6e-7 | + + case118 beats the cuDSS bar (3.7e-13 with matching + 5 steps) without matching; case1354 is still 1e4–1e6 above + its bar (4.6e-11) and needs T21's scaling (repeat this table then). On a badly scaled SPD matrix + (`badly_scaled_spd(Float64, 300, 0.02)`) one step changes relres by 2×–7× only (3e-16 → 6e-17); on the KKT + generator with perturbed pivots (`kkt_matrix(150, 60, 0)`, `pivot_pairs = "none"`) relres goes 1.5e-4 → 2.4e-9 + → … and `ir_tol = 1e-14` stops after 3 steps. `test_refinement` takes 9 s inside the full run, ~6 min alone + (compilation of the kernel variants over 4 element types × structures × flags). +- Deviations from PLAN.md / this task: + - "On a badly row-scaled SPD matrix one refinement step reduces relres by ≥ 100×" cannot be observed: Cholesky + is backward stable and invariant under symmetric diagonal scaling (the only scaling that keeps the matrix + SPD). Measured on 3 generators × 3 scalings × 5 right-hand-side choices: the unrefined relres is either at + rounding level (`b = A x`) or limited by `cond(A)` (random `b`), and one step gains ≤ 7× (once 42×). The SPD + assertion is `@test_broken` with that comment; the 100× reduction is asserted on a KKT matrix whose + factorization carries static pivot perturbations (the case refinement is meant for, issue #71), together with + the `ir_tol`/early-exit/steps-performed checks (also run on the SPD matrix, where they pass). + - `ir_tol` is the largest `‖Rₖ‖₂/‖Bₖ‖₂` over the right-hand sides; checking it costs one host synchronization + per step, so it runs only when `ir_tol > 0`. The LinearAlgebra layer sets `ir_n_steps = 2` but keeps + `ir_tol = 0` (PLAN §3.1 "early exit on ir_tol" is available once the user sets `ir_tol`), so `ldiv!` adds no + synchronizations beyond its final one. + - The refinement workspace (map ≈ 2·nnz(full) `INT`s plus `2·n·nrhs` values) is allocated by the first solve + that refines, not at analysis, so the handle layer with `ir_n_steps = 0` (MadNLP) pays nothing. + - The interrupt is polled between launch groups, not per level inside a group, and not inside the solve sweeps. +- Open issues / follow-ups: + - `"solve_refinement"` and the residual read the solver's current values: `update!` without a refactorization + refines towards the new matrix (tested), which MadNLP could use to skip refactorizations. +- Suggested plan changes: + - TASKS T16 test 1: phrase the 100× criterion on a factorization with active static perturbation (or with + Float32 factors, §3.8), not on a scaled SPD matrix. + - PLAN §1.4: document `getparam(solver, "ir_n_steps")` as "steps performed after a solve, configured value + otherwise" (one name for config and data in CUDSS.jl). + --- ## T17 — Uniform batch (v1 cut line) `[ ]` diff --git a/bench/refinement.jl b/bench/refinement.jl new file mode 100644 index 0000000..b4f7891 --- /dev/null +++ b/bench/refinement.jl @@ -0,0 +1,68 @@ +# Iterative refinement on the MadNLP K2 dumps (T16, issue #71): relres = +# ‖b − A x‖ / ‖b‖ after `ir_n_steps` ∈ {0, 2, 5} with the handle-layer solver +# (structure "S", view 'L', default pivoting), b = A * ones. Prints a Markdown +# table; dumps come from bench/dump_madnlp_kkt.jl (bench/data/kkt_*_k2_*.mtx). +# Runs on the KA CPU backend, or on CUDA with --backend=cuda (needs CUDA). +# +# julia --project=. bench/refinement.jl [--only=substring,...] [--steps=0,2,5] [--backend=cpu|cuda] + +using LinearAlgebra +using SparseArrays +using SparseDirectSolver + +include(joinpath(@__DIR__, "matrices.jl")) + +function parse_args(args) + o = Dict{String, String}("only" => "", "steps" => "0,2,5", "backend" => "cpu") + for a in args + m = match(r"^--(only|steps|backend)=(.*)$", a) + m === nothing && error("unknown argument \"$a\"; see the header of bench/refinement.jl") + o[m[1]] = m[2] + end + return o +end + +const OPTS = parse_args(ARGS) +const STEPS = parse.(Int, split(OPTS["steps"], ',')) + +if OPTS["backend"] == "cuda" + @eval using CUDA + to_dev(x::AbstractVector) = CuArray(x) + to_dev(A::SparseMatrixCSC) = CUDA.CUSPARSE.CuSparseMatrixCSR(A) +else + to_dev(x) = x +end + +function relres_table(name, A::SparseMatrixCSC{T}) where {T} + L = tril(A) + solver = DirectSolver(to_dev(L), "S", 'L') + analyze!(solver) + factorize!(solver; asynchronous = false) + b = A * ones(T, size(A, 1)) + bd = to_dev(b) + xd = similar(bd) + res = Float64[] + for k in STEPS + setparam!(solver, "ir_n_steps", k) + solve!(solver, xd, bd; asynchronous = false) + x = Array(xd) + push!(res, norm(b - A * x) / norm(b)) + end + println("| ", name, " | ", size(A, 1), " | ", getparam(solver, "npivots"), " | ", + join((string(round(r; sigdigits = 2)) for r in res), " | "), " |") + return res +end + +selected(name) = isempty(OPTS["only"]) || any(s -> occursin(s, name), split(OPTS["only"], ',')) + +dumps = isdir(BenchMatrices.DATA_DIR) ? + sort([f for f in readdir(BenchMatrices.DATA_DIR) if (d = BenchMatrices.parse_dump_name(f)) !== nothing && + d.kind == "k2" && d.ext == "mtx"]) : String[] +isempty(dumps) && println("no K2 dumps in ", BenchMatrices.DATA_DIR, " (run bench/dump_madnlp_kkt.jl)") +println("| K2 dump | n | npivots | ", join(("relres, ir_n_steps = $k" for k in STEPS), " | "), " |") +println("| --- | ---: | ---: | ", join(("---:" for _ in STEPS), " | "), " |") +for f in dumps + selected(f) || continue + A = BenchMatrices.symmetrize_triangle(tril(BenchMatrices.read_mtx(joinpath(BenchMatrices.DATA_DIR, f)))) + relres_table(f, A) +end diff --git a/src/SparseDirectSolver.jl b/src/SparseDirectSolver.jl index 5c77ab7..16ab9ad 100644 --- a/src/SparseDirectSolver.jl +++ b/src/SparseDirectSolver.jl @@ -26,6 +26,7 @@ export CSR, csr_of_transpose, to_backend, nbatch, MatrixDescriptor, update! export AbstractDirectSolver, DirectSolver, execute!, analyze!, factorize!, refactorize!, solve!, getparam! include("errors.jl") +include("logging.jl") include("types.jl") include("options.jl") include("matrix.jl") @@ -66,9 +67,12 @@ include("numeric/extract.jl") # solve phase on the device (PLAN §2.5) include("solve/permute.jl") include("solve/sweeps.jl") +include("solve/refinement.jl") # public API (PLAN §3.1): handle-style layer and LinearAlgebra layer include("solver.jl") include("generic.jl") +__init__() = _init_log_level() + end # module SparseDirectSolver diff --git a/src/generic.jl b/src/generic.jl index 5c9cfbc..1fa5d8b 100644 --- a/src/generic.jl +++ b/src/generic.jl @@ -4,23 +4,34 @@ # `Symmetric`/`Hermitian` wrappers, as CUDSS.jl does. A `SparseMatrixCSC` is not # accepted by `cholesky` (that method belongs to CHOLMOD): wrap it with `CSR(A)`. # -# PLAN §3.1 turns iterative refinement on in this layer (`ir_n_steps = 2` with -# early exit). Refinement arrives in T16; until then the layer keeps the -# handle-layer default `ir_n_steps = 0`, and a value set by the user is stored -# and ignored (the solve warns once). +# PLAN §3.1 turns iterative refinement on in this layer: the solvers created +# by `cholesky`/`ldlt` start with `ir_n_steps = 2` (the handle layer keeps the +# cuDSS default 0); `ir_tol` (default 0: no early exit, no synchronization per +# step) stops it early when the user sets it. Both stay settable on the solver. + +# the defaults of this layer on a new solver +function _linear_algebra_defaults!(solver::DirectSolver) + setparam!(solver, "ir_n_steps", LINEAR_ALGEBRA_IR_STEPS) + return solver +end + +"Refinement steps of the solvers created by `cholesky` and `ldlt` (PLAN §3.1; the handle layer uses 0)." +const LINEAR_ALGEBRA_IR_STEPS = 2 """ cholesky(A::CSR, NoPivot(); view = 'F', check = false) -> DirectSolver LLᵀ (real `T`, structure `"SPD"`) or LLᴴ (complex `T`, `"HPD"`) factorization of the sparse matrix `A` on its backend: a [`DirectSolver`](@ref) after -`"analysis"` and `"factorization"` (synchronized). `view` selects the triangle +`"analysis"` and `"factorization"` (synchronized), with `ir_n_steps = 2` +steps of iterative refinement in its solves (the handle layer's default is 0; +change it with [`setparam!`](@ref)). `view` selects the triangle of `A` that is read (`'L'`, `'U'`, `'F'`). As in CUDSS.jl a failed factorization does not throw unless `check = true` (then [`FactorizationError`](@ref)); otherwise check `getparam(solver, "info")`. """ function LinearAlgebra.cholesky(A::CSR{T}, ::NoPivot = NoPivot(); view::Char = 'F', check::Bool = false) where {T} - solver = DirectSolver(A, T <: Real ? "SPD" : "HPD", view) + solver = _linear_algebra_defaults!(DirectSolver(A, T <: Real ? "SPD" : "HPD", view)) execute!("analysis", solver, nothing, nothing) execute!("factorization", solver, nothing, nothing; asynchronous = false) _check_info(solver, check) @@ -39,13 +50,13 @@ end 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 +(synchronized, `ir_n_steps = 2` as [`cholesky`](@ref)), 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) + solver = _linear_algebra_defaults!(DirectSolver(A, T <: Real ? "S" : "H", view)) execute!("analysis", solver, nothing, nothing) execute!("factorization", solver, nothing, nothing; asynchronous = false) _check_info(solver, check) diff --git a/src/logging.jl b/src/logging.jl new file mode 100644 index 0000000..64cd18a --- /dev/null +++ b/src/logging.jl @@ -0,0 +1,63 @@ +# Logging and the host interrupt flag (PLAN §1.4 `user_host_interrupt`). +# +# Messages go through Julia's logging: phase summaries and refinement details +# are emitted with `@info` when the environment variable `SDS_LOG_LEVEL` (read +# when the package is loaded, or set later with `set_log_level!`) asks for +# them, and with `@debug` otherwise (visible with `JULIA_DEBUG=SparseDirectSolver`). +# Messages are built only when they are emitted. + +const LOG_NONE = 0 +const LOG_INFO = 1 +const LOG_DEBUG = 2 + +const LOG_LEVEL = Ref(LOG_NONE) + +function _parse_log_level(value) + s = lowercase(strip(string(value))) + s in ("", "0", "none", "off", "warn") && return LOG_NONE + s in ("1", "info") && return LOG_INFO + s in ("2", "debug") && return LOG_DEBUG + throw(InvalidValueError("invalid log level $(repr(value)); expected 0/\"none\", 1/\"info\" or 2/\"debug\"")) +end + +""" + SparseDirectSolver.set_log_level!(level) -> Int + +Set the verbosity of the solver's log messages: `0`/`"none"` (default; the +messages are `@debug` records), `1`/`"info"` (phase summaries: analysis, +factorization, refinement steps performed, as `@info`), `2`/`"debug"` (also the +relative residual of every refinement step). The initial value comes from the +environment variable `SDS_LOG_LEVEL`. Returns the previous level. +""" +function set_log_level!(level) + old = LOG_LEVEL[] + LOG_LEVEL[] = level isa Integer ? _parse_log_level(string(level)) : _parse_log_level(level) + return old +end + +function _init_log_level() + try + LOG_LEVEL[] = _parse_log_level(get(ENV, "SDS_LOG_LEVEL", "")) + catch err + err isa InvalidValueError || rethrow() + @warn "SDS_LOG_LEVEL: $(err.msg); logging stays off" + LOG_LEVEL[] = LOG_NONE + end + return nothing +end + +# emit `msg()` at `level` (LOG_INFO or LOG_DEBUG) +function _log(level::Int, msg) + if LOG_LEVEL[] >= level + @info msg() + else + @debug msg() + end + return nothing +end + +# raise InterruptedError when the user's flag is set (a host read, no device synchronization) +@inline function _poll_interrupt(flag::Union{Nothing, Threads.Atomic{Bool}}) + flag !== nothing && flag[] && throw(InterruptedError()) + return nothing +end diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index ce1bbf1..58585b6 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -145,7 +145,9 @@ the analysis used `factorization_alg = "algo1"`, no vendor calls). Then Returns `info = 0` on success, else the original column of the smallest non-positive pivot of the factor (as [`ref_factorize!`](@ref)); fronts above a -failed one hold garbage. Structure `"SPD"` (real) or `"HPD"`. Nothing is +failed one hold garbage. `opts.user_host_interrupt` is polled before every +launch group (a host read): when set, [`InterruptedError`](@ref) is raised and +the factor is incomplete. Structure `"SPD"` (real) or `"HPD"`. Nothing is allocated on the device; the panels are bitwise reproducible for a fixed `impl` and backend (regimes A and B are deterministic by construction). """ @@ -156,11 +158,14 @@ function factorize!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; impl::Sym p = _front_impls(N, S, impl) plan = N.plan nodes = S.schedule.group_nodes + flag = opts.user_host_interrupt for k in eachindex(plan.sub_first) + _poll_interrupt(flag) a, b = plan.sub_first[k], plan.sub_last[k] factorize_subtrees!(N, S, nzval, a, b - a + 1, plan.sub_local[k]) end for k in eachindex(plan.group_first) + _poll_interrupt(flag) a, b = plan.group_first[k], plan.group_last[k] W = plan.group_width[k] if W > 0 diff --git a/src/numeric/ldlt.jl b/src/numeric/ldlt.jl index bdd6e3e..32e364c 100644 --- a/src/numeric/ldlt.jl +++ b/src/numeric/ldlt.jl @@ -809,15 +809,17 @@ function _launch_subtrees_ldlt!(N::Numeric{T}, S::Symbolic, nzval, first, count, return nothing end -function _factorize_ldlt_groups!(N::Numeric, S::Symbolic, nzval::AbstractVector, prm) +function _factorize_ldlt_groups!(N::Numeric, S::Symbolic, nzval::AbstractVector, prm, flag) plan = N.plan for k in eachindex(plan.sub_first) + _poll_interrupt(flag) 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) + _poll_interrupt(flag) 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 @@ -836,7 +838,8 @@ 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"` +T15 report); then [`reduce_stats!`](@ref). `opts.user_host_interrupt` is polled +before every launch group ([`InterruptedError`](@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 @@ -856,9 +859,9 @@ function factorize_ldlt!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; opts end prm.scaled && abs_max!(N.aux, nzval) if herm - _factorize_ldlt_groups!(N, S, nzval, prm) + _factorize_ldlt_groups!(N, S, nzval, prm, opts.user_host_interrupt) else - _factorize_ldlt_groups!(N, S, nzval, _ldlt_device_params(S, T, opts, Val(false))) + _factorize_ldlt_groups!(N, S, nzval, _ldlt_device_params(S, T, opts, Val(false)), opts.user_host_interrupt) end reduce_stats!(N, S) return 0 diff --git a/src/solve/permute.jl b/src/solve/permute.jl index 519957e..e9f4672 100644 --- a/src/solve/permute.jl +++ b/src/solve/permute.jl @@ -9,6 +9,15 @@ # `n * nrhs` entries (column-major), or, with `transposed = true`, row-major # data, i.e. a column-major `nrhs × n` matrix (or the strided vector of one). +# call `f(Val(a), Val(b))` with compile-time flags (a runtime `Val(flag)` would dispatch dynamically) +@inline function _with_flags(f, a::Bool, b::Bool) + if a + return b ? f(Val(true), Val(true)) : f(Val(true), Val(false)) + else + return b ? f(Val(false), Val(true)) : f(Val(false), Val(false)) + end +end + "Workgroup size of the permutation kernels." const PERMUTE_WORKGROUP = 256 @@ -45,21 +54,23 @@ function rhs_count(B::AbstractVecOrMat, n::Integer; transposed::Bool = false) return nrhs end -@kernel function _permute_rhs_kernel!(Y, B, perm, n, nrhs, ::Val{TR}) where {TR} +@kernel function _permute_rhs_kernel!(Y, B, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 @inbounds if r <= nrhs - Y[k, r] = _rhs_get(B, perm[k], r, n, nrhs, Val(TR)) + v = _rhs_get(B, perm[k], r, n, nrhs, Val(TR)) + Y[k, r] = CJ ? conj(v) : v end end -@kernel function _unpermute_solution_kernel!(X, Y, perm, n, nrhs, ::Val{TR}) where {TR} +@kernel function _unpermute_solution_kernel!(X, Y, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} q = @index(Global, Linear) k = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 @inbounds if r <= nrhs - _rhs_set!(X, Y[k, r], perm[k], r, n, nrhs, Val(TR)) + v = Y[k, r] + _rhs_set!(X, CJ ? conj(v) : v, perm[k], r, n, nrhs, Val(TR)) end end @@ -72,40 +83,39 @@ function _check_permute(Y, B, perm, transposed) end """ - permute_rhs!(Y, B, perm; transposed = false) -> Y + permute_rhs!(Y, B, perm; transposed = false, conjugate = false) -> Y -`Y[k, r] = B[perm[k], r]` for the `nrhs` right-hand sides of the user array +`Y[k, r] = B[perm[k], r]` (conjugated when `conjugate`, for the solves with +`conj(A)` of `solve_mode`) for the `nrhs` right-hand sides of the user array `B` ([`rhs_count`](@ref): vector, matrix, strided vector; row-major when `transposed`) into the first `nrhs` columns of the `n × ≥ nrhs` device matrix `Y`. One launch; asynchronous. """ -function permute_rhs!(Y::AbstractMatrix, B::AbstractVecOrMat, perm::AbstractVector; transposed::Bool = false) +function permute_rhs!(Y::AbstractMatrix, B::AbstractVecOrMat, perm::AbstractVector; transposed::Bool = false, + conjugate::Bool = false) n, nrhs = _check_permute(Y, B, perm, transposed) n * nrhs > 0 || return Y kernel! = _permute_rhs_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) - if transposed - kernel!(Y, B, perm, n, nrhs, Val(true); ndrange = n * nrhs) - else - kernel!(Y, B, perm, n, nrhs, Val(false); ndrange = n * nrhs) + _with_flags(transposed, conjugate) do tr, cj + kernel!(Y, B, perm, n, nrhs, tr, cj; ndrange = n * nrhs) end return Y end """ - unpermute_solution!(X, Y, perm; transposed = false) -> X + unpermute_solution!(X, Y, perm; transposed = false, conjugate = false) -> X -`X[perm[k], r] = Y[k, r]`: the inverse of [`permute_rhs!`](@ref), from the +`X[perm[k], r] = Y[k, r]` (conjugated when `conjugate`): the inverse of [`permute_rhs!`](@ref), from the first `nrhs` columns of `Y` into the user array `X` (same layouts). One launch; asynchronous. """ -function unpermute_solution!(X::AbstractVecOrMat, Y::AbstractMatrix, perm::AbstractVector; transposed::Bool = false) +function unpermute_solution!(X::AbstractVecOrMat, Y::AbstractMatrix, perm::AbstractVector; transposed::Bool = false, + conjugate::Bool = false) n, nrhs = _check_permute(Y, X, perm, transposed) n * nrhs > 0 || return X kernel! = _unpermute_solution_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) - if transposed - kernel!(X, Y, perm, n, nrhs, Val(true); ndrange = n * nrhs) - else - kernel!(X, Y, perm, n, nrhs, Val(false); ndrange = n * nrhs) + _with_flags(transposed, conjugate) do tr, cj + kernel!(X, Y, perm, n, nrhs, tr, cj; ndrange = n * nrhs) end return X end diff --git a/src/solve/refinement.jl b/src/solve/refinement.jl new file mode 100644 index 0000000..4ad23f7 --- /dev/null +++ b/src/solve/refinement.jl @@ -0,0 +1,331 @@ +# Iterative refinement of the solve phase (PLAN §2.5, §1.3 `ir_n_steps`, +# `ir_tol`): the residual `R = B - op(A) X` by a KA CSR SpMV over the +# full-pattern map of the stored matrix (`full_pattern_map`: the symmetric +# expansion of the triangle given by the view, mirrored entries conjugated for +# `"H"`/`"HPD"`), its column norms reduced on the device, and the correction +# `X += A⁻¹ R` through the solve sweeps. Everything runs in the input +# precision; the residual uses the solver's *current* values (`update!` without +# a refactorization refines towards the new matrix with the old factor). +# +# `op(A)` is `A`, `Aᵀ` or `Aᴴ` (`solve_mode`) of the user matrix, which is the +# stored CSR matrix `M` or, for a CSC input, `Mᵀ`. `M` is symmetric or +# Hermitian, so every `op(A)` is `M` or `conj(M)`; `conj(M) x = b` is solved as +# `x = conj(M⁻¹ conj(b))` ([`solve_conjugated`](@ref)). +# +# User arrays come in the layouts of `src/solve/permute.jl` (vector, matrix, +# strided vector, row-major when transposed); the residual `R` is a column-major +# `n × nrhs` device matrix in the original numbering. + +"Workgroup size of the residual-norm reduction (one workgroup per right-hand side)." +const REFINE_WORKGROUP = 256 + +""" + RefinementWorkspace{T, R, VI, MT, VR} + +Device storage of iterative refinement, allocated at the first solve that +refines ([`allocate_refinement`](@ref)) and grown with the number of right-hand +sides: + +* `rowptr`, `colval`, `src` (`INT` vectors): the full matrix `M` as a CSR over + the contributions of the user's `nzval` (`src[c] > 0`: `nzval[src[c]]`, + `src[c] < 0`: `conj(nzval[-src[c]])`), from [`refinement_map`](@ref); +* `R` (`n × nrhs`): the residual, then the correction; +* `Bc` (`n × nrhs`): a copy of the right-hand side when `X` and `B` alias; +* `norms` (`2 nrhs`, real): `‖Rₖ‖₂²` and `‖Bₖ‖₂²` per right-hand side, and + `norms_host`, its host copy (the early-exit test of `ir_tol`). +""" +struct RefinementWorkspace{T, R <: Real, VI <: AbstractVector, MT <: AbstractMatrix{T}, VR <: AbstractVector{R}} + rowptr::VI + colval::VI + src::VI + R::MT + Bc::MT + norms::VR + norms_host::Vector{R} +end + +max_rhs(W::RefinementWorkspace) = size(W.R, 2) + +""" + refinement_map(F::FullPatternMap) -> (rowptr, colval, src) + +Host CSR of the full matrix of `F` over its contributions (duplicates are not +merged; their products are summed by the SpMV): row `i` holds the entries +`rowptr[i]:rowptr[i+1]-1`, column `colval[c]`, value `nzval[src[c]]`, conjugated +when `src[c] < 0`. 1-based. +""" +function refinement_map(F::FullPatternMap) + n = F.n + rowptr = Vector{Int}(undef, n + 1) + rowptr[1] = 1 + colval = Int[] + src = Int[] + for i in 1:n + for e in F.rowptr[i]:(F.rowptr[i + 1] - 1), s in F.srcptr[e]:(F.srcptr[e + 1] - 1) + push!(colval, F.colval[e]) + push!(src, F.conjflag[s] ? -F.src[s] : F.src[s]) + end + rowptr[i + 1] = length(colval) + 1 + end + return rowptr, colval, src +end + +""" + allocate_refinement(map, T, INT, backend, nrhs) -> RefinementWorkspace + +Move the host [`refinement_map`](@ref) `map` to `backend` as `INT` vectors +(overflow-checked) and allocate the residual storage for `nrhs` right-hand +sides of element type `T`. +""" +function allocate_refinement(map::NTuple{3, Vector{Int}}, ::Type{T}, ::Type{INT}, + backend::KernelAbstractions.Backend, nrhs::Integer) where {T, INT} + rowptr, colval, src = map + n = length(rowptr) - 1 + dev(v) = (v_ = _to_index_type(INT, v); d = KernelAbstractions.allocate(backend, INT, length(v_)); copyto!(d, v_); d) + R = KernelAbstractions.zeros(backend, T, n, nrhs) + Bc = KernelAbstractions.zeros(backend, T, n, nrhs) + norms = KernelAbstractions.zeros(backend, real(T), 2 * nrhs) + return RefinementWorkspace{T, real(T), typeof(dev(Int[])), typeof(R), typeof(norms)}( + dev(rowptr), dev(colval), dev(src), R, Bc, norms, zeros(real(T), 2 * nrhs)) +end + +function _to_index_type(::Type{INT}, v::Vector{Int}) where {INT} + all(x -> typemin(INT) <= x <= typemax(INT), v) || + throw(InvalidValueError("the refinement map does not fit index type $INT")) + return convert(Vector{INT}, v) +end + +# a fresh workspace for `nrhs` right-hand sides that reuses the device map of `W` +function _grow_refinement(W::RefinementWorkspace{T, R, VI, MT, VR}, nrhs::Integer) where {T, R, VI, MT, VR} + backend = KernelAbstractions.get_backend(W.R) + n = size(W.R, 1) + Rm = KernelAbstractions.zeros(backend, T, n, nrhs) + Bc = KernelAbstractions.zeros(backend, T, n, nrhs) + norms = KernelAbstractions.zeros(backend, R, 2 * nrhs) + return RefinementWorkspace{T, R, VI, MT, VR}(W.rowptr, W.colval, W.src, Rm, Bc, norms, zeros(R, 2 * nrhs)) +end + +@kernel function _residual_kernel!(R, rowptr, colval, src, nzval, X, B, n, nrhs, ::Val{XT}, ::Val{BT}, + ::Val{CJ}) where {XT, BT, CJ} + q = @index(Global, Linear) + i = (q - 1) % n + 1 + r = (q - 1) ÷ n + 1 + @inbounds if r <= nrhs + acc = zero(eltype(R)) + for c in Int(rowptr[i]):(Int(rowptr[i + 1]) - 1) + s = Int(src[c]) + v = nzval[abs(s)] + v = xor(s < 0, CJ) ? conj(v) : v + acc += v * _rhs_get(X, Int(colval[c]), r, n, nrhs, Val(XT)) + end + R[i, r] = _rhs_get(B, i, r, n, nrhs, Val(BT)) - acc + end +end + +@kernel function _residual_norms_kernel!(norms, R, B, n, nrhs, ::Val{BT}, ::Val{WG}, ::Val{LOG2WG}) where {BT, WG, LOG2WG} + @uniform RT = eltype(norms) + li = @index(Local, Linear) + r = @index(Group, Linear) + nr = @localmem RT (WG,) + nb = @localmem RT (WG,) + @inbounds begin + sr = zero(RT) + sb = zero(RT) + for i in li:WG:n + sr += abs2(R[i, r]) + sb += abs2(_rhs_get(B, i, r, n, nrhs, Val(BT))) + end + nr[li] = sr + nb[li] = sb + end + @synchronize + for lev in 1:LOG2WG + @inbounds begin + h = WG >> lev + if li <= h + nr[li] += nr[li + h] + nb[li] += nb[li + h] + end + end + @synchronize + end + if li == 1 + @inbounds norms[2 * r - 1] = nr[1] + @inbounds norms[2 * r] = nb[1] + end +end + +@kernel function _add_correction_kernel!(X, Y, perm, n, nrhs, ::Val{TR}, ::Val{CJ}) where {TR, CJ} + q = @index(Global, Linear) + k = (q - 1) % n + 1 + r = (q - 1) ÷ n + 1 + @inbounds if r <= nrhs + i = Int(perm[k]) + y = Y[k, r] + _rhs_set!(X, _rhs_get(X, i, r, n, nrhs, Val(TR)) + (CJ ? conj(y) : y), i, r, n, nrhs, Val(TR)) + end +end + +@kernel function _copy_rhs_kernel!(C, B, n, nrhs, ::Val{TR}) where {TR} + q = @index(Global, Linear) + i = (q - 1) % n + 1 + r = (q - 1) ÷ n + 1 + @inbounds if r <= nrhs + C[i, r] = _rhs_get(B, i, r, n, nrhs, Val(TR)) + end +end + +""" + residual!(W, nzval, X, B; nrhs, transposed = false, b_transposed = transposed, conjugate = false) -> W.R + +`W.R[:, 1:nrhs] = B - M X` (`conj(M)` when `conjugate`) with the full matrix +`M` of the map of `W` and the user values `nzval`; `X`, `B` in the layouts of +[`rhs_count`](@ref) (row-major when `transposed`/`b_transposed`). One launch, +one work item per (row, right-hand side), gather-based (no atomics). +Asynchronous. +""" +function residual!(W::RefinementWorkspace, nzval::AbstractVector, X::AbstractVecOrMat, B::AbstractVecOrMat; + nrhs::Integer, transposed::Bool = false, b_transposed::Bool = transposed, + conjugate::Bool = false) + n = size(W.R, 1) + n * nrhs > 0 || return W.R + kernel! = _residual_kernel!(KernelAbstractions.get_backend(W.R), PERMUTE_WORKGROUP) + _with_flags(transposed, b_transposed) do xt, bt + if conjugate + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), xt, bt, Val(true); ndrange = n * Int(nrhs)) + else + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), xt, bt, Val(false); ndrange = n * Int(nrhs)) + end + end + return W.R +end + +""" + residual_norms!(W, B; nrhs, transposed = false) -> W.norms_host + +`‖W.R[:, k]‖₂²` and `‖B[:, k]‖₂²` for `k = 1:nrhs` reduced on the device (one +workgroup per right-hand side, `@localmem` tree reduction) into `W.norms`, then +copied to `W.norms_host` (one host synchronization). +""" +function residual_norms!(W::RefinementWorkspace, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false) + n = size(W.R, 1) + WG = REFINE_WORKGROUP + kernel! = _residual_norms_kernel!(KernelAbstractions.get_backend(W.R), WG) + if transposed + kernel!(W.norms, W.R, B, n, Int(nrhs), Val(true), Val(WG), Val(_ilog2(WG)); ndrange = WG * Int(nrhs)) + else + kernel!(W.norms, W.R, B, n, Int(nrhs), Val(false), Val(WG), Val(_ilog2(WG)); ndrange = WG * Int(nrhs)) + end + copyto!(W.norms_host, 1, W.norms, 1, 2 * Int(nrhs)) + return W.norms_host +end + +# largest relative residual ‖Rₖ‖/‖Bₖ‖ over the right-hand sides (‖Rₖ‖ when Bₖ = 0) +function _max_relative_residual(norms::Vector{R}, nrhs::Int) where {R} + m = zero(R) + for k in 1:nrhs + rk, bk = sqrt(norms[2k - 1]), sqrt(norms[2k]) + m = max(m, bk > 0 ? rk / bk : rk) + end + return m +end + +""" + copy_rhs!(C, B; nrhs, transposed = false) -> C + +`C[:, 1:nrhs]` = the right-hand sides of the user array `B` (column-major copy; +the refinement keeps `B` when the solution overwrites it). One launch. +""" +function copy_rhs!(C::AbstractMatrix, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false) + n = size(C, 1) + n * nrhs > 0 || return C + kernel! = _copy_rhs_kernel!(KernelAbstractions.get_backend(C), PERMUTE_WORKGROUP) + if transposed + kernel!(C, B, n, Int(nrhs), Val(true); ndrange = n * Int(nrhs)) + else + kernel!(C, B, n, Int(nrhs), Val(false); ndrange = n * Int(nrhs)) + end + return C +end + +""" + add_correction!(X, Y, perm; nrhs, transposed = false, conjugate = false) -> X + +`X[perm[k], r] += Y[k, r]` (`conj(Y[k, r])` when `conjugate`): the solution of +the correction system, in factor order in the workspace `Y`, added to the user +array `X`. One launch, conflict-free (`perm` is a permutation). +""" +function add_correction!(X::AbstractVecOrMat, Y::AbstractMatrix, perm::AbstractVector; nrhs::Integer, + transposed::Bool = false, conjugate::Bool = false) + n = length(perm) + n * nrhs > 0 || return X + kernel! = _add_correction_kernel!(KernelAbstractions.get_backend(Y), PERMUTE_WORKGROUP) + _with_flags(transposed, conjugate) do tr, cj + kernel!(X, Y, perm, n, Int(nrhs), tr, cj; ndrange = n * Int(nrhs)) + end + return X +end + +""" + refine!(X, B, W, ws, symbolic, numeric, nzval; nsteps, tol = 0, transposed = false, + b_transposed = transposed, conjugate = false, deterministic = false, + interrupt = nothing, progress = Ref(0)) -> steps + +Plain iterative refinement of the solution `X` of `op(A) X = B` (`op(A) = M`, +or `conj(M)` when `conjugate`), at most `nsteps` steps of + + R = B - op(A) X (residual!) + stop if tol > 0 and maxₖ ‖Rₖ‖₂ / ‖Bₖ‖₂ ≤ tol (residual_norms!, one host synchronization) + X += op(A)⁻¹ R (permute, forward / diagonal / backward sweeps, add_correction!) + +with the factor in `numeric`, the solve workspace `ws` and the refinement +workspace `W`. `tol = 0` (the default of `ir_tol`) never stops early and never +synchronizes. `interrupt` (a `Threads.Atomic{Bool}` or `nothing`) is polled +before every step ([`InterruptedError`](@ref); `X` then holds the last +completed iterate). Returns the number of corrections applied, which is also +kept in `progress[]` after every step (so it survives an interrupt). +""" +function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspace, ws::SolveWorkspace, S::Symbolic, + N::Numeric, nzval::AbstractVector; nsteps::Integer, tol::Real = 0, transposed::Bool = false, + b_transposed::Bool = transposed, conjugate::Bool = false, deterministic::Bool = false, + interrupt::Union{Nothing, Threads.Atomic{Bool}} = nothing, + progress::Base.RefValue{Int} = Ref(0)) + nrhs = rhs_count(X, S.n; transposed) + max_rhs(W) >= nrhs && max_rhs(ws) >= nrhs || + throw(DimensionMismatch("the refinement workspace holds $(max_rhs(W)) right-hand sides, need $nrhs")) + steps = 0 + progress[] = 0 + for _ in 1:nsteps + _poll_interrupt(interrupt) + residual!(W, nzval, X, B; nrhs, transposed, b_transposed, conjugate) + if tol > 0 + norms = residual_norms!(W, B; nrhs, transposed = b_transposed) + rel = _max_relative_residual(norms, nrhs) + _log(LOG_DEBUG, () -> "refinement: step $steps, relative residual $rel") + rel <= tol && break + end + permute_rhs!(ws.Y, W.R, S.perm; conjugate) + forward_sweep!(ws, S, N; nrhs, deterministic) + diagonal_sweep!(ws, S, N; nrhs) + backward_sweep!(ws, S, N; nrhs) + add_correction!(X, ws.Y, S.perm; nrhs, transposed, conjugate) + steps += 1 + progress[] = steps + end + return steps +end + +""" + solve_conjugated(structure, T, csc_input::Bool, solve_mode) -> Bool + +Whether `op(A)` (`solve_mode` 0: `A`, 1: `Aᵀ`, 2: `Aᴴ`) of the user matrix is +the conjugate of the stored, factorized matrix `M` (`A = M`, or `A = Mᵀ` for a +CSC input). Real `T`: never. Complex symmetric `"S"` (`M = Mᵀ`): for `Aᴴ`. +Hermitian `"H"`/`"HPD"` (`Mᵀ = conj(M)`): for `Aᵀ` of a CSR input, and for `A` +and `Aᴴ` of a CSC input. +""" +function solve_conjugated(structure::Structure, ::Type{T}, csc_input::Bool, solve_mode::Integer) where {T} + T <: Complex || return false + _is_hermitian(structure) && return xor(csc_input, solve_mode == 1) + return solve_mode == 2 +end diff --git a/src/solver.jl b/src/solver.jl index d71ff2f..a983bd6 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -31,9 +31,13 @@ _numeric_type(backend, ::Type{T}) where {T} = Numeric{T, typeof(KernelAbstractions.allocate(backend, T, 0)), typeof(KernelAbstractions.allocate(backend, Int64, 0)), typeof(KernelAbstractions.allocate(backend, Int32, 0)), typeof(KernelAbstractions.allocate(backend, Int8, 0))} _workspace_type(backend, ::Type{T}) where {T} = SolveWorkspace{T, typeof(KernelAbstractions.allocate(backend, T, 0, 0))} +_refinement_type(backend, ::Type{T}, ::Type{INT}) where {T, INT} = + RefinementWorkspace{T, real(T), typeof(KernelAbstractions.allocate(backend, INT, 0)), + typeof(KernelAbstractions.allocate(backend, T, 0, 0)), + typeof(KernelAbstractions.allocate(backend, real(T), 0))} """ - DirectSolver{T, INT, M, B, SY, NU, WS} <: AbstractDirectSolver{T, INT} + DirectSolver{T, INT, M, B, SY, NU, WS, RF} <: AbstractDirectSolver{T, INT} Sparse direct solver handle (≅ `CudssSolver`), PLAN §3.1–§3.2. @@ -51,16 +55,18 @@ CSR arrays are wrapped without copies and live on a KernelAbstractions backend host CSR arrays once. A CSC matrix (`CuSparseMatrixCSC`, or a [`CSR`](@ref) with `transposed = true`, e.g. from [`csr_of_transpose`](@ref)) is read as the CSR 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. +matrices that transpose is the conjugate, and the solve phases conjugate the +right-hand side and the solution (as for `solve_mode`, see +[`solve_conjugated`](@ref)). Implemented at this point: structures `"SPD"` (real `T`), `"HPD"`, `"S"` (LDLᵀ; complex symmetric for complex `T`) and `"H"` (LDLᴴ), a single matrix (no uniform batch), the phases `"reordering"`, `"symbolic_factorization"`, `"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. +`"solve_fwd_perm"`, `"solve_fwd"`, `"solve_diag"`, `"solve_bwd"`, +`"solve_bwd_perm"` and `"solve_refinement"` ([`execute!`](@ref)). Structure +`"G"` and the Schur phases raise [`NotSupportedError`](@ref) when they are +executed. Fields: `A` (the current [`CSR`](@ref), re-pointed by [`update!`](@ref)), `structure`, `view` (as given), `options` ([`Options`](@ref), set through @@ -68,9 +74,12 @@ Fields: `A` (the current [`CSR`](@ref), re-pointed by [`update!`](@ref)), the first `"factorization"` after an analysis; the `LinearAlgebra` layer then switches `cholesky!` to `"refactorization"`), `info` (the `"info"` data parameter), the analysis (`ordering`, host `host_symbolic`, device `symbolic`), -the numeric storage `numeric` and the solve `workspace`. +the numeric storage `numeric`, the solve `workspace`, the refinement storage +`refinement` (allocated by the first solve that refines) and `ir_steps` (the +`"ir_n_steps"` data parameter: refinement steps of the last solve, `-1` before +one). """ -mutable struct DirectSolver{T, INT, M <: CSR{T, INT}, B <: KernelAbstractions.Backend, SY, NU, WS} <: +mutable struct DirectSolver{T, INT, M <: CSR{T, INT}, B <: KernelAbstractions.Backend, SY, NU, WS, RF} <: AbstractDirectSolver{T, INT} A::M structure::Structure @@ -88,15 +97,18 @@ mutable struct DirectSolver{T, INT, M <: CSR{T, INT}, B <: KernelAbstractions.Ba symbolic::Union{Nothing, SY} numeric::Union{Nothing, NU} workspace::Union{Nothing, WS} - # explicit parameters only: the default outer constructor would leave SY, NU, WS unbound + refinement::Union{Nothing, RF} + ir_steps::Int + # explicit parameters only: the default outer constructor would leave SY, NU, WS, RF unbound # (they occur only in `Union{Nothing, …}` fields; Aqua on Julia 1.10) - function DirectSolver{T, INT, M, B, SY, NU, WS}(A, structure, view, options, backend, nbatch, - fresh_factorization, info, stage, host_rowptr, - host_colval, ordering, host_symbolic, symbolic, - numeric, workspace) where {T, INT, M, B, SY, NU, WS} - return new{T, INT, M, B, SY, NU, WS}(A, structure, view, options, backend, nbatch, fresh_factorization, - info, stage, host_rowptr, host_colval, ordering, host_symbolic, - symbolic, numeric, workspace) + function DirectSolver{T, INT, M, B, SY, NU, WS, RF}(A, structure, view, options, backend, nbatch, + fresh_factorization, info, stage, host_rowptr, + host_colval, ordering, host_symbolic, symbolic, + numeric, workspace, refinement, + ir_steps) where {T, INT, M, B, SY, NU, WS, RF} + return new{T, INT, M, B, SY, NU, WS, RF}(A, structure, view, options, backend, nbatch, fresh_factorization, + info, stage, host_rowptr, host_colval, ordering, host_symbolic, + symbolic, numeric, workspace, refinement, ir_steps) end end @@ -113,17 +125,14 @@ function DirectSolver(A::CSR{T, INT}, structure, view; index = A.index) where {T if base != A.index A = CSR(A.rowptr, A.colval, A.nzval, A.nrows, A.ncols; index = base, transposed = A.transposed) end - if A.transposed && T <: Complex && _is_hermitian(s) - throw(NotSupportedError("a complex Hermitian matrix given as CSC (the CSR of its transpose) needs the " * - "conjugated solve of solve_mode (T16); pass the CSR arrays instead")) - end nb = nbatch(A) nb == 1 || throw(NotSupportedError("uniform batches (nbatch = $nb) are not implemented yet (T17)")) backend = KernelAbstractions.get_backend(A) SY, NU, WS = _symbolic_type(backend, INT), _numeric_type(backend, T), _workspace_type(backend, T) - return DirectSolver{T, INT, typeof(A), typeof(backend), SY, NU, WS}( + RF = _refinement_type(backend, T, INT) + return DirectSolver{T, INT, typeof(A), typeof(backend), SY, NU, WS, RF}( A, s, v, Options(), backend, nb, true, 0, STAGE_NONE, INT[], INT[], nothing, nothing, nothing, nothing, - nothing) + nothing, nothing, -1) end function DirectSolver(rowptr::AbstractVector{<:Integer}, colval::AbstractVector{<:Integer}, nzval::AbstractVecOrMat, @@ -228,24 +237,43 @@ Execute `phase` on `solver` (≅ `cudss(phase, solver, x, b)`), PLAN §1.2: `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_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`. +* `"solve"`: `X = op(A)⁻¹ B` (`solve_mode` 0: `A`, 1: `Aᵀ`, 2: `Aᴴ`), followed by + `ir_n_steps` steps of iterative refinement; the sub-phases `"solve_fwd_perm"` + (permute `B` into the workspace), `"solve_fwd"` (forward sweep), + `"solve_diag"` (`D⁻¹` of LDLᵀ/LDLᴴ, the identity for Cholesky), `"solve_bwd"` + (backward sweep), `"solve_bwd_perm"` (inverse permutation into `X`) and + `"solve_refinement"` (refine `X` in place, [`refine!`](@ref)) do the same in + six calls with the same `X`, `B` (bitwise equal to `"solve"`). `X` and `B` are only read by the solve phases (pass anything, e.g. `nothing`, to the others): vectors, `n × nrhs` matrices, strided vectors or [`MatrixDescriptor`](@ref)s (row-major when `transposed`) on the solver's -backend; `X === B` is allowed. A failed factorization does not throw (as in -cuDSS): check `getparam(solver, "info")`; solving with it gives garbage. - -`"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). +backend; `X === B` is allowed in every phase but `"solve_refinement"` (it needs +the original `B`; `"solve"` keeps a copy of `B` when they alias). A failed +factorization does not throw (as in cuDSS): check `getparam(solver, "info")`; +solving with it gives garbage. + +Refinement (PLAN §2.5): each step computes the residual `R = B - op(A) X` with +the solver's current matrix values (KA SpMV over the full pattern), and, when +`ir_tol > 0`, stops once `‖Rₖ‖₂ ≤ ir_tol ‖Bₖ‖₂` for every right-hand side `k` +(this costs one host synchronization per step; `ir_tol = 0`, the default, never +stops early and never synchronizes), then adds the correction `op(A)⁻¹ R`. The +steps performed are the data parameter `"ir_n_steps"` ([`getparam`](@ref)). + +The flag `"user_host_interrupt"` (a `Threads.Atomic{Bool}`, [`setparam!`](@ref)) +is polled at the start of the analysis phases, between the launch groups of the +factorization and between refinement steps: when set, the phase raises +[`InterruptedError`](@ref). An interrupted factorization leaves the solver +analyzed (the next phase must be `"factorization"`); an interrupted refinement +leaves the last completed iterate in `X` and reports the steps completed in +`"ir_n_steps"`. + +`"solve_fwd_schur"` and `"solve_bwd_schur"` raise [`NotSupportedError`](@ref) +until their task (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`). +(`KernelAbstractions.synchronize`). Phase summaries are logged (see +[`SparseDirectSolver.set_log_level!`](@ref)). """ function execute!(phase::AbstractString, solver::DirectSolver, X, B; asynchronous::Bool = true) p = _parse_phase(phase) @@ -259,11 +287,10 @@ function execute!(phase::AbstractString, solver::DirectSolver, X, B; asynchronou 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_DIAG || - p == PHASE_SOLVE_BWD || p == PHASE_SOLVE_BWD_PERM + p == PHASE_SOLVE_BWD || p == PHASE_SOLVE_BWD_PERM || p == PHASE_SOLVE_REFINEMENT _solve!(solver, p, X, B) else - throw(NotSupportedError("phase \"$phase\" is not implemented yet " * - "($(p == PHASE_SOLVE_REFINEMENT ? "T16" : "T20"))")) + throw(NotSupportedError("phase \"$phase\" is not implemented yet (T20)")) end asynchronous || KernelAbstractions.synchronize(solver.backend) return nothing @@ -292,6 +319,8 @@ end function _reorder!(solver::DirectSolver{T}) where {T} _check_analysis_supported(solver) + _poll_interrupt(solver.options.user_host_interrupt) + tic = time_ns() A = solver.A solver.host_rowptr = Array(A.rowptr) solver.host_colval = Array(A.colval) @@ -301,8 +330,9 @@ function _reorder!(solver::DirectSolver{T}) where {T} pp = analysis_pairs(P, solver.host_rowptr, solver.host_colval, A.nzval, A.nrows, solver.structure, solver.options; view = _stored_view(solver), index = A.index) solver.ordering = compute_ordering(P, solver.options; T, pp.pairs, pp.candidates) - solver.host_symbolic = solver.symbolic = solver.numeric = solver.workspace = nothing + solver.host_symbolic = solver.symbolic = solver.numeric = solver.workspace = solver.refinement = nothing solver.stage = STAGE_REORDERED + _log(LOG_INFO, () -> "reordering: n = $(A.nrows), nnz = $(nnz(A)), $(_elapsed(tic))") return solver end @@ -310,6 +340,8 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} solver.stage >= STAGE_REORDERED || throw(_phase_error("symbolic_factorization", "needs \"reordering\" (or run \"analysis\")")) _check_analysis_supported(solver) + _poll_interrupt(solver.options.user_host_interrupt) + tic = time_ns() opts = solver.options A = solver.A ord = solver.ordering @@ -326,9 +358,12 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} solver.symbolic = Sd solver.numeric = allocate_numeric(Sd, T, solver.backend) solver.workspace = allocate_solve(Sd, T, solver.backend, nrhs) + solver.refinement = nothing solver.stage = STAGE_ANALYZED solver.fresh_factorization = true solver.info = 0 + _log(LOG_INFO, () -> "symbolic_factorization: $(nsupernodes(Sh)) supernodes, nnz(L) = $(sp.nnz_L), " * + "flops = $(sp.flops), $(_elapsed(tic))") return solver end @@ -338,12 +373,27 @@ 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.options) + tic = time_ns() + try + solver.info = _numeric_phase!(solver.numeric, solver.symbolic, solver.A.nzval, solver.options) + catch err + if err isa InterruptedError + # the panels are partly overwritten: back to "analyzed", a "factorization" must follow + solver.stage = STAGE_ANALYZED + solver.fresh_factorization = true + solver.info = 0 + _log(LOG_INFO, () -> "$name: interrupted") + end + rethrow() + end solver.stage = STAGE_FACTORIZED p == PHASE_FACTORIZATION && (solver.fresh_factorization = false) + _log(LOG_INFO, () -> "$name: info = $(solver.info), $(_elapsed(tic)) (launches, asynchronous)") return solver end +_elapsed(tic) = string(round((time_ns() - tic) / 1.0e6; digits = 3), " ms") + # function barrier: concrete storage types for the numeric phase _numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractVector, opts::Options) = factorize!(N, S, nzval; opts) _numeric_phase!(N::Numeric, S::Symbolic, nzval::AbstractMatrix, opts::Options) = factorize!(N, S, vec(nzval); opts) @@ -375,9 +425,19 @@ function _workspace!(solver::DirectSolver{T}, nrhs::Int) where {T} return ws end -function _warn_ir_ignored(n) - @warn "ir_n_steps = $n: iterative refinement is not implemented yet (T16); the solve runs without it" maxlog = 1 - return nothing +# the refinement workspace for `nrhs` right-hand sides (the device map is built at the first use after an +# analysis, the residual storage reallocated only when nrhs grows) +function _refinement!(solver::DirectSolver{T, INT}, nrhs::Int) where {T, INT} + W = solver.refinement + if W === nothing + F = full_pattern_map(solver.host_rowptr, solver.host_colval, solver.A.nrows, solver.structure; + view = _stored_view(solver), index = solver.A.index) + W = allocate_refinement(refinement_map(F), T, INT, solver.backend, max(nrhs, max_rhs(solver.workspace))) + elseif max_rhs(W) < nrhs + W = _grow_refinement(W, nrhs) + end + solver.refinement = W + return W end function _solve!(solver::DirectSolver{T}, p::Phase, X, B) where {T} @@ -386,10 +446,8 @@ function _solve!(solver::DirectSolver{T}, p::Phase, X, B) where {T} opts = solver.options opts.solve_alg == SOLVE_DEFAULT || throw(NotSupportedError("solve_alg = \"$(convert(String, opts.solve_alg))\" is not implemented yet (M11)")) - # SPD/HPD: A = Aᴴ, and A = Aᵀ for real T; the transposed solve of a complex matrix needs T16 - opts.solve_mode == 1 && T <: Complex && - throw(NotSupportedError("solve_mode = 1 (Aᵀ) for complex matrices is not implemented yet (T16)")) - opts.ir_n_steps > 0 && _warn_ir_ignored(opts.ir_n_steps) + opts.ir_mode == IR_PLAIN || + throw(NotSupportedError("ir_mode = \"fgmres\" is not implemented yet (T18)")) Bd, bt = _rhs_data(B) Xd, xt = _rhs_data(X) xt == bt || throw(InvalidValueError("X and B must have the same layout (transposed or not)")) @@ -403,25 +461,60 @@ end function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::AbstractVecOrMat, transposed::Bool) S, N = solver.symbolic, solver.numeric + opts = solver.options n = S.n nrhs = rhs_count(B, n; transposed) rhs_count(X, n; transposed) == nrhs || throw(DimensionMismatch("X has $(rhs_count(X, n; transposed)) right-hand sides, B has $nrhs")) ws = _workspace!(solver, nrhs) - det = solver.options.deterministic_mode == 1 + det = opts.deterministic_mode == 1 + cj = solve_conjugated(solver.structure, eltype(ws), solver.A.transposed, opts.solve_mode) + refine = opts.ir_n_steps > 0 && (p == PHASE_SOLVE || p == PHASE_SOLVE_REFINEMENT) + if p == PHASE_SOLVE_REFINEMENT && refine && Base.mightalias(X, B) + throw(InvalidValueError("\"solve_refinement\" needs the original right-hand side: X and B must not alias")) + end + W = refine ? _refinement!(solver, nrhs) : nothing if p == PHASE_SOLVE - sweep_solve!(X, ws, S, N, B; transposed, deterministic = det) + Bs, bt = B, transposed + if refine && Base.mightalias(X, B) # the solve overwrites B: keep a copy for the residual + copy_rhs!(W.Bc, B; nrhs, transposed) + Bs, bt = W.Bc, false + end + permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj) + forward_sweep!(ws, S, N; nrhs, deterministic = det) + diagonal_sweep!(ws, S, N; nrhs) + backward_sweep!(ws, S, N; nrhs) + unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj) + refine ? _refine_phase!(solver, W, ws, X, Bs, transposed, bt, cj, det) : (solver.ir_steps = 0) elseif p == PHASE_SOLVE_FWD_PERM - permute_rhs!(ws.Y, B, S.perm; transposed) + permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj) 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 - unpermute_solution!(X, ws.Y, S.perm; transposed) + elseif p == PHASE_SOLVE_BWD_PERM + unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj) + else # PHASE_SOLVE_REFINEMENT + refine ? _refine_phase!(solver, W, ws, X, B, transposed, transposed, cj, det) : (solver.ir_steps = 0) + end + return nothing +end + +function _refine_phase!(solver::DirectSolver, W::RefinementWorkspace, ws::SolveWorkspace, X, B, xt::Bool, bt::Bool, + cj::Bool, det::Bool) + opts = solver.options + done = Ref(0) # the corrections applied, also when the refinement is interrupted + try + refine!(X, B, W, ws, solver.symbolic, solver.numeric, vec(solver.A.nzval); nsteps = opts.ir_n_steps, + tol = opts.ir_tol, transposed = xt, b_transposed = bt, conjugate = cj, deterministic = det, + interrupt = opts.user_host_interrupt, progress = done) + finally + solver.ir_steps = done[] end + steps = done[] + _log(LOG_INFO, () -> "solve_refinement: $steps of $(opts.ir_n_steps) steps") return nothing end @@ -490,7 +583,9 @@ everything [`setparam!`](@ref)`(::Options, …)` accepts, stored in 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. +`"factorization"`/`"refactorization"` copies it to the device. Setting +`"ir_n_steps"` (the number of refinement steps requested) makes +[`getparam`](@ref) report that value again until the next solve. 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. @@ -506,6 +601,7 @@ function setparam!(solver::DirectSolver, name::AbstractString, value) throw(InvalidValueError("pivot_sign has $(length(value)) entries, the matrix has $(size(solver, 1)) rows")) else setparam!(solver.options, name, value) + name == "ir_n_steps" && (solver.ir_steps = -1) end return nothing end @@ -527,6 +623,7 @@ The data parameters computed by the solver: | name | value | available after | | --- | --- | --- | | `"info"` | `Int`: `0`, or the original column (1-based) of the first non-positive pivot | always | +| `"ir_n_steps"` | `Int`: refinement steps performed by the last `"solve"`/`"solve_refinement"` (≤ the configured `ir_n_steps`, fewer when `ir_tol` stopped early); the configured value before the first solve and after `setparam!(solver, "ir_n_steps", k)` | always | | `"lu_nnz"` | `Int64`: nonzeros of `L` (diagonal included, amalgamation zeros excluded) | analysis | | `"flops"` | `Float64`: factorization flops of the stored panels | analysis | | `"nsuperpanels"` | `Int`: supernodes after amalgamation | analysis | @@ -549,6 +646,7 @@ one before the phase that computes it raises [`FactorizationError`](@ref). function getparam(solver::DirectSolver, name::AbstractString) spec = parameter_spec(name) name == "info" && return solver.info + name == "ir_n_steps" && solver.ir_steps >= 0 && return solver.ir_steps spec.status === :output || spec.status === :solver || return getparam(solver.options, name) name in SOLVER_OUTPUTS || throw(NotSupportedError("the data parameter \"$name\" is not implemented yet ($(_output_task(name)))")) diff --git a/test/matrices.jl b/test/matrices.jl index 72a843e..2e7b7a2 100644 --- a/test/matrices.jl +++ b/test/matrices.jl @@ -67,6 +67,22 @@ function random_spd(::Type{T}, n::Integer, density::Real; rng::AbstractRNG = Ran end random_spd(n::Integer, density::Real; kwargs...) = random_spd(Float64, n, density; kwargs...) +""" + badly_scaled_spd(T, n, density; exponent = 4, rng) + +`D A D` with `A = random_spd(T, n, density)` and `D` diagonal with entries +`10^(±exponent)` (random signs): symmetric (Hermitian) positive definite, row +`i` scaled by `dᵢ` on top of the column scaling, so the entries span +`10^(±2 exponent)` relative to `A` (the default: rows scaled by up to `10^(±8)`). +""" +function badly_scaled_spd(::Type{T}, n::Integer, density::Real; exponent::Real = 4, + rng::AbstractRNG = Random.default_rng()) where {T} + A = random_spd(T, n, density; rng) + D = Diagonal([real(T)(10)^(exponent * rand(rng, (-1, 1))) for _ in 1:n]) + S = D * A * D + return (S + S') / 2 +end + """ random_hpd(T, n, density; rng) diff --git a/test/test_api.jl b/test/test_api.jl index 34a634c..05671e7 100644 --- a/test/test_api.jl +++ b/test/test_api.jl @@ -122,7 +122,8 @@ 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_refinement", "solve_fwd_schur", "solve_bwd_schur") + @test execute!("solve_refinement", solver, x, b) === nothing # T16; ir_n_steps = 0: no-op + for phase in ("solve_fwd_schur", "solve_bwd_schur") @test thrown(() -> execute!(phase, solver, x, b)) isa NotSupportedError end # "solve_diag" (T15) is the identity for Cholesky @@ -311,20 +312,16 @@ end @test x1 == api_solve(backend, solver, b) @test relres(A, x1, b) <= tol(T) setparam!(solver, "deterministic_mode", 0) - # solve_mode: A = Aᴴ, and A = Aᵀ for real matrices + # solve_mode: A = Aᴴ, and Aᵀ = conj(A) (T16) setparam!(solver, "solve_mode", 2) @test relres(A, api_solve(backend, solver, b), b) <= tol(T) setparam!(solver, "solve_mode", 1) - if T <: Real - @test relres(A, api_solve(backend, solver, b), b) <= tol(T) - else - @test thrown(() -> api_solve(backend, solver, b)) isa NotSupportedError - end + @test relres(sparse(transpose(A)), api_solve(backend, solver, b), b) <= tol(T) setparam!(solver, "solve_mode", 0) - # ir_n_steps is stored, refinement comes in T16 (warns once) + # iterative refinement (T16) setparam!(solver, "ir_n_steps", 2) - @test relres(A, (@test_logs (:warn, r"iterative refinement") match_mode = :any api_solve(backend, solver, b)), b) <= - tol(T) + @test relres(A, api_solve(backend, solver, b), b) <= tol(T) + @test getparam(solver, "ir_n_steps") == 2 setparam!(solver, "ir_n_steps", 0) end @@ -335,15 +332,12 @@ end Ct = csr_of_transpose(Lh) C = CSR(to_device(backend, Ct.rowptr), to_device(backend, Ct.colval), to_device(backend, Ct.nzval), 70, 70; transposed = true) - if T <: Real - solver = DirectSolver(C, spd_structure(T), 'L') - execute!("analysis", solver, nothing, nothing) - execute!("factorization", solver, nothing, nothing) - b = rand(T, 70) - @test relres(A, api_solve(backend, solver, b), b) <= tol(T) - else - @test thrown(() -> DirectSolver(C, spd_structure(T), 'L')) isa NotSupportedError - end + # complex Hermitian: the stored matrix is conj(A), the solve conjugates (T16) + solver = DirectSolver(C, spd_structure(T), 'L') + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + b = rand(T, 70) + @test relres(A, api_solve(backend, solver, b), b) <= tol(T) end @testset "LinearAlgebra layer ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, diff --git a/test/test_refinement.jl b/test/test_refinement.jl new file mode 100644 index 0000000..a6cbf97 --- /dev/null +++ b/test/test_refinement.jl @@ -0,0 +1,302 @@ +# T16: iterative refinement ("solve_refinement", `ir_n_steps`, `ir_tol`), the +# solve sub-phases composed, `solve_mode`, `user_host_interrupt` and logging. + +# a solver on `backend` after "analysis" and "factorization" of the `view` triangle of `A` +function ir_solver(backend, A::SparseMatrixCSC{T}, structure, ::Type{INT} = Int32; view = 'L', + params = ()) where {T, INT} + solver = DirectSolver(api_matrix(backend, triangle_view(A, view), INT), structure, view) + for (name, value) in params + setparam!(solver, name, value) + end + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + return solver +end + +function ir_solve(backend, solver, b; steps = nothing, tol = nothing) + steps === nothing || setparam!(solver, "ir_n_steps", steps) + tol === nothing || setparam!(solver, "ir_tol", tol) + bd = to_device(backend, b) + xd = similar(bd) + execute!("solve", solver, xd, bd; asynchronous = false) + return to_host(xd) +end + +# a logger that sets `flag` when the refinement logs the residual after `step` corrections (an interrupt +# raised between two refinement steps, deterministically) +struct InterruptAtStepLogger <: Base.CoreLogging.AbstractLogger + flag::Threads.Atomic{Bool} + step::Int +end +Base.CoreLogging.min_enabled_level(::InterruptAtStepLogger) = Base.CoreLogging.Debug +Base.CoreLogging.shouldlog(::InterruptAtStepLogger, args...) = true +Base.CoreLogging.catch_exceptions(::InterruptAtStepLogger) = false +function Base.CoreLogging.handle_message(L::InterruptAtStepLogger, level, message, args...; kwargs...) + startswith(string(message), "refinement: step $(L.step),") && (L.flag[] = true) + return nothing +end + +const SOLVE_SUBPHASES = ("solve_fwd_perm", "solve_fwd", "solve_diag", "solve_bwd", "solve_bwd_perm", "solve_refinement") + +@testset "refinement on a badly scaled SPD matrix ($(backend_name(backend)))" for backend in BACKENDS + T = Float64 + A = badly_scaled_spd(T, 300, 0.02) # rows scaled by up to 10^(±8) + b = A * rand(T, 300) + solver = ir_solver(backend, A, "SPD") + r0 = relres(A, ir_solve(backend, solver, b; steps = 0), b) + @test getparam(solver, "ir_n_steps") == 0 + r1 = relres(A, ir_solve(backend, solver, b; steps = 1), b) + @test getparam(solver, "ir_n_steps") == 1 + @test r1 <= tol(T) + # Cholesky is backward stable under symmetric scaling: the unrefined residual is already at + # rounding level (≈ 3e-16, the gain of one step is 2× to 7×, see the T16 report), so the + # 100× reduction of the task text can't be observed on an SPD matrix; it is checked on the + # perturbed LDLᵀ factorization below. + @test_broken r1 <= r0 / 100 + # ir_tol: early exit, the data parameter reports the steps performed + x = ir_solve(backend, solver, b; steps = 10, tol = 1.0e-14) + @test getparam(solver, "ir_n_steps") < 10 + @test relres(A, x, b) <= 1.0e-14 + @test getparam(solver.options, "ir_n_steps") == 10 +end + +@testset "refinement with static pivot perturbation ($(backend_name(backend)), $INT)" for backend in BACKENDS, + INT in INTTYPES + # KKT matrix without 2×2 pairs: the zero (2,2) block is perturbed (pivot_epsilon), and refinement + # removes the perturbation error (the case of the MadNLP K2 systems, issue #71) + T = Float64 + K = kkt_matrix(T, 150, 60, 0.0) + b = K * rand(T, size(K, 1)) + solver = ir_solver(backend, K, "S", INT; params = (("pivot_pairs", "none"),)) + @test getparam(solver, "npivots") > 0 + r0 = relres(K, ir_solve(backend, solver, b; steps = 0), b) + r1 = relres(K, ir_solve(backend, solver, b; steps = 1), b) + @test r0 > 1.0e-8 + @test r1 <= r0 / 100 + x = ir_solve(backend, solver, b; steps = 10, tol = 1.0e-14) + steps = getparam(solver, "ir_n_steps") + @test 1 < steps < 10 + @test relres(K, x, b) <= 1.0e-14 + # the same steps without the early exit reach the same residual level + @test relres(K, ir_solve(backend, solver, b; steps, tol = 0.0), b) <= 1.0e-14 + @test getparam(solver, "ir_n_steps") == steps + # setting ir_n_steps reports the requested value until the next solve + setparam!(solver, "ir_n_steps", 7) + @test getparam(solver, "ir_n_steps") == 7 +end + +@testset "refinement: layouts, aliasing, element types ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + n = 200 + for (A, structure) in ((random_spd(T, n, 0.02), spd_structure(T)), (random_symindef(T, n, 0.02), sym_structure(T))) + # bitwise comparisons between solves: the atomic forward sweep (default on GPUs for real T) + # sums in a run-dependent order, so these solves use the deterministic variant + solver = ir_solver(backend, A, structure; params = (("ir_n_steps", 2), ("deterministic_mode", 1))) + for nrhs in (1, 3) + b = nrhs == 1 ? rand(T, n) : rand(T, n, nrhs) + x = ir_solve(backend, solver, b) + @test relres(A, x, b) <= tol(T) + @test getparam(solver, "ir_n_steps") == 2 + # X === B: "solve" keeps a copy of B for the residual, the result is the same + bd = to_device(backend, b) + execute!("solve", solver, bd, bd; asynchronous = false) + @test to_host(bd) == x + # "solve_refinement" needs the original B + @test thrown(() -> execute!("solve_refinement", solver, bd, bd)) isa InvalidValueError + nrhs == 1 && continue + # row-major and strided layouts + bt = to_device(backend, permutedims(b)) + Xt = MatrixDescriptor(similar(bt); transposed = true) + execute!("solve", solver, Xt, MatrixDescriptor(bt; transposed = true); asynchronous = false) + @test permutedims(to_host(Xt.data)) == x + bs = to_device(backend, vec(b)) + xs = similar(bs) + execute!("solve", solver, xs, bs; asynchronous = false) + @test reshape(to_host(xs), n, nrhs) == x + end + @test SDS.max_rhs(solver.refinement) >= 3 + # a row-major strided vector with fewer right-hand sides than the workspace, ir_tol > 0: + # the norm reduction indexes B with the solve's nrhs, not the workspace capacity + b2 = rand(T, n, 2) + Bd = MatrixDescriptor(T, n, 2; transposed = true) + update!(Bd, to_device(backend, vec(permutedims(b2)))) + Xd = MatrixDescriptor(T, n, 2; transposed = true) + update!(Xd, similar(Bd.data)) + setparam!(solver, "ir_tol", 1.0e-30) + execute!("solve", solver, Xd, Bd; asynchronous = false) + x2 = permutedims(reshape(to_host(Xd.data), 2, n)) + @test relres(A, x2, b2) <= tol(T) + W = solver.refinement + SDS.residual!(W, vec(solver.A.nzval), Xd.data, Bd.data; nrhs = 2, transposed = true) + nh = SDS.residual_norms!(W, Bd.data; nrhs = 2, transposed = true) + @test [nh[2], nh[4]] ≈ [sum(abs2, b2[:, 1]), sum(abs2, b2[:, 2])] + setparam!(solver, "ir_tol", 0.0) + # "solve_refinement" with ir_n_steps = 0 is a no-op + setparam!(solver, "ir_n_steps", 0) + b = rand(T, n) + x = ir_solve(backend, solver, b) + xd = to_device(backend, x) + execute!("solve_refinement", solver, xd, to_device(backend, b); asynchronous = false) + @test to_host(xd) == x && getparam(solver, "ir_n_steps") == 0 + # the residual uses the current values: update! without refactorization refines towards them + if structure == spd_structure(T) + A2 = A + T(n / 100) * I + M2 = api_matrix(backend, triangle_view(A2, 'L'), Int32) + update!(solver, M2) + x2 = ir_solve(backend, solver, b; steps = 20, tol = 0.0) + @test relres(A2, x2, b) <= tol(T) + end + end +end + +@testset "solve sub-phases compose to \"solve\" ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + n = 150 + for (A, structure) in ((random_spd(T, n, 0.03), spd_structure(T)), (random_symindef(T, n, 0.03), sym_structure(T))) + # bitwise comparison: deterministic forward sweep (the atomic one is run-dependent on GPUs) + solver = ir_solver(backend, A, structure; params = (("deterministic_mode", 1),)) + for steps in (0, 2), nrhs in (1, 2) + setparam!(solver, "ir_n_steps", steps) + b = to_device(backend, nrhs == 1 ? rand(T, n) : rand(T, n, nrhs)) + x1 = similar(b) + execute!("solve", solver, x1, b; asynchronous = false) + x2 = similar(b) + fill!(x2, zero(T)) + for phase in SOLVE_SUBPHASES + execute!(phase, solver, x2, b) + end + KernelAbstractions.synchronize(backend) + @test to_host(x2) == to_host(x1) + @test getparam(solver, "ir_n_steps") == steps + @test relres(A, to_host(x1), to_host(b)) <= tol(T) + end + end +end + +@testset "solve_mode ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + n = 120 + b = rand(T, n) + cases = Any[(random_spd(T, n, 0.03), spd_structure(T)), (random_symindef(T, n, 0.03), sym_structure(T))] + T <: Complex && push!(cases, (random_symindef(T, n, 0.03; hermitian = false), "S")) # complex symmetric + for (A, structure) in cases, steps in (0, 1) + solver = ir_solver(backend, A, structure; params = (("ir_n_steps", steps),)) + for (mode, op) in ((0, A), (1, transpose(A)), (2, adjoint(A))) + setparam!(solver, "solve_mode", mode) + @test relres(sparse(op), ir_solve(backend, solver, b), b) <= tol(T) + end + end + # the task's case: complex symmetric "S", solve_mode = 2 solves Aᴴ x = b (Aᴴ ≠ A) + if T <: Complex + A = random_symindef(T, n, 0.03; hermitian = false) + solver = ir_solver(backend, A, "S"; params = (("solve_mode", 2),)) + x = ir_solve(backend, solver, b) + @test relres(sparse(A'), x, b) <= tol(T) + @test relres(A, x, b) > 100 * tol(T) + end + # CSC input (the CSR of the transpose): for Hermitian matrices that is the conjugate + A = random_spd(T, n, 0.03) + Ct = csr_of_transpose(tril(A)) + C = CSR(to_device(backend, Ct.rowptr), to_device(backend, Ct.colval), to_device(backend, Ct.nzval), n, n; + transposed = true) + for structure in (spd_structure(T), sym_structure(T)), mode in 0:2, steps in (0, 2) + solver = DirectSolver(C, structure, 'L') + setparam!(solver, "solve_mode", mode) + setparam!(solver, "ir_n_steps", steps) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + op = mode == 1 ? sparse(transpose(A)) : A + @test relres(op, ir_solve(backend, solver, b), b) <= tol(T) + end +end + +@testset "user_host_interrupt ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + A = random_symindef(T, 300, 0.02) + b = rand(T, 300) + for structure in (sym_structure(T), spd_structure(T)) + M = structure == spd_structure(T) ? random_spd(T, 300, 0.02) : A + flag = Threads.Atomic{Bool}(false) + solver = DirectSolver(api_matrix(backend, tril(M), Int32), structure, 'L') + setparam!(solver, "user_host_interrupt", flag) + @test getparam(solver, "user_host_interrupt") === flag + flag[] = true + @test thrown(() -> execute!("analysis", solver, nothing, nothing)) isa InterruptedError + flag[] = false + execute!("analysis", solver, nothing, nothing) + flag[] = true + @test thrown(() -> factorize!(solver)) isa InterruptedError + @test thrown(() -> execute!("solve", solver, to_device(backend, b), to_device(backend, b))) isa + FactorizationError + flag[] = false + factorize!(solver) + @test getparam(solver, "info") == 0 + @test relres(M, ir_solve(backend, solver, b), b) <= tol(T) + # an interrupted refactorization: the solver needs a "factorization" again + flag[] = true + @test thrown(() -> refactorize!(solver)) isa InterruptedError + @test thrown(() -> refactorize!(solver)) isa FactorizationError + flag[] = false + factorize!(solver) + refactorize!(solver) + # the LinearAlgebra layer recovers too: cholesky!/ldlt! pick "factorization" again + flag[] = true + @test thrown(() -> refactorize!(solver)) isa InterruptedError + @test solver.fresh_factorization + flag[] = false + structure == spd_structure(T) ? cholesky!(solver, api_matrix(backend, tril(M), Int32)) : + ldlt!(solver, api_matrix(backend, tril(M), Int32)) + @test !solver.fresh_factorization + @test relres(M, ir_solve(backend, solver, b), b) <= tol(T) + # an interrupt after two refinement steps: X holds that iterate, ir_n_steps reports 2 + setparam!(solver, "ir_n_steps", 5) + setparam!(solver, "ir_tol", 1.0e-30) + @test Base.CoreLogging.with_logger(() -> thrown(() -> ir_solve(backend, solver, b)), + InterruptAtStepLogger(flag, 1)) isa InterruptedError + @test getparam(solver, "ir_n_steps") == 2 + flag[] = false + setparam!(solver, "ir_tol", 0.0) + # refinement polls the flag between steps + setparam!(solver, "ir_n_steps", 2) + flag[] = true + @test thrown(() -> ir_solve(backend, solver, b)) isa InterruptedError + flag[] = false + @test relres(M, ir_solve(backend, solver, b), b) <= tol(T) + setparam!(solver, "user_host_interrupt", nothing) + end +end + +@testset "LinearAlgebra layer refines ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + A = random_spd(T, 100, 0.03) + b = rand(T, 100) + for F in (cholesky(api_matrix(backend, tril(A), Int32); view = 'L'), ldlt(api_matrix(backend, tril(A), Int32); view = 'L')) + @test getparam(F, "ir_n_steps") == 2 + @test relres(A, to_host(F \ to_device(backend, b)), b) <= tol(T) + @test getparam(F, "ir_n_steps") == 2 + bd = to_device(backend, b) + ldiv!(F, bd) + @test relres(A, to_host(bd), b) <= tol(T) + end + # the handle layer keeps the cuDSS default + @test getparam(DirectSolver(api_matrix(backend, tril(A), Int32), spd_structure(T), 'L'), "ir_n_steps") == 0 +end + +@testset "logging" begin + backend = first(BACKENDS) + A = random_spd(Float64, 80, 0.05) + b = rand(80) + old = SDS.set_log_level!("info") + try + solver = DirectSolver(api_matrix(backend, tril(A), Int32), "SPD", 'L') + @test_logs (:info, r"^reordering") (:info, r"^symbolic_factorization") match_mode = :any analyze!(solver) + @test_logs (:info, r"^factorization: info = 0") match_mode = :any factorize!(solver) + setparam!(solver, "ir_n_steps", 3) + setparam!(solver, "ir_tol", 1.0e-30) + @test_logs (:info, r"^solve_refinement: 3 of 3 steps") match_mode = :any ir_solve(backend, solver, b) + SDS.set_log_level!(2) + @test_logs (:info, r"^refinement: step 0, relative residual") match_mode = :any ir_solve(backend, solver, b) + SDS.set_log_level!(0) + @test_logs min_level = Base.CoreLogging.Info ir_solve(backend, solver, b) # silent + finally + SDS.set_log_level!(old) + end + @test thrown(() -> SDS.set_log_level!("verbose")) isa InvalidValueError + @test SDS._parse_log_level("DEBUG") == SDS.LOG_DEBUG && SDS._parse_log_level("") == SDS.LOG_NONE +end