Skip to content

[bicoloring] Use Br and Bc directly for the decompression #194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
amontoison wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[bicoloring] Use Br and Bc directly for the decompression #194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93459cb
[bicoloring] Use Br and Bc directly for the decompression
amontoison Sep 18, 2026
1f98c32
Support bicoloring on GPUs
amontoison Sep 19, 2026
c07b879
Improve star_bicoloring_csc_indices and star_bicoloring_csr_indices
amontoison Sep 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion docs/src/dev.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ SparseMatrixColorings.RowColoringResult
SparseMatrixColorings.StarSetColoringResult
SparseMatrixColorings.TreeSetColoringResult
SparseMatrixColorings.LinearSystemColoringResult
SparseMatrixColorings.BicoloringResult
SparseMatrixColorings.StarSetBicoloringResult
SparseMatrixColorings.TreeSetBicoloringResult
SparseMatrixColorings.remap_colors
SparseMatrixColorings.decompress_csc!
```
Expand Down
126 changes: 126 additions & 0 deletions ext/SparseMatrixColoringsAMDGPUExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
126 changes: 126 additions & 0 deletions ext/SparseMatrixColoringsCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
37 changes: 37 additions & 0 deletions ext/SparseMatrixColoringsGPUArraysExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion src/adtypes.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 5 additions & 13 deletions src/check.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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]
Expand Down
Loading
Loading