T16: iterative refinement, solve_mode, interrupt, logging

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)   `[ ]`

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

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

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)

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

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