Support decompression in one triangle

Author: amontoison

ext/SparseMatrixColoringsCUDAExt.jl

@@ -48,12 +48,13 @@ function SMC.StarSetColoringResult(
     ag::SMC.AdjacencyGraph{T},
     color::Vector{<:Integer},
     star_set::SMC.StarSet{<:Integer},
+    decompression_uplo::Symbol,
 ) where {T<:Integer}
     group = SMC.group_by_color(T, color)
-    compressed_indices = SMC.star_csc_indices(ag, color, star_set)
+    compressed_indices = SMC.star_csc_indices(ag, color, star_set, decompression_uplo)
     additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
     return SMC.StarSetColoringResult(
-        A, ag, color, group, compressed_indices, additional_info
+        A, ag, color, group, compressed_indices, decompression_uplo, additional_info
     )
 end
 
@@ -86,12 +87,14 @@ function SMC.StarSetColoringResult(
     ag::SMC.AdjacencyGraph{T},
     color::Vector{<:Integer},
     star_set::SMC.StarSet{<:Integer},
+    decompression_uplo::Symbol
 ) where {T<:Integer}
     group = SMC.group_by_color(T, color)
-    compressed_indices = SMC.star_csc_indices(ag, color, star_set)
+    reversed_uplo = (decompression_uplo == :L) ? :U : (decompression_uplo == :U ? :L : decompression_uplo)
+    compressed_indices = SMC.star_csc_indices(ag, color, star_set, reversed_uplo)
     additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices))
     return SMC.StarSetColoringResult(
-        A, ag, color, group, compressed_indices, additional_info
+        A, ag, color, group, compressed_indices, decompression_uplo, additional_info
     )
 end
 
@@ -120,15 +123,7 @@ function SMC.decompress!(
     A::CuSparseMatrixCSC,
     B::CuMatrix,
     result::SMC.StarSetColoringResult{<:CuSparseMatrixCSC},
-    uplo::Symbol=:F,
 )
-    if uplo != :F
-        throw(
-            SMC.UnsupportedDecompressionError(
-                "Single-triangle decompression is not supported on GPU matrices"
-            ),
-        )
-    end
     compressed_indices = result.additional_info.compressed_indices_gpu_csc
     copyto!(A.nzVal, view(B, compressed_indices))
     return A
@@ -138,15 +133,7 @@ function SMC.decompress!(
     A::CuSparseMatrixCSR,
     B::CuMatrix,
     result::SMC.StarSetColoringResult{<:CuSparseMatrixCSR},
-    uplo::Symbol=:F,
 )
-    if uplo != :F
-        throw(
-            SMC.UnsupportedDecompressionError(
-                "Single-triangle decompression is not supported on GPU matrices"
-            ),
-        )
-    end
     compressed_indices = result.additional_info.compressed_indices_gpu_csr
     copyto!(A.nzVal, view(B, compressed_indices))
     return A

src/SparseMatrixColorings.jl

@@ -26,7 +26,9 @@ using LinearAlgebra:
     issymmetric,
     ldiv!,
     parent,
-    transpose
+    transpose,
+    tril,
+    triu
 using PrecompileTools: @compile_workload
 using Random: Random, AbstractRNG, default_rng, randperm
 using SparseArrays:

src/decompression.jl

@@ -175,19 +175,34 @@ function decompress(B::AbstractMatrix, result::AbstractColoringResult)
     return decompress!(A, B, result)
 end
 
+function decompress(
+    B::AbstractMatrix,
+    result::AbstractColoringResult{:symmetric,:column})
+    A = respectful_similar(result.A, eltype(B))
+    if A isa SparseMatrixCSC && result.uplo != :F
+        (result.uplo == :L) && (A = tril(A))
+        (result.uplo == :U) && (A = triu(A))
+    end
+    return decompress!(A, B, result)
+end
+
 function decompress(
     Br::AbstractMatrix,
     Bc::AbstractMatrix,
     result::AbstractColoringResult{structure,:bidirectional},
 ) where {structure}
     A = respectful_similar(result.A, Base.promote_eltype(Br, Bc))
+    if A isa SparseMatrixCSC && result.symmetric_result.uplo != :F
+        (result.symmetric_result.uplo == :L) && (A = tril(A))
+        (result.symmetric_result.uplo == :U) && (A = triu(A))
+    end
     return decompress!(A, Br, Bc, result)
 end
 
 """
     decompress!(
         A::AbstractMatrix, B::AbstractMatrix,
-        result::AbstractColoringResult{_,:column/:row}, [uplo=:F]
+        result::AbstractColoringResult{_,:column/:row},
     )
 
     decompress!(
@@ -204,9 +219,6 @@ The out-of-place alternative is [`decompress`](@ref).
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
 
-For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
-When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
-
 !!! warning
     For some coloring variants, the `result` object is mutated during decompression.
 
@@ -260,7 +272,7 @@ function decompress! end
 """
     decompress_single_color!(
         A::AbstractMatrix, b::AbstractVector, c::Integer,
-        result::AbstractColoringResult, [uplo=:F]
+        result::AbstractColoringResult,
     )
 
 Decompress the vector `b` corresponding to color `c` in-place into `A`, given a `:direct` coloring `result` of the sparsity pattern of `A` (it will not work with a `:substitution` coloring).
@@ -272,9 +284,6 @@ Decompress the vector `b` corresponding to color `c` in-place into `A`, given a
 !!! warning
     This function will only update some coefficients of `A`, without resetting the rest to zero.
 
