diff --git a/Project.toml b/Project.toml index cf31f67..26ee3f1 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ 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" @@ -23,6 +24,7 @@ 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" diff --git a/README.md b/README.md index d678ae6..d297971 100644 --- a/README.md +++ b/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. diff --git a/TASKS.md b/TASKS.md index 8e18ba4..a2584f5 100644 --- a/TASKS.md +++ b/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 diff --git a/bench/Project.toml b/bench/Project.toml index 921478f..efcfe2a 100644 --- a/bench/Project.toml +++ b/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" diff --git a/bench/refinement.jl b/bench/refinement.jl index b4f7891..5eccc55 100644 --- a/bench/refinement.jl +++ b/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 diff --git a/ext/SparseDirectSolverKrylovExt.jl b/ext/SparseDirectSolverKrylovExt.jl new file mode 100644 index 0000000..bc66c1f --- /dev/null +++ b/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 diff --git a/src/solve/refinement.jl b/src/solve/refinement.jl index ccd57e7..6cdbe2b 100644 --- a/src/solve/refinement.jl +++ b/src/solve/refinement.jl @@ -34,7 +34,9 @@ sides: * `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`). + `norms_host`, its host copy (the early-exit test of `ir_tol`); +* `krylov`: the storage of FGMRES-IR (`ir_mode = "fgmres"`), owned by the + Krylov.jl extension (`nothing` until its first solve). """ struct RefinementWorkspace{T, R <: Real, VI <: AbstractVector, MT <: AbstractMatrix{T}, VR <: AbstractVector{R}} rowptr::VI @@ -44,6 +46,7 @@ struct RefinementWorkspace{T, R <: Real, VI <: AbstractVector, MT <: AbstractMat Bc::MT norms::VR norms_host::Vector{R} + krylov::Base.RefValue{Any} end max_rhs(W::RefinementWorkspace) = size(W.R, 2) @@ -88,7 +91,7 @@ function allocate_refinement(map::NTuple{3, Vector{Int}}, ::Type{T}, ::Type{INT} 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)) + dev(rowptr), dev(colval), dev(src), R, Bc, norms, zeros(real(T), 2 * nrhs), Ref{Any}(nothing)) end function _to_index_type(::Type{INT}, v::Vector{Int}) where {INT} @@ -104,7 +107,8 @@ function _grow_refinement(W::RefinementWorkspace{T, R, VI, MT, VR}, nrhs::Intege 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)) + return RefinementWorkspace{T, R, VI, MT, VR}(W.rowptr, W.colval, W.src, Rm, Bc, norms, zeros(R, 2 * nrhs), + Ref{Any}(nothing)) end # compact column r (user column `_bm_ucol(bm, r)` of the `nrhs` user columns, values of member @@ -348,6 +352,198 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac return steps end +# --- FGMRES-IR (`ir_mode = "fgmres"`, Krylov.jl extension) ---------------------- +# +# FGMRES on the correction system `op(A) D = R₀` (`R₀ = B - op(A) X₀`, `X₀` the +# solution of the sweeps), right-preconditioned by the factorization, then +# `X = X₀ + D`. Krylov.jl works on vectors: the compact columns of all +# right-hand sides (and active batch members) are stacked into one vector of +# length `n ncols`, i.e. FGMRES runs on the block-diagonal system `I ⊗ op(A)` +# with the operators below; both apply all columns per launch. + +""" + FGMRES_PROVIDER + +`Ref` to the FGMRES driver of the Krylov.jl extension +(`SparseDirectSolverKrylovExt`), set in its `__init__`; `nothing` while +Krylov.jl is not loaded (`ir_mode = "fgmres"` then raises +[`NotSupportedError`](@ref)). Called as +`provider(cache, A, P, R, len; atol, itmax) -> (D, iterations)`: FGMRES for +`A D = R[1:len]` with right preconditioner `P` (`mul!(y, P, x)`), absolute +residual tolerance `atol`, at most `itmax` iterations, storage kept in `cache` +(a `Ref{Any}`). +""" +const FGMRES_PROVIDER = Ref{Any}(nothing) + +""" + fgmres_available() -> Bool + +Whether the Krylov.jl extension is loaded (`ir_mode = "fgmres"` works). +""" +fgmres_available() = FGMRES_PROVIDER[] !== nothing + +""" + RefinementOperator(W, nzval, n, ncols, bm, conjugate) + +`op(A)` (the full matrix `M` of the [`RefinementWorkspace`](@ref) `W` with +the values `nzval`, `conj(M)` when `conjugate`) acting on the `n × ncols` +compact columns stacked into vectors of length `n ncols` (column `r` uses +the values of the member of compact column `r` of the [`BatchMap`](@ref) +`bm`). Supports `size`, `eltype` and `mul!(y, op, x)` (one launch, +asynchronous): the matrix operator of FGMRES-IR. +""" +struct RefinementOperator{T, RW <: RefinementWorkspace{T}, V <: AbstractVector, BM <: BatchMap} + W::RW + nzval::V + n::Int + ncols::Int + bm::BM + conjugate::Bool +end + +Base.size(op::RefinementOperator) = (op.n * op.ncols, op.n * op.ncols) +Base.size(op::RefinementOperator, d::Integer) = d <= 2 ? op.n * op.ncols : 1 +Base.eltype(::RefinementOperator{T}) where {T} = T + +@kernel function _spmv_kernel!(y, rowptr, colval, src, nzval_all, x, n, ncols, bm, ::Val{CJ}) where {CJ} + q = @index(Global, Linear) + i = (q - 1) % n + 1 + r = (q - 1) ÷ n + 1 + @inbounds if r <= ncols + nzval = _mview(nzval_all, _bm_cmember(bm, r), bm.nbatch) + off = (r - 1) * n + acc = zero(eltype(y)) + 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 * x[Int(colval[c]) + off] + end + y[i + off] = acc + end +end + +function LinearAlgebra.mul!(y::AbstractVector, op::RefinementOperator, x::AbstractVector) + m = op.n * op.ncols + length(x) == m && length(y) == m || + throw(DimensionMismatch("the operator has size $m, x has $(length(x)) and y $(length(y)) entries")) + m > 0 || return y + W = op.W + kernel! = _spmv_kernel!(KernelAbstractions.get_backend(W.R), PERMUTE_WORKGROUP) + if op.conjugate + kernel!(y, W.rowptr, W.colval, W.src, op.nzval, x, op.n, op.ncols, op.bm, Val(true); ndrange = m) + else + kernel!(y, W.rowptr, W.colval, W.src, op.nzval, x, op.n, op.ncols, op.bm, Val(false); ndrange = m) + end + return y +end + +""" + FactorPreconditioner(ws, S, N, nrhs, ncols, conjugate, deterministic, interrupt) + +`op(A)⁻¹` through the factorization in `N` (permutation, forward, diagonal +and backward sweeps in the solve workspace `ws`; `conj(M⁻¹ conj(x))` when +`conjugate`) on the stacked compact columns (`nrhs` per active member, +`ncols` in all) of the [`RefinementOperator`](@ref). `mul!(y, P, x)` polls +`interrupt` first ([`InterruptedError`](@ref)); asynchronous otherwise. The +right preconditioner of FGMRES-IR. +""" +struct FactorPreconditioner{T, WS <: SolveWorkspace{T}, SY <: Symbolic, NU <: Numeric} + ws::WS + S::SY + N::NU + nrhs::Int + ncols::Int + conjugate::Bool + deterministic::Bool + interrupt::Union{Nothing, Threads.Atomic{Bool}} +end + +Base.size(P::FactorPreconditioner) = (P.S.n * P.ncols, P.S.n * P.ncols) +Base.size(P::FactorPreconditioner, d::Integer) = d <= 2 ? P.S.n * P.ncols : 1 +Base.eltype(::FactorPreconditioner{T}) where {T} = T + +function LinearAlgebra.mul!(y::AbstractVector, P::FactorPreconditioner, x::AbstractVector) + _poll_interrupt(P.interrupt) + perm = P.S.perm + permute_rhs!(P.ws.Y, x, perm; conjugate = P.conjugate) + forward_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs, deterministic = P.deterministic) + diagonal_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs) + backward_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs) + unpermute_solution!(y, P.ws.Y, perm; conjugate = P.conjugate) + return y +end + +# smallest ‖Bₖ‖ over the right-hand sides (1 for Bₖ = 0, as in `_max_relative_residual`) +function _min_rhs_norm(norms::Vector{R}, nrhs::Int) where {R} + m = typemax(R) + for k in 1:nrhs + bk = sqrt(norms[2k]) + m = min(m, bk > 0 ? bk : one(R)) + end + return m +end + +""" + fgmres_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)) -> iterations + +FGMRES-IR (`ir_mode = "fgmres"`, needs Krylov.jl) of the solution `X` of +`op(A) X = B`, with the arguments of [`refine!`](@ref): + + R₀ = B - op(A) X (residual!) + stop if tol > 0 and maxₖ ‖R₀ₖ‖₂ / ‖Bₖ‖₂ ≤ tol + D = FGMRES(op(A), R₀; right preconditioner op(A)⁻¹ by the factor, + at most nsteps iterations, ‖R₀ - op(A) D‖₂ ≤ tol minₖ ‖Bₖ‖₂) + X += D (add_correction!) + +on the stacked compact columns ([`RefinementOperator`](@ref), +[`FactorPreconditioner`](@ref)). Each iteration costs one SpMV and one solve, +as a step of plain refinement. `tol = 0` runs `nsteps` iterations (FGMRES +stops earlier only on an exact residual). The joint criterion bounds every +column's relative residual by `tol`. Returns the iterations performed (also +in `progress[]`); an interrupt (polled before every preconditioner +application) leaves `X` unchanged and `progress[] = 0`. FGMRES synchronizes +with the host every iteration (Krylov.jl's dot products and norms). +""" +function fgmres_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)) + provider = FGMRES_PROVIDER[] + provider === nothing && + throw(NotSupportedError("ir_mode = \"fgmres\" needs Krylov.jl: load it with `using Krylov`")) + nu = rhs_count(X, S.n; transposed) + nu % N.nbatch == 0 || throw(DimensionMismatch("$nu right-hand sides for a batch of $(N.nbatch) members")) + nrhs = nu ÷ N.nbatch + ncols = _ncols(N, nrhs) + max_rhs(W) >= ncols && max_rhs(ws) >= ncols || + throw(DimensionMismatch("the refinement workspace holds $(max_rhs(W)) right-hand sides, need $ncols")) + progress[] = 0 + nsteps > 0 && S.n * ncols > 0 || return 0 + bm = batch_map(N; nrhs) + _poll_interrupt(interrupt) + residual!(W, nzval, X, B; nrhs = nu, transposed, b_transposed, conjugate, bm) + R = real(eltype(W.R)) + atol = zero(R) + if tol > 0 + norms = residual_norms!(W, B; nrhs = nu, transposed = b_transposed, bm) + rel = _max_relative_residual(norms, ncols) + _log(LOG_DEBUG, () -> "fgmres refinement: initial relative residual $rel") + rel <= tol && return 0 + atol = R(tol) * _min_rhs_norm(norms, ncols) + end + A = RefinementOperator(W, nzval, S.n, ncols, bm, conjugate) + P = FactorPreconditioner(ws, S, N, nrhs, ncols, conjugate, deterministic, interrupt) + D, iters = provider(W.krylov, A, P, W.R, S.n * ncols; atol, itmax = Int(nsteps)) + permute_rhs!(ws.Y, D, S.perm) + add_correction!(X, ws.Y, S.perm; nrhs = nu, transposed, bm) + progress[] = iters + return iters +end + """ solve_conjugated(structure, T, csc_input::Bool, solve_mode) -> Bool diff --git a/src/solver.jl b/src/solver.jl index e0ba86f..291e34d 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -275,6 +275,12 @@ the solver's current matrix values (KA SpMV over the full pattern), and, when (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)). +With `ir_mode = "fgmres"` (needs `using Krylov`, else [`NotSupportedError`](@ref)) +the refinement is FGMRES on the correction system, right-preconditioned by the +factorization ([`fgmres_refine!`](@ref)): at most `ir_n_steps` iterations (one +SpMV and one solve each, as a plain step), stopping once every relative +residual is below `ir_tol` (`ir_tol = 0`: all `ir_n_steps` iterations); +`"ir_n_steps"` then reports the FGMRES iterations. 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 @@ -498,8 +504,6 @@ 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)")) - 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)")) @@ -567,14 +571,16 @@ function _refine_phase!(solver::DirectSolver, W::RefinementWorkspace, ws::SolveW 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) + driver = opts.ir_mode == IR_FGMRES ? fgmres_refine! : refine! + driver(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") + unit = opts.ir_mode == IR_FGMRES ? "FGMRES iterations" : "steps" + _log(LOG_INFO, () -> "solve_refinement: $steps of $(opts.ir_n_steps) $unit") return nothing end diff --git a/test/Project.toml b/test/Project.toml index 414925d..a0d9f1f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -16,6 +17,7 @@ Adapt = "4" Aqua = "0.8" DelimitedFiles = "1" KernelAbstractions = "0.9" +Krylov = "0.10" LinearAlgebra = "1.13" Metis = "1" Printf = "1.13" diff --git a/test/runtests.jl b/test/runtests.jl index 01e23fc..ba60b4c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ using Aqua using KernelAbstractions using SparseDirectSolver using Metis # activates SparseDirectSolverMetisExt (nested-dissection orderings, T05) +using Krylov # activates SparseDirectSolverKrylovExt (ir_mode = "fgmres", T18) const SDS = SparseDirectSolver diff --git a/test/test_fgmres.jl b/test/test_fgmres.jl new file mode 100644 index 0000000..65a1c87 --- /dev/null +++ b/test/test_fgmres.jl @@ -0,0 +1,129 @@ +# T18: FGMRES-IR (`ir_mode = "fgmres"`, Krylov.jl extension) with the +# factorization as right preconditioner. + +# KKT matrix with 5 duals whose static pivot perturbation (`pivot_epsilon = 1`, no 2×2 pairs) is far larger +# than the pivots: the factor is that of `A + E`, `E` of rank 5 and `ρ((A + E)⁻¹ E) ≈ 1`, so plain refinement +# stalls, while `(A + E)⁻¹ A` is the identity plus a rank-5 term and FGMRES converges in a few iterations +const FGMRES_STALL_PARAMS = (("pivot_pairs", "none"), ("pivot_epsilon", 1.0)) +fgmres_stall_matrix() = kkt_matrix(Float64, 150, 5, 0.0) + +@testset "FGMRES-IR where plain IR stalls ($(backend_name(backend)), $INT)" for backend in BACKENDS, + INT in INTTYPES + Random.seed!(666) + K = fgmres_stall_matrix() + b = K * rand(size(K, 1)) + solver = ir_solver(backend, K, "S", INT; params = FGMRES_STALL_PARAMS) + @test getparam(solver, "npivots") > 0 + # plain IR: no step of the first 5 gains a digit, and 20 steps stay far from 1e-12 + r = [relres(K, ir_solve(backend, solver, b; steps), b) for steps in 0:5] + @test r[1] > 1.0e-8 + @test minimum(r[2:end]) > r[1] / 10 + @test relres(K, ir_solve(backend, solver, b; steps = 20), b) > 1.0e-8 + # FGMRES-IR: relres ≤ 1e-12 within 20 iterations (early exit on ir_tol) + setparam!(solver, "ir_mode", "fgmres") + x = ir_solve(backend, solver, b; steps = 20, tol = 1.0e-12) + iters = getparam(solver, "ir_n_steps") + @test 1 <= iters <= 20 + @test relres(K, x, b) <= 1.0e-12 + # without ir_tol, all 20 iterations are run (or FGMRES ends on an exact residual) + x = ir_solve(backend, solver, b; steps = 20, tol = 0.0) + @test iters <= getparam(solver, "ir_n_steps") <= 20 + @test relres(K, x, b) <= 1.0e-12 + @test getparam(solver, "ir_mode") == "fgmres" + # more iterations than the stored Krylov basis: the workspace is reallocated + x = ir_solve(backend, solver, b; steps = 30, tol = 1.0e-14) + @test relres(K, x, b) <= 1.0e-14 +end + +@testset "FGMRES-IR: element types, multiple right-hand sides, solve_mode ($(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))) + solver = ir_solver(backend, A, structure; params = (("ir_mode", "fgmres"), ("deterministic_mode", 1))) + for nrhs in (1, 3) + b = nrhs == 1 ? rand(T, n) : rand(T, n, nrhs) + x = ir_solve(backend, solver, b; steps = 3, tol = 0.0) + @test relres(A, x, b) <= tol(T) + @test 1 <= getparam(solver, "ir_n_steps") <= 3 + # ir_tol: early exit, every column below the tolerance + x = ir_solve(backend, solver, b; steps = 10, tol = 100 * eps(real(T))) + @test getparam(solver, "ir_n_steps") < 10 + @test all(k -> relres(A, x[:, k], b[:, k]) <= tol(T), 1:nrhs) + # X === B is allowed with "solve"; the six sub-phases compose to "solve" + bd = to_device(backend, b) + execute!("solve", solver, bd, bd; asynchronous = false) + xs = to_host(bd) + @test relres(A, xs, b) <= tol(T) + bd = to_device(backend, b) + xd = similar(bd) + for phase in SOLVE_SUBPHASES + execute!(phase, solver, xd, bd) + end + @test to_host(xd) == xs + end + # transposed and adjoint solves (complex "S": Aᴴ x = b goes through the conjugated operators) + b = rand(T, n) + for (mode, op) in ((1, transpose), (2, adjoint)) + setparam!(solver, "solve_mode", mode) + x = ir_solve(backend, solver, b; steps = 3, tol = 0.0) + @test relres(op(A), x, b) <= tol(T) + end + setparam!(solver, "solve_mode", 0) + end +end + +@testset "FGMRES-IR: uniform batch ($(backend_name(backend)), $T)" for backend in BACKENDS, T in (Float64, ComplexF32) + n, nb, nrhs = 120, 3, 2 + members = batch_members(random_symindef(T, n, 0.03), nb) + solver = DirectSolver(api_batch_matrix(backend, members, 'L'), sym_structure(T), 'L') + setparam!(solver, "ir_mode", "fgmres") + setparam!(solver, "ir_n_steps", 3) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + B = rand(T, n, nrhs, nb) + Bd = to_device(backend, B) + Xd = similar(Bd) + execute!("solve", solver, Xd, Bd; asynchronous = false) + @test all(batch_relres(members, to_host(Xd), B) .<= tol(T)) + @test 1 <= getparam(solver, "ir_n_steps") <= 3 + # ubatch_index: only member 2 is solved and refined, the others' columns are untouched + setparam!(solver, "ubatch_index", 1) + X2 = to_device(backend, zeros(T, n, nrhs, nb)) + execute!("solve", solver, X2, Bd; asynchronous = false) + Xh = to_host(X2) + @test relres(members[2], Xh[:, :, 2], B[:, :, 2]) <= tol(T) + @test iszero(Xh[:, :, 1]) && iszero(Xh[:, :, 3]) +end + +@testset "FGMRES-IR: interrupt and missing Krylov ($(backend_name(backend)))" for backend in BACKENDS + Random.seed!(666) + K = fgmres_stall_matrix() + b = K * rand(size(K, 1)) + flag = Threads.Atomic{Bool}(false) + solver = ir_solver(backend, K, "S"; params = (FGMRES_STALL_PARAMS..., ("deterministic_mode", 1), + ("user_host_interrupt", flag))) + x0 = ir_solve(backend, solver, b; steps = 0) + setparam!(solver, "ir_mode", "fgmres") + setparam!(solver, "ir_n_steps", 10) + # an interrupt during FGMRES leaves the unrefined solution in X and reports 0 iterations + flag[] = true + bd = to_device(backend, b) + xd = similar(bd) + @test thrown(() -> execute!("solve", solver, xd, bd; asynchronous = false)) isa InterruptedError + @test to_host(xd) == x0 + @test getparam(solver, "ir_n_steps") == 0 + flag[] = false + @test relres(K, ir_solve(backend, solver, b; tol = 1.0e-12), b) <= 1.0e-12 + # without Krylov.jl, ir_mode = "fgmres" raises NotSupportedError at the refinement (not at setparam!) + provider = SDS.FGMRES_PROVIDER[] + @test SDS.fgmres_available() + try + SDS.FGMRES_PROVIDER[] = nothing + @test !SDS.fgmres_available() + @test thrown(() -> execute!("solve", solver, xd, bd; asynchronous = false)) isa NotSupportedError + # no refinement requested: nothing needs Krylov.jl + @test ir_solve(backend, solver, b; steps = 0) == x0 + finally + SDS.FGMRES_PROVIDER[] = provider + end +end diff --git a/test/test_refinement.jl b/test/test_refinement.jl index a6cbf97..a32f85e 100644 --- a/test/test_refinement.jl +++ b/test/test_refinement.jl @@ -1,27 +1,6 @@ # 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 @@ -36,8 +15,6 @@ function Base.CoreLogging.handle_message(L::InterruptAtStepLogger, level, messag 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) diff --git a/test/utils.jl b/test/utils.jl index 68874ff..10d867f 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -465,3 +465,42 @@ function batch_relres(members, X::AbstractArray, B::AbstractArray) Br = reshape(Array(B), n, :, nb) return [relres(members[k], Xr[:, :, k], Br[:, :, k]) for k in 1:nb] end + +# --------------------------------------------------------------------------- +# refinement (T16, T18) + +"The six solve sub-phases in order; composed they equal `\"solve\"` (T16)." +const SOLVE_SUBPHASES = ("solve_fwd_perm", "solve_fwd", "solve_diag", "solve_bwd", "solve_bwd_perm", "solve_refinement") + +""" + ir_solver(backend, A, structure, INT = Int32; view = 'L', params = ()) + +A [`DirectSolver`](@ref) on `backend` after `"analysis"` and `"factorization"` +of the `view` triangle of `A`, with the parameters `params` (name/value pairs) +set first (T16, T18). +""" +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 + +""" + ir_solve(backend, solver, b; steps = nothing, tol = nothing) -> x + +`"solve"` of `b` (host array, moved to `backend`) with `ir_n_steps = steps` +and `ir_tol = tol` when given; returns the solution on the host (T16, T18). +""" +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