Store CSR indices as additional info

Author: gdalle

ext/SparseMatrixColoringsCUDAExt.jl

@@ -22,24 +22,26 @@ function SMC.compress(
     return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
 end
 
-## CSC
+## CSC Result
 
 function SMC.ColumnColoringResult(
     A::CuSparseMatrixCSC, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
 ) where {T<:Integer}
-    A_cpu = SparseMatrixCSC(A)
-    result_cpu = SMC.ColumnColoringResult(A_cpu, bg, color)
-    compressed_indices = CuVector(result_cpu.compressed_indices)
-    return SMC.ColumnColoringResult(A, bg, color, result_cpu.group, compressed_indices)
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.column_csc_indices(bg, color)
+    additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
+    return SMC.ColumnColoringResult(
+        A, bg, color, group, compressed_indices, additional_info
+    )
 end
 
 function SMC.RowColoringResult(
     A::CuSparseMatrixCSC, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
 ) where {T<:Integer}
-    A_cpu = SparseMatrixCSC(A)
-    result_cpu = SMC.RowColoringResult(A_cpu, bg, color)
-    compressed_indices = CuVector(result_cpu.compressed_indices)
-    return SMC.RowColoringResult(A, bg, color, result_cpu.group, compressed_indices)
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.row_csc_indices(bg, color)
+    additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
+    return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
 end
 
 function SMC.StarSetColoringResult(
@@ -48,12 +50,54 @@ function SMC.StarSetColoringResult(
     color::Vector{<:Integer},
     star_set::SMC.StarSet{<:Integer},
 ) where {T<:Integer}
-    A_cpu = SparseMatrixCSC(A)
-    result_cpu = SMC.StarSetColoringResult(A_cpu, ag, color, star_set)
-    compressed_indices = CuVector(result_cpu.compressed_indices)
-    return SMC.StarSetColoringResult(A, ag, color, result_cpu.group, compressed_indices)
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.star_csc_indices(ag, color, star_set)
+    additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
+    return SMC.StarSetColoringResult(
+        A, ag, color, group, compressed_indices, additional_info
+    )
+end
+
+## CSR Result
+
+function SMC.ColumnColoringResult(
+    A::CuSparseMatrixCSR, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
+) where {T<:Integer}
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.column_csc_indices(bg, color)
+    compressed_indices_csr = SMC.column_csr_indices(bg, color)
+    additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices_csr))
+    return SMC.ColumnColoringResult(
+        A, bg, color, group, compressed_indices, additional_info
+    )
+end
+
+function SMC.RowColoringResult(
+    A::CuSparseMatrixCSR, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
+) where {T<:Integer}
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.row_csc_indices(bg, color)
+    compressed_indices_csr = SMC.row_csr_indices(bg, color)
+    additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices_csr))
+    return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
+end
+
+function SMC.StarSetColoringResult(
+    A::CuSparseMatrixCSR,
+    ag::SMC.AdjacencyGraph{T},
+    color::Vector{<:Integer},
+    star_set::SMC.StarSet{<:Integer},
+) where {T<:Integer}
+    group = SMC.group_by_color(T, color)
+    compressed_indices = SMC.star_csc_indices(ag, color, star_set)
+    additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices))
+    return SMC.StarSetColoringResult(
+        A, ag, color, group, compressed_indices, additional_info
+    )
 end
 
+## Decompression
+
 function update_nzval_from_matrix!(
     nzVal::AbstractVector, B::AbstractMatrix, compressed_indices::AbstractVector{<:Integer}
 )
@@ -70,14 +114,30 @@ for R in (:ColumnColoringResult, :RowColoringResult, :StarSetColoringResult)
     @eval function SMC.decompress!(
         A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSC}
     )