-For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
-When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
-
 !!! warning
     For some coloring variants, the `result` object is mutated during decompression.
 
@@ -446,95 +455,68 @@ end
 ## StarSetColoringResult
 
 function decompress!(
-    A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
-)
-    (; ag, compressed_indices) = result
+    A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult)
+    (; ag, compressed_indices, uplo) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     fill!(A, zero(eltype(A)))
 
-    rvS = rowvals(S)
+    l = 0
+    rvS = rowvals(A)
     for j in axes(S, 2)
         for k in nzrange(S, j)
             i = rvS[k]
             if in_triangle(i, j, uplo)
-                A[i, j] = B[compressed_indices[k]]
+                l += 1
+                A[i, j] = B[compressed_indices[l]]
             end
         end
     end
     return A
 end
 
 function decompress_single_color!(
-    A::AbstractMatrix,
+    A::SparseMatrixCSC,
     b::AbstractVector,
     c::Integer,
     result::StarSetColoringResult,
-    uplo::Symbol=:F,
 )
-    (; ag, compressed_indices, group) = result
+    (; ag, compressed_indices, group, uplo) = result
     (; S) = ag
+    println(uplo)
     uplo == :F && check_same_pattern(A, S)
 
     lower_index = (c - 1) * S.n + 1
     upper_index = c * S.n
-    rvS = rowvals(S)
+    rvA = rowvals(A)
+    nzA = nonzeros(A)
     for j in group[c]
-        for k in nzrange(S, j)
-            # Check if the color c is used to recover A[i,j] / A[j,i]
+        for k in nzrange(A, j)
+            # Check if the color c is used to recover A[i,j]
             if lower_index <= compressed_indices[k] <= upper_index
-                i = rvS[k]
-                if i == j
-                    # Recover the diagonal coefficients of A
-                    A[i, i] = b[i]
-                else
-                    # Recover the off-diagonal coefficients of A
-                    if in_triangle(i, j, uplo)
-                        A[i, j] = b[i]
-                    end
-                    if in_triangle(j, i, uplo)
-                        A[j, i] = b[i]
-                    end
-                end
+                i = rvA[k]
+                nzA[k] = b[i]
             end
         end
     end
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
-)
-    (; ag, compressed_indices) = result
+function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult)
+    (; ag, compressed_indices, uplo) = result
     (; S) = ag
     nzA = nonzeros(A)
-    if uplo == :F
-        check_same_pattern(A, S)
-        for k in eachindex(nzA, compressed_indices)
-            nzA[k] = B[compressed_indices[k]]
-        end
-    else
-        rvS = rowvals(S)
-        l = 0  # assume A has the same pattern as the triangle
-        for j in axes(S, 2)
-            for k in nzrange(S, j)
-                i = rvS[k]
-                if in_triangle(i, j, uplo)
-                    l += 1
-                    nzA[l] = B[compressed_indices[k]]
-                end
-            end
-        end
+    uplo == :F && check_same_pattern(A, S)
+    for k in eachindex(nzA, compressed_indices)
+        nzA[k] = B[compressed_indices[k]]
     end
     return A
 end
 
 ## TreeSetColoringResult
 
-function decompress!(
-    A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
-)
-    (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, buffer) = result
+function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult)
+    (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, diagonal_indices, buffer, uplo) = result
     (; S) = ag
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
@@ -548,10 +530,8 @@ function decompress!(
 
     # Recover the diagonal coefficients of A
     if !augmented_graph(ag)
-        for i in axes(S, 1)
-            if !iszero(S[i, i])
-                A[i, i] = B[i, color[i]]
-            end
+        for i in diagonal_indices
+            A[i, i] = B[i, color[i]]
         end
     end
 
@@ -590,7 +570,6 @@ function decompress!(
     A::SparseMatrixCSC{R},
     B::AbstractMatrix{R},
     result::TreeSetColoringResult,
-    uplo::Symbol=:F,
 ) where {R<:Real}
     (;
         ag,
@@ -603,6 +582,7 @@ function decompress!(
         lower_triangle_offsets,
         upper_triangle_offsets,
         buffer,
+        uplo,
     ) = result
     (; S) = ag
     A_colptr = A.colptr
@@ -706,9 +686,8 @@ function decompress!(
     A::AbstractMatrix,
     B::AbstractMatrix,
     result::LinearSystemColoringResult,
-    uplo::Symbol=:F,
 )
-    (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = result
+    (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A, uplo) = result
     S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
 
@@ -770,10 +749,20 @@ end
 function decompress!(
     A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
 )
+    (; large_colptr, large_rowval, symmetric_result) = result
     m, n = size(A)
     Br_and_Bc = _join_compressed!(result, Br, Bc)
-    A_and_Aᵀ = decompress(Br_and_Bc, result.symmetric_result)
-    copyto!(A, A_and_Aᵀ[(n + 1):(n + m), 1:n])  # original matrix in bottom left corner
+    nzval = Vector{eltype(A)}(undef, length(large_rowval))
+    A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, nzval)
+    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result)
+    rvA = rowvals(A_and_noAᵀ)
+    nzA = nonzeros(A_and_noAᵀ)
+    for j in axes(A_and_noAᵀ, 2)
+        for k in nzrange(A_and_noAᵀ, j)
+            i = rvA[k]
+            A[i-n, j] = nzA[k]
+        end
+    end
     return A
 end
 
@@ -786,6 +775,6 @@ function decompress!(
     # pretend A is larger
     A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval)
     # decompress lower triangle only
-    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
+    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result)
     return A
 end