Add pivoting to KA sparse LU fallback

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 <olav.moyner@gmail.com>"]
 
 [deps]

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

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