+        compressed_indices = result.additional_info.compressed_indices_gpu_csc
+        A.nnz == 0 && return A
+        kernel = @cuda launch = false update_nzval_from_matrix!(
+            A.nzVal, B, compressed_indices
+        )
+        config = launch_configuration(kernel.fun)
+        threads = min(A.nnz, config.threads)
+        blocks = cld(A.nnz, threads)
+        kernel(A.nzVal, B, compressed_indices; threads, blocks)
+        return A
+    end
+
+    @eval function SMC.decompress!(
+        A::CuSparseMatrixCSR, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSR}
+    )
+        compressed_indices = result.additional_info.compressed_indices_gpu_csr
         A.nnz == 0 && return A
         kernel = @cuda launch = false update_nzval_from_matrix!(
-            A.nzVal, B, result.compressed_indices
+            A.nzVal, B, compressed_indices
         )
         config = launch_configuration(kernel.fun)
         threads = min(A.nnz, config.threads)
         blocks = cld(A.nnz, threads)
-        kernel(A.nzVal, B, result.compressed_indices; threads, blocks)
+        kernel(A.nzVal, B, compressed_indices; threads, blocks)
         return A
     end
 end

src/graph.jl

@@ -345,8 +345,6 @@ end
 
 Base.eltype(::BipartiteGraph{T}) where {T} = T
 
-Base.transpose(bg::BipartiteGraph) = BipartiteGraph(bg.S2, bg.S1)
-
 function BipartiteGraph(A::AbstractMatrix; symmetric_pattern::Bool=false)
     return BipartiteGraph(SparseMatrixCSC(A); symmetric_pattern)
 end

src/result.jl

@@ -152,6 +152,7 @@ struct ColumnColoringResult{
     CT<:AbstractVector{T},
     GT<:AbstractGroups{T},
     VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:nonsymmetric,:column,:direct}
     "matrix that was colored"
     A::M
