From e9819172bb01e57ec2ef0bb7ba0077d5b3f794c7 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Sun, 4 Jan 2026 23:31:42 +0100 Subject: [PATCH 1/4] Add an option decompression_uplo for symmetric results --- ext/SparseMatrixColoringsCUDAExt.jl | 16 +++--- src/decompression.jl | 6 ++- src/interface.jl | 8 +-- src/result.jl | 83 +++++++++++++++++++---------- 4 files changed, 75 insertions(+), 38 deletions(-) diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl index ed33eece..265fcb93 100644 --- a/ext/SparseMatrixColoringsCUDAExt.jl +++ b/ext/SparseMatrixColoringsCUDAExt.jl @@ -31,13 +31,15 @@ function SMC.StarSetColoringResult( A::CuSparseMatrixCSC, ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, - star_set::SMC.StarSet{<:Integer}, + star_set::SMC.StarSet{<:Integer}; + decompression_uplo::Symbol=:F, ) where {T<:Integer} + @assert decompression_uplo == :F 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 @@ -69,13 +71,15 @@ function SMC.StarSetColoringResult( A::CuSparseMatrixCSR, ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, - star_set::SMC.StarSet{<:Integer}, + star_set::SMC.StarSet{<:Integer}; + decompression_uplo::Symbol=:F, ) where {T<:Integer} + @assert decompression_uplo == :F 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_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 diff --git a/src/decompression.jl b/src/decompression.jl index 2dcf1847..c4ce5acf 100644 --- a/src/decompression.jl +++ b/src/decompression.jl @@ -448,6 +448,7 @@ end function decompress!( A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F ) + @assert result.decompression_uplo == :F (; ag, compressed_indices) = result (; S) = ag check_compatible_pattern(A, ag, uplo) @@ -472,6 +473,7 @@ function decompress_single_color!( result::StarSetColoringResult, uplo::Symbol=:F, ) + @assert result.decompression_uplo == :F (; ag, compressed_indices, group) = result (; S) = ag check_compatible_pattern(A, ag, uplo) @@ -509,11 +511,12 @@ function decompress!( (; S) = ag nzA = nonzeros(A) check_compatible_pattern(A, ag, uplo) - if uplo == :F + if result.decompression_uplo == uplo for k in eachindex(nzA, compressed_indices) nzA[k] = B[compressed_indices[k]] end else + @assert result.decompression_uplo == :F rvS = rowvals(S) l = 0 # assume A has the same pattern as the triangle for j in axes(S, 2) @@ -534,6 +537,7 @@ end function decompress!( A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F ) + @assert result.decompression_uplo == :F (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, diagonal_indices, buffer) = result (; S) = ag diff --git a/src/interface.jl b/src/interface.jl index 9480e07c..49331f20 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -294,7 +294,7 @@ function _coloring( end color, star_set = argmin(maximum ∘ first, color_and_star_set_by_order) if speed_setting isa WithResult - return StarSetColoringResult(A, ag, color, star_set) + return StarSetColoringResult(A, ag, color, star_set, :F) else return color end @@ -321,7 +321,7 @@ function _coloring( color, tree_set = argmin(maximum ∘ first, color_and_tree_set_by_order) end if speed_setting isa WithResult - return TreeSetColoringResult(A, ag, color, tree_set, R) + return TreeSetColoringResult(A, ag, color, tree_set, R, :F) else return color end @@ -366,7 +366,7 @@ 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) + symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set, :L) return BicoloringResult( A, ag, @@ -416,7 +416,7 @@ 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) + symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R, :L) return BicoloringResult( A, ag, diff --git a/src/result.jl b/src/result.jl index 489426c9..22956a64 100644 --- a/src/result.jl +++ b/src/result.jl @@ -312,6 +312,7 @@ struct StarSetColoringResult{ color::CT group::GT compressed_indices::VT + decompression_uplo::Symbol additional_info::A end @@ -320,43 +321,58 @@ function StarSetColoringResult( ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} group = group_by_color(T, color) - compressed_indices = star_csc_indices(ag, color, star_set) - return StarSetColoringResult(A, ag, color, group, compressed_indices, nothing) + compressed_indices = star_csc_indices(ag, color, star_set, decompression_uplo) + return StarSetColoringResult( + A, ag, color, group, compressed_indices, decompression_uplo, nothing + ) end function star_csc_indices( - ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set + ag::AdjacencyGraph{T}, + color::Vector{<:Integer}, + star_set::StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T} (; star, hub) = star_set S = pattern(ag) edge_to_index = edge_indices(ag) n = S.n rvS = rowvals(S) - compressed_indices = zeros(T, nnz(S)) # needs to be independent from the storage in the graph, in case the graph gets reused + nb_indices = nnz(S) + if decompression_uplo != :F + nb_indices = nb_edges(ag) + ag.nb_self_loops + end + compressed_indices = zeros(T, nb_indices) # needs to be independent from the storage in the graph, in case the graph gets reused + l = 0 for j in axes(S, 2) for k in nzrange(S, j) i = rvS[k] if i == j # diagonal coefficients + l += 1 c = color[i] - compressed_indices[k] = (c - 1) * n + i + compressed_indices[l] = (c - 1) * n + i else - # off-diagonal coefficients - index_ij = edge_to_index[k] - s = star[index_ij] - h = abs(hub[s]) - - # Assign the non-hub vertex (spoke) to the correct position in spokes - if i == h - # i is the hub and j is the spoke - c = color[i] - compressed_indices[k] = (c - 1) * n + j - else # j == h - # j is the hub and i is the spoke - c = color[j] - compressed_indices[k] = (c - 1) * n + i + if in_triangle(i, j, decompression_uplo) + # off-diagonal coefficients + l += 1 + index_ij = edge_to_index[k] + s = star[index_ij] + h = abs(hub[s]) + + # Assign the non-hub vertex (spoke) to the correct position in spokes + if i == h + # i is the hub and j is the spoke + c = color[i] + compressed_indices[l] = (c - 1) * n + j + else # j == h + # j is the hub and i is the spoke + c = color[j] + compressed_indices[l] = (c - 1) * n + i + end end end end @@ -394,6 +410,7 @@ struct TreeSetColoringResult{ lower_triangle_offsets::Vector{T} upper_triangle_offsets::Vector{T} buffer::Vector{R} + decompression_uplo::Symbol end function TreeSetColoringResult( @@ -402,6 +419,7 @@ function TreeSetColoringResult( color::Vector{<:Integer}, tree_set::TreeSet{<:Integer}, decompression_eltype::Type{R}, + decompression_uplo::Symbol, ) where {T<:Integer,R} (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set (; S, nb_self_loops) = ag @@ -411,7 +429,7 @@ function TreeSetColoringResult( # Vector for the decompression of the diagonal coefficients diagonal_indices = Vector{T}(undef, nb_self_loops) - diagonal_nzind = Vector{T}(undef, nb_self_loops) + diagonal_nzind = (decompression_uplo == :F) ? Vector{T}(undef, nb_self_loops) : T[] if !augmented_graph(ag) l = 0 @@ -421,7 +439,9 @@ function TreeSetColoringResult( if i == j l += 1 diagonal_indices[l] = i - diagonal_nzind[l] = k + if decompression_uplo == :F + diagonal_nzind[l] = k + end end end end @@ -429,8 +449,8 @@ function TreeSetColoringResult( # Vectors for the decompression of the off-diagonal coefficients nedges = nb_edges(ag) - lower_triangle_offsets = Vector{T}(undef, nedges) - upper_triangle_offsets = Vector{T}(undef, nedges) + lower_triangle_offsets = decompression_uplo == :U ? T[] : Vector{T}(undef, nedges) + upper_triangle_offsets = decompression_uplo == :L ? T[] : Vector{T}(undef, nedges) # Index in lower_triangle_offsets and upper_triangle_offsets index_offsets = 0 @@ -454,21 +474,29 @@ function TreeSetColoringResult( if in_triangle(i, j, :L) # uplo = :L or uplo = :F # S[i,j] is stored at index_ij = (S.colptr[j+1] - offset_L) in S.nzval - lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1 + if decompression_uplo != :U + lower_triangle_offsets[index_offsets] = length(col_j) - searchsortedfirst(col_j, i) + 1 + end # uplo = :U or uplo = :F # S[j,i] is stored at index_ji = (S.colptr[i] + offset_U) in S.nzval - upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1 + if decompression_uplo != :L + upper_triangle_offsets[index_offsets] = searchsortedfirst(col_i, j)::Int - 1 + end # S[i,j] is in the upper triangular part of S else # uplo = :U or uplo = :F # S[i,j] is stored at index_ij = (S.colptr[j] + offset_U) in S.nzval - upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1 + if decompression_uplo != :L + upper_triangle_offsets[index_offsets] = searchsortedfirst(col_j, i)::Int - 1 + end # uplo = :L or uplo = :F # S[j,i] is stored at index_ji = (S.colptr[i+1] - offset_L) in S.nzval - lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1 + if decompression_uplo != :U + lower_triangle_offsets[index_offsets] = length(col_i) - searchsortedfirst(col_i, j) + 1 + end end #! format: on end @@ -491,6 +519,7 @@ function TreeSetColoringResult( lower_triangle_offsets, upper_triangle_offsets, buffer, + decompression_uplo, ) end From fa20b33f80037ee156d037d344893bddb2557b24 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 5 Jan 2026 00:07:38 +0100 Subject: [PATCH 2/4] Fix tests --- ext/SparseMatrixColoringsCUDAExt.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl index 265fcb93..0f6face0 100644 --- a/ext/SparseMatrixColoringsCUDAExt.jl +++ b/ext/SparseMatrixColoringsCUDAExt.jl @@ -31,8 +31,8 @@ function SMC.StarSetColoringResult( A::CuSparseMatrixCSC, ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, - star_set::SMC.StarSet{<:Integer}; - decompression_uplo::Symbol=:F, + star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} @assert decompression_uplo == :F group = SMC.group_by_color(T, color) @@ -71,8 +71,8 @@ function SMC.StarSetColoringResult( A::CuSparseMatrixCSR, ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, - star_set::SMC.StarSet{<:Integer}; - decompression_uplo::Symbol=:F, + star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} @assert decompression_uplo == :F group = SMC.group_by_color(T, color) From bbd9261f6c0285ecf9b84fe71a639f8bebebd678 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Mon, 7 Sep 2026 14:37:27 -0600 Subject: [PATCH 3/4] Add a missing assert --- src/decompression.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/decompression.jl b/src/decompression.jl index c4ce5acf..c18ed323 100644 --- a/src/decompression.jl +++ b/src/decompression.jl @@ -721,6 +721,7 @@ function decompress!( uplo::Symbol=:F, ) where {R<:Real} check_compatible_pattern(A, result.ag, uplo) + @assert result.decompression_uplo == uplo || result.decompression_uplo == :F decompress_csc!(nonzeros(A), A.colptr, B, result, uplo) return A end From dd445bd350cf4a04cdbb0ce511d647e8608ec289 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Tue, 8 Sep 2026 21:58:38 -0600 Subject: [PATCH 4/4] Update the extension for AMD GPUs --- ext/SparseMatrixColoringsAMDGPUExt.jl | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/ext/SparseMatrixColoringsAMDGPUExt.jl b/ext/SparseMatrixColoringsAMDGPUExt.jl index 58714ad2..d79d5a14 100644 --- a/ext/SparseMatrixColoringsAMDGPUExt.jl +++ b/ext/SparseMatrixColoringsAMDGPUExt.jl @@ -32,12 +32,14 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + @assert decompression_uplo == :F 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=ROCVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end @@ -70,12 +72,14 @@ function SMC.StarSetColoringResult( ag::SMC.AdjacencyGraph{T}, color::Vector{<:Integer}, star_set::SMC.StarSet{<:Integer}, + decompression_uplo::Symbol, ) where {T<:Integer} + @assert decompression_uplo == :F 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_csr=ROCVector(compressed_indices)) return SMC.StarSetColoringResult( - A, ag, color, group, compressed_indices, additional_info + A, ag, color, group, compressed_indices, decompression_uplo, additional_info ) end