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] diff --git a/src/KAPreconditioners/sparse_lu.jl b/src/KAPreconditioners/sparse_lu.jl index 3cbff0d3..8d93f7df 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,39 @@ 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..b243c4ec 100644 --- a/test/ka_preconditioners.jl +++ b/test/ka_preconditioners.jl @@ -81,6 +81,39 @@ 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