Skip to content

Support decompression in one triangle #286

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

Draft
amontoison wants to merge 2 commits into
base: am/postprocessing_v2
Choose a base branch
from
Draft

Support decompression in one triangle #286

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
87c8016
Revisit the post-processing
amontoison Dec 31, 2025
cc29f75
Support decompression in one triangle
amontoison Dec 31, 2025
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
31 changes: 11 additions & 20 deletions ext/SparseMatrixColoringsCUDAExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -86,12 +87,18 @@ 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 = if (decompression_uplo == :L)
:U
else
(decompression_uplo == :U ? :L : decompression_uplo)
end
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

Expand Down Expand Up @@ -120,15 +127,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
Expand All @@ -138,15 +137,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
Expand Down
4 changes: 3 additions & 1 deletion src/SparseMatrixColorings.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
15 changes: 10 additions & 5 deletions src/coloring.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ end
"""
star_coloring(
g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
forced_colors::Union{AbstractVector,Nothing}=nothing
postprocessing_minimizes::Symbol=:all_colors, forced_colors::Union{AbstractVector,Nothing}=nothing
)

Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where
Expand Down Expand Up @@ -110,6 +110,7 @@ function star_coloring(
g::AdjacencyGraph{T},
vertices_in_order::AbstractVector{<:Integer},
postprocessing::Bool;
postprocessing_minimizes::Symbol=:all_colors,
forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
) where {T<:Integer}
# Initialize data structures
Expand Down Expand Up @@ -168,7 +169,7 @@ function star_coloring(
if postprocessing
# Reuse the vector forbidden_colors to compute offsets during post-processing
offsets = forbidden_colors
postprocess!(color, star_set, g, offsets)
postprocess!(color, star_set, g, offsets, postprocessing_minimizes)
end
return color, star_set
end
Expand Down Expand Up @@ -250,7 +251,8 @@ struct StarSet{T}
end

"""
acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
postprocessing_minimizes::Symbol=:all_colors)

Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a tuple `(color, tree_set)`, where

Expand All @@ -273,7 +275,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 3.1
"""
function acyclic_coloring(
g::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
g::AdjacencyGraph{T},
vertices_in_order::AbstractVector{<:Integer},
postprocessing::Bool;
postprocessing_minimizes::Symbol=:all_colors,
) where {T<:Integer}
# Initialize data structures
nv = nb_vertices(g)
Expand Down Expand Up @@ -345,7 +350,7 @@ function acyclic_coloring(
if postprocessing
# Reuse the vector forbidden_colors to compute offsets during post-processing
offsets = forbidden_colors
postprocess!(color, tree_set, g, offsets)
postprocess!(color, tree_set, g, offsets, postprocessing_minimizes)
end
return color, tree_set
end
Expand Down
118 changes: 45 additions & 73 deletions src/decompression.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,19 +175,32 @@ 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!(
Expand All @@ -204,9 +217,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.

Expand Down Expand Up @@ -260,7 +270,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).
Expand All @@ -272,9 +282,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.

Expand Down Expand Up @@ -445,97 +452,66 @@ end

## StarSetColoringResult

function decompress!(
A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
)
(; ag, compressed_indices) = result
function decompress!(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,
b::AbstractVector,
c::Integer,
result::StarSetColoringResult,
uplo::Symbol=:F,
A::SparseMatrixCSC, b::AbstractVector, c::Integer, result::StarSetColoringResult
)
(; ag, compressed_indices, group) = result
(; ag, compressed_indices, group, uplo) = result
(; S) = ag
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, diagonal_indices, 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)
Expand Down Expand Up @@ -586,10 +562,7 @@ function decompress!(
end

function decompress!(
A::SparseMatrixCSC{R},
B::AbstractMatrix{R},
result::TreeSetColoringResult,
uplo::Symbol=:F,
A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
) where {R<:Real}
(;
ag,
Expand All @@ -602,6 +575,7 @@ function decompress!(
lower_triangle_offsets,
upper_triangle_offsets,
buffer,
uplo,
) = result
(; S) = ag
A_colptr = A.colptr
Expand Down Expand Up @@ -702,12 +676,10 @@ end
## MatrixInverseColoringResult

function decompress!(
A::AbstractMatrix,
B::AbstractMatrix,
result::LinearSystemColoringResult,
uplo::Symbol=:F,
A::AbstractMatrix, B::AbstractMatrix, result::LinearSystemColoringResult
)
(; 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)

Expand Down Expand Up @@ -776,7 +748,7 @@ function decompress!(
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)
decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result)
rvA = rowvals(A_and_noAᵀ)
nzA = nonzeros(A_and_noAᵀ)
for j in 1:n
Expand All @@ -797,6 +769,6 @@ function decompress!(
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)
decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result)
return A
end
Loading
Loading