@@ -161,15 +162,22 @@ struct ColumnColoringResult{
     color::CT
     "color groups for columns or rows (depending on `partition`)"
     group::GT
-    "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A`. When `A isa SparseMatrixCSC`, they satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`."
+    "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A` when `A isa SparseMatrixCSC`. They satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
     compressed_indices::VT
+    "optional data used for decompressing into specific matrix types"
+    additional_info::A
 end
 
 function ColumnColoringResult(
     A::AbstractMatrix, bg::BipartiteGraph{T}, color::Vector{<:Integer}
 ) where {T<:Integer}
-    S = bg.S2
     group = group_by_color(T, color)
+    compressed_indices = column_csc_indices(bg, color)
+    return ColumnColoringResult(A, bg, color, group, compressed_indices, nothing)
+end
+
+function column_csc_indices(bg::BipartiteGraph{T}, color::Vector{<:Integer}) where {T}
+    S = bg.S2
     n = size(S, 1)
     rv = rowvals(S)
     compressed_indices = zeros(T, nnz(S))
@@ -181,7 +189,23 @@ function ColumnColoringResult(
             compressed_indices[k] = (c - 1) * n + i
         end
     end
-    return ColumnColoringResult(A, bg, color, group, compressed_indices)
+    return compressed_indices
+end
+
+function column_csr_indices(bg::BipartiteGraph{T}, color::Vector{<:Integer}) where {T}
+    Sᵀ = bg.S1  # CSC storage of transpose(A)
+    n = size(Sᵀ, 2)
+    rv = rowvals(Sᵀ)
+    compressed_indices = zeros(T, nnz(Sᵀ))
+    for i in axes(Sᵀ, 2)
+        for k in nzrange(Sᵀ, i)
+            j = rv[k]
+            c = color[j]
+            # A[i, j] = B[i, c]
+            compressed_indices[k] = (c - 1) * n + i
+        end
+    end
+    return compressed_indices
 end
 
 """
@@ -206,20 +230,27 @@ struct RowColoringResult{
     CT<:AbstractVector{T},
     GT<:AbstractGroups{T},
     VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:nonsymmetric,:row,:direct}
     A::M
     bg::G
     color::CT
     group::GT
     compressed_indices::VT
+    additional_info::A
 end
 
 function RowColoringResult(
     A::AbstractMatrix, bg::BipartiteGraph{T}, color::Vector{<:Integer}
 ) where {T<:Integer}
-    S = bg.S2
     group = group_by_color(T, color)
-    C = length(group)  # ncolors
+    compressed_indices = row_csc_indices(bg, color)
+    return RowColoringResult(A, bg, color, group, compressed_indices, nothing)
+end
+
+function row_csc_indices(bg::BipartiteGraph{T}, color::Vector{<:Integer}) where {T}
+    S = bg.S2
+    C = maximum(color)  # ncolors
     rv = rowvals(S)
     compressed_indices = zeros(T, nnz(S))
     for j in axes(S, 2)
@@ -230,7 +261,23 @@ function RowColoringResult(
             compressed_indices[k] = (j - 1) * C + c
         end
     end
-    return RowColoringResult(A, bg, color, group, compressed_indices)
+    return compressed_indices
+end
+
+function row_csr_indices(bg::BipartiteGraph{T}, color::Vector{<:Integer}) where {T}
+    Sᵀ = bg.S1  # CSC storage of transpose(A)
+    C = maximum(color)  # ncolors
+    rv = rowvals(Sᵀ)
+    compressed_indices = zeros(T, nnz(Sᵀ))
+    for i in axes(Sᵀ, 2)
+        for k in nzrange(Sᵀ, i)
+            j = rv[k]
+            c = color[i]
+            # A[i, j] = B[c, j]
+            compressed_indices[k] = (j - 1) * C + c
+        end
+    end
+    return compressed_indices
 end
 
 """
@@ -255,12 +302,14 @@ struct StarSetColoringResult{
     CT<:AbstractVector{T},
     GT<:AbstractGroups{T},
     VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:symmetric,:column,:direct}
     A::M
     ag::G
     color::CT
     group::GT
     compressed_indices::VT
+    additional_info::A
 end
 
 function StarSetColoringResult(
@@ -269,11 +318,18 @@ function StarSetColoringResult(
     color::Vector{<:Integer},
     star_set::StarSet{<:Integer},
 ) 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)
+end
+
+function star_csc_indices(
+    ag::AdjacencyGraph{T}, color::Vector{<:Integer}, star_set
+) where {T}
     (; star, hub) = star_set
     S = pattern(ag)
     edge_to_index = edge_indices(ag)
     n = S.n
-    group = group_by_color(T, color)
     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
     for j in axes(S, 2)
@@ -302,8 +358,7 @@ function StarSetColoringResult(
             end
         end
     end
-
-    return StarSetColoringResult(A, ag, color, group, compressed_indices)
+    return compressed_indices
 end
 
 """

test/cuda.jl

@@ -5,6 +5,8 @@ using SparseMatrixColorings
 using StableRNGs
 using Test
 
+include("utils.jl")
+
 rng = StableRNG(63)
 
 asymmetric_params = vcat(
@@ -19,32 +21,32 @@ symmetric_params = vcat(
     [(100, p) for p in (0.01:0.02:0.05)],
 )
 
-@testset "Column coloring & decompression" begin
+@testset verbose = true "Column coloring & decompression" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
     algo = GreedyColoringAlgorithm(; decompression=:direct)
-    @testset for T in (CuSparseMatrixCSC,)
+    @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_coloring_decompression(A0, problem, algo; gpu=true)
         end
     end
 end;
 
-@testset "Row coloring & decompression" begin
+@testset verbose = true "Row coloring & decompression" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
     algo = GreedyColoringAlgorithm(; decompression=:direct)
-    @testset for T in (CuSparseMatrixCSC,)
+    @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_coloring_decompression(A0, problem, algo; gpu=true)
         end
     end
 end;
 
-@testset "Symmetric coloring & direct decompression" begin
+@testset verbose = true "Symmetric coloring & direct decompression" begin
     problem = ColoringProblem(; structure=:symmetric, partition=:column)
     algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:direct)
-    @testset for T in (CuSparseMatrixCSC,)
+    @testset for T in (CuSparseMatrixCSC, CuSparseMatrixCSR)
         @testset "$((; n, p))" for (n, p) in symmetric_params
             A0 = T(sparse(Symmetric(sprand(rng, n, n, p))))
             test_coloring_decompression(A0, problem, algo; gpu=true)