T18: FGMRES-IR extension (Krylov.jl), ir_mode = "fgmres"

Project.toml

@@ -16,13 +16,15 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
 cuBLAS = "182d3088-87b7-4494-8cad-fc6afaa545bc"
 cuSOLVER = "887afef0-6a32-4de5-add4-7827692ba8fc"
 cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
 
 [extensions]
 SparseDirectSolverCUDAExt = ["CUDACore", "cuBLAS", "cuSOLVER", "cuSPARSE"]
+SparseDirectSolverKrylovExt = "Krylov"
 SparseDirectSolverMetisExt = "Metis"
 
 [compat]
@@ -34,6 +36,7 @@ CliqueTrees = "1"
 GPUArrays = "11"
 GPUArraysCore = "0.2"
 KernelAbstractions = "0.9"
+Krylov = "0.10"
 LinearAlgebra = "1.13"
 Metis = "1"
 SparseArrays = "1.13"

README.md

@@ -28,7 +28,9 @@ CUDA:
 * GPU triangular solves with multiple right-hand sides, forward, diagonal and
   backward sub-phases, permutations, `solve_mode` (transposed and conjugated
   systems), and iterative refinement (`ir_n_steps`, `ir_tol`), allocation-free
-  after the analysis; `user_host_interrupt` is polled between launch groups;
+  after the analysis, or FGMRES-IR with the factorization as preconditioner
+  (`ir_mode = "fgmres"`, after `using Krylov`); `user_host_interrupt` is
+  polled between launch groups;
 * the public API: `DirectSolver`, `execute!` with cuDSS phase strings, named
   phase wrappers, `update!`, `setparam!`/`getparam`, and the `LinearAlgebra`
   layer (`cholesky`, `cholesky!`, `ldlt`, `ldlt!`, `ldiv!`, `\`, `logabsdet`),
@@ -37,7 +39,7 @@ CUDA:
 
 Not there yet: LU (`"G"`), batches, Schur complements, matching and scaling
 (badly scaled MadNLP K2 systems still need them, see the refinement table in
-`TASKS.md`), FGMRES refinement, mixed precision, and the AMDGPU, oneAPI and
+`TASKS.md`), mixed precision, and the AMDGPU, oneAPI and
 Metal extensions. Unsupported structures, phases and parameters raise
 `NotSupportedError` rather than falling back silently.
 

TASKS.md

@@ -2382,13 +2382,91 @@ tasks below are refined.
 
 Each keeps the same structure; the test criteria are the minimum.
 
-### T18 — FGMRES-IR extension (Krylov.jl)   `[ ]`
+### T18 — FGMRES-IR extension (Krylov.jl)   `[!]`
 
 `ext/SparseDirectSolverKrylovExt.jl`: `ir_mode = "fgmres"` runs FGMRES with
 the factorization as right preconditioner; test: on an ill-conditioned
 Float64 system where plain IR stalls (`relres` not improving over 5 steps),
 FGMRES-IR reaches `relres ≤ 1e-12` within 20 iterations on CPU and CUDA.
 
+#### Report
+
+- Status: [!] (done on the KA CPU backend; CUDA from CI; deviations below)
+- What was built:
+  - `ext/SparseDirectSolverKrylovExt.jl` (new, weak dependency `Krylov = "0.10"`): registers
+    `fgmres_correction` in `SparseDirectSolver.FGMRES_PROVIDER` (same provider pattern as the Metis extension),
+    which runs `Krylov.fgmres!` (no restart, `rtol = 0`, `atol`, `itmax`) with the factorization as right
+    preconditioner; the `FgmresWorkspace` and its right-hand side vector are cached in the refinement workspace
+    and reallocated only when the length, the array type or a larger `itmax` require it.
+  - `src/solve/refinement.jl`: `RefinementOperator` (`op(A)` on the stacked compact columns, new KA kernel
+    `_spmv_kernel!` over the existing refinement map, member values per column for uniform batches, conjugated
+    for `conj(M)`), `FactorPreconditioner` (`op(A)⁻¹`: permute, forward/diagonal/backward sweeps, unpermute,
+    conjugating when `op(A) = conj(M)`; polls `user_host_interrupt`), both with `size`/`eltype`/`mul!`;
+    `fgmres_refine!` (residual `R₀ = B − op(A)X₀`, optional early exit on `ir_tol`, FGMRES on `op(A) D = R₀`,
+    `X += D`), `FGMRES_PROVIDER`, `fgmres_available()`. `RefinementWorkspace` gains the field `krylov`.
+  - `src/solver.jl`: `ir_mode = "fgmres"` selects `fgmres_refine!` for `"solve"` and `"solve_refinement"` (the
+    T16 `NotSupportedError` is gone); without `using Krylov` the refinement raises `NotSupportedError` (only
+    when refinement runs, `ir_n_steps > 0`). `getparam(solver, "ir_n_steps")` reports the FGMRES iterations.
+  - `test/test_fgmres.jl` (new); `ir_solver`, `ir_solve`, `SOLVE_SUBPHASES` moved from `test/test_refinement.jl`
+    to `test/utils.jl` (shared by both files); `test/runtests.jl` loads Krylov; Krylov added to
+    `test/Project.toml` and `bench/Project.toml`. `bench/refinement.jl --mode=fgmres`. README feature list.
+- Tests: `SDS_TEST_GPU=0 SDS_TEST_ONLY=test_fgmres`: 148 pass (CPU, 5.6 min, mostly compilation).
+  Full `SDS_TEST_GPU=0 julia --project=. -e 'using Pkg; Pkg.test()'` (CPU): 62532 pass, 1 fail, 1 broken (the
+  T16 `@test_broken`); the failure was the T16 logging test matching the `"solve_refinement: k of n steps"`
+  message I had changed. After restoring that message (FGMRES logs `"… FGMRES iterations"`),
+  `SDS_TEST_ONLY=test_refinement,test_aqua,test_fgmres`: 675 pass, 1 broken. Combined: 62533 pass, 0 fail,
+  1 broken. CUDA/AMDGPU: pending CI on the PR.
+  The criterion test: `kkt_matrix(Float64, 150, 5, 0)` with `pivot_pairs = "none"`, `pivot_epsilon = 1`
+  (5 duals perturbed far above their pivots: the factor is that of `A + E`, rank-5 `E`, `ρ((A+E)⁻¹E) ≈ 1`).
+  Plain IR: relres 7.8e-5 → 7.2e-5, 6.8e-5, …, 5.9e-5 after 5 steps, 1.0e-4 after 20 (asserted: no step of the
+  first 5 gains a digit, 20 steps stay > 1e-8); FGMRES-IR: ≤ 1e-12 with `ir_tol = 1e-12` in ≤ 20 iterations
+  (2.7e-16 after 10 iterations), for `INT ∈ (Int32, Int64)`. Also tested for all four element types: SPD/HPD
+  and `"S"`/`"H"` with 1 and 3 right-hand sides, `ir_tol` early exit, `X === B`, sub-phases composed equal
+  `"solve"` bitwise, `solve_mode` 1/2; a uniform batch (3 members × 2 rhs, `ubatch_index`); the interrupt (X
+  keeps the unrefined solution, 0 iterations reported, next solve works); `NotSupportedError` without the
+  provider.
+- Measurements: relres `‖b − Ax‖/‖b‖` on the MadNLP K2 dumps (regenerated in this session with
+  `bench/dump_madnlp_kkt.jl`, current MadNLP/ExaModelsPower), handle layer `"S"`, default pivoting, no matching,
+  KA CPU backend, `b = A·1`, `ir_tol = 0`, `bench/refinement.jl --steps=0,2,5,10,20 --mode=ir|fgmres`:
+
+  | K2 dump | n | npivots | ir 2 | ir 5 | ir 10 | ir 20 | fgmres 2 | fgmres 5 | fgmres 10 | fgmres 20 |
+  | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+  | case118 k2_1 (relres₀ 0.031) | 3150 | 5 | 4.2e-10 | 7.5e-16 | 7.6e-16 | 9.9e-16 | 4.7e-11 | 1.3e-15 | 1.3e-15 | 1.3e-15 |
+  | case118 k2_10 (4.8e-5) | 3150 | 3 | 1.5e-11 | 3.3e-17 | 3.4e-17 | 3.1e-17 | 2.3e-12 | 2.9e-16 | 2.9e-16 | 2.9e-16 |
+  | case118 k2_20 (6.8e-7) | 3150 | 3 | 3.6e-14 | 8.8e-19 | 3.8e-18 | 3.8e-18 | 3.4e-14 | 8.9e-18 | 8.9e-18 | 8.9e-18 |
+  | case1354 k2_1 (0.12) | 33811 | 34 | 7.1e-4 | 5.1e-7 | 1.4e-12 | 8.1e-16 | 4.3e-4 | 9.0e-10 | 4.1e-15 | 4.0e-15 |
+  | case1354 k2_10 (0.038) | 33811 | 46 | 5.0e-3 | 4.2e-4 | 5.3e-6 | 9.2e-10 | 1.5e-3 | 5.1e-7 | 2.6e-12 | 1.6e-15 |
+  | case1354 k2_20 (4.9e-3) | 33811 | 32 | 1.3e-4 | 8.3e-7 | 2.8e-10 | 1.4e-15 | 7.3e-5 | 2.5e-8 | 5.9e-15 | 1.3e-15 |
+
+  FGMRES-IR with 10 iterations beats the cuDSS bar of issue #71 (4.6e-11 on case1354 with `matching_alg =
+  "algo5"` and 5 IR steps) on all three case1354 dumps without matching; plain IR needs 20 steps and still
+  misses it on k2_10. Each FGMRES iteration costs one SpMV and one solve (as one IR step) plus Krylov's
+  orthogonalization (`O(k n)` per iteration) and one host synchronization per iteration.
+- Deviations from PLAN.md / this task:
+  - FGMRES runs on one vector holding every right-hand side and every active batch member (the block-diagonal
+    system `I ⊗ op(A)`), so the SpMV and the sweeps stay one launch group per application and no single-column
+    variants of the kernels are needed. The stopping test is joint: `‖R₀ − op(A)D‖₂ ≤ ir_tol · minₖ‖Bₖ‖₂`,
+    which bounds every column's relative residual by `ir_tol`. With several right-hand sides this can take
+    more iterations than separate FGMRES runs per column (MadNLP solves one right-hand side).
+  - `ir_n_steps` is the maximum number of FGMRES iterations (no restart; the Krylov basis holds `2·ir_n_steps`
+    vectors of `n·ncols` entries, allocated at the first FGMRES solve and kept). `ir_tol = 0` runs all
+    `ir_n_steps` iterations (Krylov.jl stops earlier only on an exact residual).
+  - FGMRES is not synchronization-free: Krylov.jl reads its dot products and norms on the host every iteration,
+    and allocates small host vectors. The "no host synchronization inside solve" rule holds only for
+    `ir_mode = "ir"` with `ir_tol = 0`.
+  - An interrupt during FGMRES leaves `X` at the unrefined solution (the correction is added only at the end)
+    and reports 0 iterations; plain IR keeps the last completed step.
+  - The stalling test matrix is `kkt_matrix` with parameters (no new generator).
+- Open issues / follow-ups:
+  - The CUDA path (Krylov.jl on `CuVector`, the new SpMV kernel) is exercised only by CI.
+  - Mixed precision (`factor_precision = Float32`, M13) can reuse `FactorPreconditioner` unchanged: the operators
+    already work in the input precision.
+- Suggested plan changes:
+  - PLAN §1.3 `ir_mode`/§2.5: document "`ir_n_steps` = maximum FGMRES iterations, joint stopping over the
+    stacked right-hand sides, one host synchronization per iteration".
+  - Given the K2 table, consider recommending `ir_mode = "fgmres"` with `ir_n_steps ≈ 10`, `ir_tol ≈ 1e-12` for
+    MadNLP's K2 systems, independent of T21.
+
 ### T19 — General LU (`"G"`): CPU reference + GPU   `[ ]`
 
 Symmetric-pattern multifrontal LU with `U` panels, in-block threshold partial

bench/Project.toml

@@ -3,6 +3,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixDepot = "b51810bb-c9f3-55da-ae3c-350fc1fbce05"
 Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
@@ -19,6 +20,7 @@ BenchmarkTools = "1"
 CUDA = "6"
 CUDSS = "0.8"
 DelimitedFiles = "1"
+Krylov = "0.10"
 MatrixDepot = "1"
 Metis = "1"
 julia = "1.13"

bench/refinement.jl

@@ -3,8 +3,11 @@
 # (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).
+# --mode=fgmres runs FGMRES-IR (`ir_mode = "fgmres"`, T18; needs Krylov in the
+# environment, e.g. --project=bench) with up to `ir_n_steps` iterations.
 #
 #   julia --project=. bench/refinement.jl [--only=substring,...] [--steps=0,2,5] [--backend=cpu|cuda]
+#                                         [--mode=ir|fgmres]
 
 using LinearAlgebra
 using SparseArrays
@@ -13,9 +16,9 @@ using SparseDirectSolver
 include(joinpath(@__DIR__, "matrices.jl"))
 
 function parse_args(args)
-    o = Dict{String, String}("only" => "", "steps" => "0,2,5", "backend" => "cpu")
+    o = Dict{String, String}("only" => "", "steps" => "0,2,5", "backend" => "cpu", "mode" => "ir")
     for a in args
-        m = match(r"^--(only|steps|backend)=(.*)$", a)
+        m = match(r"^--(only|steps|backend|mode)=(.*)$", a)
         m === nothing && error("unknown argument \"$a\"; see the header of bench/refinement.jl")
         o[m[1]] = m[2]
     end
@@ -24,6 +27,8 @@ end
 
 const OPTS = parse_args(ARGS)
 const STEPS = parse.(Int, split(OPTS["steps"], ','))
+OPTS["mode"] in ("ir", "fgmres") || error("--mode must be ir or fgmres")
+OPTS["mode"] == "fgmres" && @eval using Krylov
 
 if OPTS["backend"] == "cuda"
     @eval using CUDA
@@ -36,6 +41,7 @@ end
 function relres_table(name, A::SparseMatrixCSC{T}) where {T}
     L = tril(A)
     solver = DirectSolver(to_dev(L), "S", 'L')
+    setparam!(solver, "ir_mode", OPTS["mode"])
     analyze!(solver)
     factorize!(solver; asynchronous = false)
     b = A * ones(T, size(A, 1))
@@ -59,7 +65,7 @@ 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("| K2 dump | n | npivots | ", join(("relres, $(OPTS["mode"]) $k" for k in STEPS), " | "), " |")
 println("| --- | ---: | ---: | ", join(("---:" for _ in STEPS), " | "), " |")
 for f in dumps
     selected(f) || continue

ext/SparseDirectSolverKrylovExt.jl

@@ -0,0 +1,49 @@
+"""
+    SparseDirectSolverKrylovExt
+
+FGMRES-IR for SparseDirectSolver (PLAN §2.5, `ir_mode = "fgmres"`). Loading
+Krylov.jl registers [`fgmres_correction`](@ref) in
+`SparseDirectSolver.FGMRES_PROVIDER`; the operators (SpMV and the
+factorization as right preconditioner) and the refinement loop live in the
+core (`src/solve/refinement.jl`), this extension only runs `Krylov.fgmres!`.
+"""
+module SparseDirectSolverKrylovExt
+
+using SparseDirectSolver
+using Krylov
+
+"""
+    fgmres_correction(cache, A, P, R, len; atol, itmax) -> (D, iterations)
+
+Solve `A D = R[1:len]` (the first `len` entries of the device array `R`,
+copied into the Krylov storage) by `Krylov.fgmres!` with right
+preconditioner `P` (`mul!`), no restart, absolute tolerance `atol`, relative
+tolerance 0 and at most `itmax` iterations. The `FgmresWorkspace` and the
+right-hand side vector are kept in `cache` (a `Ref{Any}`) and reallocated only
+when `len` or the element/array type change or `itmax` exceeds the stored
+Krylov basis. `D` is the workspace's solution vector (valid until the next
+call).
+"""
+function fgmres_correction(cache::Base.RefValue{Any}, A, P, R::AbstractArray{FC}, len::Int; atol::Real,
+                           itmax::Int) where {FC}
+    S = typeof(similar(R, FC, 0))
+    c = cache[]
+    if c isa Tuple{Krylov.FgmresWorkspace{real(FC), FC, S}, S, Int} && c[1].m == len && c[3] >= itmax
+        workspace, b, _ = c
+    else
+        b = similar(R, FC, len)
+        workspace = Krylov.FgmresWorkspace(len, len, S; memory = itmax)
+        cache[] = (workspace, b, itmax)
+    end
+    copyto!(b, 1, R, 1, len)
+    Krylov.fgmres!(workspace, A, b; N = P, ldiv = false, restart = false, atol = real(FC)(atol),
+                   rtol = zero(real(FC)), itmax)
+    return workspace.x, workspace.stats.niter
+end
+
+function __init__()
+    SparseDirectSolver.FGMRES_PROVIDER[] = fgmres_correction
+    return nothing
+end
+
+end # module SparseDirectSolverKrylovExt