From e6d84470cd3fa20206dde6e53d45c1f7ceaf5f61 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olav=20M=C3=B8yner?= Date: Thu, 1 Oct 2026 13:13:29 +0200 Subject: [PATCH 1/3] Pivoting version of sparse lu for KA --- src/KAPreconditioners/sparse_lu.jl | 162 ++++++++++++++++++++++++++--- test/ka_preconditioners.jl | 32 ++++++ 2 files changed, 179 insertions(+), 15 deletions(-) diff --git a/src/KAPreconditioners/sparse_lu.jl b/src/KAPreconditioners/sparse_lu.jl index 3cbff0d3..13ac1ded 100644 --- a/src/KAPreconditioners/sparse_lu.jl +++ b/src/KAPreconditioners/sparse_lu.jl @@ -1,8 +1,9 @@ function setup_preferred_sparse_lu end """ -The portable sparse LU uses a fixed symbolic fill pattern and no pivoting. It is -intended for eliminated blocks that admit LU without pivoting. +The portable sparse LU accepts a fixed CSR pattern. With pivoting enabled it +uses a dense factor buffer so row swaps can introduce arbitrary fill. Set +`pivoting=false` to use the sparse symbolic fill pattern without row swaps. """ struct KASparseLUFactor{V, I, B, R, C} rowptr::I @@ -16,6 +17,9 @@ struct KASparseLUFactor{V, I, B, R, C} pattern_colval::Vector{Int} source_rowptr::R source_colval::C + pivots::I + status::I + pivoting::Bool end function sparse_lu_pattern(A::StaticSparsityMatrixCSR) @@ -63,10 +67,12 @@ function sparse_lu_pattern(A::StaticSparsityMatrixCSR) end @kernel function sparse_lu_factor_kernel!( - factors, source, rowptr, colval, diagonal, source_to_factor, nnz_source, n + factors, source, rowptr, colval, diagonal, source_to_factor, + status, nnz_source, n ) index = @index(Global) if index == 1 + status[1] = zero(eltype(status)) @inbounds for k in eachindex(factors) factors[k] = zero(eltype(factors)) end @@ -76,7 +82,12 @@ end @inbounds for i in 1:n for ij in rowptr[i]:(diagonal[i] - 1) j = colval[ij] - lij = factors[ij] / factors[diagonal[j]] + pivot = factors[diagonal[j]] + if iszero(pivot) || !isfinite(pivot) + status[1] = j + break + end + lij = factors[ij] / pivot factors[ij] = lij for jk in (diagonal[j] + 1):(rowptr[j + 1] - 1) col = colval[jk] @@ -88,6 +99,67 @@ end end end end + if !iszero(status[1]) + break + end + pivot = factors[diagonal[i]] + if iszero(pivot) || !isfinite(pivot) + status[1] = i + break + end + end + end +end + +@kernel function pivoted_sparse_lu_factor_kernel!( + factors, pivots, status, source, rowptr, colval, n + ) + index = @index(Global) + if index == 1 + status[1] = zero(eltype(status)) + @inbounds for k in eachindex(factors) + factors[k] = zero(eltype(factors)) + end + @inbounds for row in 1:n + for position in rowptr[row]:(rowptr[row + 1] - 1) + col = colval[position] + factors[(row - 1)*n + col] = source[position] + end + end + @inbounds for k in 1:n + pivot_row = k + pivot_size = abs(factors[(k - 1)*n + k]) + for row in (k + 1):n + candidate = abs(factors[(row - 1)*n + k]) + if candidate > pivot_size + pivot_size = candidate + pivot_row = row + end + end + if iszero(pivot_size) || !isfinite(pivot_size) + status[1] = k + break + end + pivots[k] = pivot_row + if pivot_row != k + for col in 1:n + first = (k - 1)*n + col + second = (pivot_row - 1)*n + col + value = factors[first] + factors[first] = factors[second] + factors[second] = value + end + end + pivot = factors[(k - 1)*n + k] + for row in (k + 1):n + lower = (row - 1)*n + k + multiplier = factors[lower] / pivot + factors[lower] = multiplier + for col in (k + 1):n + factors[(row - 1)*n + col] -= + multiplier*factors[(k - 1)*n + col] + end + end end end end @@ -114,16 +186,63 @@ end end end -function setup_sparse_lu(A::StaticSparsityMatrixCSR) +@kernel function pivoted_sparse_lu_solve_kernel!(output, rhs, factors, pivots, n) + index = @index(Global) + if index == 1 + @inbounds for i in 1:n + output[i] = rhs[i] + end + @inbounds for k in 1:n + pivot_row = pivots[k] + value = output[k] + output[k] = output[pivot_row] + output[pivot_row] = value + end + @inbounds for i in 1:n + value = output[i] + for j in 1:(i - 1) + value -= factors[(i - 1)*n + j]*output[j] + end + output[i] = value + end + @inbounds for i in n:-1:1 + value = output[i] + for j in (i + 1):n + value -= factors[(i - 1)*n + j]*output[j] + end + output[i] = value / factors[(i - 1)*n + i] + end + end +end + +function setup_sparse_lu(A::StaticSparsityMatrixCSR; pivoting::Bool = true) + size(A, 1) == size(A, 2) || + throw(DimensionMismatch("sparse LU requires a square matrix")) backend = matrix_backend(A) - rowptr, colval, diagonal, map, pattern_rowptr, pattern_colval = sparse_lu_pattern(A) + if pivoting + host = host_csr(A) + rowptr = Int.(host.rowptr) + colval = Int.(host.colval) + diagonal = zeros(Int, size(A, 1)) + map = collect(1:nnz(A)) + pattern_rowptr = rowptr + pattern_colval = colval + nfactors = Base.checked_mul(size(A, 1), size(A, 1)) + else + rowptr, colval, diagonal, map, pattern_rowptr, pattern_colval = + sparse_lu_pattern(A) + nfactors = length(colval) + end Ti = eltype(A.rowptr) - factors = KernelAbstractions.allocate(backend, eltype(A.nzval), length(colval)) + factors = KernelAbstractions.allocate(backend, eltype(A.nzval), nfactors) + pivots = KernelAbstractions.allocate(backend, Ti, size(A, 1)) + status = KernelAbstractions.allocate(backend, Ti, 1) state = KASparseLUFactor( backend_copy(backend, Ti.(rowptr)), backend_copy(backend, Ti.(colval)), backend_copy(backend, Ti.(diagonal)), backend_copy(backend, Ti.(map)), factors, backend, size(A, 1), - pattern_rowptr, pattern_colval, A.rowptr, A.colval + pattern_rowptr, pattern_colval, A.rowptr, A.colval, + pivots, status, pivoting ) return resetup_sparse_lu!(SparseLU(state), A) end @@ -143,20 +262,33 @@ function resetup_sparse_lu!(S::SparseLU{<:KASparseLUFactor}, A::StaticSparsityMa F = S.factorization sparse_lu_same_pattern(S, A) || throw(ArgumentError("sparse LU resetup requires the same CSR pattern")) - kernel! = sparse_lu_factor_kernel!(F.backend, 1) - event = kernel!( - F.factors, A.nzval, F.rowptr, F.colval, F.diagonal, - F.source_to_factor, nnz(A), F.n; ndrange = 1 - ) + if F.pivoting + kernel! = pivoted_sparse_lu_factor_kernel!(F.backend, 1) + event = kernel!(F.factors, F.pivots, F.status, + A.nzval, F.rowptr, F.colval, F.n; ndrange = 1) + else + kernel! = sparse_lu_factor_kernel!(F.backend, 1) + event = kernel!(F.factors, A.nzval, F.rowptr, F.colval, + F.diagonal, F.source_to_factor, F.status, nnz(A), F.n; + ndrange = 1) + end isnothing(event) || wait(event) + singular = only(Array(F.status)) + iszero(singular) || throw(LinearAlgebra.SingularException(Int(singular))) return S end function LinearAlgebra.ldiv!(x, S::SparseLU{<:KASparseLUFactor}, b) F = S.factorization length(x) == F.n && length(b) == F.n || throw(DimensionMismatch()) - kernel! = sparse_lu_solve_kernel!(F.backend, 1) - event = kernel!(x, b, F.factors, F.rowptr, F.colval, F.diagonal, F.n; ndrange = 1) + if F.pivoting + kernel! = pivoted_sparse_lu_solve_kernel!(F.backend, 1) + event = kernel!(x, b, F.factors, F.pivots, F.n; ndrange = 1) + else + kernel! = sparse_lu_solve_kernel!(F.backend, 1) + event = kernel!(x, b, F.factors, F.rowptr, F.colval, + F.diagonal, F.n; ndrange = 1) + end isnothing(event) || wait(event) return x end diff --git a/test/ka_preconditioners.jl b/test/ka_preconditioners.jl index 3f2e6c22..f3cd0ec5 100644 --- a/test/ka_preconditioners.jl +++ b/test/ka_preconditioners.jl @@ -81,6 +81,38 @@ end @test_throws ArgumentError KAPreconditioners.resetup_sparse_lu!( F, csr_matrix(different; backend) ) + + unpivoted = KAPreconditioners.setup_sparse_lu(matrix; pivoting = false) + @test !unpivoted.factorization.pivoting + ldiv!(x, unpivoted, rhs) + @test Array(x) ≈ A \ b + end +end + +@testset "SparseLU row pivoting" begin + # Resetup changes which row is selected as the first pivot while + # preserving the CSR pattern, including its explicit zero diagonal. + rows = [1, 1, 2, 2] + cols = [1, 2, 1, 2] + A = sparse(rows, cols, [0.0, 2.0, 3.0, 4.0], 2, 2) + B = sparse(rows, cols, [5.0, 2.0, 3.0, 4.0], 2, 2) + C = sparse([1, 2, 2], [2, 1, 2], [2.0, 3.0, 4.0], 2, 2) + b = [1.0, 2.0] + for backend in (KernelAbstractions.CPU(), JLBackend()) + matrix = csr_matrix(A; backend) + F = KAPreconditioners.setup_sparse_lu(matrix) + @test F.factorization.pivoting + rhs = backend isa KernelAbstractions.CPU ? b : JLArray(b) + x = similar(rhs) + ldiv!(x, F, rhs) + @test Array(x) ≈ A \ b + @test KAPreconditioners.resetup_sparse_lu!( + F, csr_matrix(B; backend)) === F + ldiv!(x, F, rhs) + @test Array(x) ≈ B \ b + missing_diagonal = KAPreconditioners.setup_sparse_lu(csr_matrix(C; backend)) + ldiv!(x, missing_diagonal, rhs) + @test Array(x) ≈ C \ b end end From 013259bf06a43f29d1a97243f4b2ac1d250901b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olav=20M=C3=B8yner?= Date: Thu, 1 Oct 2026 16:32:39 +0200 Subject: [PATCH 2/3] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index e35c3b54..31930d18 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Jutul" uuid = "2b460a1a-8a2b-45b2-b125-b5c536396eb9" -version = "0.5.1" +version = "0.5.2" authors = ["Olav Møyner "] [deps] From 86e3b2eb239af9c63b3dd7077cef39183fe6c0aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olav=20M=C3=B8yner?= Date: Thu, 1 Oct 2026 16:41:20 +0200 Subject: [PATCH 3/3] Appease runic --- src/KAPreconditioners/sparse_lu.jl | 42 +++++++++++++++++------------- test/ka_preconditioners.jl | 3 ++- 2 files changed, 26 insertions(+), 19 deletions(-) diff --git a/src/KAPreconditioners/sparse_lu.jl b/src/KAPreconditioners/sparse_lu.jl index 13ac1ded..8d93f7df 100644 --- a/src/KAPreconditioners/sparse_lu.jl +++ b/src/KAPreconditioners/sparse_lu.jl @@ -123,14 +123,14 @@ end @inbounds for row in 1:n for position in rowptr[row]:(rowptr[row + 1] - 1) col = colval[position] - factors[(row - 1)*n + col] = source[position] + factors[(row - 1) * n + col] = source[position] end end @inbounds for k in 1:n pivot_row = k - pivot_size = abs(factors[(k - 1)*n + k]) + pivot_size = abs(factors[(k - 1) * n + k]) for row in (k + 1):n - candidate = abs(factors[(row - 1)*n + k]) + candidate = abs(factors[(row - 1) * n + k]) if candidate > pivot_size pivot_size = candidate pivot_row = row @@ -143,21 +143,21 @@ end pivots[k] = pivot_row if pivot_row != k for col in 1:n - first = (k - 1)*n + col - second = (pivot_row - 1)*n + col + first = (k - 1) * n + col + second = (pivot_row - 1) * n + col value = factors[first] factors[first] = factors[second] factors[second] = value end end - pivot = factors[(k - 1)*n + k] + pivot = factors[(k - 1) * n + k] for row in (k + 1):n - lower = (row - 1)*n + k + lower = (row - 1) * n + k multiplier = factors[lower] / pivot factors[lower] = multiplier for col in (k + 1):n - factors[(row - 1)*n + col] -= - multiplier*factors[(k - 1)*n + col] + factors[(row - 1) * n + col] -= + multiplier * factors[(k - 1) * n + col] end end end @@ -201,16 +201,16 @@ end @inbounds for i in 1:n value = output[i] for j in 1:(i - 1) - value -= factors[(i - 1)*n + j]*output[j] + value -= factors[(i - 1) * n + j] * output[j] end output[i] = value end @inbounds for i in n:-1:1 value = output[i] for j in (i + 1):n - value -= factors[(i - 1)*n + j]*output[j] + value -= factors[(i - 1) * n + j] * output[j] end - output[i] = value / factors[(i - 1)*n + i] + output[i] = value / factors[(i - 1) * n + i] end end end @@ -264,13 +264,17 @@ function resetup_sparse_lu!(S::SparseLU{<:KASparseLUFactor}, A::StaticSparsityMa throw(ArgumentError("sparse LU resetup requires the same CSR pattern")) if F.pivoting kernel! = pivoted_sparse_lu_factor_kernel!(F.backend, 1) - event = kernel!(F.factors, F.pivots, F.status, - A.nzval, F.rowptr, F.colval, F.n; ndrange = 1) + event = kernel!( + F.factors, F.pivots, F.status, + A.nzval, F.rowptr, F.colval, F.n; ndrange = 1 + ) else kernel! = sparse_lu_factor_kernel!(F.backend, 1) - event = kernel!(F.factors, A.nzval, F.rowptr, F.colval, + event = kernel!( + F.factors, A.nzval, F.rowptr, F.colval, F.diagonal, F.source_to_factor, F.status, nnz(A), F.n; - ndrange = 1) + ndrange = 1 + ) end isnothing(event) || wait(event) singular = only(Array(F.status)) @@ -286,8 +290,10 @@ function LinearAlgebra.ldiv!(x, S::SparseLU{<:KASparseLUFactor}, b) event = kernel!(x, b, F.factors, F.pivots, F.n; ndrange = 1) else kernel! = sparse_lu_solve_kernel!(F.backend, 1) - event = kernel!(x, b, F.factors, F.rowptr, F.colval, - F.diagonal, F.n; ndrange = 1) + event = kernel!( + x, b, F.factors, F.rowptr, F.colval, + F.diagonal, F.n; ndrange = 1 + ) end isnothing(event) || wait(event) return x diff --git a/test/ka_preconditioners.jl b/test/ka_preconditioners.jl index f3cd0ec5..b243c4ec 100644 --- a/test/ka_preconditioners.jl +++ b/test/ka_preconditioners.jl @@ -107,7 +107,8 @@ end ldiv!(x, F, rhs) @test Array(x) ≈ A \ b @test KAPreconditioners.resetup_sparse_lu!( - F, csr_matrix(B; backend)) === F + F, csr_matrix(B; backend) + ) === F ldiv!(x, F, rhs) @test Array(x) ≈ B \ b missing_diagonal = KAPreconditioners.setup_sparse_lu(csr_matrix(C; backend))