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,155 @@
+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)
+    # 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
+        copyto!(A.nzVal, view(B, 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
+        copyto!(A.nzVal, view(B, compressed_indices))
+        return A
+    end
+end
+
+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
+end
+
+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
+end
+
+end

src/decompression.jl

@@ -106,6 +106,14 @@
     return Br, Bc
 end
 
+struct UnsupportedDecompressionError
+    msg::String
+end
+
+function Base.showerror(io::IO, err::UnsupportedDecompressionError)
+    return print(io, "UnsupportedDecompressionError: $(err.msg)")
+end
+
 """
     decompress(B::AbstractMatrix, result::AbstractColoringResult{_,:column/:row})
     decompress(Br::AbstractMatrix, Bc::AbstractMatrix, result::AbstractColoringResult{_,:bidirectional})

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 = [