diff --git a/docs/src/dev.md b/docs/src/dev.md index 0f7b86fa..16ab34d3 100644 --- a/docs/src/dev.md +++ b/docs/src/dev.md @@ -39,7 +39,8 @@ SparseMatrixColorings.RowColoringResult SparseMatrixColorings.StarSetColoringResult SparseMatrixColorings.TreeSetColoringResult SparseMatrixColorings.LinearSystemColoringResult -SparseMatrixColorings.BicoloringResult +SparseMatrixColorings.StarSetBicoloringResult +SparseMatrixColorings.TreeSetBicoloringResult SparseMatrixColorings.remap_colors SparseMatrixColorings.decompress_csc! ``` diff --git a/ext/SparseMatrixColoringsAMDGPUExt.jl b/ext/SparseMatrixColoringsAMDGPUExt.jl index 58714ad2..4189f062 100644 --- a/ext/SparseMatrixColoringsAMDGPUExt.jl +++ b/ext/SparseMatrixColoringsAMDGPUExt.jl @@ -41,6 +41,48 @@ function SMC.StarSetColoringResult( ) end +function SMC.StarSetBicoloringResult( + A::ROCSparseMatrixCSC, + S::SMC.SparsityPatternCSC{T}, + ag::SMC.AdjacencyGraph{T}, + symmetric_color::Vector{<:Integer}, + star_set::SMC.StarSet{<:Integer}, + row_color::Vector{T}, + column_color::Vector{T}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, +) where {T<:Integer} + column_group = SMC.group_by_color(T, column_color) + row_group = SMC.group_by_color(T, row_color) + num_row_colors = length(row_group) + A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br = SMC.star_bicoloring_csc_indices( + S, symmetric_color, star_set, symmetric_to_row, symmetric_to_column, num_row_colors + ) + additional_info = (; + A_indices_gpu_bc_csc=ROCVector(A_indices_bc), + compressed_indices_gpu_bc_csc=ROCVector(compressed_indices_bc), + A_indices_gpu_br_csc=ROCVector(A_indices_br), + compressed_indices_gpu_br_csc=ROCVector(compressed_indices_br), + ) + return SMC.StarSetBicoloringResult( + A, + S, + ag, + symmetric_color, + column_color, + row_color, + column_group, + row_group, + symmetric_to_column, + symmetric_to_row, + A_indices_bc, + compressed_indices_bc, + A_indices_br, + compressed_indices_br, + additional_info, + ) +end + ## CSR Result function SMC.ColumnColoringResult( @@ -79,6 +121,54 @@ function SMC.StarSetColoringResult( ) end +function SMC.StarSetBicoloringResult( + A::ROCSparseMatrixCSR, + S::SMC.SparsityPatternCSC{T}, + ag::SMC.AdjacencyGraph{T}, + symmetric_color::Vector{<:Integer}, + star_set::SMC.StarSet{<:Integer}, + row_color::Vector{T}, + column_color::Vector{T}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, +) where {T<:Integer} + column_group = SMC.group_by_color(T, column_color) + row_group = SMC.group_by_color(T, row_color) + num_row_colors = length(row_group) + A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br = SMC.star_bicoloring_csr_indices( + ag, + S, + symmetric_color, + star_set, + symmetric_to_row, + symmetric_to_column, + num_row_colors, + ) + additional_info = (; + A_indices_gpu_bc_csr=ROCVector(A_indices_bc), + compressed_indices_gpu_bc_csr=ROCVector(compressed_indices_bc), + A_indices_gpu_br_csr=ROCVector(A_indices_br), + compressed_indices_gpu_br_csr=ROCVector(compressed_indices_br), + ) + return SMC.StarSetBicoloringResult( + A, + S, + ag, + symmetric_color, + column_color, + row_color, + column_group, + row_group, + symmetric_to_column, + symmetric_to_row, + A_indices_bc, + compressed_indices_bc, + A_indices_br, + compressed_indices_br, + additional_info, + ) +end + ## Decompression for R in (:ColumnColoringResult, :RowColoringResult) @@ -136,4 +226,40 @@ function SMC.decompress!( return A end +function SMC.decompress!( + A::ROCSparseMatrixCSC, + Br::ROCMatrix, + Bc::ROCMatrix, + result::SMC.StarSetBicoloringResult{<:ROCSparseMatrixCSC}, +) + (; + A_indices_gpu_bc_csc, + compressed_indices_gpu_bc_csc, + A_indices_gpu_br_csc, + compressed_indices_gpu_br_csc, + ) = result.additional_info + nzA = A.nzVal + view(nzA, A_indices_gpu_bc_csc) .= view(Bc, compressed_indices_gpu_bc_csc) + view(nzA, A_indices_gpu_br_csc) .= view(Br, compressed_indices_gpu_br_csc) + return A +end + +function SMC.decompress!( + A::ROCSparseMatrixCSR, + Br::ROCMatrix, + Bc::ROCMatrix, + result::SMC.StarSetBicoloringResult{<:ROCSparseMatrixCSR}, +) + (; + A_indices_gpu_bc_csr, + compressed_indices_gpu_bc_csr, + A_indices_gpu_br_csr, + compressed_indices_gpu_br_csr, + ) = result.additional_info + nzA = A.nzVal + view(nzA, A_indices_gpu_bc_csr) .= view(Bc, compressed_indices_gpu_bc_csr) + view(nzA, A_indices_gpu_br_csr) .= view(Br, compressed_indices_gpu_br_csr) + return A +end + end diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl index ed33eece..8fc9eb7a 100644 --- a/ext/SparseMatrixColoringsCUDAExt.jl +++ b/ext/SparseMatrixColoringsCUDAExt.jl @@ -41,6 +41,48 @@ function SMC.StarSetColoringResult( ) end +function SMC.StarSetBicoloringResult( + A::CuSparseMatrixCSC, + S::SMC.SparsityPatternCSC{T}, + ag::SMC.AdjacencyGraph{T}, + symmetric_color::Vector{<:Integer}, + star_set::SMC.StarSet{<:Integer}, + row_color::Vector{T}, + column_color::Vector{T}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, +) where {T<:Integer} + column_group = SMC.group_by_color(T, column_color) + row_group = SMC.group_by_color(T, row_color) + num_row_colors = length(row_group) + A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br = SMC.star_bicoloring_csc_indices( + S, symmetric_color, star_set, symmetric_to_row, symmetric_to_column, num_row_colors + ) + additional_info = (; + A_indices_gpu_bc_csc=CuVector(A_indices_bc), + compressed_indices_gpu_bc_csc=CuVector(compressed_indices_bc), + A_indices_gpu_br_csc=CuVector(A_indices_br), + compressed_indices_gpu_br_csc=CuVector(compressed_indices_br), + ) + return SMC.StarSetBicoloringResult( + A, + S, + ag, + symmetric_color, + column_color, + row_color, + column_group, + row_group, + symmetric_to_column, + symmetric_to_row, + A_indices_bc, + compressed_indices_bc, + A_indices_br, + compressed_indices_br, + additional_info, + ) +end + ## CSR Result function SMC.ColumnColoringResult( @@ -79,6 +121,54 @@ function SMC.StarSetColoringResult( ) end +function SMC.StarSetBicoloringResult( + A::CuSparseMatrixCSR, + S::SMC.SparsityPatternCSC{T}, + ag::SMC.AdjacencyGraph{T}, + symmetric_color::Vector{<:Integer}, + star_set::SMC.StarSet{<:Integer}, + row_color::Vector{T}, + column_color::Vector{T}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, +) where {T<:Integer} + column_group = SMC.group_by_color(T, column_color) + row_group = SMC.group_by_color(T, row_color) + num_row_colors = length(row_group) + A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br = SMC.star_bicoloring_csr_indices( + ag, + S, + symmetric_color, + star_set, + symmetric_to_row, + symmetric_to_column, + num_row_colors, + ) + additional_info = (; + A_indices_gpu_bc_csr=CuVector(A_indices_bc), + compressed_indices_gpu_bc_csr=CuVector(compressed_indices_bc), + A_indices_gpu_br_csr=CuVector(A_indices_br), + compressed_indices_gpu_br_csr=CuVector(compressed_indices_br), + ) + return SMC.StarSetBicoloringResult( + A, + S, + ag, + symmetric_color, + column_color, + row_color, + column_group, + row_group, + symmetric_to_column, + symmetric_to_row, + A_indices_bc, + compressed_indices_bc, + A_indices_br, + compressed_indices_br, + additional_info, + ) +end + ## Decompression for R in (:ColumnColoringResult, :RowColoringResult) @@ -136,4 +226,40 @@ function SMC.decompress!( return A end +function SMC.decompress!( + A::CuSparseMatrixCSC, + Br::CuMatrix, + Bc::CuMatrix, + result::SMC.StarSetBicoloringResult{<:CuSparseMatrixCSC}, +) + (; + A_indices_gpu_bc_csc, + compressed_indices_gpu_bc_csc, + A_indices_gpu_br_csc, + compressed_indices_gpu_br_csc, + ) = result.additional_info + nzA = A.nzVal + view(nzA, A_indices_gpu_bc_csc) .= view(Bc, compressed_indices_gpu_bc_csc) + view(nzA, A_indices_gpu_br_csc) .= view(Br, compressed_indices_gpu_br_csc) + return A +end + +function SMC.decompress!( + A::CuSparseMatrixCSR, + Br::CuMatrix, + Bc::CuMatrix, + result::SMC.StarSetBicoloringResult{<:CuSparseMatrixCSR}, +) + (; + A_indices_gpu_bc_csr, + compressed_indices_gpu_bc_csr, + A_indices_gpu_br_csr, + compressed_indices_gpu_br_csr, + ) = result.additional_info + nzA = A.nzVal + view(nzA, A_indices_gpu_bc_csr) .= view(Bc, compressed_indices_gpu_bc_csr) + view(nzA, A_indices_gpu_br_csr) .= view(Br, compressed_indices_gpu_br_csr) + return A +end + end diff --git a/ext/SparseMatrixColoringsGPUArraysExt.jl b/ext/SparseMatrixColoringsGPUArraysExt.jl index 37574e73..9db952b9 100644 --- a/ext/SparseMatrixColoringsGPUArraysExt.jl +++ b/ext/SparseMatrixColoringsGPUArraysExt.jl @@ -26,4 +26,41 @@ function SMC.compress( return B end +function SMC.compress( + A::AbstractGPUSparseMatrix, result::SMC.AbstractColoringResult{structure,:bidirectional} +) where {structure} + A_cpu = SparseMatrixCSC(A) + Br_cpu, Bc_cpu = SMC.compress(A_cpu, result) + M = dense_array_type(A) + return M(Br_cpu), M(Bc_cpu) +end + +## Decompression + +function SMC.decompress!( + A::AbstractGPUSparseMatrix, + B::AbstractMatrix, + result::SMC.TreeSetColoringResult, + uplo::Symbol=:F, +) + return throw( + SMC.UnsupportedDecompressionError( + "Symmetric decompression by substitution is not supported on GPU matrices" + ), + ) +end + +function SMC.decompress!( + A::AbstractGPUSparseMatrix, + Br::AbstractMatrix, + Bc::AbstractMatrix, + result::SMC.TreeSetBicoloringResult, +) + return throw( + SMC.UnsupportedDecompressionError( + "Bidirectional decompression by substitution is not supported on GPU matrices" + ), + ) +end + end diff --git a/src/adtypes.jl b/src/adtypes.jl index 7c464433..12a162f0 100644 --- a/src/adtypes.jl +++ b/src/adtypes.jl @@ -15,7 +15,7 @@ function coloring( forced_colors = ADTypes.row_coloring(A, algo) else # TODO: improve once https://github.com/SciML/ADTypes.jl/issues/69 is done - A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern) + _, A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern) forced_colors = ADTypes.symmetric_coloring(A_and_Aᵀ, algo) end else diff --git a/src/check.jl b/src/check.jl index 3ad11133..6afda69e 100644 --- a/src/check.jl +++ b/src/check.jl @@ -503,7 +503,7 @@ end """ rank_nonzeros_from_trees(result::TreeSetColoringResult) - rank_nonzeros_from_trees(result::BicoloringResult) + rank_nonzeros_from_trees(result::TreeSetBicoloringResult) Construct a sparse matrix `rank_nonzeros` that assigns a unique recovery rank to each nonzero coefficient associated with an acyclic coloring or bicoloring. @@ -543,19 +543,11 @@ function rank_nonzeros_from_trees(result::TreeSetColoringResult) return rank_nonzeros end -function rank_nonzeros_from_trees(result::BicoloringResult) - (; A, abg, row_color, column_color, symmetric_result, large_colptr, large_rowval) = - result - @assert symmetric_result isa TreeSetColoringResult - (; ag, reverse_bfs_orders, tree_edge_indices, nt) = symmetric_result - (; S) = ag +function rank_nonzeros_from_trees(result::TreeSetBicoloringResult) + (; A, S, reverse_bfs_orders, tree_edge_indices, nt) = result m, n = size(A) - nnzA = nnz(S) ÷ 2 - nzval = zeros(Int, nnzA) - colptr = large_colptr[1:(n + 1)] - rowval = large_rowval[1:nnzA] - rowval .-= n - rank_nonzeros = SparseMatrixCSC(m, n, colptr, rowval, nzval) + nzval = zeros(Int, nnz(S)) + rank_nonzeros = SparseMatrixCSC(m, n, S.colptr, S.rowval, nzval) counter = 0 for k in 1:nt first = tree_edge_indices[k] diff --git a/src/decompression.jl b/src/decompression.jl index 2dcf1847..23d51c6c 100644 --- a/src/decompression.jl +++ b/src/decompression.jl @@ -753,73 +753,181 @@ function decompress!( return A end -## BicoloringResult - -function _join_compressed!(result::BicoloringResult, Br::AbstractMatrix, Bc::AbstractMatrix) - #= - Say we have an original matrix `A` of size `(n, m)` and we build an augmented matrix `A_and_Aᵀ = [zeros(n, n) Aᵀ; A zeros(m, m)]`. - Its first `1:n` columns have the form `[zeros(n); A[:, j]]` and its following `n+1:n+m` columns have the form `[A[i, :]; zeros(m)]`. - The symmetric column coloring is performed on `A_and_Aᵀ` and the column-wise compression of `A_and_Aᵀ` should return a matrix `Br_and_Bc`. - But in reality, `Br_and_Bc` is computed as two partial compressions: the row-wise compression `Br` (corresponding to `Aᵀ`) and the columnwise compression `Bc` (corresponding to `A`). - Before symmetric decompression, we must reconstruct `Br_and_Bc` from `Br` and `Bc`, knowing that the symmetric colors (those making up `Br_and_Bc`) are present in either a row of `Br`, a column of `Bc`, or both. - Therefore, the column indices in `Br_and_Bc` don't necessarily match with the row indices in `Br` or the column indices in `Bc` since some colors may be missing in the partial compressions. - The columns of the top part of `Br_and_Bc` (rows `1:n`) are the rows of `Br`, interlaced with zero columns whenever the current color hasn't been used to color any row. - The columns of the bottom part of `Br_and_Bc` (rows `n+1:n+m`) are the columns of `Bc`, interlaced with zero columns whenever the current color hasn't been used to color any column. - We use the vectors `symmetric_to_row` and `symmetric_to_column` to map from symmetric colors to row and column colors. - =# - (; A, symmetric_to_column, symmetric_to_row) = result - m, n = size(A) - R = Base.promote_eltype(Br, Bc) - if eltype(result.Br_and_Bc) == R - Br_and_Bc = result.Br_and_Bc - else - Br_and_Bc = similar(result.Br_and_Bc, R) - end - fill!(Br_and_Bc, zero(R)) - for c in axes(Br_and_Bc, 2) - if symmetric_to_row[c] > 0 # some rows were colored with the symmetric color c - copyto!(view(Br_and_Bc, 1:n, c), view(Br, symmetric_to_row[c], :)) - end - if symmetric_to_column[c] > 0 # some columns were colored with the symmetric color c - copyto!( - view(Br_and_Bc, (n + 1):(n + m), c), view(Bc, :, symmetric_to_column[c]) - ) +## StarSetBicoloringResult + +#= +Each nonzero `A[i, j]` is recovered from a single coefficient of `Bc` or of `Br`, depending on +whether the hub of its star in the augmented graph is the column vertex `j` or the row vertex `i + n`. +The constructor therefore splits the nonzeros of `A` into two groups, each described by its own +pair of index vectors. Both `A_indices_bc` and `A_indices_br` are increasing and together they +partition `1:nnz(A)`, which lets the dense method below merge them in a single traversal. +=# + +function decompress!( + A::AbstractMatrix, + Br::AbstractMatrix, + Bc::AbstractMatrix, + result::StarSetBicoloringResult, +) + (; S, A_indices_bc, compressed_indices_bc, compressed_indices_br) = result + fill!(A, zero(eltype(A))) + rvS = rowvals(S) + nb_bc = length(A_indices_bc) + ind_bc = 1 + ind_br = 1 + for j in axes(S, 2) + for k in nzrange(S, j) + i = rvS[k] + if ind_bc <= nb_bc && A_indices_bc[ind_bc] == k + A[i, j] = Bc[compressed_indices_bc[ind_bc]] + ind_bc += 1 + else + A[i, j] = Br[compressed_indices_br[ind_br]] + ind_br += 1 + end end end - return Br_and_Bc + return A +end + +function decompress!( + A::SparseMatrixCSC, + Br::AbstractMatrix, + Bc::AbstractMatrix, + result::StarSetBicoloringResult, +) + (; A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br) = result + nzA = nonzeros(A) + for t in eachindex(A_indices_bc) + nzA[A_indices_bc[t]] = Bc[compressed_indices_bc[t]] + end + for t in eachindex(A_indices_br) + nzA[A_indices_br[t]] = Br[compressed_indices_br[t]] + end + return A end +## TreeSetBicoloringResult + function decompress!( - A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult + A::AbstractMatrix, + Br::AbstractMatrix, + Bc::AbstractMatrix, + result::TreeSetBicoloringResult, ) - (; large_colptr, large_rowval, symmetric_result) = result + (; + symmetric_color, + symmetric_to_row, + symmetric_to_column, + reverse_bfs_orders, + tree_edge_indices, + nt, + buffer, + ) = result m, n = size(A) R = eltype(A) fill!(A, zero(R)) - nzval = Vector{R}(undef, length(large_rowval)) - A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, nzval) - Br_and_Bc = _join_compressed!(result, Br, Bc) - decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L) - rvA = rowvals(A_and_noAᵀ) - nzA = nonzeros(A_and_noAᵀ) - for j in 1:n - for k in nzrange(A_and_noAᵀ, j) - i = rvA[k] - A[i - n, j] = nzA[k] + + if eltype(buffer) == R + buffer_right_type = buffer + else + buffer_right_type = similar(buffer, R) + end + + for k in 1:nt + # Positions of the first and last edges of the tree + first = tree_edge_indices[k] + last = tree_edge_indices[k + 1] - 1 + + # Reset the buffer to zero for all vertices in the tree (except the root) + for pos in first:last + (vertex, _) = reverse_bfs_orders[pos] + buffer_right_type[vertex] = zero(R) + end + # Reset the buffer to zero for the root vertex + (_, root) = reverse_bfs_orders[last] + buffer_right_type[root] = zero(R) + + for pos in first:last + (i, j) = reverse_bfs_orders[pos] + cj = symmetric_color[j] + if in_triangle(i, j, :L) + val = Bc[i - n, symmetric_to_column[cj]] - buffer_right_type[i] + buffer_right_type[j] = buffer_right_type[j] + val + A[i - n, j] = val + else + val = Br[symmetric_to_row[cj], i] - buffer_right_type[i] + buffer_right_type[j] = buffer_right_type[j] + val + A[j - n, i] = val + end end end return A end function decompress!( - A::SparseMatrixCSC, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult + A::SparseMatrixCSC, + Br::AbstractMatrix, + Bc::AbstractMatrix, + result::TreeSetBicoloringResult, ) - (; large_colptr, large_rowval, symmetric_result) = result + (; + symmetric_color, + symmetric_to_column, + symmetric_to_row, + reverse_bfs_orders, + tree_edge_indices, + nt, + A_indices, + buffer, + ) = result m, n = size(A) - # pretend A is larger - A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval) - # decompress lower triangle only - Br_and_Bc = _join_compressed!(result, Br, Bc) - decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L) + R = eltype(A) + nzA = nonzeros(A) + + if eltype(buffer) == R + buffer_right_type = buffer + else + buffer_right_type = similar(buffer, R) + end + + counter = 0 + for k in 1:nt + # Positions of the first and last edges of the tree + first = tree_edge_indices[k] + last = tree_edge_indices[k + 1] - 1 + + # Reset the buffer to zero for all vertices in the tree (except the root) + for pos in first:last + (vertex, _) = reverse_bfs_orders[pos] + buffer_right_type[vertex] = zero(R) + end + # Reset the buffer to zero for the root vertex + (_, root) = reverse_bfs_orders[last] + buffer_right_type[root] = zero(R) + + for pos in first:last + (i, j) = reverse_bfs_orders[pos] + cj = symmetric_color[j] + counter += 1 + + #! format: off + if in_triangle(i, j, :L) + val = Bc[i - n, symmetric_to_column[cj]] - buffer_right_type[i] + buffer_right_type[j] = buffer_right_type[j] + val + + # A[i-n,j] is stored at A_indices[counter] in nonzeros(A) + nzA[A_indices[counter]] = val + + else + val = Br[symmetric_to_row[cj], i] - buffer_right_type[i] + buffer_right_type[j] = buffer_right_type[j] + val + + # A[j-n,i] is stored at A_indices[counter] in nonzeros(A) + nzA[A_indices[counter]] = val + end + #! format: on + end + end return A end diff --git a/src/graph.jl b/src/graph.jl index eef76b92..e217a86e 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -101,10 +101,13 @@ end """ bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) -Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations. +Return the [`SparsityPatternCSC`](@ref) `S` of `A`, the [`SparsityPatternCSC`](@ref) `S_and_Sᵀ` corresponding to the matrix `[0 Aᵀ; A 0]`, and the mapping `edge_to_index`, with a minimum of allocations. + +`S` is returned alongside `S_and_Sᵀ` because bidirectional decompression indexes directly into the nonzeros of `A`. """ function bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) - return bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern) + S = SparsityPatternCSC(SparseMatrixCSC(A)) + return bidirectional_pattern(S; symmetric_pattern) end function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T} @@ -176,7 +179,7 @@ function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool # Create the SparsityPatternCSC of the augmented adjacency matrix S_and_Sᵀ = SparsityPatternCSC{T}(p, p, colptr, rowval) - return S_and_Sᵀ, edge_to_index + return S, S_and_Sᵀ, edge_to_index end function build_edge_to_index(S::SparsityPatternCSC{T}) where {T} diff --git a/src/interface.jl b/src/interface.jl index a219186d..8e697847 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -338,7 +338,7 @@ function _coloring( symmetric_pattern::Bool; forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing, ) where {R} - A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern) + S, A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern) ag = AdjacencyGraph( A_and_Aᵀ, edge_to_index, 0; augmented_graph=true, original_size=size(A) ) @@ -368,16 +368,16 @@ function _coloring( t -> maximum(t[3]) + maximum(t[4]), outputs_by_order ) # can't use ncolors without computing the full result if speed_setting isa WithResult - symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set) - return BicoloringResult( + return StarSetBicoloringResult( A, + S, ag, - symmetric_result, + color, + star_set, row_color, column_color, symmetric_to_row, symmetric_to_column, - R, ) else return row_color, column_color @@ -392,7 +392,7 @@ function _coloring( decompression_eltype::Type{R}, symmetric_pattern::Bool, ) where {R} - A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern) + S, A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern) ag = AdjacencyGraph( A_and_Aᵀ, edge_to_index, 0; augmented_graph=true, original_size=size(A) ) @@ -418,11 +418,12 @@ function _coloring( t -> maximum(t[3]) + maximum(t[4]), outputs_by_order ) # can't use ncolors without computing the full result if speed_setting isa WithResult - symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R) - return BicoloringResult( + return TreeSetBicoloringResult( A, + S, ag, - symmetric_result, + color, + tree_set, row_color, column_color, symmetric_to_row, diff --git a/src/result.jl b/src/result.jl index 489426c9..4b1325f4 100644 --- a/src/result.jl +++ b/src/result.jl @@ -631,10 +631,20 @@ function remap_colors( return row_color, column_color, symmetric_to_row, symmetric_to_column end +function column_colors(result::AbstractColoringResult{:nonsymmetric,:bidirectional}) + return result.column_color +end +function column_groups(result::AbstractColoringResult{:nonsymmetric,:bidirectional}) + return result.column_group +end + +row_colors(result::AbstractColoringResult{:nonsymmetric,:bidirectional}) = result.row_color +row_groups(result::AbstractColoringResult{:nonsymmetric,:bidirectional}) = result.row_group + """ $TYPEDEF -Storage for the result of a bidirectional coloring with direct or substitution decompression, based on the symmetric coloring of a 2x2 block matrix. +Storage for the result of a bidirectional coloring with direct decompression, based on the symmetric star coloring of a 2 x 2 block matrix. # Fields @@ -644,19 +654,22 @@ $TYPEDFIELDS - [`AbstractColoringResult`](@ref) """ -struct BicoloringResult{ +struct StarSetBicoloringResult{ M<:AbstractMatrix, T<:Integer, G<:AdjacencyGraph{T}, - decompression, GT<:AbstractGroups{T}, - SR<:AbstractColoringResult{:symmetric,:column,decompression}, - R, -} <: AbstractColoringResult{:nonsymmetric,:bidirectional,decompression} + VT<:AbstractVector{T}, + A, +} <: AbstractColoringResult{:nonsymmetric,:bidirectional,:direct} "matrix that was colored" A::M + "sparsity pattern of the matrix that was colored" + S::SparsityPatternCSC{T} "augmented adjacency graph that was used for bicoloring" abg::G + "one integer color for each vertex of the augmented adjacency graph" + symmetric_color::Vector{T} "one integer color for each column" column_color::Vector{T} "one integer color for each row" @@ -665,57 +678,297 @@ struct BicoloringResult{ column_group::GT "color groups for rows" row_group::GT - "result for the coloring of the symmetric 2 x 2 block matrix" - symmetric_result::SR "maps symmetric colors to column colors" symmetric_to_column::Vector{T} "maps symmetric colors to row colors" symmetric_to_row::Vector{T} - "combination of `Br` and `Bc` (almost a concatenation up to color remapping)" - Br_and_Bc::Matrix{R} - "CSC storage of `A_and_noAᵀ - `colptr`" - large_colptr::Vector{T} - "CSC storage of `A_and_noAᵀ - `rowval`" - large_rowval::Vector{T} + "increasing positions in `nonzeros(A)` of the coefficients recovered from `Bc`" + A_indices_bc::VT + "linear indices in `Bc` of those same coefficients" + compressed_indices_bc::VT + "increasing positions in `nonzeros(A)` of the coefficients recovered from `Br`" + A_indices_br::VT + "linear indices in `Br` of those same coefficients" + compressed_indices_br::VT + "optional data used for decompressing into specific matrix types" + additional_info::A +end + +function StarSetBicoloringResult( + A::AbstractMatrix, + S::SparsityPatternCSC{T}, + ag::AdjacencyGraph{T}, + symmetric_color::Vector{<:Integer}, + star_set::StarSet{<:Integer}, + row_color::Vector{T}, + column_color::Vector{T}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, +) where {T<:Integer} + column_group = group_by_color(T, column_color) + row_group = group_by_color(T, row_color) + A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br = star_bicoloring_csc_indices( + S, + symmetric_color, + star_set, + symmetric_to_row, + symmetric_to_column, + length(row_group), + ) + return StarSetBicoloringResult( + A, + S, + ag, + symmetric_color, + column_color, + row_color, + column_group, + row_group, + symmetric_to_column, + symmetric_to_row, + A_indices_bc, + compressed_indices_bc, + A_indices_br, + compressed_indices_br, + nothing, + ) end -column_colors(result::BicoloringResult) = result.column_color -column_groups(result::BicoloringResult) = result.column_group +#= +Return `(A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br)` for a direct +bidirectional decompression where the nonzero coefficients of `A` are stored in CSC order. +They satisfy `nonzeros(A)[A_indices_bc[t]] = vec(Bc)[compressed_indices_bc[t]]` and +`nonzeros(A)[A_indices_br[t]] = vec(Br)[compressed_indices_br[t]]`. +`A_indices_bc` and `A_indices_br` are increasing and partition `1:nnz(S)`. +=# +function star_bicoloring_csc_indices( + S::SparsityPatternCSC{T}, + symmetric_color::Vector{<:Integer}, + star_set::StarSet{<:Integer}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, + num_row_colors::Integer, +) where {T<:Integer} + m, n = size(S) + (; star, hub) = star_set + rvS = rowvals(S) + nnzA = nnz(S) + A_indices_bc = Vector{T}(undef, nnzA) + compressed_indices_bc = Vector{T}(undef, nnzA) + A_indices_br = Vector{T}(undef, nnzA) + compressed_indices_br = Vector{T}(undef, nnzA) + + nb_bc = 0 + nb_br = 0 + for j in 1:n + for k in nzrange(S, j) + i = rvS[k] + # the first nnzA edges of the augmented graph are the nonzeros of A in CSC order + s = star[k] + h = abs(hub[s]) + if j == h + # j is the hub and (i + n) is the spoke + c = symmetric_color[j] + # A[i, j] = Bc[i, symmetric_to_column[c]] + nb_bc += 1 + A_indices_bc[nb_bc] = k + compressed_indices_bc[nb_bc] = (symmetric_to_column[c] - 1) * m + i + else # i + n == h + # (i + n) is the hub and j is the spoke + c = symmetric_color[i + n] + # A[i, j] = Br[symmetric_to_row[c], j] + nb_br += 1 + A_indices_br[nb_br] = k + compressed_indices_br[nb_br] = + (j - 1) * num_row_colors + symmetric_to_row[c] + end + end + end + resize!(A_indices_bc, nb_bc) + resize!(compressed_indices_bc, nb_bc) + resize!(A_indices_br, nb_br) + resize!(compressed_indices_br, nb_br) + return A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br +end -row_colors(result::BicoloringResult) = result.row_color -row_groups(result::BicoloringResult) = result.row_group +#= +Same as `star_bicoloring_csc_indices`, except that `A_indices` refers to the nonzero coefficients +of `A` stored in CSR order. +The columns `n+1:n+m` of the augmented adjacency graph hold `Aᵀ` in CSC order, which is exactly +`A` in CSR order, and `edge_indices(ag)` maps each of those positions back to the CSC position of +the same coefficient. +=# +function star_bicoloring_csr_indices( + ag::AdjacencyGraph{T}, + S::SparsityPatternCSC{T}, + symmetric_color::Vector{<:Integer}, + star_set::StarSet{<:Integer}, + symmetric_to_row::Vector{T}, + symmetric_to_column::Vector{T}, + num_row_colors::Integer, +) where {T<:Integer} + m, n = size(S) + (; star, hub) = star_set + S_aug = pattern(ag) + edge_to_index = edge_indices(ag) + rv_aug = rowvals(S_aug) + nnzA = nnz(S) + A_indices_bc = Vector{T}(undef, nnzA) + compressed_indices_bc = Vector{T}(undef, nnzA) + A_indices_br = Vector{T}(undef, nnzA) + compressed_indices_br = Vector{T}(undef, nnzA) + + nb_bc = 0 + nb_br = 0 + for i in 1:m + for t in nzrange(S_aug, n + i) + j = rv_aug[t] + k = edge_to_index[t] # position of A[i, j] in CSC order + csr_position = t - nnzA # position of A[i, j] in CSR order + s = star[k] + h = abs(hub[s]) + if j == h + # j is the hub and (i + n) is the spoke + c = symmetric_color[j] + # A[i, j] = Bc[i, symmetric_to_column[c]] + nb_bc += 1 + A_indices_bc[nb_bc] = csr_position + compressed_indices_bc[nb_bc] = (symmetric_to_column[c] - 1) * m + i + else # i + n == h + # (i + n) is the hub and j is the spoke + c = symmetric_color[i + n] + # A[i, j] = Br[symmetric_to_row[c], j] + nb_br += 1 + A_indices_br[nb_br] = csr_position + compressed_indices_br[nb_br] = + (j - 1) * num_row_colors + symmetric_to_row[c] + end + end + end + resize!(A_indices_bc, nb_bc) + resize!(compressed_indices_bc, nb_bc) + resize!(A_indices_br, nb_br) + resize!(compressed_indices_br, nb_br) + return A_indices_bc, compressed_indices_bc, A_indices_br, compressed_indices_br +end + +""" +$TYPEDEF + +Storage for the result of a bidirectional coloring with decompression by substitution, based on the symmetric acyclic coloring of a 2 x 2 block matrix. + +# Fields -function BicoloringResult( +$TYPEDFIELDS + +# See also + +- [`AbstractColoringResult`](@ref) +""" +struct TreeSetBicoloringResult{ + M<:AbstractMatrix,T<:Integer,G<:AdjacencyGraph{T},GT<:AbstractGroups{T},R +} <: AbstractColoringResult{:nonsymmetric,:bidirectional,:substitution} + "matrix that was colored" + A::M + "sparsity pattern of the matrix that was colored" + S::SparsityPatternCSC{T} + "augmented adjacency graph that was used for bicoloring" + abg::G + "one integer color for each vertex of the augmented adjacency graph" + symmetric_color::Vector{T} + "one integer color for each column" + column_color::Vector{T} + "one integer color for each row" + row_color::Vector{T} + "color groups for columns" + column_group::GT + "color groups for rows" + row_group::GT + "maps symmetric colors to column colors" + symmetric_to_column::Vector{T} + "maps symmetric colors to row colors" + symmetric_to_row::Vector{T} + "position in `nonzeros(A)` of the coefficient recovered at each step of the substitution" + A_indices::Vector{T} + "storage for the edges of each tree in reverse BFS order" + reverse_bfs_orders::Vector{Tuple{T,T}} + "internal storage for the positions of the trees inside `reverse_bfs_orders`" + tree_edge_indices::Vector{T} + "number of trees" + nt::T + "buffer needed during decompression by substitution" + buffer::Vector{R} +end + +function TreeSetBicoloringResult( A::AbstractMatrix, + S::SparsityPatternCSC{T}, ag::AdjacencyGraph{T}, - symmetric_result::AbstractColoringResult{:symmetric,:column}, + symmetric_color::Vector{<:Integer}, + tree_set::TreeSet{<:Integer}, row_color::Vector{T}, column_color::Vector{T}, symmetric_to_row::Vector{T}, symmetric_to_column::Vector{T}, decompression_eltype::Type{R}, -) where {T,R} +) where {T<:Integer,R} + (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set m, n = size(A) - symmetric_color = column_colors(symmetric_result) - num_sym_colors = maximum(symmetric_color) column_group = group_by_color(T, column_color) row_group = group_by_color(T, row_color) - Br_and_Bc = Matrix{R}(undef, n + m, num_sym_colors) - large_colptr = copy(ag.S.colptr) - large_colptr[(n + 2):end] .= large_colptr[n + 1] # last few columns are empty - large_rowval = ag.S.rowval[1:(end ÷ 2)] # forget the second half of nonzeros - return BicoloringResult( + + rvS = rowvals(S) + A_indices = Vector{T}(undef, nnz(S)) + + index = 0 + for k in 1:nt + # Positions of the edges for each tree + first = tree_edge_indices[k] + last = tree_edge_indices[k + 1] - 1 + + for pos in first:last + (leaf, neighbor) = reverse_bfs_orders[pos] + i = leaf + j = neighbor + index += 1 + + #! format: off + # The vertices 1:n are the columns of A, the vertices n+1:n+m are its rows + if in_triangle(i, j, :L) + # (i - n, j) is a nonzero of A, stored at (S.colptr[j] + offset) in nonzeros(A) + col_j = view(rvS, nzrange(S, j)) + offset = searchsortedfirst(col_j, i - n)::Int - 1 + A_indices[index] = S.colptr[j] + offset + + else + # (j - n, i) is a nonzero of A, stored at (S.colptr[i] + offset) in nonzeros(A) + col_i = view(rvS, nzrange(S, i)) + offset = searchsortedfirst(col_i, j - n)::Int - 1 + A_indices[index] = S.colptr[i] + offset + end + #! format: on + end + end + + # buffer holds the sum of edge values for subtrees in a tree. + # For each vertex i, buffer[i] is the sum of edge values in the subtree rooted at i. + buffer = Vector{R}(undef, n + m) + + return TreeSetBicoloringResult( A, + S, ag, + symmetric_color, column_color, row_color, column_group, row_group, - symmetric_result, symmetric_to_column, symmetric_to_row, - Br_and_Bc, - large_colptr, - large_rowval, + A_indices, + reverse_bfs_orders, + tree_edge_indices, + nt, + buffer, ) end diff --git a/test/allocations.jl b/test/allocations.jl index 31305e63..7bdee088 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -35,7 +35,7 @@ function test_noallocs_sparse_decompression( bench1_full = @be similar(A) decompress!(_, Br, Bc, result) evals = 1 bench2_full = @be similar(Matrix(A)) decompress!(_, Br, Bc, result) evals = 1 @test minimum(bench1_full).allocs == 0 - @test_broken minimum(bench2_full).allocs == 0 + @test minimum(bench2_full).allocs == 0 end else B = compress(A, result) diff --git a/test/cuda.jl b/test/cuda.jl index 6ae73060..2d17083a 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -60,3 +60,27 @@ end; ) end end; + +@testset verbose = true "Bidirectional coloring & direct decompression" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional) + algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:direct) + @testset for T in (CuSparseMatrixCSC, CuSparseMatrixCSR) + @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params + A0 = T(sprand(rng, m, n, p)) + test_bicoloring_decompression(A0, problem, algo; gpu=true) + end + end +end; + +@testset verbose = true "Bidirectional decompression by substitution is unsupported" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional) + algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:substitution) + @testset for T in (CuSparseMatrixCSC, CuSparseMatrixCSR) + A0 = T(sprand(rng, 20, 10, 0.3)) + result = coloring(A0, problem, algo) + Br, Bc = compress(A0, result) + @test_throws SMC.UnsupportedDecompressionError decompress!( + similar(A0), Br, Bc, result + ) + end +end; diff --git a/test/graph.jl b/test/graph.jl index 2a6211c4..7d56d2d5 100644 --- a/test/graph.jl +++ b/test/graph.jl @@ -36,7 +36,11 @@ using Test p = 0.05 * rand() A = sprand(Bool, m, n, p) A_and_Aᵀ = [spzeros(Bool, n, n) transpose(A); A spzeros(Bool, m, m)] - S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=false) + S, S_and_Sᵀ, edge_to_index = bidirectional_pattern( + A; symmetric_pattern=false + ) + @test S.colptr == A.colptr + @test S.rowval == A.rowval @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval M = SparseMatrixCSC( @@ -51,7 +55,11 @@ using Test p = 0.05 * rand() A = sparse(Symmetric(sprand(Bool, m, m, p))) A_and_Aᵀ = [spzeros(Bool, m, m) transpose(A); A spzeros(Bool, m, m)] - S_and_Sᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern=true) + S, S_and_Sᵀ, edge_to_index = bidirectional_pattern( + A; symmetric_pattern=true + ) + @test S.colptr == A.colptr + @test S.rowval == A.rowval @test S_and_Sᵀ.colptr == A_and_Aᵀ.colptr @test S_and_Sᵀ.rowval == A_and_Aᵀ.rowval M = SparseMatrixCSC( diff --git a/test/rocm.jl b/test/rocm.jl index e2aa7840..f212eec0 100644 --- a/test/rocm.jl +++ b/test/rocm.jl @@ -60,3 +60,27 @@ end; ) end end; + +@testset verbose = true "Bidirectional coloring & direct decompression" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional) + algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:direct) + @testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR) + @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params + A0 = T(sprand(rng, m, n, p)) + test_bicoloring_decompression(A0, problem, algo; gpu=true) + end + end +end; + +@testset verbose = true "Bidirectional decompression by substitution is unsupported" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional) + algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:substitution) + @testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR) + A0 = T(sprand(rng, 20, 10, 0.3)) + result = coloring(A0, problem, algo) + Br, Bc = compress(A0, result) + @test_throws SMC.UnsupportedDecompressionError decompress!( + similar(A0), Br, Bc, result + ) + end +end; diff --git a/test/utils.jl b/test/utils.jl index 98c12afa..4e4898d8 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -207,6 +207,7 @@ function test_bicoloring_decompression( problem::ColoringProblem{:nonsymmetric,:bidirectional}, algo::GreedyColoringAlgorithm{decompression}; test_fast=false, + gpu=false, ) where {decompression} @testset "$(typeof(A))" for A in matrix_versions(A0) yield() @@ -241,6 +242,10 @@ function test_bicoloring_decompression( ) ≈ A0 end + if gpu + continue + end + if decompression == :direct @testset "Recoverability" begin @test structurally_biorthogonal(A0, row_color, column_color) @@ -257,6 +262,10 @@ function test_bicoloring_decompression( end end + if gpu + return nothing + end + @testset "More orders is better" begin more_orders = (algo.orders..., _ALL_ORDERS...) better_algo = GreedyColoringAlgorithm{decompression}(