feat: decompression for cuSPARSE matrices

.github/workflows/Test-GPU.yml

@@ -0,0 +1,46 @@
+name: Test-GPU
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+# needed to allow julia-actions/cache to delete old caches that it has created
+permissions:
+  actions: write
+  contents: read
+
+jobs:
+  test:
+    runs-on: self-hosted
+    env:
+      CUDA_VISIBLE_DEVICES: 1
+      JULIA_DEPOT_PATH: /scratch/github-actions/julia_depot_smc
+      JULIA_SMC_TEST_GROUP: "GPU"
+    strategy:
+      matrix:
+        julia-version: ['1.10', '1']
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.julia-version }}
+          arch: x64
+      - uses: julia-actions/julia-downgrade-compat@v1
+        if: ${{ matrix.version == '1.10' }}
+        with:
+          skip: LinearAlgebra, Random, SparseArrays
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v5
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle", "Alexis Montoison"]
-version = "0.4.20"
+version = "0.4.21"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -12,15 +12,18 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 
 [extensions]
+SparseMatrixColoringsCUDAExt = "CUDA"
 SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
 SparseMatrixColoringsColorsExt = "Colors"
 
 [compat]
 ADTypes = "1.2.1"
+CUDA = "5.8.2"
 CliqueTrees = "1"
 Colors = "0.12.11, 0.13"
 DocStringExtensions = "0.8,0.9"

ext/SparseMatrixColoringsCUDAExt.jl

@@ -0,0 +1,119 @@
+module SparseMatrixColoringsCUDAExt
+
+import SparseMatrixColorings as SMC
+using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
+using CUDA: CuVector, CuMatrix
+using CUDA.CUSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
+
+SMC.matrix_versions(A::AbstractCuSparseMatrix) = (A,)
+
+## Compression (slow, through CPU)
+
+function SMC.compress(
+    A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:column}
+) where {structure}
+    return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
+end
+
+function SMC.compress(
+    A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:row}
+) where {structure}
+    return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
+end
+
+## CSC Result
+
+function SMC.ColumnColoringResult(
+    A::CuSparseMatrixCSC, 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)
+    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}
+    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(
+    A::CuSparseMatrixCSC,
+    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_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
+
+for R in (:ColumnColoringResult, :RowColoringResult, :StarSetColoringResult)
+    # loop to avoid method ambiguity
+    @eval function SMC.decompress!(
+        A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSC}
+    )
+        compressed_indices = result.additional_info.compressed_indices_gpu_csc
+        map!(Base.Fix1(getindex, B), A.nzVal, compressed_indices)
+        return A
+    end
+
+    @eval function SMC.decompress!(
+        A::CuSparseMatrixCSR, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSR}
+    )
+        compressed_indices = result.additional_info.compressed_indices_gpu_csr
+        map!(Base.Fix1(getindex, B), A.nzVal, compressed_indices)
+        return A
+    end
+end
+
+end

src/graph.jl

@@ -100,7 +100,7 @@
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
 function bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
-    bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
+    return bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
 end
 
 function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T}

src/matrices.jl

@@ -10,7 +10,7 @@
 
 Used for internal testing.
 """
-function matrix_versions(A)
+function matrix_versions(A::AbstractMatrix)
     A_dense = Matrix(A)
     A_sparse = sparse(A)
     versions = [

src/result.jl

@@ -146,25 +146,38 @@
 - [`AbstractColoringResult`](@ref)
 """
 struct ColumnColoringResult{
-    M<:AbstractMatrix,T<:Integer,G<:BipartiteGraph{T},GT<:AbstractGroups{T}
+    M<:AbstractMatrix,
+    T<:Integer,
+    G<:BipartiteGraph{T},
+    CT<:AbstractVector{T},
+    GT<:AbstractGroups{T},
+    VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:nonsymmetric,:column,:direct}
     "matrix that was colored"
     A::M
     "bipartite graph that was used for coloring"
     bg::G
     "one integer color for each column or row (depending on `partition`)"
-    color::Vector{T}
+    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]]`"
-    compressed_indices::Vector{T}
+    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))
@@ -176,7 +189,23 @@
             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
 
 """
@@ -195,21 +224,33 @@
 - [`AbstractColoringResult`](@ref)
 """
 struct RowColoringResult{
-    M<:AbstractMatrix,T<:Integer,G<:BipartiteGraph{T},GT<:AbstractGroups{T}
+    M<:AbstractMatrix,
+    T<:Integer,
+    G<:BipartiteGraph{T},
+    CT<:AbstractVector{T},
+    GT<:AbstractGroups{T},
+    VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:nonsymmetric,:row,:direct}
     A::M
     bg::G
-    color::Vector{T}
+    color::CT
     group::GT
-    compressed_indices::Vector{T}
+    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)
@@ -220,7 +261,23 @@
             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
 
 """
@@ -239,13 +296,20 @@
 - [`AbstractColoringResult`](@ref)
 """
 struct StarSetColoringResult{
-    M<:AbstractMatrix,T<:Integer,G<:AdjacencyGraph{T},GT<:AbstractGroups{T}
+    M<:AbstractMatrix,
+    T<:Integer,
+    G<:AdjacencyGraph{T},
+    CT<:AbstractVector{T},
+    GT<:AbstractGroups{T},
+    VT<:AbstractVector{T},
+    A,
 } <: AbstractColoringResult{:symmetric,:column,:direct}
     A::M
     ag::G
-    color::Vector{T}
+    color::CT
     group::GT
-    compressed_indices::Vector{T}
+    compressed_indices::VT
+    additional_info::A
 end
 
 function StarSetColoringResult(
@@ -254,11 +318,18 @@
     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)
@@ -287,8 +358,7 @@
             end
         end
     end
-
-    return StarSetColoringResult(A, ag, color, group, compressed_indices)
+    return compressed_indices
 end
 
 """