diff --git a/AGENTS.md b/AGENTS.md index fae6a6a..e969ae1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -218,7 +218,11 @@ a split file, write every testset either over `ELTYPES` or behind `RUN_SHARED`. tolerances inside individual test files. * Every numeric test loops over `BACKENDS` and over `ELTYPES` unless the task says otherwise. Seed with `Random.seed!(666)`. Tolerance is `tol(T) = - sqrt(eps(real(T)))` on well-conditioned generators. + sqrt(eps(real(T)))` on well-conditioned generators. Elementwise + device-vs-reference panel checks use `panel_tol(T)` only on well-conditioned + generators; on generators with element growth (weak pivots, static pivoting) + use `growth_tol(T, Nr)`, since GPUs fuse multiply-adds and the two roundings + differ by O(eps · growth). The pivot sequence is always compared exactly. * Tests must not need the network. SuiteSparse downloads belong to `bench/`. * A task is done only when its listed tests pass on both the CPU backend and CUDA on this machine, and the Report records the counts. diff --git a/README.md b/README.md index d297971..4633f99 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,11 @@ CUDA: with a user-chosen sign per row (`pivot_sign`, which cuDSS cannot do), and `inertia`, `npivots` and `pivot_stats` read back from the device; the same pivot sequence as the CPU reference LDLᵀ, which serves as its oracle; +* GPU LU (`"G"`, `L D U` on the symmetric pattern of `A + Aᵀ`) with row + interchanges inside the fully-summed block of each front (threshold partial + pivoting, `pivot_threshold`), static perturbation of tiny pivots, `perm_row` + and `perm_col`, solves with `A`, `Aᵀ` and `Aᴴ`, and uniform batches; the same + pivot sequence as the CPU reference LU; * 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 @@ -33,11 +38,11 @@ CUDA: 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`), + layer (`cholesky`, `cholesky!`, `ldlt`, `ldlt!`, `lu`, `lu!`, `ldiv!`, `\`, `logabsdet`), checked by the test suite of CUDSS.jl ported to this package; phase logging through `SDS_LOG_LEVEL`. -Not there yet: LU (`"G"`), batches, Schur complements, matching and scaling +Not there yet: non-uniform batches, Schur complements, matching and scaling (badly scaled MadNLP K2 systems still need them, see the refinement table in `TASKS.md`), mixed precision, and the AMDGPU, oneAPI and Metal extensions. Unsupported structures, phases and parameters raise @@ -78,7 +83,7 @@ dissection ordering; without it the ordering is AMD. The handle API mirrors CUDSS.jl: a `DirectSolver` is created from a CSR matrix living on a KernelAbstractions backend, with a structure string (`"SPD"`, -`"HPD"`, `"S"`, `"H"`; `"G"` is reserved for the LU milestone) and the +`"HPD"`, `"S"`, `"H"`, or `"G"` for LU, which reads the full matrix, view `'F'`) and the triangle that is read (`'L'`, `'U'` or `'F'`). Phases are run with `execute!`. ```julia diff --git a/TASKS.md b/TASKS.md index a6f759b..b8c856d 100644 --- a/TASKS.md +++ b/TASKS.md @@ -2467,13 +2467,131 @@ FGMRES-IR reaches `relres ≤ 1e-12` within 20 iterations on CPU and CUDA. - 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 `[ ]` +### T19 — General LU (`"G"`): CPU reference + GPU `[!]` Symmetric-pattern multifrontal LU with `U` panels, in-block threshold partial pivoting (A/B) and `getrf!` + `laswp!` + post-check (C), `perm_row/col`, `solve_mode` 1/2, `lu`/`lu!`; ported CUDSS.jl unsymmetric tests and the uniform-batch LU tests pass; `random_general` residuals `≤ tol(T)`. +#### Report + +- Status: [!] (done on the KA CPU backend; regime C runs the fused KA kernel, no vendor `getrf`/`laswp`, see + deviations; CUDA from CI) +- What was built: + - Storage (PLAN §2.3 step 8): an LU front is two lower-triangular structures with the Cholesky layout, `L` + (lower triangle with the diagonal: `numeric.factor`, `stack`) and `Uᵀ` (strict upper triangle transposed: + new `Numeric` fields `ufactor`, `ustack`, same offsets, empty for the other structures). `F[i, j]` is + `L[i, j]` for `i ≥ j`, `Uᵀ[j, i]` otherwise. Panels, update-stack blocks, regime-A local layout, extend-add + maps, batch interleaving and the solve gather lists are reused unchanged on both structures. `d` = D of + `L D U` (unit `L` and `U`), `piv` = local row order, `pivot_kind` 1×1 / perturbed. + - Symbolic: `assembly_map` for `"G"` (upper-triangle entries of `P A Pᵀ` → `Uᵀ` structure, encoded as + negative offsets; `_group_amap` sorts L before Uᵀ entries of an offset so runs are one structure); + `schedule_elsize(structure, T)` and `build_schedule(…; elsize)`: `2 sizeof(T)` per packed entry for LU (two + triangles in local memory and on the stack); `memory_estimates` slots 7/9 doubled for `"G"`. + - `src/reference/lu.jl` (new): `ref_lu!` (full `f×f` front on the host, in-block threshold partial pivoting: + keep the diagonal if `|F[k,k]| ≥ u maxᵢ≥ₖ|F[i,k]|` over all rows of the front, else the first fully-summed + row with the largest entry, taken even below the threshold; `|d| < ε` → `sign(d) ε`, GESP style; + `F₂₂ −= (L₂₁ D) U₁₂`), `lu_pivoting(pivot_type)` (`'A'`/`'B'`/`'L'` pivot, `'N'`/`'D'` diagonal), + `extract_lu` (`A[p, q] = L D U` with `p = perm[piv]`, `q = perm`), `ref_solve_lu!` (`A` or `Aᵀ`); + `ref_solve!` dispatches to it. + - `src/numeric/lu.jl` (new): `front_lu_kernel!` (regimes B and C: one workgroup per front, fused + zero/scatter/extend-add on both structures, per pivot step: choose row (work item 1), swap rows `k`, `r` + over all columns, pivot/perturb, rank-1 update of the fully-summed columns and rows, scale; then the two + packed contribution blocks, unit diagonals), `subtree_lu_kernel!` (regime A: both structures in two + `@localmem` buffers, 32 kernel arguments), `factorize_lu!` (driver, `abs_max!` for the scaled ε, + `reduce_stats!`, interrupt polling; no allocation, no host sync, deterministic). `factorize!` dispatches + `"G"` to it. + - Solve (`src/solve/sweeps.jl`): the sweep kernels take the panel buffer; kinds `_SV_LU` (forward `L` with the + local row order, diagonal, backward `U = (Uᵀ)ᵀ` without reordering) and `_SV_LU_T` (forward `Uᵀ`, diagonal, + backward `Lᵀ` restoring the order), also on the regime-C dense path (`'T'` calls on the chosen buffer); + `forward_sweep!`/`backward_sweep!`/`sweep_solve!` gain `transpose`. `solve_transposed(structure, csc, mode)`: + `Mᵀ` for `Aᵀ`/`Aᴴ` of CSR input and `A` of CSC input; `solve_conjugated` unchanged (`Aᴴ` conjugates). + - Refinement: the `"G"` map also holds the rows of `Mᵀ` (`refinement_map(F; transpose = true)`, `2n + 1` row + pointers), selected by `transpose_matrix` in `residual!`, `refine!`, `fgmres_refine!` + (`RefinementOperator.roff`, `FactorPreconditioner.transpose`); `allocate_refinement(…; n)`. + - API: `"G"` accepted (view `'F'`); `getparam` `"perm_row"` = `perm[piv]` after a factorization (member 1), + `"perm_col"` = `perm`, `"diag"` = D, `"npivots"`, `"pivot_stats"`, `"inertia"` = `(0, 0)`; + `solve_mode` 0/1/2 incl. sub-phases and refinement; `lu(::CSR; check)`, `lu!(solver, A; check)` (refinement + on, 2 steps, as `ldlt`), `logabsdet` (parity of the local row order × Π d), uniform batches through the + existing batch machinery; CUDA extension `lu(::CuSparseMatrixCSR)`. README and `pivot_type` docs. + - Tests: `test/test_numeric_lu.jl` (new): reference (`A[p,q] = LDU`, `A`/`Aᵀ` solves), device vs reference + under three regime mixes on five cases (`piv`, kinds, stats, totals `==`; D, `L`, `Uᵀ` within + `panel_tol`), solves `nrhs ∈ {1,5}`, atomic/deterministic, `A` and `Aᵀ`; regime-A groups with Int64 maps; + perturbation (zero row/column, default and regime-C root, `pivot_epsilon`, scaled); determinism and + allocation budgets; sweeps one at a time against `extract_lu` and every dense implementation; uniform + batch (3 members × 2 rhs, `solve_mode = 1`, per-member reference equality, `ubatch_index`); public API + for `INT ∈ INTTYPES` (perm_row/col vs `extract_lu`, `solve_mode` 0/1/2 with sub-phases `==` `"solve"`, + IR and FGMRES-IR, `logabsdet`, CSC input, view refusal, `lu`/`lu!`). Ported: `cudss_execution` + "Unsymmetric -- Non-Hermitian", `cudss_generic` `lu`/`lu!`/`ldiv!`/`\`, `cudss_solver` with `"G"`, + `cudss_uniform_batch` `uniform_batch_lu` (both APIs, strided or not). Helpers: `weak_diagonal_general` + (test/matrices.jl), `lu_setup`, `upanel_error`, `lu_error` (test/utils.jl). + - Test adaptations required by the feature: `test_api.jl` no longer expects `"G"` to raise + `NotSupportedError` (asserts the view-'L' refusal instead); `test_symbolic_schedule.jl` checks the `"G"` + analysis (elsize, doubled regime-A peak, half the off-diagonal entries in `Uᵀ`) and its memory-estimate + check counts two structures for `"G"`. No assertion was loosened. +- Tests: `SDS_TEST_GPU=0 julia --project=. -e 'using Pkg; Pkg.test()'` (Julia 1.13.1, ubuntu-latest, KA CPU + backend): 65224 pass / 0 fail / 1 broken (the T16 `@test_broken`), 18.5 min; `test_numeric_lu` 1756 pass + (≈ 9.5 min standalone, mostly kernel compilation; 28 s inside the full run), `test_ported` 9464 pass (with the + LU ports). The first full `test_numeric_lu` run found one bug (the refinement workspace took `n` from the + `2n`-row `"G"` map) and the schedule memory-estimate check needed the two-structure count; both fixed before + the final run. CUDA/AMDGPU: pending CI on the PR. + After the `growth_tol` fix (owner's machine, Julia 1.13.1, CUDA.jl 6.4.2, RTX 4080): + `SDS_TEST_ONLY=test_numeric_lu` 1816 pass / 0 fail on the KA CPU backend; with `SDS_TEST_CPU=0` (CUDA only) + 1788 pass / 0 fail, including the 12 checks that failed before (9 Float32 'weak diagonal', 3 Float64 + 'weak diagonal, pivot_type N'). +- Measurements (ubuntu-latest KA CPU backend, Float64, `pivot_threshold = 1`, best of 3; "device" is + `factorize!` on the CPU backend, so only equality and launch counts are meaningful, GPU numbers are owed): + + ```text + matrix fronts A/B/C launches ref_lu! device solve A solve Aᵀ swaps same piv + random_general(2000,0.002) 181/46/23 127 413 ms 567 ms 1.04 ms 1.08 ms 0 yes + weak_diagonal_general(3000,0.002) 175/84/32 160 2.26 s 6.41 s 2.44 ms 2.30 ms 515 yes + laplacian2d(100,100) + 0.5 triu 2881/451/1 57 25 ms 60 ms 4.66 ms 4.55 ms 1630 yes + ``` + + Residuals: `relres(Aᵀ)` 2.2e-16 and 1.8e-13 on the first two; the third matrix is numerically singular + (cond₂ ≈ 2.8e18; UMFPACK's `Aᵀ` solve gives relres 16, ours 6.1; on the 30×30 version, cond₁ 2e8, the + reference factor has backward error 7e-14 and relres 6e-10). The device/reference gap on the second matrix + is the per-front serial structure of one workgroup per front on the KA CPU backend (wide fronts with a + pivot loop of `w` steps and 5 barriers each). +- Deviations from PLAN.md / this task: + - **No vendor `getrf!` + `laswp!` + post-check in regime C.** As for LDLᵀ in T15, regime-B and regime-C groups + run the fused KA `front_lu_kernel!` (one workgroup per front). `getrf` picks its own pivots over the whole + `F₁₁` (no threshold against the rows of `F₂₁`, its own tie-breaking), so the device factor would not equal the + reference, and the split `L`/`Uᵀ` storage would need a dense copy of `F₁₁`/`F₁₂` per front. Vendor BLAS-3 for + `L₂₁`, `U₁₂` and the contribution block is the same follow-up as #75 (T25). + - `U` is stored transposed in a second buffer with the `L` layout (`ufactor`, `ustack`) instead of separate + `w×f` U panels: every map, the layout, the regime-A local layout and the sweeps are reused, and the solves + with `A` and `Aᵀ` are the existing sweeps with another buffer and pivot-order flags. Memory: the factor and + the update stack are twice the Cholesky size (as for any LU); regime-A budgets count `2 sizeof(T)` per entry. + - `pivot_type`: `'C'`/`'R'` stay `NotSupportedError` (PLAN §3.3), so the ported tests assert that, as the + symmetric ports do; `'N'` and `'D'` mean diagonal pivots. `pivot_sign` is ignored by LU (no inertia); + `"inertia"` is `(0, 0)`. + - Static pivoting limits accuracy when no fully-summed row can replace a small pivot: on + `weak_diagonal_general` with scale `1e-3` (cond ≈ 7e4) the reference reaches only relres 1e-8 (Float64) and + 0.7 (Float32) with up to 130 row interchanges, `max|L|` ≈ 3e5; with scale 0.1 (the test matrix) it is at + rounding level. This is PLAN R9 (matching, T21; APTP/delayed pivots, T27). + - `"G"` keeps T05's view rule (`'F'` only); cuDSS accepts `'L'`/`'U'` for `"G"` (#84). +- Open issues / follow-ups: + - #84 (found-by-agent): views `'L'`/`'U'` with `"G"`. + - Regime-C performance (vendor BLAS-3, parallel pivot search) with #75 in T25. The pivot search is serial on + work item 1 (`O(f)` per column, cheaper than the LDLᵀ fallback scan). + - `perm_row` reports batch member 1 for a uniform batch (each member has its own local row order). + - CUDA path exercised only by CI. + - Device arithmetic fuses multiply-adds (GPUCompiler `-nvptx-fma-level=1`), so device-vs-reference equality is + bitwise on the CPU backend only. On `weak_diagonal_general` with `pivot_threshold = 0.01` (growth ≈ 3e2) CUDA + deviated by up to 900 eps from `ref_lu!` (at least as accurate against a Float64 oracle). The elementwise + panel checks now use `growth_tol(T, Nr)` (`panel_tol` scaled by `max(|L|, |Uᵀ|) / min |d|`), the pivot + sequence stays exact, and the device testset also checks `lu_error ≤ tol(T)`. `test_numeric_ldlt.jl` has the + same latent fragility (see #86) and should move to `growth_tol` when touched. + - The GitHub App token expired before the last push: the final commit is local and published by the + workflow's "Publish the branch" step (which also opens the PR). +- Suggested plan changes: + - PLAN §2.3 step 8 / §3.2: LU stores `Uᵀ` in a second buffer with the `L` layout (`ufactor`, `ustack`). + - PLAN §2.4 / §3.3: regime C of `"G"` uses the KA in-front kernel (same pivot sequence as the oracle); vendor + calls only for the BLAS-3 parts (T25). + ### T20 — Schur complement mode `[ ]` Constrained ordering, unfactorized root front, `schur_shape`, dense and CSR diff --git a/ext/SparseDirectSolverCUDAExt.jl b/ext/SparseDirectSolverCUDAExt.jl index 89ec4c9..d7454aa 100644 --- a/ext/SparseDirectSolverCUDAExt.jl +++ b/ext/SparseDirectSolverCUDAExt.jl @@ -123,6 +123,14 @@ LinearAlgebra.ldlt(A::Symmetric{T, <:CuSparseMatrixCSR{T}}; check::Bool = false) LinearAlgebra.ldlt(A::Hermitian{T, <:CuSparseMatrixCSR{T}}; check::Bool = false) where {T} = ldlt(CSR(A.data); view = A.uplo, check) +""" + lu(A::CuSparseMatrixCSR; check = false) -> DirectSolver + +`L D U` factorization (structure `"G"`) of `A` on the GPU (≅ CUDSS.jl's `lu`). +See `lu(::CSR)`. +""" +LinearAlgebra.lu(A::CuSparseMatrixCSR; check::Bool = false) = lu(CSR(A); check) + # --------------------------------------------------------------------------- # vendor dense bindings (see `src/dense/vendor.jl` for the contracts) diff --git a/src/SparseDirectSolver.jl b/src/SparseDirectSolver.jl index bacea50..ee1ed55 100644 --- a/src/SparseDirectSolver.jl +++ b/src/SparseDirectSolver.jl @@ -56,6 +56,7 @@ include("numeric/batch.jl") include("numeric/storage.jl") include("reference/cholesky.jl") include("reference/ldlt.jl") +include("reference/lu.jl") # numeric phase on the device (PLAN §2.4) include("numeric/assembly.jl") @@ -64,6 +65,7 @@ include("numeric/subtree.jl") include("numeric/factorize.jl") include("numeric/ldlt.jl") include("numeric/ldlt_c.jl") +include("numeric/lu.jl") include("numeric/extract.jl") # solve phase on the device (PLAN §2.5) diff --git a/src/generic.jl b/src/generic.jl index f0bab5d..a485863 100644 --- a/src/generic.jl +++ b/src/generic.jl @@ -1,5 +1,5 @@ # Julia-native layer (PLAN §1.5, §3.1) on top of the handle layer: `cholesky`, -# `cholesky!`, `ldlt`, `ldlt!`, `ldiv!`, `\`, `logabsdet`, `diag`, `nnz`. Methods take the +# `cholesky!`, `ldlt`, `ldlt!`, `lu`, `lu!`, `ldiv!`, `\`, `logabsdet`, `diag`, `nnz`. Methods take the # in-package `CSR`; the CUDA extension adds `CuSparseMatrixCSR` and its # `Symmetric`/`Hermitian` wrappers, as CUDSS.jl does. A `SparseMatrixCSC` is not # accepted by `cholesky` (that method belongs to CHOLMOD): wrap it with `CSR(A)`. @@ -11,7 +11,7 @@ # or strided vectors). # # 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 +# by `cholesky`/`ldlt`/`lu` 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. @@ -21,7 +21,7 @@ function _linear_algebra_defaults!(solver::DirectSolver) return solver end -"Refinement steps of the solvers created by `cholesky` and `ldlt` (PLAN §3.1; the handle layer uses 0)." +"Refinement steps of the solvers created by `cholesky`, `ldlt` and `lu` (PLAN §3.1; the handle layer uses 0)." const LINEAR_ALGEBRA_IR_STEPS = 2 """ @@ -73,6 +73,41 @@ function LinearAlgebra.ldlt(A::CSR{T}; view::Char = 'F', check::Bool = false) wh return solver end +""" + lu(A::CSR; check = false) -> DirectSolver + +`L D U` factorization (structure `"G"`) of the general sparse matrix `A` on its +backend (≅ CUDSS.jl's `lu`): a [`DirectSolver`](@ref) after `"analysis"` and +`"factorization"` (synchronized, `ir_n_steps = 2` as [`cholesky`](@ref)), on the +symmetric pattern of `A + Aᵀ` with row pivoting inside the fully-summed block of +each front and static perturbation of tiny pivots (read +`getparam(solver, "npivots")`). The factorization does not fail (`"info"` stays +0), so `check` has no effect. `ldiv!`/`\` solve with `A`; `setparam!(solver, +"solve_mode", 1)` (or `2`) solves with `Aᵀ` (`Aᴴ`). +""" +function LinearAlgebra.lu(A::CSR{T}; check::Bool = false) where {T} + solver = _linear_algebra_defaults!(DirectSolver(A, "G", 'F')) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing; asynchronous = false) + _check_info(solver, check) + return solver +end + +""" + lu!(solver::DirectSolver, A; check = false) -> solver + +Factorize the new values `A` reusing the analysis of `solver`, as +[`cholesky!`](@ref) (`"factorization"` the first time after an analysis, +`"refactorization"` afterwards, synchronized). +""" +function LinearAlgebra.lu!(solver::DirectSolver, A; check::Bool = false) + update!(solver, A) + phase = solver.fresh_factorization ? "factorization" : "refactorization" + execute!(phase, solver, nothing, nothing; asynchronous = false) + _check_info(solver, check) + return solver +end + """ ldlt!(solver::DirectSolver, A; check = false) -> solver @@ -166,6 +201,8 @@ determinants of the 2×2 blocks; `sign` is `±1` (real or Hermitian) or `det D / the determinant of `A + E`. Raises [`FactorizationError`](@ref) when the factorization failed (`"info" ≠ 0`) and [`NotSupportedError`](@ref) for a uniform batch (read the members' diagonals with `getparam(solver, "diag")`). +LU (`"G"`): `det A = ± det D`, the sign of the local row interchanges times the +product of the pivots (`sign` is `±1` for real `T`, a unit complex number otherwise). """ function LinearAlgebra.logabsdet(solver::DirectSolver{T}) where {T} solver.nbatch == 1 || @@ -173,6 +210,7 @@ function LinearAlgebra.logabsdet(solver::DirectSolver{T}) where {T} solver.stage >= STAGE_FACTORIZED && solver.info[1] == 0 || throw(FactorizationError(solver.info[1], "logabsdet needs a successful factorization")) _is_ldlt_structure(solver.structure) && return _logabsdet_ldlt(solver) + solver.structure == STRUCTURE_GENERAL && return _logabsdet_lu(solver) d = Array(getparam(solver, "diag")) return 2 * sum(x -> log(abs(x)), d; init = zero(real(T))), one(T) end @@ -200,6 +238,31 @@ function _logabsdet_ldlt(solver::DirectSolver{T}) where {T} return la, sgn end +function _logabsdet_lu(solver::DirectSolver{T}) where {T} + n = size(solver, 1) + d = Array(view(solver.numeric.d, 1:n)) + piv = Array(view(solver.numeric.piv, 1:n)) + la, sgn = zero(real(T)), one(T) + for x in d + la += log(abs(x)) + sgn *= T <: Real ? T(sign(x)) : x / abs(x) + end + # parity of the local row interchanges (P A Pᵀ: the symmetric permutation cancels) + seen = falses(n) + for k in 1:n + seen[k] && continue + len = 0 + j = k + while !seen[j] + seen[j] = true + j = Int(piv[j]) + len += 1 + end + iseven(len) && (sgn = -sgn) + end + return la, sgn +end + """ logdet(solver::DirectSolver) diff --git a/src/numeric/extract.jl b/src/numeric/extract.jl index 1fb3cbe..890c4ae 100644 --- a/src/numeric/extract.jl +++ b/src/numeric/extract.jl @@ -5,7 +5,7 @@ host_numeric(numeric) -> Numeric{T, Vector{T}} Host copy of `numeric` (panels, D, update stack, workspace, statistics and totals, status, pivot order, kinds and sign -requests, `aux`; the +requests, `aux`, the `Uᵀ` panels and stack of LU; the plan is shared), usable by [`ref_solve!`](@ref) and [`extract_L`](@ref). A host `numeric` is returned as is. """ @@ -16,7 +16,8 @@ function host_numeric(N::Numeric{T}) where {T} return Numeric{T, Vector{T}, Vector{Int64}, Vector{Int32}, Vector{Int8}}(factor, d, stack, work, stats, info, piv, pivot_kind, Array(N.totals), Array(N.psign), Array(N.aux), N.plan, - Array(N.members), N.nbatch) + Array(N.members), N.nbatch, + Array(N.ufactor), Array(N.ustack)) end """ @@ -48,10 +49,12 @@ function member_numeric(N::Numeric{T}, S::Symbolic, k::Integer) where {T} L = S.layout ns = nsupernodes(S) factor = zeros(T, L.factor_len) + ufactor = zeros(T, isempty(H.ufactor) ? 0 : L.factor_len) for s in 1:ns len = L.panel_ptr[s + 1] - L.panel_ptr[s] p0 = panel_offset(L.panel_ptr, s, k, nb) copyto!(factor, L.panel_ptr[s], H.factor, p0, len) + isempty(ufactor) || copyto!(ufactor, L.panel_ptr[s], H.ufactor, p0, len) end mv(A) = A[((k - 1) * (length(A) ÷ nb) + 1):(k * (length(A) ÷ nb))] info = H.info[k:nb:end] @@ -59,5 +62,6 @@ function member_numeric(N::Numeric{T}, S::Symbolic, k::Integer) where {T} mv(H.stats), info, mv(H.piv), mv(H.pivot_kind), mv(H.totals), copy(H.psign), mv(H.aux), - NumericPlan(S), Int32[1], 1) + NumericPlan(S), Int32[1], 1, ufactor, + isempty(H.ustack) ? T[] : mv(H.ustack)) end diff --git a/src/numeric/factorize.jl b/src/numeric/factorize.jl index 5358a07..476f66f 100644 --- a/src/numeric/factorize.jl +++ b/src/numeric/factorize.jl @@ -66,8 +66,9 @@ function cholesky_stats!(N::Numeric, S::Symbolic) end function _check_numeric(N::Numeric{T}, S::Symbolic, nzval::AbstractVector) where {T} - _is_ldlt_structure(S.structure) ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) - sizeof(T) <= S.elsize || + lu = S.structure == STRUCTURE_GENERAL + _is_ldlt_structure(S.structure) || lu ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) + (lu ? 2 : 1) * sizeof(T) <= S.elsize || throw(InvalidValueError("the analysis was built for $(S.elsize)-byte elements, the numeric storage has " * "$(sizeof(T))-byte $T")) nb = N.nbatch @@ -79,6 +80,8 @@ function _check_numeric(N::Numeric{T}, S::Symbolic, nzval::AbstractVector) where throw(InvalidValueError("the numeric storage was not allocated for this analysis")) length(N.work) == nb * S.layout.work_len && length(N.piv) == nb * S.n && length(N.psign) == S.n || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) + lu && (length(N.ufactor) != length(N.factor) || length(N.ustack) != length(N.stack)) && + throw(InvalidValueError("the numeric storage was not allocated for structure \"G\"")) backend = KernelAbstractions.get_backend(N.factor) for x in (nzval, S.amap, N.info) typeof(KernelAbstractions.get_backend(x)) == typeof(backend) || @@ -194,7 +197,7 @@ _check_ldlt_eltype(::Type{T}) where {T} = T <: LinearAlgebra.BlasFloat || factorize!(numeric, symbolic, nzval; impl = :auto, opts = Options()) -> info::Int Structures `"S"`/`"H"`: [`factorize_ldlt!`](@ref) with the pivoting options -`opts` (`impl` is not used). Structures `"SPD"`/`"HPD"`: multifrontal Cholesky `P A Pᵀ = L Lᴴ` on the device: `nzval` are the stored +`opts` (`impl` is not used); structure `"G"`: [`factorize_lu!`](@ref) likewise. Structures `"SPD"`/`"HPD"`: multifrontal Cholesky `P A Pᵀ = L Lᴴ` on the device: `nzval` are the stored values of A (same pattern, view and index base as the analysis) on the backend of `numeric`, and `symbolic` has its maps on that backend (`adapt(backend, symbolic, INT)`). For every launch group of the schedule, @@ -230,6 +233,7 @@ allocated on the device; the panels are bitwise reproducible for a fixed """ function factorize!(N::Numeric, S::Symbolic, nzval::AbstractVector; impl::Symbol = :auto, opts::Options = Options()) _is_ldlt_structure(S.structure) && return factorize_ldlt!(N, S, nzval; impl, opts) + S.structure == STRUCTURE_GENERAL && return factorize_lu!(N, S, nzval; opts) return factorize_cholesky!(N, S, nzval; impl, opts) end @@ -286,11 +290,11 @@ _original_column(S::Symbolic, k::Integer) = k == 0 ? 0 : S.partition.perm[k] Write the `info` of the last [`factorize!`](@ref) of every active batch member `k` into `info[k]` (the original column of the member's first non-positive pivot, 0 = success; inactive members keep their entries). Host data only: the -statuses were read by `factorize!` (Cholesky) or are 0 (LDLᵀ/LDLᴴ). +statuses were read by `factorize!` (Cholesky) or are 0 (LDLᵀ/LDLᴴ, LU). """ function member_info!(info::AbstractVector{<:Integer}, N::Numeric, S::Symbolic) plan = N.plan - chol = !_is_ldlt_structure(S.structure) + chol = !_is_ldlt_structure(S.structure) && S.structure != STRUCTURE_GENERAL for j in 1:plan.nact[] k = Int(plan.members_host[j]) info[k] = chol ? _original_column(S, plan.info_host[k]) : 0 diff --git a/src/numeric/lu.jl b/src/numeric/lu.jl new file mode 100644 index 0000000..570ad6a --- /dev/null +++ b/src/numeric/lu.jl @@ -0,0 +1,575 @@ +# Device LU (PLAN §2.2, §2.4, §3.3, M7): the in-front `L D U` factorization of +# the CPU reference (`src/reference/lu.jl`) as KA kernels, with the same pivot +# sequence: row interchanges inside the fully-summed rows of each front +# (threshold partial pivoting against the whole remaining front column), +# static perturbation `±ε` of tiny pivots, and the contribution block updated +# once per front with `F₂₂ ← F₂₂ − (L₂₁ D) U₁₂`. +# +# A front is kept as two lower-triangular structures with the layout of the +# Cholesky factor: `L` (the lower triangle of the front, diagonal included: the +# panel in `factor`, the contribution block in `stack`) and `Uᵀ` (the strict +# upper triangle, transposed: `ufactor`, `ustack`; its diagonal stays zero +# until the panel gets its unit diagonal). Entry `F[i, j]` is `L[i, j]` for +# `i ≥ j` and `Uᵀ[j, i]` otherwise (`_lu_get`), so the assembly, the extend-add, +# the update-stack layout, the regime-A local-memory layout and the solve +# sweeps of the symmetric structures are reused on both structures. +# +# * Regime A: `subtree_lu_kernel!`, one workgroup per subtree, the serial stack +# of the two packed structures in two `@localmem` buffers (the schedule +# budgets `2 sizeof(T)` bytes per packed entry, `schedule_elsize`). +# * Regimes B and C: `front_lu_kernel!`, one workgroup per front of a launch +# group, fused assembly, the pivoted factorization of the two panels in +# global memory and the update of the front's two packed contribution blocks. +# Regime C uses this KA kernel too (no vendor `getrf`: see the T19 report). +# +# Per pivot step `k` (a loop over the `w` columns, uniform trip count): work item +# 1 chooses the pivot row, the workgroup swaps rows `k` and `r` (every column of +# the front), work item 1 stores D (perturbing a tiny pivot) and counts, the +# workgroup applies the rank-1 update to the rest of the fully-summed columns +# and rows, then scales column `k` into L and row `k` into U. No atomics: +# every destination has one writer per segment, so the result is deterministic. + +"Workgroup size of the LU front kernel (regimes B and C)." +const LU_WORKGROUP = 128 + +# resolved pivoting parameters passed to the kernels (isbits) +struct _LUDevice{R} + pivot::Bool # in-block threshold partial pivoting + u::R # pivot_threshold + eps::R # pivot_epsilon (multiplied by max |aᵢⱼ| = real(aux[1]) when `scaled`) + scaled::Bool # pivot_epsilon_alg = "algo1" +end + +function _lu_device_params(S::Symbolic, ::Type{T}, opts::Options) where {T} + S.structure == STRUCTURE_GENERAL || + throw(InvalidValueError("LU needs structure \"G\", got \"$(convert(String, S.structure))\"")) + R = real(T) + return _LUDevice{R}(lu_pivoting(opts.pivot_type), R(opts.pivot_threshold), R(resolved_pivot_epsilon(opts, R)), + opts.pivot_epsilon_alg == PIVOT_EPSILON_SCALED) +end + +@inline function _lu_eps(p::_LUDevice, aux) + p.scaled || return p.eps + @inbounds a = real(aux[1]) + return a > 0 ? p.eps * a : p.eps +end + +# F[i, j] of the front from its two structures (`Fl`: L, `Fu`: Uᵀ; `_PackedFront` or `_PanelFront`) +@inline _lu_get(Fl, Fu, i, j) = i >= j ? _fget(Fl, i, j) : _fget(Fu, j, i) +@inline _lu_set!(Fl, Fu, i, j, v) = i >= j ? _fset!(Fl, i, j, v) : _fset!(Fu, j, i, v) + +# --------------------------------------------------------------------------- +# phases of one pivot step k (work item `li` of `WG`; `pk`: packed local fronts) + +# work item 1: the pivot row of column k (`_lu_choose_row` of the reference) +@inline function _lu_choose!(fa, ua, ctl, k, p::_LUDevice, pk::Val) + @inbounds begin + IT = eltype(ctl) + r = k + if p.pivot + Fl, Fu = _lt_front(fa, ctl, pk), _lt_front(ua, ctl, pk) + f = Int(ctl[_ST_F]) + γ = zero(real(eltype(fa))) + for i in k:f + γ = max(γ, abs(_fget(Fl, i, k))) + end + best = abs(_fget(Fl, k, k)) + if !(best >= p.u * γ) + for i in (k + 1):Int(ctl[_ST_W]) + a = abs(_fget(Fl, i, k)) + if a > best + r = i + best = a + end + end + end + end + ctl[_LT_C] = r % IT + end + return nothing +end + +# interchange of the rows k and r ≤ w of the front (every column), and of their pivot order entries +@inline function _lu_swap!(fa, ua, ctl, piv, k, li, ::Val{WG}, pk::Val) where {WG} + @inbounds begin + r = Int(ctl[_LT_C]) + if r != k + Fl, Fu = _lt_front(fa, ctl, pk), _lt_front(ua, ctl, pk) + for j in li:WG:Int(ctl[_ST_F]) + x = _lu_get(Fl, Fu, k, j) + _lu_set!(Fl, Fu, k, j, _lu_get(Fl, Fu, r, j)) + _lu_set!(Fl, Fu, r, j, x) + end + if li == 1 + c0 = Int(ctl[_LT_C0]) + x = piv[c0 + k - 1] + piv[c0 + k - 1] = piv[c0 + r - 1] + piv[c0 + r - 1] = x + end + end + end + return nothing +end + +# work item 1: store the pivot in D (perturbing a tiny one), its kind and the statistics; `pv[1] = d` +@inline function _lu_pivot!(fa, ctl, pv, d, pivot_kind, aux, k, p::_LUDevice, pk::Val) + @inbounds begin + IT = eltype(ctl) + T = eltype(fa) + Fl = _lt_front(fa, ctl, pk) + g = Int(ctl[_LT_C0]) + k - 1 + ε = _lu_eps(p, aux) + dk = _fget(Fl, k, k) + kind = PIVOT_KIND_1X1 + if !(abs(dk) >= ε) # tiny (or NaN): perturb + iszero(dk) && (ctl[_LT_STAT + STAT_NZERO] += one(IT)) + dk = _perturbation_sign(dk, 0, !(T <: Complex)) * ε + kind = PIVOT_KIND_PERTURBED + ctl[_LT_STAT + STAT_NPERTURBED] += one(IT) + end + _fset!(Fl, k, k, dk) + d[g] = dk + pivot_kind[g] = kind + pv[1] = dk + end + return nothing +end + +# rank-1 update with the unscaled pivot column and row: F[i, j] -= F[i, k] / d F[k, j] for i, j > k with +# j ≤ w (columns k+1:w, rows to f) or i ≤ w (rows k+1:w, columns w+1:f) +@inline function _lu_update!(fa, ua, ctl, pv, k, li, ::Val{WG}, pk::Val) where {WG} + @inbounds begin + Fl, Fu = _lt_front(fa, ctl, pk), _lt_front(ua, ctl, pk) + f = Int(ctl[_ST_F]) + w = Int(ctl[_ST_W]) + nr = f - k + nc = w - k + n1 = nr * nc + d = pv[1] + for q in (li - 1):WG:(n1 + nc * (f - w) - 1) + if q < n1 + i = k + 1 + q % nr + j = k + 1 + q ÷ nr + else + q2 = q - n1 + i = k + 1 + q2 % nc + j = w + 1 + q2 ÷ nc + end + _lu_set!(Fl, Fu, i, j, _lu_get(Fl, Fu, i, j) - _lu_get(Fl, Fu, i, k) / d * _lu_get(Fl, Fu, k, j)) + end + end + return nothing +end + +# scale column k below the pivot into L and row k right of the pivot into U +@inline function _lu_scale!(fa, ua, ctl, pv, k, li, ::Val{WG}, pk::Val) where {WG} + @inbounds begin + Fl, Fu = _lt_front(fa, ctl, pk), _lt_front(ua, ctl, pk) + nr = Int(ctl[_ST_F]) - k + d = pv[1] + for q in (li - 1):WG:(2 * nr - 1) + i = k + 1 + q % nr + if q < nr + _fset!(Fl, i, k, _fget(Fl, i, k) / d) # L[i, k] + else + _fset!(Fu, i, k, _fget(Fu, i, k) / d) # U[k, i] + end + end + end + return nothing +end + +# F₂₂ ← F₂₂ − (L₂₁ D) U₁₂ on the packed m×m blocks at `coff` of `ca` (lower triangle, diagonal included) and +# `cu` (strict upper triangle, transposed); `coff < 0`: no block +@inline function _lu_cb_update!(fa, ua, ctl, ca, cu, coff, d, li, ::Val{WG}, pk::Val) where {WG} + @inbounds begin + T = eltype(fa) + f = Int(ctl[_ST_F]) + w = Int(ctl[_ST_W]) + m = f - w + if coff >= 0 && m > 0 + Fl, Fu = _lt_front(fa, ctl, pk), _lt_front(ua, ctl, pk) + c0 = Int(ctl[_LT_C0]) + for q in (li - 1):WG:(m * m - 1) + jj = q ÷ m + 1 + ii = q - (jj - 1) * m + 1 + acc = zero(T) + for k in 1:w + acc += _fget(Fl, w + ii, k) * d[c0 + k - 1] * _fget(Fu, w + jj, k) + end + if ii >= jj + pos = coff + _packed(ii, jj, m) + ca[pos] -= acc + else + pos = coff + _packed(jj, ii, m) + cu[pos] -= acc + end + end + end + end + return nothing +end + +# --------------------------------------------------------------------------- +# regimes B/C: one workgroup per front, panels in global memory + +# scatter the values of A into the zeroed panels of front `s`: positive `amap` offsets into `factor` (L), +# negative ones into `ufactor` (Uᵀ); runs of one destination are summed by one work item in `nzval` order +@inline function _scatter_front_lu!(factor, ufactor, nzval, amap, amap_ptr, amap_src, s, shift, li, + ::Val{WG}) where {WG} + @inbounds begin + T = eltype(factor) + a = amap_ptr[s] + b = amap_ptr[s + 1] - 1 + for k in (a + li - 1):WG:b + dest = amap[amap_src[k]] + if k == a || amap[amap_src[k - 1]] != dest + v = zero(T) + kk = k + while kk <= b + p = amap_src[kk] + amap[p] == dest || break + v += T(nzval[p]) + kk += 1 + end + if dest > 0 + factor[dest + shift] = v + else + ufactor[-dest + shift] = v + end + end + end + end + return nothing +end + +# unit diagonals of both panels, statistics and status +@inline function _lu_front_finish!(factor, ufactor, stats, info, ctl, li, ::Val{WG}) where {WG} + @inbounds begin + Fl, Fu = _lt_front(factor, ctl, Val(false)), _lt_front(ufactor, ctl, Val(false)) + for j in li:WG:Int(ctl[_ST_W]) + _fset!(Fl, j, j, one(eltype(factor))) + _fset!(Fu, j, j, one(eltype(factor))) + end + li == 1 && _lt_stats!(stats, info, ctl, ctl[_ST_NODE]) + end + return nothing +end + +""" + front_lu_kernel!(backend, WG)(factor, ufactor, stack, ustack, info, stats, d, piv, pivot_kind, aux, nzval, amap, + amap_ptr, amap_src, nodes, bm, super_ptr, front_ptr, front_nrows, front_ncols, + cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, prm, Val(WG); + ndrange = WG * count * bm.nact) + +LU kernel of regimes B and C: workgroup `G` takes front +`s = nodes[bm.first + g - 1]` of batch member `k` (`g`, `k` from the +[`BatchMap`](@ref) `bm`; the per-member arrays are the member's), zeroes and +assembles its `L` and `Uᵀ` structures (A through the `amap`, then the +`maxchild` children's two packed contribution blocks in `child_list` order), +factors its `w` fully-summed columns in place in the two panels with in-block +row pivoting, threshold test and perturbation (`prm`, an `_LUDevice`), writes +D, the local row order `piv`, the pivot kinds and the front's statistics, +updates its two packed contribution blocks (`cb_ptr[s] > 0`) with +`F₂₂ − (L₂₁ D) U₁₂` and leaves unit-lower panels of `L` and `Uᵀ`. +""" +@kernel function front_lu_kernel!(factor, ufactor, stack, ustack, info, stats, d, piv, pivot_kind, aux, nzval, amap, + amap_ptr, amap_src, nodes, bm, super_ptr, front_ptr, front_nrows, front_ncols, + cb_ptr, child_ptr, child_list, relind_ptr, relind, maxchild, prm::_LUDevice, + ::Val{WG}) where {WG} + @uniform TT = eltype(factor) + @uniform IT = eltype(front_ptr) + li = @index(Local, Linear) + G = @index(Group, Linear) + ctl = @localmem IT (_LT_CTL,) + pv = @localmem TT (1,) + if li == 1 + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + _lt_front_setup!(ctl, s, super_ptr, member_panels(front_ptr, _bm_gmember(bm, G), bm.nbatch), front_nrows, + front_ncols, cb_ptr) + end + @synchronize + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + nb = bm.nbatch + P = member_panels(front_ptr, k, nb) + _zero_front!(factor, _mview(stack, k, nb), s, li, P, front_nrows, front_ncols, cb_ptr, Val(WG)) + _zero_front!(ufactor, _mview(ustack, k, nb), s, li, P, front_nrows, front_ncols, cb_ptr, Val(WG)) + _lt_init_piv!(_mview(piv, k, nb), ctl, li, Val(WG)) + @synchronize + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + _scatter_front_lu!(factor, ufactor, _mview(nzval, k, bm.nbatch), amap, amap_ptr, amap_src, s, + _member_shift(front_ptr, s, k, bm.nbatch), li, Val(WG)) + @synchronize + for kc in 1:maxchild + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + k = _bm_gmember(bm, G) + nb = bm.nbatch + P = member_panels(front_ptr, k, nb) + _extend_add_child!(factor, _mview(stack, k, nb), s, kc, li, P, front_nrows, front_ncols, cb_ptr, child_ptr, + child_list, relind_ptr, relind, Val(WG)) + _extend_add_child!(ufactor, _mview(ustack, k, nb), s, kc, li, P, front_nrows, front_ncols, cb_ptr, child_ptr, + child_list, relind_ptr, relind, Val(WG)) + @synchronize + end + for it in 1:ctl[_ST_W] + if li == 1 + _lu_choose!(factor, ufactor, ctl, it, prm, Val(false)) + end + @synchronize + _lu_swap!(factor, ufactor, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), it, li, Val(WG), Val(false)) + @synchronize + if li == 1 + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lu_pivot!(factor, ctl, pv, _mview(d, k, nb), _mview(pivot_kind, k, nb), _mview(aux, k, nb), it, prm, + Val(false)) + end + @synchronize + _lu_update!(factor, ufactor, ctl, pv, it, li, Val(WG), Val(false)) + @synchronize + _lu_scale!(factor, ufactor, ctl, pv, it, li, Val(WG), Val(false)) + @synchronize + end + @inbounds s = nodes[bm.first + _bm_node(bm, G) - 1] + @inbounds coff = Int(cb_ptr[s]) - 1 + k = _bm_gmember(bm, G) + nb = bm.nbatch + _lu_cb_update!(factor, ufactor, ctl, _mview(stack, k, nb), _mview(ustack, k, nb), coff, _mview(d, k, nb), li, + Val(WG), Val(false)) + _lu_front_finish!(factor, ufactor, _mview(stats, k, nb), _iview(info, k, nb), ctl, li, Val(WG)) +end + +# --------------------------------------------------------------------------- +# regime A: one workgroup per subtree, the two structures in local memory + +# scatter A into the zeroed packed structures of the current node (`buf`: L, `ubuf`: Uᵀ) +@inline function _subtree_scatter_lu!(buf, ubuf, ctl, nzval, amap, amap_ptr, amap_src, front_ptr, li, + ::Val{WG}) where {WG} + @inbounds begin + T = eltype(buf) + v = ctl[_ST_NODE] + f = ctl[_ST_F] + lf = ctl[_ST_LF] + p0 = front_ptr[v] + a = amap_ptr[v] + b = amap_ptr[v + 1] - 1 + for k in (a + li - 1):WG:b + dest = amap[amap_src[k]] + if k == a || amap[amap_src[k - 1]] != dest + x = zero(T) + kk = k + while kk <= b + p = amap_src[kk] + amap[p] == dest || break + x += T(nzval[p]) + kk += 1 + end + loc = abs(dest) - p0 # 0-based position in the f×w panel + j = loc ÷ f + 1 + i = loc - (j - 1) * f + 1 + if dest > 0 + buf[lf + _packed(i, j, f) - 1] = x + else + ubuf[lf + _packed(i, j, f) - 1] = x + end + end + end + end + return nothing +end + +# write the unit-lower panels of L and Uᵀ, statistics and status; the subtree root writes its two +# contribution blocks to the update stacks +@inline function _lu_subtree_write!(factor, ufactor, stack, ustack, info, stats, buf, ubuf, ctl, front_ptr, cb_ptr, li, + ::Val{WG}) where {WG} + @inbounds begin + v = ctl[_ST_NODE] + f = Int(ctl[_ST_F]) + w = Int(ctl[_ST_W]) + lf = Int(ctl[_ST_LF]) - 1 + p0 = Int(front_ptr[v]) + T = eltype(factor) + for q in (li - 1):WG:(f * w - 1) + j = q ÷ f + 1 + i = q - (j - 1) * f + 1 + if i > j + factor[p0 + q] = buf[lf + _packed(i, j, f)] + ufactor[p0 + q] = ubuf[lf + _packed(i, j, f)] + else + x = i == j ? one(T) : zero(T) + factor[p0 + q] = x + ufactor[p0 + q] = x + end + end + li == 1 && _lt_stats!(stats, info, ctl, v) + c0 = cb_ptr[v] + if c0 > 0 + m = f - w + src = _lt_local_cb_offset(ctl) + 1 + for q in (li - 1):WG:(m * (m + 1) ÷ 2 - 1) + stack[c0 + q] = buf[src + q] + ustack[c0 + q] = ubuf[src + q] + end + end + end + return nothing +end + +""" + subtree_lu_kernel!(backend, WG)(factor, ufactor, stack, ustack, info, stats, d, piv, pivot_kind, aux, nzval, amap, + amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, super_ptr, front_ptr, + front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, child_list, + relind_ptr, relind, prm, Val(NE), Val(WG); ndrange = WG * count * bm.nact) + +LU kernel of regime A: workgroup `G` takes subtree `t = trees[bm.first + g - 1]` +of batch member `k` and processes its supernodes in order with the serial +stacks of the packed `L` and `Uᵀ` structures in two `@localmem` buffers of +`NE` entries each (the layout of [`subtree_local_layout`](@ref), shared by +both): zero, scatter A, extend-add the children, factor the front's first `w` +columns with the pivoted `L D U` of [`front_lu_kernel!`](@ref), update the +trailing contribution blocks, write the unit-lower panels, D, `piv`, the +pivot kinds and the statistics, then move both blocks down or (subtree root) +write them to the update stacks. +""" +@kernel function subtree_lu_kernel!(factor, ufactor, stack, ustack, info, stats, d, piv, pivot_kind, aux, nzval, amap, + amap_ptr, amap_src, trees, bm, subtree_ptr, subtree_nodes, super_ptr, front_ptr, + front_nrows, front_ncols, cb_ptr, local_front, local_cb, child_ptr, child_list, + relind_ptr, relind, prm::_LUDevice, ::Val{NE}, ::Val{WG}) where {NE, WG} + @uniform TT = eltype(factor) + @uniform IT = eltype(subtree_nodes) + li = @index(Local, Linear) + G = @index(Group, Linear) + buf = @localmem TT (NE,) + ubuf = @localmem TT (NE,) + ctl = @localmem IT (_LT_CTL,) + pv = @localmem TT (1,) + if li == 1 + @inbounds t = trees[bm.first + _bm_node(bm, G) - 1] + @inbounds ctl[_ST_FIRST] = subtree_ptr[t] + @inbounds ctl[_ST_COUNT] = subtree_ptr[t + 1] - subtree_ptr[t] + end + @synchronize + for k in 1:ctl[_ST_COUNT] + if li == 1 + _subtree_setup!(ctl, k, subtree_nodes, front_nrows, front_ncols, local_front, local_cb, child_ptr) + @inbounds _lt_reset!(ctl, super_ptr[ctl[_ST_NODE]]) + end + @synchronize + _subtree_zero!(buf, ctl, li, Val(WG)) + _subtree_zero!(ubuf, ctl, li, Val(WG)) + _lt_init_piv!(_mview(piv, _bm_gmember(bm, G), bm.nbatch), ctl, li, Val(WG)) + @synchronize + _subtree_scatter_lu!(buf, ubuf, ctl, _mview(nzval, _bm_gmember(bm, G), bm.nbatch), amap, amap_ptr, amap_src, + front_ptr, li, Val(WG)) + @synchronize + for kc in 1:ctl[_ST_NCHILD] + _subtree_extend_add!(buf, ctl, kc, child_ptr, child_list, front_nrows, front_ncols, local_cb, relind_ptr, + relind, li, Val(WG)) + _subtree_extend_add!(ubuf, ctl, kc, child_ptr, child_list, front_nrows, front_ncols, local_cb, relind_ptr, + relind, li, Val(WG)) + @synchronize + end + for it in 1:ctl[_ST_W] + if li == 1 + _lu_choose!(buf, ubuf, ctl, it, prm, Val(true)) + end + @synchronize + _lu_swap!(buf, ubuf, ctl, _mview(piv, _bm_gmember(bm, G), bm.nbatch), it, li, Val(WG), Val(true)) + @synchronize + if li == 1 + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lu_pivot!(buf, ctl, pv, _mview(d, mb, nb), _mview(pivot_kind, mb, nb), _mview(aux, mb, nb), it, prm, + Val(true)) + end + @synchronize + _lu_update!(buf, ubuf, ctl, pv, it, li, Val(WG), Val(true)) + @synchronize + _lu_scale!(buf, ubuf, ctl, pv, it, li, Val(WG), Val(true)) + @synchronize + end + mb = _bm_gmember(bm, G) + off = _lt_local_cb_offset(ctl) + _lu_cb_update!(buf, ubuf, ctl, buf, ubuf, off, _mview(d, mb, bm.nbatch), li, Val(WG), Val(true)) + @synchronize + mb = _bm_gmember(bm, G) + nb = bm.nbatch + _lu_subtree_write!(factor, ufactor, _mview(stack, mb, nb), _mview(ustack, mb, nb), _iview(info, mb, nb), + _mview(stats, mb, nb), buf, ubuf, ctl, member_panels(front_ptr, mb, nb), cb_ptr, li, + Val(WG)) + @synchronize + for r in 1:ctl[_ST_ROUNDS] + _subtree_move!(buf, ctl, r, li, Val(WG)) + _subtree_move!(ubuf, ctl, r, li, Val(WG)) + @synchronize + end + end +end + +# --------------------------------------------------------------------------- +# driver + +function _launch_front_lu!(N::Numeric, S::Symbolic, nzval, first, count, maxchild, prm, ::Val{WG}) where {WG} + bm = batch_map(N; first) + kernel! = front_lu_kernel!(KernelAbstractions.get_backend(N.factor), WG) + kernel!(N.factor, N.ufactor, N.stack, N.ustack, N.info, N.stats, N.d, N.piv, N.pivot_kind, N.aux, nzval, S.amap, + S.amap_ptr, S.amap_src, S.group_nodes, bm, S.super_ptr, S.front_ptr, S.front_nrows, S.front_ncols, + S.cb_ptr, S.child_ptr, S.child_list, S.relind_ptr, S.relind, Int(maxchild), prm, Val(WG); + ndrange = WG * count * bm.nact) + return nothing +end + +function _launch_subtrees_lu!(N::Numeric{T}, S::Symbolic, nzval, first, count, prm, ::Val{LB}, + ::Val{WG}) where {T, LB, WG} + bm = batch_map(N; first) + kernel! = subtree_lu_kernel!(KernelAbstractions.get_backend(N.factor), WG) + kernel!(N.factor, N.ufactor, N.stack, N.ustack, N.info, N.stats, N.d, N.piv, N.pivot_kind, N.aux, nzval, S.amap, + S.amap_ptr, S.amap_src, S.group_nodes, bm, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.front_ptr, + S.front_nrows, S.front_ncols, S.cb_ptr, S.local_front, S.local_cb, S.child_ptr, S.child_list, + S.relind_ptr, S.relind, prm, Val((LB - SUBTREE_LOCAL_RESERVE) ÷ (2 * sizeof(T))), Val(WG); + ndrange = WG * count * bm.nact) + return nothing +end + +""" + factorize_lu!(numeric, symbolic, nzval; opts = Options()) -> 0 + +Multifrontal `P_r P A Pᵀ = L D U` of a general matrix (structure `"G"`) on the +device, with the in-front row pivoting and the static perturbation of the +reference [`ref_lu!`](@ref) (same pivot choice, same storage): +`opts.pivot_type` ([`lu_pivoting`](@ref)), `pivot_threshold`, `pivot_epsilon` +and `pivot_epsilon_alg`. Regime-A groups are one [`subtree_lu_kernel!`](@ref) +launch per budget class; every regime-B and regime-C launch group is one +[`front_lu_kernel!`](@ref) launch (fused assembly and factorization, one +workgroup per front; no vendor calls); 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` (`L`), +`numeric.ufactor` (`Uᵀ`), `d`, `piv` (local row order), `pivot_kind`, `stats` +and `totals`. The factorization always completes (`info = 0`); the phase +allocates nothing on the device, never synchronizes with the host, and is +deterministic. Uniform batch as [`factorize!`](@ref). +""" +function factorize_lu!(N::Numeric{T}, S::Symbolic, nzval::AbstractVector; opts::Options = Options()) where {T} + _check_numeric(N, S, nzval) + prm = _lu_device_params(S, T, opts) + prm.scaled && abs_max!(N.aux, nzval, batch_map(N)) + plan = N.plan + 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] + _with_local_bytes(plan.sub_local[k]) do lb + _launch_subtrees_lu!(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_lu!(N, S, nzval, a, b - a + 1, plan.group_maxchild[k], prm, Val(LU_WORKGROUP)) + end + reduce_stats!(N, S) + return 0 +end diff --git a/src/numeric/storage.jl b/src/numeric/storage.jl index a5d426a..b84bcc5 100644 --- a/src/numeric/storage.jl +++ b/src/numeric/storage.jl @@ -130,7 +130,13 @@ Numeric storage of a factorization (PLAN §3.2), laid out by the * `plan`: the host [`NumericPlan`](@ref); * `members` (`nbatch` `Int32`): the active batch members of the last phase (first `plan.nact[]` entries, see [`BatchMap`](@ref)); -* `nbatch`: the number of matrices of the uniform batch (PLAN §3.5). +* `nbatch`: the number of matrices of the uniform batch (PLAN §3.5); +* `ufactor`, `ustack` (structure `"G"` only, empty otherwise): the `Uᵀ` + structure of LU (`src/numeric/lu.jl`), laid out exactly as `factor` and + `stack`: the panel of `s` holds `U[1:w, :]ᵀ` (unit diagonal stored), the + packed contribution block the transposed strict upper triangle of the `m×m` + block (its diagonal stays zero). For LU, `factor` holds the unit-lower `L`, + `d` the diagonal `D` of `L D U` and `piv` the local row pivot order. A uniform batch of `nbatch` members multiplies every length above by `nbatch`, in the layouts of `src/numeric/batch.jl`: interleaved per front for @@ -154,6 +160,8 @@ struct Numeric{T, VT <: AbstractVector{T}, VS <: AbstractVector{Int64}, VI <: Ab plan::NumericPlan members::VI nbatch::Int + ufactor::VT + ustack::VT end # `pivot_kind` codes of LDLᵀ/LDLᴴ factor columns @@ -172,7 +180,8 @@ Base.show(io::IO, N::Numeric{T, VT}) where {T, VT} = allocate_numeric(symbolic, T, backend = CPU(); nbatch = 1) -> Numeric{T} Allocate (zero-filled) the factor panels, D, update stack, regime-C workspace, per-front -statistics and their totals, status vector, pivot order, pivot kinds, pivot sign requests and `aux` of `symbolic`'s [`Layout`](@ref) for element type +statistics and their totals, status vector, pivot order, pivot kinds, pivot sign requests, `aux` (and for +`"G"` the `Uᵀ` panels and update stack) of `symbolic`'s [`Layout`](@ref) for element type `T` on the KernelAbstractions `backend`, for a uniform batch of `nbatch` matrices (every length times `nbatch`, all members active), and build its [`NumericPlan`](@ref). This is the only allocation of the numeric phase. @@ -198,11 +207,15 @@ function allocate_numeric(S::Symbolic{INT}, ::Type{T}, backend::KernelAbstractio aux = KernelAbstractions.zeros(backend, T, nb) members = KernelAbstractions.allocate(backend, Int32, nb) copyto!(members, Int32.(1:nb)) + lu = S.structure == STRUCTURE_GENERAL + ufactor = KernelAbstractions.zeros(backend, T, lu ? nb * L.factor_len : 0) + ustack = KernelAbstractions.zeros(backend, T, lu ? nb * L.stack_len : 0) plan = NumericPlan(S, nb) nb > 1 && _vendor_batch_pointers!(plan.vendor_ptrs, factor, S, nb) return Numeric{T, typeof(factor), typeof(stats), typeof(info), typeof(pivot_kind)}(factor, d, stack, work, stats, info, piv, pivot_kind, totals, - psign, aux, plan, members, nb) + psign, aux, plan, members, nb, + ufactor, ustack) end # member pointers of the F11 and F21 blocks of the `nb` panels of every regime-C front, when the backend has batched @@ -210,6 +223,7 @@ end function _vendor_batch_pointers!(ptrs::Vector{Any}, factor::AbstractVector{T}, S::Symbolic, nb::Int) where {T} sc, L = S.schedule, S.layout _is_ldlt_structure(S.structure) && return ptrs # LDLᵀ/LDLᴴ: KA kernels in every regime + S.structure == STRUCTURE_GENERAL && return ptrs # LU: likewise sc.vendor_c || return ptrs caps = capabilities(KernelAbstractions.get_backend(factor), T) caps.vendor_potrf_batched || caps.vendor_trsm_batched || return ptrs diff --git a/src/options.jl b/src/options.jl index 14691a1..19d5826 100644 --- a/src/options.jl +++ b/src/options.jl @@ -524,6 +524,8 @@ Accepted values: * algorithms (`"reordering_alg"`, `"factorization_alg"`, `"solve_alg"`, `"matching_alg"`, `"pivot_epsilon_alg"`): `"default"`, `"algoN"` or the integer `N`; * `"pivot_type"`: `'A'`, `'N'`, `'D'`, `'L'`, `'B'` (`'C'`/`'R'` raise `NotSupportedError`); + LDLᵀ/LDLᴴ: `'A'`, `'B'`, `'L'` Bunch–Kaufman, `'D'` 1×1 only, `'N'` none; LU (`"G"`, [`lu_pivoting`](@ref)): + `'A'`, `'B'`, `'L'` in-block threshold row pivoting, `'N'`, `'D'` diagonal pivots; * integer and flag parameters: an `Integer` in the documented range; * `"ir_tol"`, `"pivot_threshold"`, `"pivot_epsilon"`: a finite real `≥ 0` (`nothing` resets `"pivot_epsilon"` to [`default_pivot_epsilon`](@ref)); diff --git a/src/reference/cholesky.jl b/src/reference/cholesky.jl index 3a2376d..291e41e 100644 --- a/src/reference/cholesky.jl +++ b/src/reference/cholesky.jl @@ -117,10 +117,12 @@ supernodes in `snpost` order (`trsm` on the diagonal block, `gemm` update of the rows below), backward sweep `Lᴴ W = Z` in reverse order, `X[perm, :] = W`. `B` and `X` are host vectors or `n × nrhs` matrices (`X === B` is allowed). For structures `"S"`/`"H"` this is the LDLᵀ/LDLᴴ solve of -[`ref_solve_ldlt!`](@ref) with a factor of [`ref_ldlt!`](@ref). +[`ref_solve_ldlt!`](@ref) with a factor of [`ref_ldlt!`](@ref), for `"G"` the +LU solve of [`ref_solve_lu!`](@ref) with a factor of [`ref_lu!`](@ref). """ function ref_solve!(X::AbstractVecOrMat, S::Symbolic, N::Numeric{T, Vector{T}}, B::AbstractVecOrMat) where {T} _is_ldlt_structure(S.structure) && return ref_solve_ldlt!(X, S, N, B) + S.structure == STRUCTURE_GENERAL && return ref_solve_lu!(X, S, N, B) n = S.n size(B, 1) == n && size(X, 1) == n && size(X, 2) == size(B, 2) || throw(DimensionMismatch("A is $n×$n, X is $(size(X)), B is $(size(B))")) diff --git a/src/reference/lu.jl b/src/reference/lu.jl new file mode 100644 index 0000000..07a812c --- /dev/null +++ b/src/reference/lu.jl @@ -0,0 +1,290 @@ +# CPU reference multifrontal LU (PLAN §2.2, §3.3, M7; the oracle of the device +# LU of `src/numeric/lu.jl`): plain Julia on host arrays, the traversal, +# assembly and extend-add of the reference Cholesky on the symmetric pattern of +# `A + Aᵀ`, with a dense in-front `L D U` factorization that pivots on rows +# inside the fully-summed block of each front (threshold partial pivoting) and +# replaces pivots that are too small by `±ε` (static perturbation, as +# SuperLU_DIST's GESP). No delayed pivots: the symbolic structure never changes. +# +# Storage (see `Numeric`): two lower-triangular structures with the layout of +# the Cholesky factor. `factor` holds the unit-lower `L` (front rows after the +# local row interchanges), `ufactor` holds `Uᵀ` (unit diagonal; the panel of `s` +# is `U[c0:c0+w-1, rows of s]ᵀ`, columns not permuted), `d[k]` the pivot `D[k, k]` +# and `piv[k]` the row of `P A Pᵀ` (supernodal numbering) eliminated at factor +# column `k`. Contribution blocks of the device path are two packed triangles +# (the lower triangle with the diagonal in `stack`, the transposed strict upper +# triangle in `ustack`); the reference keeps full host matrices. + +# resolved pivoting parameters of one LU factorization +struct _LUParams{R} + pivot::Bool # in-block threshold partial pivoting (row interchanges) + u::R # pivot_threshold + eps::R # effective perturbation ε (scaled when pivot_epsilon_alg = "algo1") +end + +""" + lu_pivoting(pivot_type::PivotType) -> Bool + +Whether the LU of structure `"G"` interchanges rows inside the fully-summed +block of a front: `'A'` (default), `'L'` and `'B'` mean in-block threshold +partial pivoting; `'N'` and `'D'` keep the diagonal pivots (perturbed when +tiny). Global pivoting (`'C'`, `'R'`) is not planned (PLAN §3.3). +""" +function lu_pivoting(p::PivotType) + (p == PIVOT_NONE || p == PIVOT_DIAGONAL) && return false + p == PIVOT_AUTO || p == PIVOT_BUNCH_KAUFMAN || p == PIVOT_LOCAL_BLOCK || + throw(NotSupportedError("pivot_type = $(repr(convert(Char, p))) is not supported by LU")) + return true +end + +function _lu_params(S::Symbolic, ::Type{T}, opts::Options, nz::AbstractVector) where {T} + S.structure == STRUCTURE_GENERAL || + throw(InvalidValueError("LU needs structure \"G\", got \"$(convert(String, S.structure))\"")) + R = real(T) + eps = R(resolved_pivot_epsilon(opts, R)) + if opts.pivot_epsilon_alg == PIVOT_EPSILON_SCALED + amax = maximum(abs, nz; init = zero(R)) + amax > 0 && (eps *= R(amax)) + end + return _LUParams{R}(lu_pivoting(opts.pivot_type), R(opts.pivot_threshold), eps) +end + +""" + ref_lu!(numeric::Numeric{T, Vector{T}}, symbolic, nzval; opts = Options()) -> 0 + +Reference multifrontal `L D U` of the general matrix whose stored values are +`nzval` (structure `"G"`, same CSR pattern and index base as the analysis) +into the host `numeric`. The supernodes of the symmetric pattern of `A + Aᵀ` +are processed in postorder: the full `f×f` front is assembled (A through the +`amap`, whose negative offsets are the entries above the diagonal; the +children's full contribution blocks through `relind`), its `w` fully-summed +columns are factored with row interchanges inside the fully-summed rows, and +the contribution block is updated once with `F₂₂ ← F₂₂ − (L₂₁ D) U₁₂`. + +Pivot choice at step `k` (`opts.pivot_type`, see [`lu_pivoting`](@ref)): with +`γ = maxᵢ |F[i, k]|` over all remaining rows of the front (fully-summed or +not), the diagonal is kept when `|F[k, k]| ≥ u γ` (`u = pivot_threshold`), +otherwise the fully-summed row `r ≥ k` with the largest `|F[r, k]|` (the first +one) is swapped into row `k`; it is taken even when it fails the threshold +(static pivoting: no delayed pivots). A pivot with `|d| < ε` becomes +`d ← sign(d) ε` (`+ε` for `d = 0`; `d/|d|` for complex `T`), with `ε` as in +[`ref_ldlt!`](@ref) (`pivot_epsilon`, scaled by `max |aᵢⱼ|` with +`pivot_epsilon_alg = "algo1"`); `pivot_sign` is not used by LU. + +Fills `numeric.factor` (unit-lower `L` panels), `numeric.ufactor` (unit-lower +`Uᵀ` panels), `numeric.d`, `numeric.piv` (local row order), `numeric.pivot_kind` +(`PIVOT_KIND_1X1` or `PIVOT_KIND_PERTURBED`) and the per-front `numeric.stats` +(`nzero`, `nperturbed`; `npos`, `nneg` and `n2x2` stay 0: a general matrix has +no inertia). Returns 0: with static perturbation the factorization always +completes. +""" +function ref_lu!(N::Numeric{T, Vector{T}}, S::Symbolic, nzval::AbstractVector; opts::Options = Options()) where {T} + length(nzval) == S.nnz || + throw(InvalidValueError("nzval has $(length(nzval)) entries, the analysis expects $(S.nnz)")) + nz = _host_vector(nzval) + prm = _lu_params(S, T, opts, nz) + amap = _host_vector(S.amap) + amap_ptr = _host_vector(S.amap_ptr) + amap_src = _host_vector(S.amap_src) + child_ptr = _host_vector(S.child_ptr) + child_list = _host_vector(S.child_list) + relind_ptr = _host_vector(S.relind_ptr) + relind = _host_vector(S.relind) + sp, L = S.partition, S.layout + ns = nsupernodes(sp) + n = S.n + length(N.ufactor) == length(N.factor) || + throw(InvalidValueError("the numeric storage was not allocated for structure \"G\"")) + fill!(N.factor, zero(T)) + fill!(N.ufactor, zero(T)) + fill!(N.d, zero(T)) + fill!(N.stats, 0) + fill!(N.piv, Int32(0)) + fill!(N.pivot_kind, Int8(0)) + cbs = Vector{Matrix{T}}(undef, ns) + for s in sp.snpost + f = sp.rowptr[s + 1] - sp.rowptr[s] + w = snwidth(sp, s) + m = f - w + c0 = sp.super_ptr[s] + base = L.panel_ptr[s] + F = zeros(T, f, f) + for k in amap_ptr[s]:(amap_ptr[s + 1] - 1) + p = amap_src[k] + off = Int(amap[p]) + loc = abs(off) - base + i, j = loc % f + 1, loc ÷ f + 1 + off > 0 ? (F[i, j] += T(nz[p])) : (F[j, i] += T(nz[p])) + end + for k in child_ptr[s]:(child_ptr[s + 1] - 1) + c = child_list[k] + isassigned(cbs, c) || continue + C = cbs[c] + ri = view(relind, relind_ptr[c]:(relind_ptr[c + 1] - 1)) + for jj in axes(C, 2), ii in axes(C, 1) + F[ri[ii], ri[jj]] += C[ii, jj] + end + cbs[c] = Matrix{T}(undef, 0, 0) + end + lp, kind, dd, st = _lu_front!(F, w, prm) + for k in 1:w + g = c0 + k - 1 + N.piv[g] = Int32(c0 + lp[k] - 1) + N.pivot_kind[g] = kind[k] + N.d[g] = dd[k] + end + for q in 1:5 + N.stats[(s - 1) * FRONT_STATS_FIELDS + q] = st[q] + end + if m > 0 + # F₂₂ ← F₂₂ − (L₂₁ D) U₁₂, summed in the order of the device kernels + for jj in 1:m, ii in 1:m + acc = zero(T) + for k in 1:w + acc += F[w + ii, k] * dd[k] * F[k, w + jj] + end + F[w + ii, w + jj] -= acc + end + sp.snparent[s] != 0 && (cbs[s] = F[(w + 1):f, (w + 1):f]) + end + # unit-lower panels of L and Uᵀ + for j in 1:w + N.factor[base + (j - 1) * f + j - 1] = one(T) + N.ufactor[base + (j - 1) * f + j - 1] = one(T) + for i in (j + 1):f + N.factor[base + (j - 1) * f + i - 1] = F[i, j] + N.ufactor[base + (j - 1) * f + i - 1] = F[j, i] + end + end + end + _host_totals!(N) + return 0 +end + +ref_lu!(N::Numeric, S::Symbolic, A::CSR; kwargs...) = ref_lu!(N, S, vec(A.nzval); kwargs...) + +# Dense L D U of the fully-summed columns 1:w of the full front F with row interchanges among the rows 1:w +# (columns 1:f of the fully-summed rows and rows 1:f of the fully-summed columns are read and updated; the +# contribution block is left to the caller). On return the strict lower part of F[:, 1:w] is L, the strict +# upper part of F[1:w, :] is U (both unit, scaled by D), F[k, k] = D[k]. Returns the local row order `lp` +# (`lp[k]` = assembled row eliminated at k), the pivot kinds, D and the statistics. +function _lu_front!(F::AbstractMatrix{T}, w::Int, prm::_LUParams) where {T} + f = size(F, 1) + lp = collect(1:w) + kind = fill(PIVOT_KIND_1X1, w) + dd = zeros(T, w) + st = zeros(Int64, 5) + for k in 1:w + r = prm.pivot ? _lu_choose_row(F, k, w, prm.u) : k + if r != k + for j in 1:f + F[k, j], F[r, j] = F[r, j], F[k, j] + end + lp[k], lp[r] = lp[r], lp[k] + end + d = F[k, k] + if !(abs(d) >= prm.eps) # tiny (or NaN): perturb + iszero(d) && (st[STAT_NZERO] += 1) + d = _perturbation_sign(d, 0, !(T <: Complex)) * prm.eps + kind[k] = PIVOT_KIND_PERTURBED + st[STAT_NPERTURBED] += 1 + end + F[k, k] = d + dd[k] = d + for j in (k + 1):w, i in (k + 1):f + F[i, j] -= F[i, k] / d * F[k, j] + end + for j in (w + 1):f, i in (k + 1):w + F[i, j] -= F[i, k] / d * F[k, j] + end + for i in (k + 1):f + F[i, k] /= d + end + for j in (k + 1):f + F[k, j] /= d + end + end + return lp, kind, dd, st +end + +# in-block threshold partial pivoting: the diagonal when |F[k, k]| ≥ u maxᵢ≥ₖ |F[i, k]| (all rows of the front), +# else the first fully-summed row with the largest |F[r, k]| +function _lu_choose_row(F::AbstractMatrix{T}, k::Int, w::Int, u) where {T} + γ = zero(real(T)) + for i in k:size(F, 1) + γ = max(γ, abs(F[i, k])) + end + akk = abs(F[k, k]) + akk >= u * γ && return k + r, best = k, akk + for i in (k + 1):w + a = abs(F[i, k]) + if a > best + r, best = i, a + end + end + return r +end + +""" + extract_lu(symbolic, numeric) -> (L, D, U, p, q) + +The LU factor in `numeric` (host or device, batch member 1) as host sparse +matrices in factor order: `L` unit lower and `U` unit upper triangular, +`D` the pivots, `p` and `q` the row and column permutations (original +indices, `p[k] = perm[piv[k]]`, `q = perm`), so that +`A[p, q] ≈ L * Diagonal(D) * U` (with the perturbed pivots, `A + E`). Rows of `L` +below a diagonal block are the rows of later supernodes, which their own local +row interchanges moved to `ipiv[row]`. +""" +function extract_lu(S::Symbolic, N::Numeric{T}) where {T} + S.structure == STRUCTURE_GENERAL || throw(InvalidValueError("extract_lu needs structure \"G\"")) + H = N.nbatch == 1 ? host_numeric(N) : member_numeric(N, S, 1) + sp, L = S.partition, S.layout + n = S.n + piv = Int.(H.piv[1:n]) + ipiv = invperm(piv) + Li, Lj, Lv = Int[], Int[], T[] + Ui, Uj, Uv = Int[], Int[], T[] + for s in 1:nsupernodes(sp) + rows = snrows(sp, s) + f = length(rows) + c0 = sp.super_ptr[s] + w = snwidth(sp, s) + base = L.panel_ptr[s] + for j in 1:w, i in j:f + g = c0 + j - 1 + l = i == j ? one(T) : H.factor[base + (j - 1) * f + i - 1] + u = i == j ? one(T) : H.ufactor[base + (j - 1) * f + i - 1] + push!(Li, i <= w ? c0 + i - 1 : ipiv[rows[i]]); push!(Lj, g); push!(Lv, l) + push!(Ui, g); push!(Uj, i <= w ? c0 + i - 1 : rows[i]); push!(Uv, u) + end + end + perm = sp.perm + return sparse(Li, Lj, Lv, n, n), H.d[1:n], sparse(Ui, Uj, Uv, n, n), perm[piv], copy(perm) +end + +""" + ref_solve_lu!(X, symbolic, numeric, B; transpose = false) -> X + +Host solve of `A X = B` (`Aᵀ X = B` when `transpose`) with the LU factor in +`numeric` through [`extract_lu`](@ref) (sparse triangular solves; tests and +debugging). `B`, `X`: vectors or matrices. +""" +function ref_solve_lu!(X::AbstractVecOrMat, S::Symbolic, N::Numeric{T}, B::AbstractVecOrMat; + transpose::Bool = false) where {T} + L, D, U, p, q = extract_lu(S, N) + Bh = Array(B) + Bm = reshape(Bh, S.n, :) + Y = similar(Bm) + if transpose + # Aᵀ[q, p] = Uᵀ D Lᵀ + Z = UnitLowerTriangular(sparse(Base.transpose(U))) \ Bm[q, :] + Y[p, :] = UnitUpperTriangular(sparse(Base.transpose(L))) \ (Diagonal(D) \ Z) + else + Z = UnitLowerTriangular(L) \ Bm[p, :] + Y[q, :] = UnitUpperTriangular(U) \ (Diagonal(D) \ Z) + end + copyto!(X, vec(Y)) + return X +end diff --git a/src/solve/refinement.jl b/src/solve/refinement.jl index 6cdbe2b..7f38051 100644 --- a/src/solve/refinement.jl +++ b/src/solve/refinement.jl @@ -8,9 +8,12 @@ # 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)). +# stored CSR matrix `M` or, for a CSC input, `Mᵀ`. For symmetric or Hermitian +# `M` every `op(A)` is `M` or `conj(M)`; `conj(M) x = b` is solved as +# `x = conj(M⁻¹ conj(b))` ([`solve_conjugated`](@ref)). For a general `M` +# (`"G"`, LU) `op(A)` is `M`, `Mᵀ` or `conj(Mᵀ)` ([`solve_transposed`](@ref)): +# the refinement map then also holds the rows of `Mᵀ` (`transpose_matrix` +# selects them) and the sweeps solve with the transposed factors. # # 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 @@ -30,7 +33,8 @@ 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); + `src[c] < 0`: `conj(nzval[-src[c]])`), from [`refinement_map`](@ref); for + `"G"` the rows `n+1:2n` of `rowptr` are the rows of `Mᵀ`; * `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 @@ -52,16 +56,17 @@ end max_rhs(W::RefinementWorkspace) = size(W.R, 2) """ - refinement_map(F::FullPatternMap) -> (rowptr, colval, src) + refinement_map(F::FullPatternMap; transpose = false) -> (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. +when `src[c] < 0`. 1-based. With `transpose = true` (structure `"G"`) the rows +of `Mᵀ` follow as rows `n+1:2n` (`rowptr` has `2n + 1` entries). """ -function refinement_map(F::FullPatternMap) +function refinement_map(F::FullPatternMap; transpose::Bool = false) n = F.n - rowptr = Vector{Int}(undef, n + 1) + rowptr = Vector{Int}(undef, (transpose ? 2n : n) + 1) rowptr[1] = 1 colval = Int[] src = Int[] @@ -72,20 +77,34 @@ function refinement_map(F::FullPatternMap) end rowptr[i + 1] = length(colval) + 1 end + if transpose + # the entries of M by column: row j of Mᵀ holds the entries (i, j) of M, rows i in order + cols = [Tuple{Int, Int}[] for _ in 1:n] + for i in 1:n, e in F.rowptr[i]:(F.rowptr[i + 1] - 1), s in F.srcptr[e]:(F.srcptr[e + 1] - 1) + push!(cols[F.colval[e]], (i, F.conjflag[s] ? -F.src[s] : F.src[s])) + end + for j in 1:n + for (i, v) in cols[j] + push!(colval, i) + push!(src, v) + end + rowptr[n + j + 1] = length(colval) + 1 + end + end return rowptr, colval, src end """ - allocate_refinement(map, T, INT, backend, nrhs) -> RefinementWorkspace + allocate_refinement(map, T, INT, backend, nrhs; n = rows of map) -> 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`. +sides of element type `T` and `n` rows (pass `n` for a `"G"` map, which holds `2n` rows). """ function allocate_refinement(map::NTuple{3, Vector{Int}}, ::Type{T}, ::Type{INT}, - backend::KernelAbstractions.Backend, nrhs::Integer) where {T, INT} + backend::KernelAbstractions.Backend, nrhs::Integer; + n::Integer = length(map[1]) - 1) 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) @@ -113,8 +132,9 @@ end # compact column r (user column `_bm_ucol(bm, r)` of the `nrhs` user columns, values of member # `_bm_cmember(bm, r)`) -@kernel function _residual_kernel!(R, rowptr, colval, src, nzval_all, X, B, n, nrhs, ncols, bm, ::Val{XT}, ::Val{BT}, - ::Val{CJ}) where {XT, BT, CJ} +# (`roff = n`: the rows of Mᵀ of a "G" map) +@kernel function _residual_kernel!(R, rowptr, colval, src, nzval_all, X, B, n, nrhs, ncols, bm, roff, ::Val{XT}, + ::Val{BT}, ::Val{CJ}) where {XT, BT, CJ} q = @index(Global, Linear) i = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 @@ -122,7 +142,7 @@ end nzval = _mview(nzval_all, _bm_cmember(bm, r), bm.nbatch) u = _bm_ucol(bm, r) acc = zero(eltype(R)) - for c in Int(rowptr[i]):(Int(rowptr[i + 1]) - 1) + for c in Int(rowptr[roff + i]):(Int(rowptr[roff + i + 1]) - 1) s = Int(src[c]) v = nzval[abs(s)] v = xor(s < 0, CJ) ? conj(v) : v @@ -194,9 +214,10 @@ _user_map(::Nothing, nrhs) = single_batch(; nrhs) """ residual!(W, nzval, X, B; nrhs, transposed = false, b_transposed = transposed, conjugate = false, - bm = nothing) -> W.R + bm = nothing, transpose_matrix = false) -> W.R -`W.R[:, 1:nrhs] = B - M X` (`conj(M)` when `conjugate`) with the full matrix +`W.R[:, 1:nrhs] = B - M X` (`conj(M)` when `conjugate`; `Mᵀ` when +`transpose_matrix`, which needs a `"G"` map) 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`) with `nrhs` right-hand sides. Uniform batch ([`BatchMap`](@ref) `bm`): `nzval` holds the @@ -207,18 +228,19 @@ Asynchronous. """ function residual!(W::RefinementWorkspace, nzval::AbstractVector, X::AbstractVecOrMat, B::AbstractVecOrMat; nrhs::Integer, transposed::Bool = false, b_transposed::Bool = transposed, - conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing) + conjugate::Bool = false, bm::Union{Nothing, BatchMap} = nothing, transpose_matrix::Bool = false) n = size(W.R, 1) bm = _user_map(bm, nrhs) ncols = bm.nrhs * bm.nact n * ncols > 0 || return W.R + roff = _map_row_offset(W, transpose_matrix) 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), ncols, bm, xt, bt, Val(true); + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), ncols, bm, roff, xt, bt, Val(true); ndrange = n * ncols) else - kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), ncols, bm, xt, bt, Val(false); + kernel!(W.R, W.rowptr, W.colval, W.src, nzval, X, B, n, Int(nrhs), ncols, bm, roff, xt, bt, Val(false); ndrange = n * ncols) end end @@ -249,6 +271,15 @@ function residual_norms!(W::RefinementWorkspace, B::AbstractVecOrMat; nrhs::Inte return W.norms_host end +# first row of the map of `op(A)`: the rows of Mᵀ start at n (a "G" map) +function _map_row_offset(W::RefinementWorkspace, transpose_matrix::Bool) + transpose_matrix || return 0 + n = size(W.R, 1) + length(W.rowptr) == 2n + 1 || + throw(InvalidValueError("the refinement map has no rows of Mᵀ (only structure \"G\" solves with Aᵀ)")) + return n +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) @@ -302,10 +333,11 @@ 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 + interrupt = nothing, progress = Ref(0), transpose_matrix = false) -> 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 +or `conj(M)` when `conjugate`; `Mᵀ` or `conj(Mᵀ)` with `transpose_matrix`, LU +only), 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) @@ -322,7 +354,7 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac 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)) + progress::Base.RefValue{Int} = Ref(0), transpose_matrix::Bool = false) nu = rhs_count(X, S.n; transposed) # user columns: nrhs per member × nbatch nu % N.nbatch == 0 || throw(DimensionMismatch("$nu right-hand sides for a batch of $(N.nbatch) members")) nrhs = nu ÷ N.nbatch @@ -334,7 +366,7 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac progress[] = 0 for _ in 1:nsteps _poll_interrupt(interrupt) - residual!(W, nzval, X, B; nrhs = nu, transposed, b_transposed, conjugate, bm) + residual!(W, nzval, X, B; nrhs = nu, transposed, b_transposed, conjugate, bm, transpose_matrix) if tol > 0 norms = residual_norms!(W, B; nrhs = nu, transposed = b_transposed, bm) rel = _max_relative_residual(norms, ncols) @@ -342,9 +374,9 @@ function refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementWorkspac rel <= tol && break end permute_rhs!(ws.Y, W.R, S.perm; conjugate) - forward_sweep!(ws, S, N; nrhs, deterministic) + forward_sweep!(ws, S, N; nrhs, deterministic, transpose = transpose_matrix) diagonal_sweep!(ws, S, N; nrhs) - backward_sweep!(ws, S, N; nrhs) + backward_sweep!(ws, S, N; nrhs, transpose = transpose_matrix) add_correction!(X, ws.Y, S.perm; nrhs = nu, transposed, conjugate, bm) steps += 1 progress[] = steps @@ -383,10 +415,11 @@ Whether the Krylov.jl extension is loaded (`ir_mode = "fgmres"` works). fgmres_available() = FGMRES_PROVIDER[] !== nothing """ - RefinementOperator(W, nzval, n, ncols, bm, conjugate) + RefinementOperator(W, nzval, n, ncols, bm, conjugate, roff = 0) `op(A)` (the full matrix `M` of the [`RefinementWorkspace`](@ref) `W` with -the values `nzval`, `conj(M)` when `conjugate`) acting on the `n × ncols` +the values `nzval`, `conj(M)` when `conjugate`; `roff = n`: `Mᵀ` from the +rows of a `"G"` map) 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, @@ -399,13 +432,17 @@ struct RefinementOperator{T, RW <: RefinementWorkspace{T}, V <: AbstractVector, ncols::Int bm::BM conjugate::Bool + roff::Int end +RefinementOperator(W::RefinementWorkspace, nzval::AbstractVector, n::Integer, ncols::Integer, bm::BatchMap, + conjugate::Bool) = RefinementOperator(W, nzval, Int(n), Int(ncols), bm, conjugate, 0) + 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} +@kernel function _spmv_kernel!(y, rowptr, colval, src, nzval_all, x, n, ncols, bm, roff, ::Val{CJ}) where {CJ} q = @index(Global, Linear) i = (q - 1) % n + 1 r = (q - 1) ÷ n + 1 @@ -413,7 +450,7 @@ Base.eltype(::RefinementOperator{T}) where {T} = T 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) + for c in Int(rowptr[roff + i]):(Int(rowptr[roff + i + 1]) - 1) s = Int(src[c]) v = nzval[abs(s)] v = xor(s < 0, CJ) ? conj(v) : v @@ -431,19 +468,19 @@ function LinearAlgebra.mul!(y::AbstractVector, op::RefinementOperator, x::Abstra 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) + kernel!(y, W.rowptr, W.colval, W.src, op.nzval, x, op.n, op.ncols, op.bm, op.roff, 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) + kernel!(y, W.rowptr, W.colval, W.src, op.nzval, x, op.n, op.ncols, op.bm, op.roff, Val(false); ndrange = m) end return y end """ - FactorPreconditioner(ws, S, N, nrhs, ncols, conjugate, deterministic, interrupt) + FactorPreconditioner(ws, S, N, nrhs, ncols, conjugate, deterministic, interrupt, transpose = false) `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, +`conjugate`; the sweeps of `Mᵀ` when `transpose`, LU only) 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. @@ -457,8 +494,13 @@ struct FactorPreconditioner{T, WS <: SolveWorkspace{T}, SY <: Symbolic, NU <: Nu conjugate::Bool deterministic::Bool interrupt::Union{Nothing, Threads.Atomic{Bool}} + transpose::Bool end +FactorPreconditioner(ws::SolveWorkspace, S::Symbolic, N::Numeric, nrhs::Integer, ncols::Integer, conjugate::Bool, + deterministic::Bool, interrupt) = + FactorPreconditioner(ws, S, N, Int(nrhs), Int(ncols), conjugate, deterministic, interrupt, false) + 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 @@ -467,9 +509,9 @@ function LinearAlgebra.mul!(y::AbstractVector, P::FactorPreconditioner, x::Abstr _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) + forward_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs, deterministic = P.deterministic, transpose = P.transpose) diagonal_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs) - backward_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs) + backward_sweep!(P.ws, P.S, P.N; nrhs = P.nrhs, transpose = P.transpose) unpermute_solution!(y, P.ws.Y, perm; conjugate = P.conjugate) return y end @@ -487,7 +529,7 @@ 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 + interrupt = nothing, progress = Ref(0), transpose_matrix = false) -> iterations FGMRES-IR (`ir_mode = "fgmres"`, needs Krylov.jl) of the solution `X` of `op(A) X = B`, with the arguments of [`refine!`](@ref): @@ -511,7 +553,7 @@ function fgmres_refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementW 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)) + progress::Base.RefValue{Int} = Ref(0), transpose_matrix::Bool = false) provider = FGMRES_PROVIDER[] provider === nothing && throw(NotSupportedError("ir_mode = \"fgmres\" needs Krylov.jl: load it with `using Krylov`")) @@ -525,7 +567,7 @@ function fgmres_refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementW 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) + residual!(W, nzval, X, B; nrhs = nu, transposed, b_transposed, conjugate, bm, transpose_matrix) R = real(eltype(W.R)) atol = zero(R) if tol > 0 @@ -535,8 +577,8 @@ function fgmres_refine!(X::AbstractVecOrMat, B::AbstractVecOrMat, W::RefinementW 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) + A = RefinementOperator(W, nzval, S.n, ncols, bm, conjugate, _map_row_offset(W, transpose_matrix)) + P = FactorPreconditioner(ws, S, N, nrhs, ncols, conjugate, deterministic, interrupt, transpose_matrix) 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) @@ -551,10 +593,23 @@ Whether `op(A)` (`solve_mode` 0: `A`, 1: `Aᵀ`, 2: `Aᴴ`) of the user matrix i 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. +and `Aᴴ` of a CSC input. General `"G"`: for `Aᴴ` (`op(A) = conj(Mᵀ)` or +`conj(M)`, see [`solve_transposed`](@ref)). """ 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 + +""" + solve_transposed(structure, csc_input::Bool, solve_mode) -> Bool + +Whether the solve of `op(A)` (`solve_mode` 0: `A`, 1: `Aᵀ`, 2: `Aᴴ`) uses the +transpose `Mᵀ` of the stored, factorized matrix `M` (`A = M`, or `A = Mᵀ` for +a CSC input), conjugated or not ([`solve_conjugated`](@ref)). Only for +structure `"G"` (LU): `Mᵀ` for `Aᵀ`/`Aᴴ` of a CSR input and for `A` of a CSC +input. Symmetric and Hermitian structures never need it. +""" +solve_transposed(structure::Structure, csc_input::Bool, solve_mode::Integer) = + structure == STRUCTURE_GENERAL && xor(solve_mode != 0, csc_input) diff --git a/src/solve/sweeps.jl b/src/solve/sweeps.jl index d5655ae..b57eb86 100644 --- a/src/solve/sweeps.jl +++ b/src/solve/sweeps.jl @@ -45,6 +45,14 @@ # rows `1:w` of its gather-list buffer in `U` (unused otherwise). Complex # symmetric `"S"` (`CONJ = false`) solves with `Lᵀ` instead of `Lᴴ`. # +# LU (`"G"`, `src/numeric/lu.jl`): the factor is two unit-lower structures, `L` +# (`numeric.factor`, rows in local pivot order) and `Uᵀ` (`numeric.ufactor`), +# and the kernels take the panel buffer as an argument. `A x = b`: forward +# sweep with `L`, reordering by the local row order (`LDL = true`), diagonal, +# backward sweep with `(Uᵀ)ᵀ = U` (`CONJ = false`) without reordering. +# `Aᵀ x = b` (`transpose`): forward with `Uᵀ` without reordering, diagonal, +# backward with `Lᵀ`, restoring the order. +# # As in the numeric kernels, per-node values sit in a `@localmem` control array # written by work item 1 before a barrier, so every loop containing a barrier # has a workgroup-uniform trip count that the KA CPU backend can evaluate. @@ -609,11 +617,14 @@ end _ncols(N::Numeric, nrhs::Integer) = Int(nrhs) * N.plan.nact[] function _check_solve(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric{T}, nrhs::Integer) where {T} - _is_ldlt_structure(S.structure) ? _check_ldlt_eltype(T) : _check_reference_cholesky(S, T) + _is_ldlt_structure(S.structure) || S.structure == STRUCTURE_GENERAL ? _check_ldlt_eltype(T) : + _check_reference_cholesky(S, T) length(N.factor) == N.nbatch * S.layout.factor_len || throw(InvalidValueError("the numeric storage was not allocated for this analysis")) size(ws.Y, 1) == S.n && size(ws.U, 1) == length(S.partition.rowval) && size(ws.tmp, 1) == ws.plan.maxm || throw(InvalidValueError("the solve workspace was not allocated for this analysis")) + S.structure == STRUCTURE_GENERAL && length(N.ufactor) != length(N.factor) && + throw(InvalidValueError("the numeric storage was not allocated for structure \"G\"")) 0 <= _ncols(N, nrhs) <= max_rhs(ws) || throw(InvalidValueError("$nrhs right-hand sides for $(N.plan.nact[]) batch members, the solve workspace " * "holds $(max_rhs(ws))")) @@ -636,58 +647,72 @@ function _solve_impls(N::Numeric, S::Symbolic, impl::Symbol) end function _launch_fwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, - det::Val, ldl::Val) + det::Val, ldl::Val, factor::AbstractVector = N.factor) WG = SOLVE_WORKGROUP solve_fwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, ws.U, N.factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, + ws.Y, ws.U, factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, S.child_list, S.relind_ptr, S.relind, sub, det, ldl, Val(WG); ndrange = WG * count * _ncols(N, nrhs)) return nothing end function _launch_bwd!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Val, - ldl::Val, cj::Val) + ldl::Val, cj::Val, factor::AbstractVector = N.factor) WG = SOLVE_WORKGROUP solve_bwd_kernel!(KernelAbstractions.get_backend(ws.Y), WG)( - ws.Y, ws.U, N.factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, + ws.Y, ws.U, factor, N.piv, S.group_nodes, batch_map(N; first, nrhs), count, S.subtree_ptr, S.subtree_nodes, S.super_ptr, S.rowptr, S.rowval, S.front_ptr, S.front_nrows, S.front_ncols, S.child_ptr, sub, ldl, cj, Val(WG); ndrange = WG * count * _ncols(N, nrhs)) return nothing end -# solve kind of the analysis: Cholesky, LDLᵀ/LDLᴴ with Lᴴ (Hermitian or real), LDLᵀ complex symmetric +# solve kind of the analysis: Cholesky, LDLᵀ/LDLᴴ with Lᴴ (Hermitian or real), LDLᵀ complex symmetric, +# LU of A and LU of Aᵀ const _SV_CHOLESKY = 0 const _SV_LDL_HERM = 1 const _SV_LDL_SYM = 2 +const _SV_LU = 3 +const _SV_LU_T = 4 -_solve_kind(S::Symbolic, ::Type{T}) where {T} = +_solve_kind(S::Symbolic, ::Type{T}, transpose::Bool = false) where {T} = + S.structure == STRUCTURE_GENERAL ? (transpose ? _SV_LU_T : _SV_LU) : !_is_ldlt_structure(S.structure) ? _SV_CHOLESKY : (T <: Complex && S.structure == STRUCTURE_SYMMETRIC) ? _SV_LDL_SYM : _SV_LDL_HERM +# the panel buffer a sweep reads and whether it applies the local pivot order: LU solves `A x = b` with `L` +# (reordered) forward and `U` backward, `Aᵀ x = b` with `Uᵀ` forward and `Lᵀ` (order restored) backward +_sweep_factor(N::Numeric, kind::Int, forward::Bool) = + (kind == _SV_LU && !forward) || (kind == _SV_LU_T && forward) ? N.ufactor : N.factor +_sweep_reorders(kind::Int, forward::Bool) = + kind == _SV_LU ? forward : kind == _SV_LU_T ? !forward : kind != _SV_CHOLESKY + # the flags as compile-time constants through explicit branches (closures over runtime flags allocate) function _launch_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, sub::Bool, det::Bool, forward::Bool, kind::Int) + fac = _sweep_factor(N, kind, forward) if kind == _SV_CHOLESKY - _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(false), Val(true)) + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(false), Val(true), fac) elseif kind == _SV_LDL_HERM - _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(true)) + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(true), fac) + elseif _sweep_reorders(kind, forward) + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(false), fac) else - _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(true), Val(false)) + _launch_sweep!(ws, S, N, first, count, nrhs, sub, det, forward, Val(false), Val(false), fac) end return nothing end function _launch_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric, first::Int, count::Int, nrhs::Int, - sub::Bool, det::Bool, forward::Bool, ldl::Val, cj::Val) + sub::Bool, det::Bool, forward::Bool, ldl::Val, cj::Val, fac::AbstractVector) if !forward - sub ? _launch_bwd!(ws, S, N, first, count, nrhs, Val(true), ldl, cj) : - _launch_bwd!(ws, S, N, first, count, nrhs, Val(false), ldl, cj) + sub ? _launch_bwd!(ws, S, N, first, count, nrhs, Val(true), ldl, cj, fac) : + _launch_bwd!(ws, S, N, first, count, nrhs, Val(false), ldl, cj, fac) elseif sub - det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(true), ldl) : - _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(false), ldl) + det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(true), ldl, fac) : + _launch_fwd!(ws, S, N, first, count, nrhs, Val(true), Val(false), ldl, fac) else - det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(true), ldl) : - _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(false), ldl) + det ? _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(true), ldl, fac) : + _launch_fwd!(ws, S, N, first, count, nrhs, Val(false), Val(false), ldl, fac) end return nothing end @@ -708,25 +733,27 @@ end # the blocks L11, L21 of the panel of front `s` of the batch member in slot `j`, the front's rows of the member's # columns of Y and the m rows of its columns of tmp (the single-matrix path) -function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j::Int) +function _dense_front(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j::Int, + factor::AbstractVector = N.factor) sc = S.schedule f, w = sc.rows[s], sc.width[s] m = f - w p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j], N.nbatch) c0 = S.partition.super_ptr[s] cols = ((j - 1) * nrhs + 1):(j * nrhs) - P = reshape(view(N.factor, p0:(p0 + f * w - 1)), f, w) + P = reshape(view(factor, p0:(p0 + f * w - 1)), f, w) return view(P, 1:w, 1:w), view(P, (w + 1):f, 1:w), view(ws.Y, c0:(c0 + w - 1), cols), view(ws.tmp, 1:m, cols) end # the same for the `count` consecutive members in the slots `j0:(j0 + count - 1)`, as `×count` strided batches -function _dense_front_batch(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int) +function _dense_front_batch(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, j0::Int, count::Int, + factor::AbstractVector = N.factor) sc = S.schedule f, w = sc.rows[s], sc.width[s] m = f - w p0 = panel_offset(S.layout.panel_ptr, s, N.plan.members_host[j0], N.nbatch) c0 = S.partition.super_ptr[s] - P = reshape(view(N.factor, p0:(p0 + f * w * count - 1)), f, w, count) + P = reshape(view(factor, p0:(p0 + f * w * count - 1)), f, w, count) Y3 = _column_batch(ws.Y, (j0 - 1) * nrhs, nrhs, count) T3 = _column_batch(ws.tmp, (j0 - 1) * nrhs, nrhs, count) return view(P, 1:w, 1:w, :), view(P, (w + 1):f, 1:w, :), view(Y3, c0:(c0 + w - 1), :, :), view(T3, 1:m, :, :) @@ -742,10 +769,10 @@ end # the dense calls of front `s` for each run of consecutive active members: one strided-batched call per run, or # (single matrix, `impl = :generic`) one call per member function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs::Int, p::NamedTuple, forward::Bool, - tr::Char) + tr::Char, fac::AbstractVector = N.factor) m = S.schedule.rows[s] - S.schedule.width[s] if N.nbatch == 1 - _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, 1)...) + _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, 1, fac)...) return nothing end runs = N.plan.runs @@ -753,10 +780,10 @@ function _dense_runs!(ws::SolveWorkspace, S::Symbolic, N::Numeric, s::Int, nrhs: j0, j1 = runs[r], runs[r + 1] - 1 if p.btrsm === :none for j in j0:j1 - _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, j)...) + _dense_call!(p, forward, tr, m, _dense_front(ws, S, N, s, nrhs, j, fac)...) end else - _dense_call_batch!(p, forward, tr, m, _dense_front_batch(ws, S, N, s, nrhs, j0, j1 - j0 + 1)...) + _dense_call_batch!(p, forward, tr, m, _dense_front_batch(ws, S, N, s, nrhs, j0, j1 - j0 + 1, fac)...) end end return nothing @@ -803,8 +830,8 @@ function _fwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrh sc = S.schedule w = sc.width[s] m = sc.rows[s] - w - kind == _SV_CHOLESKY || _piv_dense!(ws, S, N, s, nrhs, true) - _dense_runs!(ws, S, N, s, nrhs, p, true, 'N') + _sweep_reorders(kind, true) && _piv_dense!(ws, S, N, s, nrhs, true) + _dense_runs!(ws, S, N, s, nrhs, p, true, 'N', _sweep_factor(N, kind, true)) m > 0 || return nothing ncols = _ncols(N, nrhs) rp = S.partition.rowptr[s] + w @@ -822,15 +849,15 @@ function _bwd_dense!(ws::SolveWorkspace{T}, S::Symbolic, N::Numeric, s::Int, nrh sc = S.schedule w = sc.width[s] m = sc.rows[s] - w - tr = kind == _SV_LDL_SYM ? 'T' : 'C' + tr = kind == _SV_LDL_HERM || kind == _SV_CHOLESKY ? 'C' : 'T' if m > 0 ncols = _ncols(N, nrhs) rp = S.partition.rowptr[s] + w kernel! = _sv_gather_kernel!(KernelAbstractions.get_backend(ws.Y), SOLVE_WORKGROUP) kernel!(ws.tmp, ws.Y, S.rowval, rp, m, ncols; ndrange = m * ncols) end - _dense_runs!(ws, S, N, s, nrhs, p, false, tr) - kind == _SV_CHOLESKY || _piv_dense!(ws, S, N, s, nrhs, false) + _dense_runs!(ws, S, N, s, nrhs, p, false, tr, _sweep_factor(N, kind, false)) + _sweep_reorders(kind, false) && _piv_dense!(ws, S, N, s, nrhs, false) return nothing end @@ -849,15 +876,17 @@ per-front buffers `ws.U` and an owner-pull of the children (one extra launch per regime-C step). Uniform batch: `nrhs` right-hand sides per active member of `numeric` (`nrhs × nact` columns of `ws.Y`, each solved with its member's factor; regime-C-path fronts use strided-batched `trsm`/`gemm` over runs of -consecutive members). Asynchronous; allocates nothing on the device. +consecutive members). LU (`"G"`): `L Z = Y` with the local row order, or +`Uᵀ Z = Y` with `transpose = true` (the solve of `Aᵀ`). Asynchronous; allocates +nothing on the device. """ function forward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[], - deterministic::Bool = false, impl::Symbol = :auto) + deterministic::Bool = false, impl::Symbol = :auto, transpose::Bool = false) _check_solve(ws, S, N, nrhs) nrhs > 0 || return ws det = deterministic || !ws.atomic p = _solve_impls(N, S, impl) - kind = _solve_kind(S, eltype(ws)) + kind = _solve_kind(S, eltype(ws), transpose) plan = ws.plan nodes = S.schedule.group_nodes for k in eachindex(plan.kind) @@ -888,10 +917,11 @@ columns of `ws.Y`: the identity for Cholesky (`"SPD"`/`"HPD"`); for LDLᵀ/LDLᴴ one launch, one work item per (factor column, right-hand side), dividing by the 1×1 pivots of `numeric.d` and solving the 2×2 blocks (`d[k]`, `d[n + k]`, `d[k + 1]`; Hermitian or complex symmetric as the -structure). Uniform batch: `nrhs` per active member, each with its member's D. -Asynchronous. +structure; LU: the pivots of `L D U`, `transpose` is accepted for symmetry). +Uniform batch: `nrhs` per active member, each with its member's D. Asynchronous. """ -function diagonal_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[]) +function diagonal_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[], + transpose::Bool = false) _check_solve(ws, S, N, nrhs) kind = _solve_kind(S, eltype(ws)) (kind == _SV_CHOLESKY || nrhs == 0 || S.n == 0) && return ws @@ -914,14 +944,16 @@ first `nrhs` columns of `ws.Y`, the forward plan in reverse: per step the regime `gemm` with `L21ᴴ`, `trsm` with `L11ᴴ`) and one launch over the regime-B fronts, then one launch over all regime-A subtrees (supernodes in reverse processing order). Gather-based and conflict-free: no atomics, deterministic. -Uniform batch as in [`forward_sweep!`](@ref). Asynchronous. +Uniform batch as in [`forward_sweep!`](@ref). LU (`"G"`): `U X = W`, or +`Lᵀ X = W` followed by the inverse local row order with `transpose = true`. +Asynchronous. """ function backward_sweep!(ws::SolveWorkspace, S::Symbolic, N::Numeric; nrhs::Integer = max_rhs(ws) ÷ N.plan.nact[], - impl::Symbol = :auto) + impl::Symbol = :auto, transpose::Bool = false) _check_solve(ws, S, N, nrhs) nrhs > 0 || return ws p = _solve_impls(N, S, impl) - kind = _solve_kind(S, eltype(ws)) + kind = _solve_kind(S, eltype(ws), transpose) plan = ws.plan nodes = S.schedule.group_nodes for k in reverse(eachindex(plan.kind)) @@ -949,10 +981,12 @@ layouts of [`rhs_count`](@ref) (vector, `n × nrhs` matrix, strided vector; row-major with `transposed = true`), with at most `max_rhs(ws)` right-hand sides; `X === B` is allowed. Uniform batch: `B` and `X` hold the right-hand sides of all `numeric.nbatch` members (member after member), only the active -members are solved. Asynchronous; allocates nothing on the device. +members are solved. LU (`"G"`): `transpose = true` solves `Aᵀ X = B`. +Asynchronous; allocates nothing on the device. """ function sweep_solve!(X::AbstractVecOrMat, ws::SolveWorkspace, S::Symbolic, N::Numeric, B::AbstractVecOrMat; - transposed::Bool = false, deterministic::Bool = false, impl::Symbol = :auto) + transposed::Bool = false, deterministic::Bool = false, impl::Symbol = :auto, + transpose::Bool = false) ncols = rhs_count(B, S.n; transposed) rhs_count(X, S.n; transposed) == ncols || throw(DimensionMismatch("X has $(rhs_count(X, S.n; transposed)) right-hand sides, B has $ncols")) @@ -962,9 +996,9 @@ function sweep_solve!(X::AbstractVecOrMat, ws::SolveWorkspace, S::Symbolic, N::N _check_solve(ws, S, N, nrhs) bm = batch_map(N; nrhs) permute_rhs!(ws.Y, B, S.perm; transposed, bm) - forward_sweep!(ws, S, N; nrhs, deterministic, impl) + forward_sweep!(ws, S, N; nrhs, deterministic, impl, transpose) diagonal_sweep!(ws, S, N; nrhs) - backward_sweep!(ws, S, N; nrhs, impl) + backward_sweep!(ws, S, N; nrhs, impl, transpose) unpermute_solution!(X, ws.Y, S.perm; transposed, bm) return X end diff --git a/src/solver.jl b/src/solver.jl index 75bcc94..ad79354 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -60,13 +60,14 @@ 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ᴴ), single matrices +(LDLᵀ; complex symmetric for complex `T`), `"H"` (LDLᴴ) and `"G"` (LU on the +symmetric pattern of `A + Aᵀ` with row pivoting inside the fully-summed block +of each front, `P_r P A Pᵀ = L D U`; view `'F'` only), single matrices and uniform batches (see below), the phases `"reordering"`, `"symbolic_factorization"`, `"analysis"`, `"factorization"`, `"refactorization"`, `"solve"`, `"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. +`"solve_bwd_perm"` and `"solve_refinement"` ([`execute!`](@ref)). The Schur +phases raise [`NotSupportedError`](@ref) when they are executed. Uniform batch (PLAN §1.6, §3.5, ≅ CUDSS.jl's uniform batch): `nbatch` matrices with the pattern of `rowptr`/`colval` and values `nzval`, either a @@ -321,9 +322,6 @@ end # options that no phase implements yet: refuse them instead of silently ignoring them function _check_analysis_supported(solver::DirectSolver{T}) where {T} s = solver.structure - s == STRUCTURE_GENERAL && - throw(NotSupportedError("structure \"G\" (LU) is not implemented yet (T19); use \"S\", \"H\", \"SPD\" or " * - "\"HPD\"")) T <: Complex && s == STRUCTURE_SPD && throw(InvalidValueError("a complex positive definite matrix needs structure \"HPD\"")) opts = solver.options @@ -376,7 +374,8 @@ function _symbolic!(solver::DirectSolver{T, INT}) where {T, INT} P = SymmetricPattern(solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) - sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(solver.structure)) + sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(solver.structure), + elsize = schedule_elsize(solver.structure, T)) layout = build_layout(sp, sc; ldlt = _is_ldlt_structure(solver.structure)) Sh = Symbolic(sp, sc, layout, solver.host_rowptr, solver.host_colval, A.nrows, solver.structure; view = _stored_view(solver), index = A.index) @@ -447,12 +446,14 @@ _info_value(solver::DirectSolver) = solver.nbatch == 1 ? solver.info[1] : copy(s _elapsed(tic) = string(round((time_ns() - tic) / 1.0e6; digits = 3), " ms") # The numeric phase of the structure, behind a barrier: the structure is a run-time value, and a static call -# would compile the Cholesky and the LDLᵀ/LDLᴴ paths with every solver, on the CPU backend all their kernels (most +# would compile the Cholesky, LDLᵀ/LDLᴴ and LU paths with every solver, on the CPU backend all their kernels (most # of the test suite's compile time). The barrier takes the mutable handle only: the dynamic call boxes nothing # (an immutable `Numeric` or `Symbolic` argument would be copied to the heap on every call). function _numeric_phase!(solver::DirectSolver) if _is_ldlt_structure(solver.structure) Base.inferencebarrier(_numeric_phase_ldlt!)(solver) + elseif solver.structure == STRUCTURE_GENERAL + Base.inferencebarrier(_numeric_phase_lu!)(solver) else Base.inferencebarrier(_numeric_phase_cholesky!)(solver) end @@ -469,6 +470,11 @@ function _numeric_phase_ldlt!(solver::DirectSolver) return nothing end +function _numeric_phase_lu!(solver::DirectSolver) + factorize_lu!(solver.numeric, solver.symbolic, _flat(solver.A.nzval); opts = solver.options) + return nothing +end + _flat(nzval::AbstractVector) = nzval _flat(nzval::AbstractMatrix) = vec(nzval) @@ -512,7 +518,8 @@ function _refinement!(solver::DirectSolver{T, INT}, nrhs::Int) where {T, INT} 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))) + W = allocate_refinement(refinement_map(F; transpose = solver.structure == STRUCTURE_GENERAL), T, INT, + solver.backend, max(nrhs, max_rhs(solver.workspace)); n = solver.A.nrows) elseif max_rhs(W) < nrhs W = _grow_refinement(W, nrhs) end @@ -555,6 +562,7 @@ function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::A bm = batch_map(N; nrhs) det = opts.deterministic_mode == 1 cj = solve_conjugated(solver.structure, eltype(ws), solver.A.transposed, opts.solve_mode) + tm = solve_transposed(solver.structure, solver.A.transposed, opts.solve_mode) # LU: solve with Mᵀ 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")) @@ -567,36 +575,36 @@ function _solve_phase!(solver::DirectSolver, p::Phase, X::AbstractVecOrMat, B::A Bs, bt = W.Bc, false end permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj, bm) - forward_sweep!(ws, S, N; nrhs, deterministic = det) + forward_sweep!(ws, S, N; nrhs, deterministic = det, transpose = tm) diagonal_sweep!(ws, S, N; nrhs) - backward_sweep!(ws, S, N; nrhs) + backward_sweep!(ws, S, N; nrhs, transpose = tm) unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj, bm) - refine ? _refine_phase!(solver, W, ws, X, Bs, transposed, bt, cj, det) : (solver.ir_steps = 0) + refine ? _refine_phase!(solver, W, ws, X, Bs, transposed, bt, cj, det, tm) : (solver.ir_steps = 0) elseif p == PHASE_SOLVE_FWD_PERM permute_rhs!(ws.Y, B, S.perm; transposed, conjugate = cj, bm) elseif p == PHASE_SOLVE_FWD - forward_sweep!(ws, S, N; nrhs, deterministic = det) + forward_sweep!(ws, S, N; nrhs, deterministic = det, transpose = tm) elseif p == PHASE_SOLVE_DIAG diagonal_sweep!(ws, S, N; nrhs) elseif p == PHASE_SOLVE_BWD - backward_sweep!(ws, S, N; nrhs) + backward_sweep!(ws, S, N; nrhs, transpose = tm) elseif p == PHASE_SOLVE_BWD_PERM unpermute_solution!(X, ws.Y, S.perm; transposed, conjugate = cj, bm) else # PHASE_SOLVE_REFINEMENT - refine ? _refine_phase!(solver, W, ws, X, B, transposed, transposed, cj, det) : (solver.ir_steps = 0) + refine ? _refine_phase!(solver, W, ws, X, B, transposed, transposed, cj, det, tm) : (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) + cj::Bool, det::Bool, tm::Bool) opts = solver.options done = Ref(0) # the corrections applied, also when the refinement is interrupted try 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) + interrupt = opts.user_host_interrupt, progress = done, transpose_matrix = tm) finally solver.ir_steps = done[] end @@ -729,10 +737,10 @@ The data parameters computed by the solver: | `"nsuperpanels"` | `Int`: supernodes after amalgamation | analysis | | `"memory_estimates"` | `Vector{Int64}` (16 entries, see [`memory_estimates`](@ref)) | analysis | | `"perm_reorder_row"`, `"perm_reorder_col"` | `Vector{Int}`: the fill-reducing permutation, 1-based (`perm[k]` = original index of the `k`-th pivot) | reordering | -| `"perm_row"`, `"perm_col"` | `Vector{Int}`: the final permutation of the factor (= the reordering for Cholesky) | analysis | -| `"diag"` | vector of `T` on the solver's backend: the diagonal of `L` (Cholesky) or of `D` (LDLᵀ/LDLᴴ; for a 2×2 block its two diagonal entries) in factor order | factorization | -| `"npivots"` | `INT`: perturbed pivots (LDLᵀ/LDLᴴ; `0` for Cholesky) | factorization | -| `"inertia"` | `Tuple{INT, INT}`: `(npos, nneg)` of D (after perturbation, so the inertia of `A + E`; read it with `"npivots"`); `(0, 0)` for complex symmetric `"S"`; Cholesky: `(number of positive pivots, 0)` | factorization | +| `"perm_row"`, `"perm_col"` | `Vector{Int}`: the final permutation of the factor (= the reordering for Cholesky and LDLᵀ/LDLᴴ); LU (`"G"`) after a factorization: `perm_row[k]` is the original row of factor row `k` (the reordering composed with the local row interchanges of batch member 1), `perm_col` the reordering, so `A[perm_row, perm_col] = L D U` | analysis | +| `"diag"` | vector of `T` on the solver's backend: the diagonal of `L` (Cholesky), of `D` (LDLᵀ/LDLᴴ; for a 2×2 block its two diagonal entries) or of `D` in `L D U` (LU, = the diagonal of `U` in `L U`) in factor order | factorization | +| `"npivots"` | `INT`: perturbed pivots (LDLᵀ/LDLᴴ, LU; `0` for Cholesky) | factorization | +| `"inertia"` | `Tuple{INT, INT}`: `(npos, nneg)` of D (after perturbation, so the inertia of `A + E`; read it with `"npivots"`); `(0, 0)` for complex symmetric `"S"` and for `"G"`; Cholesky: `(number of positive pivots, 0)` | factorization | | `"pivot_stats"` | `NamedTuple` `(npos, nneg, nzero, nperturbed, n2x2)` of `Int64` (PLAN §1.7) | factorization | Uniform batch (`nbatch > 1`): `"info"`, `"npivots"`, `"inertia"` and @@ -775,6 +783,10 @@ function getparam(solver::DirectSolver, name::AbstractString) name == "flops" && return sp.flops name == "nsuperpanels" && return nsuperpanels(sp) name == "memory_estimates" && return _memory_estimates(solver) + if name == "perm_row" && solver.structure == STRUCTURE_GENERAL && solver.stage >= STAGE_FACTORIZED + piv = Array(view(solver.numeric.piv, 1:Sh.n)) # local row order of member 1 + return sp.perm[piv] + end return copy(sp.perm) # perm_row, perm_col end @@ -824,7 +836,7 @@ end # per member for a batch function _pivot_output(solver::DirectSolver{T, INT}, name) where {T, INT} S, N = solver.symbolic, solver.numeric - if !_is_ldlt_structure(S.structure) + if !_is_ldlt_structure(S.structure) && S.structure != STRUCTURE_GENERAL _set_members!(solver) reduce_stats!(N, S) end @@ -834,13 +846,13 @@ function _pivot_output(solver::DirectSolver{T, INT}, name) where {T, INT} return [out(pivot_totals(totals, solver.nbatch, k)) for k in 1:solver.nbatch] end -# diagonal of L (Cholesky) or D (LDLᵀ/LDLᴴ) in factor order, on the solver's backend, the members of a batch one +# diagonal of L (Cholesky) or D (LDLᵀ/LDLᴴ, LU) in factor order, on the solver's backend, the members of a batch one # after the other (Cholesky: one launch, one work item per (supernode, member)) function _factor_diag(solver::DirectSolver{T}) where {T} S, N = solver.symbolic, solver.numeric n, nb = S.n, solver.nbatch d = KernelAbstractions.zeros(solver.backend, T, n * nb) - if _is_ldlt_structure(S.structure) + if _is_ldlt_structure(S.structure) || S.structure == STRUCTURE_GENERAL for k in 1:nb copyto!(d, (k - 1) * n + 1, N.d, (k - 1) * 2n + 1, n) end diff --git a/src/symbolic/maps.jl b/src/symbolic/maps.jl index c79e71f..b4a989c 100644 --- a/src/symbolic/maps.jl +++ b/src/symbolic/maps.jl @@ -45,8 +45,9 @@ Device maps (`DEVICE_MAPS`; supernodal numbering, 1-based, element offsets): * `amap` (length `nnz`): for every entry `p` of the user's `nzval`, the factor-buffer offset it is added to; `0` when the view ignores the entry, and `-offset` when the conjugate is added (a `"H"`/`"HPD"` entry that lands in the - upper triangle of `P A Pᵀ` and is mirrored). Duplicated entries map to the same - offset and are summed; + upper triangle of `P A Pᵀ` and is mirrored; for `"G"`: an upper-triangle entry + that goes to the `Uᵀ` structure). Duplicated entries map to the same offset and + are summed; * `amap_ptr`, `amap_src`: the same map grouped by supernode for owner-pull assembly: supernode `s` receives the entries `amap_src[amap_ptr[s]:(amap_ptr[s+1]-1)]` (positions in `nzval`, sorted by destination, then by position); @@ -137,19 +138,19 @@ The `amap` of [`Symbolic`](@ref): for every stored entry `p` of the user's CSR pattern (`rowptr`, `colval`, base `index`), its offset in the factor buffer of `layout`, `-offset` when its conjugate is stored instead, `0` when `view` ignores it. Entry `(r, c)` lands at `(i, j) = (iperm[r], iperm[c])` of `P A Pᵀ`, or at -`(j, i)` when `i < j`. Only symmetric structures are supported (`"G"` needs -U panels, T19). +`(j, i)` when `i < j`. Structure `"G"` (LU, two lower-triangular structures +`L` and `Uᵀ` with the same layout, see `src/numeric/lu.jl`): an entry with +`i ≥ j` goes to offset `off` of the `L` structure, an entry with `i < j` to the +offset of `(j, i)` in the `Uᵀ` structure, encoded as `-off`. """ function assembly_map(sp::SupernodePartition, layout::Layout, rowptr::AbstractVector{<:Integer}, colval::AbstractVector{<:Integer}, n::Integer, structure; view = VIEW_FULL, index = INDEX_ONE) s_ = _structure(structure) v = _matrix_view(view) - s_ == STRUCTURE_GENERAL && - throw(NotSupportedError("assembly maps for structure \"G\" (LU panels) are not implemented yet (T19)")) n == sp.n || throw(InvalidValueError("matrix size $n does not match the analysis size $(sp.n)")) rp, cv = _host_pattern(rowptr, colval, n, index) keep = _entry_filter(s_, v) - herm = _is_hermitian(s_) + neg = _is_hermitian(s_) || s_ == STRUCTURE_GENERAL # mirrored entries: conjugate, or the Uᵀ structure amap = zeros(Int, length(cv)) for r in 1:n, p in rp[r]:(rp[r + 1] - 1) c = cv[p] @@ -164,16 +165,17 @@ function assembly_map(sp::SupernodePartition, layout::Layout, rowptr::AbstractVe throw(InvalidValueError("entry ($r, $c) is outside the symbolic factor (pattern changed?)")) f = length(rows) off = layout.panel_ptr[s] + (j - sp.super_ptr[s]) * f + pos - 1 - amap[p] = flip && herm ? -off : off + amap[p] = flip && neg ? -off : off end return amap end -# owner-pull grouping of amap by supernode -function _group_amap(amap::Vector{Int}, layout::Layout, ns::Int) +# owner-pull grouping of amap by supernode; for "G" the L entries of a destination come before its Uᵀ +# entries (`signed`), so every run of equal offsets is one destination of one structure +function _group_amap(amap::Vector{Int}, layout::Layout, ns::Int, signed::Bool = false) owner(off) = searchsortedlast(layout.panel_ptr, off) order = [p for p in eachindex(amap) if amap[p] != 0] - sort!(order; by = p -> (abs(amap[p]), p)) + sort!(order; by = p -> signed ? (abs(amap[p]), amap[p] < 0, p) : (abs(amap[p]), false, p)) amap_ptr = zeros(Int, ns + 1) amap_ptr[1] = 1 for p in order @@ -194,7 +196,7 @@ function Symbolic(sp::SupernodePartition, sc::Schedule, layout::Layout, rowptr:: colval::AbstractVector{<:Integer}, n::Integer, structure; view = VIEW_FULL, index = INDEX_ONE) ns = nsupernodes(sp) amap = assembly_map(sp, layout, rowptr, colval, n, structure; view, index) - amap_ptr, amap_src = _group_amap(amap, layout, ns) + amap_ptr, amap_src = _group_amap(amap, layout, ns, _structure(structure) == STRUCTURE_GENERAL) child_ptr = zeros(Int, ns + 1) child_ptr[1] = 1 for s in 1:ns @@ -245,7 +247,8 @@ function symbolic_analysis(A::CSR, structure, view = VIEW_FULL; opts::Options = pp = analysis_pairs(P, rowptr, colval, A.nzval, A.nrows, structure, opts; view, index = A.index) ord = compute_ordering(P, opts; T, pp.pairs, pp.candidates) sp = supernode_partition(factor_pattern(P, ord), ord.perm, opts) - sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(structure)) + sc = build_schedule(sp, opts, T; reserve = subtree_local_reserve(structure), + elsize = schedule_elsize(structure, T)) layout = build_layout(sp, sc; ldlt = _is_ldlt_structure(_structure(structure))) return Symbolic(sp, sc, layout, rowptr, colval, A.nrows, structure; view, index = A.index) end @@ -301,9 +304,9 @@ element type `T` and device indices `INT` (default: the index type of | 4 | peak host memory (= slot 3) | | 5 | minimum device memory of the hybrid memory mode (= slot 2 until M12) | | 6 | maximum host memory of the hybrid memory mode (0 until M12) | -| 7 | factor panels (`Layout.factor_len` entries of `T`) | +| 7 | factor panels (`Layout.factor_len` entries of `T`; twice that for `"G"`: `L` and `Uᵀ`) | | 8 | D (`2n` entries of `T`) | -| 9 | update stack (`Layout.stack_len` entries of `T`) | +| 9 | update stack (`Layout.stack_len` entries of `T`; twice that for `"G"`) | | 10 | device maps (`INT`) | | 11 | per-front statistics (6 `Int64` per supernode) and their totals (6 `Int64`), status (`ns + 1` `Int32`), pivot order (`n` `Int32`), pivot kinds and sign requests (`n` `Int8` each), `aux` (one `T`) | | 12 | largest regime-A local memory in use (per workgroup, [`subtree_local_bytes`](@ref) of its class) | @@ -315,9 +318,10 @@ function memory_estimates(S::Symbolic{INT0}, ::Type{T}, ::Type{INT} = INT0) wher est = zeros(Int64, 16) L = S.layout sc = S.schedule - est[7] = Int64(L.factor_len) * sizeof(T) + two = S.structure == STRUCTURE_GENERAL ? 2 : 1 # LU: the L and Uᵀ structures + est[7] = two * Int64(L.factor_len) * sizeof(T) est[8] = Int64(L.d_len) * sizeof(T) - est[9] = Int64(L.stack_len) * sizeof(T) + est[9] = two * Int64(L.stack_len) * sizeof(T) est[10] = device_map_bytes(S, INT) est[11] = Int64(nsupernodes(S) + 1) * FRONT_STATS_FIELDS * sizeof(Int64) + Int64(nsupernodes(S) + 1) * sizeof(Int32) + Int64(S.n) * (sizeof(Int32) + 2 * sizeof(Int8)) + sizeof(T) diff --git a/src/symbolic/schedule.jl b/src/symbolic/schedule.jl index 7439635..a3bf29c 100644 --- a/src/symbolic/schedule.jl +++ b/src/symbolic/schedule.jl @@ -147,7 +147,7 @@ Regime assignment and launch order of the supernodes of a * `groups` and `group_nodes`: the launch groups in execution order (regime A classes first, then step by step: B bins in increasing bin id, then C); * `wclasses`, `fclasses`, `budgets`, `regime_c_width`, `regime_c_rows`, - `memory_budget`, `elsize = sizeof(T)`, `vendor_c` (whether regime C uses vendor + `memory_budget`, `elsize` (bytes per packed entry, `sizeof(T)`; `2 sizeof(T)` for LU), `vendor_c` (whether regime C uses vendor calls; `factorization_alg = "algo1"` turns them off). """ struct Schedule @@ -261,12 +261,23 @@ function front_flops(f::Integer, w::Integer) return sq(f) - sq(f - w) end +""" + schedule_elsize(structure, T) -> Int + +Bytes per packed entry of the schedule's local-memory and update-stack +budgets: `sizeof(T)`, or `2 sizeof(T)` for structure `"G"` (LU fronts and +contribution blocks are two packed triangles, `L` and `Uᵀ`). +""" +schedule_elsize(structure, ::Type{T}) where {T} = (_structure(structure) == STRUCTURE_GENERAL ? 2 : 1) * sizeof(T) + """ build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64; - reserve = SUBTREE_LOCAL_RESERVE) -> Schedule + elsize = sizeof(T), reserve = SUBTREE_LOCAL_RESERVE) -> Schedule Regime assignment, binning, levels and launch groups (PLAN §2.3 step 5) for -the supernodes of `sp`, with byte budgets for the element type `T`: +the supernodes of `sp`, with byte budgets for the element type `T` (`elsize` +bytes per packed entry: `2 sizeof(T)` for LU, whose fronts and contribution +blocks are two packed triangles, see `src/numeric/lu.jl`): 1. regime C: fronts with `w > opts.regime_c_width` or `f > opts.regime_c_rows` (every non-A front with `factorization_alg = "algo2"`); @@ -290,9 +301,9 @@ the supernodes of `sp`, with byte budgets for the element type `T`: than the budget gets its own chunk; a negative budget means no chunking). """ function build_schedule(sp::SupernodePartition, opts::Options = Options(), ::Type{T} = Float64; - reserve::Integer = SUBTREE_LOCAL_RESERVE) where {T} + elsize::Integer = sizeof(T), reserve::Integer = SUBTREE_LOCAL_RESERVE) where {T} ns = nsupernodes(sp) - elsize = sizeof(T) + elsize = Int(elsize) cw, cr = opts.regime_c_width, opts.regime_c_rows budgets = sort(opts.subtree_budgets) capacity = [subtree_capacity(b, elsize, reserve) for b in budgets] diff --git a/test/matrices.jl b/test/matrices.jl index 564ca79..59442ed 100644 --- a/test/matrices.jl +++ b/test/matrices.jl @@ -197,6 +197,22 @@ function random_general(::Type{T}, n::Integer, density::Real; rng::AbstractRNG = end random_general(n::Integer, density::Real; kwargs...) = random_general(Float64, n, density; kwargs...) +""" + weak_diagonal_general(T, n, density; every = 5, scale = 0.1, rng) + +[`random_general`](@ref) with every `every`-th diagonal entry multiplied by +`scale`: those rows are no longer diagonally dominant, so the in-front row +pivoting of LU (T19) interchanges rows where the fully-summed block allows it. +""" +function weak_diagonal_general(::Type{T}, n::Integer, density::Real; every::Integer = 5, scale::Real = 0.1, + rng::AbstractRNG = Random.default_rng()) where {T} + A = random_general(T, n, density; rng) + for i in 1:every:n + A[i, i] *= scale + end + return A +end + """ singular_block_matrix(T, n, j; stored_zero = false) diff --git a/test/ported/cudss_execution.jl b/test/ported/cudss_execution.jl index 8735b08..8a1ecd5 100644 --- a/test/ported/cudss_execution.jl +++ b/test/ported/cudss_execution.jl @@ -1,5 +1,5 @@ -# Port of `cudss_execution()`, testsets "SPD -- HPD" (T13) and "Symmetric -- Hermitian" -# (T15) (CUDSS.jl test/test_cudss.jl). +# Port of `cudss_execution()`, testsets "SPD -- HPD" (T13), "Symmetric -- Hermitian" +# (T15) and "Unsymmetric -- Non-Hermitian" (T19) (CUDSS.jl test/test_cudss.jl). # # Changes: global pivoting `pivot_type = 'C'/'R'` (with `reordering_alg = # "algo2"` in CUDSS.jl) is not planned (PLAN §3.3): setting it raises @@ -7,6 +7,9 @@ # "Symmetric -- Hermitian": the matrix is `random_symindef(T, n, 0.01)` # (symmetric/Hermitian indefinite) instead of `sprand + I` symmetrized, and # CUDSS.jl's shift `A + Diagonal(d)` with `d = rand(R, n)` is kept. +# "Unsymmetric -- Non-Hermitian": the matrix is `random_general(T, n, 0.02)` +# (diagonally dominant) instead of `sprand(T, n, n, 0.02) + I`; the shift +# `A + Diagonal(d)` with `d = rand(T, n)` is kept. function ported_execution_spd(backend, ::Type{T}, ::Type{INT}, view, pivot) where {T, INT} n, p = 100, 5 @@ -90,6 +93,51 @@ function ported_execution_sym(backend, ::Type{T}, ::Type{INT}, view, pivot) wher return nothing end +function ported_execution_lu(backend, ::Type{T}, ::Type{INT}, pivot) where {T, INT} + n = 100 + A_cpu = random_general(T, n, 0.02) + x_cpu = zeros(T, n) + b_cpu = rand(T, n) + + A_gpu = api_matrix(backend, A_cpu, INT) + x_gpu = to_device(backend, x_cpu) + b_gpu = to_device(backend, b_cpu) + + solver = DirectSolver(A_gpu, "G", 'F') + if pivot in ('C', 'R') + @test_throws NotSupportedError setparam!(solver, "pivot_type", pivot) + else + setparam!(solver, "pivot_type", pivot) + @test getparam(solver, "pivot_type") == pivot + end + + execute!("analysis", solver, x_gpu, b_gpu) + execute!("factorization", solver, x_gpu, b_gpu) + execute!("solve", solver, x_gpu, b_gpu) + + @test relres(A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + + # In-place LU + A_cpu2 = A_cpu + Diagonal(rand(T, n)) + update!(solver, api_matrix(backend, A_cpu2, INT)) + + c_cpu = rand(T, n) + c_gpu = to_device(backend, c_cpu) + + execute!("refactorization", solver, x_gpu, c_gpu) + execute!("solve", solver, x_gpu, c_gpu) + + @test relres(A_cpu2, to_host(x_gpu), c_cpu) <= tol(T) + return nothing +end + +@testset "Unsymmetric -- Non-Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + @testset "Pivoting = $pivot" for pivot in ('C', 'R', 'N') + ported_execution_lu(backend, T, INT, pivot) + end +end + @testset "Symmetric -- Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES @testset "view = $view" for view in ('F', 'L', 'U') diff --git a/test/ported/cudss_generic.jl b/test/ported/cudss_generic.jl index e723139..ef7dd81 100644 --- a/test/ported/cudss_generic.jl +++ b/test/ported/cudss_generic.jl @@ -1,9 +1,72 @@ -# Port of `cudss_generic()`, testsets "SPD -- HPD" (T13) and "Symmetric -- -# Hermitian" (`ldlt`, `ldlt!`, T15) (CUDSS.jl test/test_cudss.jl). +# Port of `cudss_generic()`, testsets "SPD -- HPD" (T13), "Symmetric -- +# Hermitian" (`ldlt`, `ldlt!`, T15) and "Unsymmetric -- Non-Hermitian" (`lu`, +# `lu!`, T19) (CUDSS.jl test/test_cudss.jl). # # Changes: on the CPU backend the matrix is a `CSR` (`cholesky(::SparseMatrixCSC)` # is CHOLMOD's); `CudssMatrix(x)` is `MatrixDescriptor(x)`. `rand(R)` scalings -# are shifted to `[1/2, 3/2)` so no scaled matrix is nearly zero. +# are shifted to `[1/2, 3/2)` so no scaled matrix is nearly zero (the complex +# `rand(T)` scalings of the LU part likewise, by `1/2` on the real part). The +# unsymmetric matrix is `random_general(T, n, 0.02)` instead of +# `sprand(T, n, n, 0.02) + I`. + +function ported_generic_lu(backend, ::Type{T}, ::Type{INT}) where {T, INT} + n = 100 + R = real(T) + A_cpu = random_general(T, n, 0.02) + b_cpu = rand(T, n) + x_cpu = zeros(T, n) + + A_gpu = api_matrix(backend, A_cpu, INT) + b_gpu = to_device(backend, b_cpu) + scaled(c) = api_matrix(backend, c * A_cpu, INT) + + @testset "lu!" begin + solver = DirectSolver(A_gpu, "G", 'F') + x_gpu = to_device(backend, x_cpu) + execute!("analysis", solver, x_gpu, b_gpu) + solver = lu!(solver, A_gpu) + @test !solver.fresh_factorization + end + + @testset "ldiv!" begin + solver = lu(A_gpu) + x_gpu = to_device(backend, x_cpu) + ldiv!(x_gpu, solver, b_gpu) + @test relres(A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + + c = rand(T) + R(0.5) + lu!(solver, scaled(c)) + x_gpu .= b_gpu + ldiv!(solver, x_gpu) + @test relres(c * A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + + c = rand(T) + R(0.5) + lu!(solver, scaled(c)) + x_gpu .= b_gpu + x_desc = MatrixDescriptor(x_gpu) + ldiv!(solver, x_desc) + @test relres(c * A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + + c = rand(T) + R(0.5) + lu!(solver, scaled(c)) + x_desc = MatrixDescriptor(x_gpu) + b_desc = MatrixDescriptor(b_gpu) + ldiv!(x_desc, solver, b_desc) + @test relres(c * A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + end + + @testset "\\" begin + solver = lu(A_gpu) + x_gpu = solver \ b_gpu + @test relres(A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + + c = rand(T) + R(0.5) + lu!(solver, scaled(c)) + x_gpu = solver \ b_gpu + @test relres(c * A_cpu, to_host(x_gpu), b_cpu) <= tol(T) + end + return nothing +end function ported_generic_spd(backend, ::Type{T}, ::Type{INT}, view) where {T, INT} n, p = 100, 5 @@ -125,6 +188,11 @@ function ported_generic_sym(backend, ::Type{T}, ::Type{INT}, view) where {T, INT return nothing end +@testset "Unsymmetric -- Non-Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + ported_generic_lu(backend, T, INT) +end + @testset "Symmetric -- Hermitian ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES @testset "view = $view" for view in ('F', 'L', 'U') diff --git a/test/ported/cudss_solver.jl b/test/ported/cudss_solver.jl index bf045f2..0c3bf0e 100644 --- a/test/ported/cudss_solver.jl +++ b/test/ported/cudss_solver.jl @@ -3,10 +3,13 @@ # loop over every data and configuration parameter. # # Changes: -# * structures "S", "H", "SPD" and "HPD" ("G" is T19; complex `T`: "SPD" is -# checked to be rejected at analysis); the matrix is SPD/HPD, so `"inertia"` -# is `(n, 0)` for real "S" and "H" (complex symmetric "S": `(0, 0)`, no -# inertia) and `"diag"` is real positive except for complex "S"; +# * structures "G", "S", "H", "SPD" and "HPD" (complex `T`: "SPD" is checked to +# be rejected at analysis); the matrix is SPD/HPD, so `"inertia"` is `(n, 0)` +# for real "S" and "H" (complex symmetric "S": `(0, 0)`, no inertia) and +# `"diag"` is real positive except for complex "S" (and has a positive real +# part for "G", whose complex `D` is not exactly real); "G" reads the full +# matrix: views 'L'/'U' are rejected at analysis (`InvalidValueError`; cuDSS +# ignores the view of a general matrix); # * `matching_alg = "algo6"` before the analysis is skipped: matching is T21 and # the analysis raises `NotSupportedError` for it (checked); # * cuDSS's set-buffer-then-get protocol for vector data is @@ -55,7 +58,11 @@ function ported_solver_data_parameters(backend, solver, structure, n, memory_est if parameter == "diag" getparam!(buffer_T, solver, parameter) csym = structure == "S" && eltype(buffer_T) <: Complex # D of a complex symmetric matrix - csym || @test all(x -> real(x) > 0 && imag(x) == 0, buffer_T) + if structure == "G" + @test all(x -> real(x) > 0, buffer_T) + else + csym || @test all(x -> real(x) > 0 && imag(x) == 0, buffer_T) + end end if parameter == "memory_estimates" getparam!(memory_estimates, solver, parameter) @@ -169,9 +176,15 @@ end INT in INTTYPES n = 20 A_cpu = random_spd(T, n, 1.0) - @testset "structure = $structure" for structure in (T <: Real ? ("S", "H", "SPD", "HPD") : ("S", "H", "HPD")) + @testset "structure = $structure" for structure in (T <: Real ? ("G", "S", "H", "SPD", "HPD") : + ("G", "S", "H", "HPD")) @testset "view = $view" for view in ('L', 'U', 'F') A_gpu = api_matrix(backend, triangle_view(A_cpu, view), INT) + if structure == "G" && view != 'F' + @test_throws InvalidValueError execute!("analysis", DirectSolver(A_gpu, structure, view), nothing, + nothing) + continue + end ported_solver(backend, T, INT, A_gpu, structure, view, n) end end diff --git a/test/ported/cudss_uniform_batch.jl b/test/ported/cudss_uniform_batch.jl index dbecc81..f525918 100644 --- a/test/ported/cudss_uniform_batch.jl +++ b/test/ported/cudss_uniform_batch.jl @@ -1,11 +1,11 @@ -# Port of CUDSS.jl test/test_uniform_batch_cudss.jl (T17): `uniform_batch_ldlt()` -# and `uniform_batch_cholesky()`, each with the cuDSS API (`DirectSolver` + +# Port of CUDSS.jl test/test_uniform_batch_cudss.jl (T17, T19): `uniform_batch_lu()`, +# `uniform_batch_ldlt()` and `uniform_batch_cholesky()`, each with the cuDSS API (`DirectSolver` + # `"ubatch_size"` + `MatrixDescriptor(T, n[, nrhs]; nbatch)`) and the generic # API (`ldlt`/`cholesky` on a matrix whose values are longer than its pattern), # strided and non-strided storage, views 'L', 'U', 'F'. # -# Changes: `uniform_batch_lu()` is skipped (structure "G" is T19). The data are -# CUDSS.jl's, written once as `nnz × nbatch` matrices (the strided storage is +# Changes: the data are CUDSS.jl's (the LU batch is `ubatch_example` of +# `test/matrices.jl`, the 3×3 example of the docs), written once as `nnz × nbatch` matrices (the strided storage is # their `vec`, exactly CUDSS.jl's strided vectors). Matrices come from # `api_csr` (a `CSR` on the CPU backend); CUDSS.jl's `As_gpu.nzVal = …` is a new # matrix on the same pattern. Residuals are checked per member with @@ -122,6 +122,60 @@ function ported_ubatch(backend, ::Type{T}, ::Type{INT}, uplo, generic::Bool, str return nothing end +# `uniform_batch_lu()`: A(λ) = [1+λ 0 3; 4 5+λ 0; 2 6 2+λ], λ ∈ (1, 10, -20), then λ ∈ (-2, -10, 30) +function ported_ubatch_lu(backend, ::Type{T}, ::Type{INT}, generic::Bool, strided::Bool) where {T, INT} + ex = ubatch_example(T) + n, nbatch = ex.n, ex.nbatch + rowptr, colval = Vector{INT}(ex.rowptr), Vector{INT}(ex.colval) + nz = reshape(ex.nzval, :, nbatch) + new_nz = reshape(T[v for λ in (-2, -10, 30) for v in (1 + λ, 3, 4, 5 + λ, 2, 6, 2 + λ)], :, nbatch) + members(z) = [host_csc(rowptr, colval, z[:, k], n, n) for k in 1:nbatch] + B = reshape(ex.b, n, nbatch) + shape = strided ? (n * nbatch,) : (n, nbatch) + nz_dev(x) = to_device(backend, strided ? vec(x) : x) + b_gpu = to_device(backend, reshape(B, shape)) + x_gpu = to_device(backend, zeros(T, shape)) + + if generic + solver = lu(api_csr(backend, rowptr, colval, vec(nz), n)) + @test solver.nbatch == nbatch + ldiv!(x_gpu, solver, b_gpu) + else + solver = DirectSolver(to_device(backend, rowptr), to_device(backend, colval), nz_dev(nz), "G", 'F') + setparam!(solver, "ubatch_size", nbatch) + b_desc = MatrixDescriptor(T, n; nbatch) + x_desc = MatrixDescriptor(T, n; nbatch) + update!(b_desc, b_gpu) + update!(x_desc, x_gpu) + execute!("analysis", solver, x_desc, b_desc) + execute!("factorization", solver, x_desc, b_desc; asynchronous = false) + execute!("solve", solver, x_desc, b_desc; asynchronous = false) + end + @test maximum(batch_relres(members(nz), to_host(x_gpu), B)) <= tol(T) + + # refactorize all matrices of the uniform batch + if generic + lu!(solver, api_csr(backend, rowptr, colval, vec(new_nz), n)) + x_gpu .= b_gpu + ldiv!(solver, x_gpu) + else + update!(solver, to_device(backend, rowptr), to_device(backend, colval), nz_dev(new_nz)) + execute!("refactorization", solver, x_desc, b_desc; asynchronous = false) + execute!("solve", solver, x_desc, b_desc; asynchronous = false) + end + @test maximum(batch_relres(members(new_nz), to_host(x_gpu), B)) <= tol(T) + return nothing +end + +@testset "uniform batch LU ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, + INT in INTTYPES + for generic in (false, true), strided in (false, true) + @testset "$(generic ? "Generic" : "cuDSS") API, strided = $strided" begin + ported_ubatch_lu(backend, T, INT, generic, strided) + end + end +end + for (name, cholesky_case) in (("LDLᵀ and LDLᴴ", false), ("Cholesky", true)) @testset "uniform batch $name ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES diff --git a/test/test_api.jl b/test/test_api.jl index 2fb2ec3..5e75d1a 100644 --- a/test/test_api.jl +++ b/test/test_api.jl @@ -167,8 +167,8 @@ end setparam!(solver, "solve_alg", "algo1") @test thrown(() -> execute!("solve", solver, x, b)) isa NotSupportedError setparam!(solver, "solve_alg", "default") - s3 = DirectSolver(api_matrix(backend, A), "G", 'F') # LU: T19 - @test thrown(() -> execute!("analysis", s3, x, b)) isa NotSupportedError + s3 = DirectSolver(api_matrix(backend, tril(A)), "G", 'L') # LU (T19) reads the full matrix: view 'F' only + @test thrown(() -> execute!("analysis", s3, x, b)) isa InvalidValueError if T <: Complex s4 = DirectSolver(api_matrix(backend, tril(A)), "SPD", 'L') @test thrown(() -> execute!("analysis", s4, x, b)) isa InvalidValueError diff --git a/test/test_numeric_lu.jl b/test/test_numeric_lu.jl new file mode 100644 index 0000000..73ea9e5 --- /dev/null +++ b/test/test_numeric_lu.jl @@ -0,0 +1,339 @@ +# T19: general LU (structure "G"): the CPU reference `ref_lu!` (multifrontal L D U on the pattern of A + Aᵀ, +# row pivoting inside the fully-summed block, static perturbation), the device kernels of regimes A/B/C against +# it, the solves with A and Aᵀ, uniform batches and the handle / LinearAlgebra layers. + +# regime mixes of the analysis (merged into each case's options) +const LU_REGIMES = (("default", (;)), + ("A+B+C, small budgets", (subtree_budgets = [8192, 16384], regime_c_width = 16, + regime_c_rows = 128)), + ("C only", (factorization_alg = "algo2", subtree_budgets = Int[]))) + +# (name, matrix, option keywords, whether the in-block pivoting must interchange rows) +function lu_cases(::Type{T}) where {T} + return (("random_general(400,0.01)", random_general(T, 400, 0.01), (;), false), + ("random_general(300,0.02), natural order", random_general(T, 300, 0.02), (reordering_alg = "algo5",), + false), + ("weak diagonal", weak_diagonal_general(T, 300, 0.02), (;), false), + ("weak diagonal, pivot_threshold = 1", weak_diagonal_general(T, 300, 0.02), (pivot_threshold = 1.0,), + true), + ("weak diagonal, pivot_type 'N'", weak_diagonal_general(T, 300, 0.02), + (pivot_type = 'N', pivot_threshold = 1.0), false)) +end + +@testset "reference LU ($T)" for T in ELTYPES + Random.seed!(666) + for (name, A, kw, swaps) in lu_cases(T) + @testset "$name" begin + opts = Options(; kw...) + S, Nr, _, _, _ = lu_setup(CPU(), A; opts) + n = size(A, 1) + @test lu_error(A, S, Nr) <= tol(T) + # the local row order permutes the columns of each supernode + sp = S.partition + @test all(s -> sort(Int.(Nr.piv[SDS.sncols(sp, s)])) == SDS.sncols(sp, s), 1:SDS.nsupernodes(sp)) + swaps && @test Nr.piv != 1:n + haskey(kw, :pivot_type) && @test Nr.piv == 1:n + @test all(==(SDS.PIVOT_KIND_1X1), Nr.pivot_kind) + @test SDS.pivot_stats(Nr) == (npos = 0, nneg = 0, nzero = 0, nperturbed = 0, n2x2 = 0) + for transpose in (false, true) + b = rand(T, n, 2) + x = similar(b) + SDS.ref_solve_lu!(x, S, Nr, b; transpose) + @test relres(transpose ? Base.transpose(A) : A, x, b) <= tol(T) + end + end + end +end + +@testset "device LU equals the reference ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + for (name, A, kw, _) in lu_cases(T), (rname, rkw) in LU_REGIMES + @testset "$name, $rname" begin + opts = Options(; kw..., rkw...) + S, Nr, Sd, Nd, nz = lu_setup(backend, A; opts) + rname == "C only" && @test all(s -> SDS.takes_c_path(S.schedule, s), 1:SDS.nsupernodes(S)) + @test SDS.factorize!(Nd, Sd, nz; opts) == 0 + Nh = SDS.host_numeric(Nd) + # the same pivot sequence: local row orders and kinds equal, D, L and Uᵀ to rounding + @test Nh.piv == Nr.piv + @test Nh.pivot_kind == Nr.pivot_kind + @test d_error(Nh, Nr) <= growth_tol(T, Nr) + @test panel_error(Nh, Nr) <= growth_tol(T, Nr) + @test upanel_error(Nh, Nr) <= growth_tol(T, Nr) + # the backward error of the device factor: A[p, q] = L D U + @test lu_error(A, S, Nh) <= tol(T) + @test Nh.stats == Nr.stats + @test SDS.pivot_totals(Nd) == SDS.pivot_stats(Nr) && Nh.totals == Nr.totals + # device solves with A (L, D, U) and Aᵀ (Uᵀ, D, Lᵀ) + ws = SDS.allocate_solve(Sd, T, backend, 5) + for nrhs in (1, 5), det in (false, true), transpose in (false, true) + b = nrhs == 1 ? rand(T, size(A, 1)) : rand(T, size(A, 1), nrhs) + x = device_solve(backend, ws, Sd, Nd, b; deterministic = det, transpose) + @test relres(transpose ? Base.transpose(A) : A, x, b) <= tol(T) + end + end + end +end + +@testset "regime-A groups and Int64 maps ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + A = weak_diagonal_general(T, 400, 0.01) + opts = Options(pivot_threshold = 1.0, subtree_budgets = [8192, 16384, 32768, 49152]) + S, Nr, Sd, Nd, nz = lu_setup(backend, A, Int64; opts) + @test !isempty(Nd.plan.sub_first) && Nr.piv != 1:400 + SDS.factorize!(Nd, Sd, nz; opts) + Nh = SDS.host_numeric(Nd) + @test Nh.piv == Nr.piv && Nh.pivot_kind == Nr.pivot_kind + gt = growth_tol(T, Nr) + @test d_error(Nh, Nr) <= gt && panel_error(Nh, Nr) <= gt && upanel_error(Nh, Nr) <= gt + ws = SDS.allocate_solve(Sd, T, backend, 2) + b = rand(T, 400, 2) + @test relres(A, device_solve(backend, ws, Sd, Nd, b), b) <= tol(T) + # the extracted factors: A[p, q] = L D U + @test lu_error(A, S, Nh) <= tol(T) +end + +@testset "perturbation ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + n, j = 60, 23 + R = real(T) + # a zero row and column j: no row of the front can replace the zero pivot, it is perturbed + for (rname, rkw) in (("default", (;)), ("regime-C root", (factorization_alg = "algo2", subtree_budgets = Int[]))), + stored_zero in (false, true) + A = singular_block_matrix(T, n, j; stored_zero) + opts = Options(; rkw...) + S, Nr, Sd, Nd, nz = lu_setup(backend, A; opts) + @test SDS.factorize!(Nd, Sd, nz; opts) == 0 + Nh = SDS.host_numeric(Nd) + @test Nh.pivot_kind == Nr.pivot_kind && Nh.piv == Nr.piv + @test d_error(Nh, Nr) <= panel_tol(T) + st = SDS.pivot_totals(Nd) + @test st.nperturbed == 1 && st.nzero == 1 + k = findfirst(==(SDS.PIVOT_KIND_PERTURBED), Nh.pivot_kind) + _, _, _, p, q = SDS.extract_lu(S, Nh) + @test k !== nothing && p[k] == j && q[k] == j + @test Nh.d[k] == R(SDS.default_pivot_epsilon(R)) + x0 = rand(T, n) + x0[j] = 0 + b = A * x0 + ws = SDS.allocate_solve(Sd, T, backend, 1) + @test relres(A, device_solve(backend, ws, Sd, Nd, b), b) <= 1.0e-6 + @test relres(Base.transpose(A), device_solve(backend, ws, Sd, Nd, b; transpose = true), b) <= 1.0e-6 + end + # pivot_epsilon and the scaled variant (ε max|aᵢⱼ| on the device) + A = singular_block_matrix(T, n, j) + for opts in (Options(pivot_epsilon = 1.0e-3), Options(pivot_epsilon = 1.0e-3, pivot_epsilon_alg = "algo1")) + S, Nr, Sd, Nd, nz = lu_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + Nh = SDS.host_numeric(Nd) + @test Nh.pivot_kind == Nr.pivot_kind && Nh.piv == Nr.piv && d_error(Nh, Nr) <= panel_tol(T) + pert = Nr.pivot_kind .== SDS.PIVOT_KIND_PERTURBED + @test Nh.d[1:n][pert] == Nr.d[1:n][pert] + scale = opts.pivot_epsilon_alg == SDS.PIVOT_EPSILON_SCALED ? maximum(abs, A) : 1 + @test Nh.d[1:n][pert] == [R(R(1.0e-3) * R(scale))] + end + # global pivoting is not planned + @test thrown(() -> Options(pivot_type = 'C')) isa NotSupportedError +end + +@testset "determinism and refactorization ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + A = weak_diagonal_general(T, 300, 0.02) + opts = Options(pivot_threshold = 1.0, subtree_budgets = [8192, 16384], regime_c_width = 16, regime_c_rows = 128) + S, Nr, Sd, Nd, nz = lu_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + factor, ufactor, d, piv = to_host(Nd.factor), to_host(Nd.ufactor), to_host(Nd.d), to_host(Nd.piv) + # new values, then the old ones again: bitwise the same factor + nz2 = to_device(backend, 2 .* to_host(nz)) + SDS.factorize!(Nd, Sd, nz2; opts) + @test to_host(Nd.d) ≈ 2 .* d + SDS.factorize!(Nd, Sd, nz; opts) + @test to_host(Nd.factor) == factor && to_host(Nd.ufactor) == ufactor && to_host(Nd.d) == d && + to_host(Nd.piv) == piv + ws = SDS.allocate_solve(Sd, T, backend, 2) + b = rand(T, 300, 2) + for transpose in (false, true) + @test device_solve(backend, ws, Sd, Nd, b; deterministic = true, transpose) == + device_solve(backend, ws, Sd, Nd, b; deterministic = true, transpose) + end + # no allocation in the numeric phase and the solve (CPU backend) + if backend isa CPU + SDS.factorize!(Nd, Sd, nz; opts) + @test (@allocated SDS.factorize!(Nd, Sd, nz; opts)) <= ldlt_alloc_budget(S) + x = similar(b) + for transpose in (false, true) + SDS.sweep_solve!(x, ws, Sd, Nd, b; transpose) + @test (@allocated SDS.sweep_solve!(x, ws, Sd, Nd, b; transpose)) <= solve_alloc_budget(S, ws; ldlt = true) + end + end +end + +@testset "solve sweeps: one at a time and dense path ($(backend_name(backend)), $T)" for backend in BACKENDS, + T in ELTYPES + Random.seed!(666) + A = weak_diagonal_general(T, 300, 0.02) + for (rname, rkw) in LU_REGIMES + opts = Options(; pivot_threshold = 1.0, rkw...) + S, Nr, Sd, Nd, nz = lu_setup(backend, A; opts) + SDS.factorize!(Nd, Sd, nz; opts) + ws = SDS.allocate_solve(Sd, T, backend, 3) + b = rand(T, 300, 3) + bd = to_device(backend, b) + L, D, U, _, _ = SDS.extract_lu(S, SDS.host_numeric(Nd)) + piv = Int.(to_host(Nd.piv)) + for transpose in (false, true) + # Y = P b; A: L Z = Y[piv] (local row orders), D W = Z, U V = W; Aᵀ: Uᵀ Z = Y, D W = Z, Lᵀ V = W + SDS.permute_rhs!(ws.Y, bd, Sd.perm) + Y = to_host(ws.Y) + SDS.forward_sweep!(ws, Sd, Nd; deterministic = true, transpose) + Z = to_host(ws.Y) + SDS.diagonal_sweep!(ws, Sd, Nd) + W = to_host(ws.Y) + SDS.backward_sweep!(ws, Sd, Nd; transpose) + V = to_host(ws.Y) + xd = similar(bd) + SDS.unpermute_solution!(xd, ws.Y, Sd.perm) + @test to_host(xd) == device_solve(backend, ws, Sd, Nd, b; deterministic = true, transpose) + @test relres(transpose ? Base.transpose(A) : A, to_host(xd), b) <= tol(T) + if transpose + @test norm(Base.transpose(U) * Z - Y) <= tol(T) * norm(Y) + @test norm(Base.transpose(L) * V[piv, :] - W) <= tol(T) * norm(W) + else + @test norm(L * Z - Y[piv, :]) <= tol(T) * norm(Y) + @test norm(U * V - W) <= tol(T) * norm(W) + end + @test norm(Diagonal(D) * W - Z) <= tol(T) * norm(Z) + # every dense implementation of the regime-C path + if rname == "C only" + for impl in SDS.dense_impls(:trsm, backend, T) + impl in SDS.dense_impls(:gemm, backend, T) || continue + x = device_solve(backend, ws, Sd, Nd, b; impl, transpose) + @test relres(transpose ? Base.transpose(A) : A, x, b) <= tol(T) + end + end + end + end +end + +@testset "uniform batch ($(backend_name(backend)), $T)" for backend in BACKENDS, T in ELTYPES + Random.seed!(666) + A = weak_diagonal_general(T, 200, 0.03) + nb, nrhs = 3, 2 + members = batch_members(A, nb) + solver = DirectSolver(api_batch_matrix(backend, members, 'F'), "G", 'F') + setparam!(solver, "pivot_threshold", 1.0) + execute!("analysis", solver, nothing, nothing) + execute!("factorization", solver, nothing, nothing) + @test solver.nbatch == nb + @test getparam(solver, "info") == zeros(Int, nb) + @test getparam(solver, "npivots") == zeros(Int32, nb) + B = rand(T, 200, nrhs, nb) + X = to_device(backend, zeros(T, 200, nrhs, nb)) + execute!("solve", solver, X, to_device(backend, B)) + @test maximum(batch_relres(members, to_host(X), B)) <= tol(T) + setparam!(solver, "solve_mode", 1) + execute!("solve", solver, X, to_device(backend, B)) + @test maximum(batch_relres(map(M -> sparse(Base.transpose(M)), members), to_host(X), B)) <= tol(T) + setparam!(solver, "solve_mode", 0) + # each member equals its own reference factorization + for k in 1:nb + _, Nr, _, _, _ = lu_setup(CPU(), members[k]; opts = solver.options) + Nk = SDS.member_numeric(solver.numeric, solver.symbolic, k) + @test Nk.piv == Nr.piv && d_error(Nk, Nr) <= growth_tol(T, Nr) && upanel_error(Nk, Nr) <= growth_tol(T, Nr) + end + # ubatch_index: only member 2 is refactorized + before = [SDS.member_numeric(solver.numeric, solver.symbolic, k) for k in 1:nb] + setparam!(solver, "ubatch_index", 1) + update!(solver, api_batch_matrix(backend, [2 .* M for M in members], 'F')) + execute!("refactorization", solver, nothing, nothing) + for k in 1:nb + Nk = SDS.member_numeric(solver.numeric, solver.symbolic, k) + if k == 2 + @test Nk.d ≈ 2 .* before[k].d + else + @test Nk.factor == before[k].factor && Nk.ufactor == before[k].ufactor && Nk.d == before[k].d + end + end + setparam!(solver, "ubatch_index", -1) +end + +@testset "public API ($(backend_name(backend)), $T, $INT)" for backend in BACKENDS, T in ELTYPES, INT in INTTYPES + Random.seed!(666) + A = weak_diagonal_general(T, 300, 0.02) + n = size(A, 1) + At = sparse(Base.transpose(A)) + solver = DirectSolver(api_matrix(backend, A, INT), "G", 'F') + setparam!(solver, "pivot_threshold", 1.0) + setparam!(solver, "deterministic_mode", 1) + b = to_device(backend, rand(T, n, 2)) + x, y = similar(b), similar(b) + execute!("analysis", solver, x, b) + @test getparam(solver, "perm_row") == getparam(solver, "perm_col") == getparam(solver, "perm_reorder_row") + execute!("factorization", solver, x, b) + @test getparam(solver, "info") == 0 + # the reference factorization of the same analysis + S = solver.host_symbolic + Nr = SDS.allocate_numeric(S, T) + SDS.ref_lu!(Nr, S, SDS.CSR(A).nzval; opts = solver.options) + @test getparam(solver, "npivots") == 0 && getparam(solver, "npivots") isa INT + @test getparam(solver, "inertia") == (0, 0) + @test getparam(solver, "pivot_stats") == SDS.pivot_stats(Nr) + d = getparam(solver, "diag") + @test norm(to_host(d) - Nr.d[1:n]) <= panel_tol(T) * norm(Nr.d) + # perm_row/perm_col: A[perm_row, perm_col] = L D U + _, _, _, p, q = SDS.extract_lu(S, Nr) + pr, pc = getparam(solver, "perm_row"), getparam(solver, "perm_col") + @test pr == p && pc == q && pr != pc + buf = zeros(INT, n) + getparam!(buf, solver, "perm_row") + @test buf == pr + # solve_mode 0, 1, 2: A, Aᵀ, Aᴴ; the sub-phases compose to "solve" + for (mode, M) in ((0, A), (1, At), (2, sparse(adjoint(A)))) + setparam!(solver, "solve_mode", mode) + execute!("solve", solver, x, b) + @test relres(M, to_host(x), to_host(b)) <= tol(T) + for phase in ("solve_fwd_perm", "solve_fwd", "solve_diag", "solve_bwd", "solve_bwd_perm") + execute!(phase, solver, y, b) + end + @test to_host(y) == to_host(x) + # refinement with the residual of op(A) + setparam!(solver, "ir_n_steps", 2) + execute!("solve", solver, x, b) + @test relres(M, to_host(x), to_host(b)) <= tol(T) + @test getparam(solver, "ir_n_steps") == 2 + if SDS.fgmres_available() + setparam!(solver, "ir_mode", "fgmres") + execute!("solve", solver, x, b) + @test relres(M, to_host(x), to_host(b)) <= tol(T) + setparam!(solver, "ir_mode", "ir") + end + setparam!(solver, "ir_n_steps", 0) + end + setparam!(solver, "solve_mode", 0) + # logabsdet = log|det A| with the sign of det A + la, sg = logabsdet(solver) + la_ref, sg_ref = logabsdet(Matrix(A)) + @test la ≈ la_ref rtol = sqrt(eps(real(T))) + @test sg ≈ sg_ref rtol = sqrt(eps(real(T))) + # CSC input: the stored CSR is Aᵀ; solve_mode 0 solves A, 1 solves Aᵀ + Ct = csr_of_transpose(SparseMatrixCSC{T, INT}(A)) + C = CSR(to_device(backend, Ct.rowptr), to_device(backend, Ct.colval), to_device(backend, Ct.nzval), n, n; + transposed = true) + cs = DirectSolver(C, "G", 'F') + execute!("analysis", cs, nothing, nothing) + execute!("factorization", cs, nothing, nothing) + for (mode, M) in ((0, A), (1, At), (2, sparse(adjoint(A)))) + setparam!(cs, "solve_mode", mode) + setparam!(cs, "ir_n_steps", mode) + execute!("solve", cs, x, b) + @test relres(M, to_host(x), to_host(b)) <= tol(T) + end + # views other than 'F' are refused for "G" + @test thrown(() -> execute!("analysis", DirectSolver(api_matrix(backend, A, INT), "G", 'L'), nothing, nothing)) isa + InvalidValueError + # the LinearAlgebra layer + F = lu(SDS.CSR(SparseMatrixCSC{T, INT}(A))) + @test F.structure == SDS.STRUCTURE_GENERAL && getparam(F, "ir_n_steps") == 2 + backend isa CPU && @test relres(A, F \ to_host(b), to_host(b)) <= tol(T) + backend isa CPU && @test relres(2 * A, lu!(F, SDS.CSR(SparseMatrixCSC{T, INT}(2 * A))) \ to_host(b), to_host(b)) <= + tol(T) +end diff --git a/test/test_symbolic_schedule.jl b/test/test_symbolic_schedule.jl index c847db6..e453c66 100644 --- a/test/test_symbolic_schedule.jl +++ b/test/test_symbolic_schedule.jl @@ -145,10 +145,11 @@ function check_schedule(S::SDS.Symbolic) for T in ELTYPES, INT in INTTYPES est = SDS.memory_estimates(S, T, INT) @test length(est) == 16 && eltype(est) == Int64 - actual = (L.factor_len + L.d_len + L.stack_len) * sizeof(T) + SDS.device_map_bytes(S, INT) + two = S.structure == SDS.STRUCTURE_GENERAL ? 2 : 1 # LU: the L and Uᵀ structures (T19) + actual = (two * L.factor_len + L.d_len + two * L.stack_len) * sizeof(T) + SDS.device_map_bytes(S, INT) @test est[2] >= actual @test est[1] + est[9] == est[2] - @test est[7] == L.factor_len * sizeof(T) && est[9] == L.stack_len * sizeof(T) + @test est[7] == two * L.factor_len * sizeof(T) && est[9] == two * L.stack_len * sizeof(T) @test all(==(0), est[13:16]) end return nothing @@ -281,8 +282,12 @@ end @test all(==(SDS.REGIME_C), S2.schedule.regime) @test SDS.nlaunches(S2.schedule) == (1 + 4 * 4) + (1 + 1) # potrf, trsm, syrk, pack_add! check_schedule(S2) - # "G" maps are not implemented yet - @test thrown(() -> SDS.symbolic_analysis(SDS.CSR(A), "G", 'F')) isa NotSupportedError + # "G" (LU, T19): the same tree, every packed entry twice (L and Uᵀ), so the regime-A peak doubles + SG = SDS.symbolic_analysis(SDS.CSR(A), "G", 'F'; opts) + @test SG.schedule.elsize == 2 * sizeof(Float64) + @test SG.schedule.subtree_peak == 2 .* sc.subtree_peak + @test count(<(0), SG.amap) == (nnz(A) - size(A, 1)) ÷ 2 # one of each symmetric pair goes to Uᵀ + check_schedule(SG) end @testset "reconstruction through amap: $T $structure '$view' index '$index'" for T in ELTYPES, diff --git a/test/utils.jl b/test/utils.jl index feefe1f..9b441cf 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -160,10 +160,37 @@ end panel_tol(T) Elementwise tolerance of device panels against the reference panels, relative -to `max |L|`: `100·eps(real(T))` (TASKS.md T09). +to `max |L|`: `100·eps(real(T))` (TASKS.md T09). This is an exact-order +comparison; use [`growth_tol`](@ref) on generators with element growth (weak +pivots, static pivoting). """ panel_tol(::Type{T}) where {T} = 100 * eps(real(T)) +""" + factor_growth(Nr) + +Element growth `max(max |L|, max |Uᵀ|) / min |d[1:n]|` of the reference factor +`Nr` (`ufactor` is empty for the symmetric structures, then only `L` counts). +""" +function factor_growth(Nr) + n = length(Nr.piv) + big = max(maximum(abs, Nr.factor; init = zero(real(eltype(Nr.factor)))), + maximum(abs, Nr.ufactor; init = zero(real(eltype(Nr.ufactor))))) + return big / minimum(abs, @view Nr.d[1:n]) +end + +""" + growth_tol(T, Nr) + +Elementwise tolerance of device LU/LDLᵀ panels against the reference panels on +matrices with element growth: [`panel_tol`](@ref) scaled by the growth +[`factor_growth`](@ref) of the reference factor, at least `panel_tol(T)`. Two +roundings of the same elimination (the device may fuse multiply-adds, the host +does not) differ by O(eps · growth), not O(eps). The pivot sequence is still +compared exactly. +""" +growth_tol(::Type{T}, Nr) where {T} = panel_tol(T) * max(one(real(T)), real(T)(factor_growth(Nr))) + """ ka_cpu_alloc_budget(launches, localmem_bytes) -> Int @@ -517,3 +544,41 @@ function ir_solve(backend, solver, b; steps = nothing, tol = nothing) execute!("solve", solver, xd, bd; asynchronous = false) return to_host(xd) end + +# --- LU (T19) --- + +""" + lu_setup(backend, A, INT = Int32; opts = Options()) -> (S, Nr, Sd, Nd, nz) + +Host analysis `S` of the general matrix `A` (structure `"G"`, view `'F'`) with +its reference LU factor `Nr` (`ref_lu!`), the analysis adapted to `backend` +with `INT` maps, device storage `Nd` and the values `nz` on `backend` (T19). +""" +function lu_setup(backend, A::SparseMatrixCSC{T}, ::Type{INT} = Int32; opts = Options()) where {T, INT} + C = SparseDirectSolver.CSR(A) + S = SparseDirectSolver.symbolic_analysis(C, "G", 'F'; opts) + Nr = SparseDirectSolver.allocate_numeric(S, T) + SparseDirectSolver.ref_lu!(Nr, S, C.nzval; opts) == 0 || error("lu_setup: the reference factorization failed") + Sd = SparseDirectSolver.adapt(backend, S, INT) + Nd = SparseDirectSolver.allocate_numeric(Sd, T, backend) + return S, Nr, Sd, Nd, to_device(backend, C.nzval) +end + +""" + upanel_error(Nh, Nr) + +`max |Uᵀ_device - Uᵀ_ref| / max |Uᵀ_ref|` of the `Uᵀ` panels (`ufactor`) of a +host copy `Nh` of a device LU factor and the reference `Nr`; compare with +[`panel_tol`](@ref) (T19). +""" +upanel_error(Nh, Nr) = maximum(abs, Nh.ufactor - Nr.ufactor) / maximum(abs, Nr.ufactor) + +""" + lu_error(A, S, N) + +`‖A[p, q] - L D U‖_F / ‖A‖_F` of the LU factor in `N` (`extract_lu`, T19). +""" +function lu_error(A, S, N) + L, D, U, p, q = SparseDirectSolver.extract_lu(S, N) + return norm(A[p, q] - L * Diagonal(D) * U) / norm(A) +end