Skip to content

Commit 108c7c1

Browse files
gdallegdalle
authored
feat: decompression for cuSPARSE matrices (#253)
1 parent f7acc3f commit 108c7c1

File tree

12 files changed

+467
-104
lines changed

12 files changed

+467
-104
lines changed

‎.github/workflows/Test-GPU.yml‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
name: Test-GPU
2+
3+
on:
4+
push:
5+
branches:
6+
- main
7+
pull_request:
8+
9+
concurrency:
10+
group: ${{ github.workflow }}-${{ github.ref }}
11+
cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
12+
13+
# needed to allow julia-actions/cache to delete old caches that it has created
14+
permissions:
15+
actions: write
16+
contents: read
17+
18+
jobs:
19+
test:
20+
runs-on: self-hosted
21+
env:
22+
CUDA_VISIBLE_DEVICES: 1
23+
JULIA_DEPOT_PATH: /scratch/github-actions/julia_depot_smc
24+
JULIA_SMC_TEST_GROUP: "GPU"
25+
strategy:
26+
matrix:
27+
julia-version: ['1.10', '1']
28+
steps:
29+
- uses: actions/checkout@v4
30+
- uses: julia-actions/setup-julia@v2
31+
with:
32+
version: ${{ matrix.julia-version }}
33+
arch: x64
34+
- uses: julia-actions/julia-downgrade-compat@v1
35+
if: ${{ matrix.version == '1.10' }}
36+
with:
37+
skip: LinearAlgebra, Random, SparseArrays
38+
- uses: julia-actions/cache@v2
39+
- uses: julia-actions/julia-buildpkg@v1
40+
- uses: julia-actions/julia-runtest@v1
41+
- uses: julia-actions/julia-processcoverage@v1
42+
- uses: codecov/codecov-action@v5
43+
with:
44+
files: lcov.info
45+
token: ${{ secrets.CODECOV_TOKEN }}
46+
fail_ci_if_error: false

‎Project.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "SparseMatrixColorings"
22
uuid = "0a514795-09f3-496d-8182-132a7b665d35"
33
authors = ["Guillaume Dalle", "Alexis Montoison"]
4-
version = "0.4.20"
4+
version = "0.4.21"
55

66
[deps]
77
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -12,15 +12,18 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
1212
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1313

1414
[weakdeps]
15+
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
1516
CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
1617
Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
1718

1819
[extensions]
20+
SparseMatrixColoringsCUDAExt = "CUDA"
1921
SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
2022
SparseMatrixColoringsColorsExt = "Colors"
2123

2224
[compat]
2325
ADTypes = "1.2.1"
26+
CUDA = "5.8.2"
2427
CliqueTrees = "1"
2528
Colors = "0.12.11, 0.13"
2629
DocStringExtensions = "0.8,0.9"

‎ext/SparseMatrixColoringsCUDAExt.jl‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
module SparseMatrixColoringsCUDAExt
2+
3+
import SparseMatrixColorings as SMC
4+
using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
5+
using CUDA: CuVector, CuMatrix
6+
using CUDA.CUSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
7+
8+
SMC.matrix_versions(A::AbstractCuSparseMatrix) = (A,)
9+
10+
## Compression (slow, through CPU)
11+
12+
function SMC.compress(
13+
A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:column}
14+
) where {structure}
15+
return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
16+
end
17+
18+
function SMC.compress(
19+
A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:row}
20+
) where {structure}
21+
return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
22+
end
23+
24+
## CSC Result
25+
26+
function SMC.ColumnColoringResult(
27+
A::CuSparseMatrixCSC, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
28+
) where {T<:Integer}
29+
group = SMC.group_by_color(T, color)
30+
compressed_indices = SMC.column_csc_indices(bg, color)
31+
additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
32+
return SMC.ColumnColoringResult(
33+
A, bg, color, group, compressed_indices, additional_info
34+
)
35+
end
36+
37+
function SMC.RowColoringResult(
38+
A::CuSparseMatrixCSC, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
39+
) where {T<:Integer}
40+
group = SMC.group_by_color(T, color)
41+
compressed_indices = SMC.row_csc_indices(bg, color)
42+
additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
43+
return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
44+
end
45+
46+
function SMC.StarSetColoringResult(
47+
A::CuSparseMatrixCSC,
48+
ag::SMC.AdjacencyGraph{T},
49+
color::Vector{<:Integer},
50+
star_set::SMC.StarSet{<:Integer},
51+
) where {T<:Integer}
52+
group = SMC.group_by_color(T, color)
53+
compressed_indices = SMC.star_csc_indices(ag, color, star_set)
54+
additional_info = (; compressed_indices_gpu_csc=CuVector(compressed_indices))
55+
return SMC.StarSetColoringResult(
56+
A, ag, color, group, compressed_indices, additional_info
57+
)
58+
end
59+
60+
## CSR Result
61+
62+
function SMC.ColumnColoringResult(
63+
A::CuSparseMatrixCSR, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
64+
) where {T<:Integer}
65+
group = SMC.group_by_color(T, color)
66+
compressed_indices = SMC.column_csc_indices(bg, color)
67+
compressed_indices_csr = SMC.column_csr_indices(bg, color)
68+
additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices_csr))
69+
return SMC.ColumnColoringResult(
70+
A, bg, color, group, compressed_indices, additional_info
71+
)
72+
end
73+
74+
function SMC.RowColoringResult(
75+
A::CuSparseMatrixCSR, bg::SMC.BipartiteGraph{T}, color::Vector{<:Integer}
76+
) where {T<:Integer}
77+
group = SMC.group_by_color(T, color)
78+
compressed_indices = SMC.row_csc_indices(bg, color)
79+
compressed_indices_csr = SMC.row_csr_indices(bg, color)
80+
additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices_csr))
81+
return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
82+
end
83+
84+
function SMC.StarSetColoringResult(
85+
A::CuSparseMatrixCSR,
86+
ag::SMC.AdjacencyGraph{T},
87+
color::Vector{<:Integer},
88+
star_set::SMC.StarSet{<:Integer},
89+
) where {T<:Integer}
90+
group = SMC.group_by_color(T, color)
91+
compressed_indices = SMC.star_csc_indices(ag, color, star_set)
92+
additional_info = (; compressed_indices_gpu_csr=CuVector(compressed_indices))
93+
return SMC.StarSetColoringResult(
94+
A, ag, color, group, compressed_indices, additional_info
95+
)
96+
end
97+
98+
## Decompression
99+
100+
for R in (:ColumnColoringResult, :RowColoringResult)
101+
# loop to avoid method ambiguity
102+
@eval function SMC.decompress!(
103+
A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSC}
104+
)
105+
compressed_indices = result.additional_info.compressed_indices_gpu_csc
106+
copyto!(A.nzVal, view(B, compressed_indices))
107+
return A
108+
end
109+
110+
@eval function SMC.decompress!(
111+
A::CuSparseMatrixCSR, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSR}
112+
)
113+
compressed_indices = result.additional_info.compressed_indices_gpu_csr
114+
copyto!(A.nzVal, view(B, compressed_indices))
115+
return A
116+
end
117+
end
118+
119+
function SMC.decompress!(
120+
A::CuSparseMatrixCSC,
121+
B::CuMatrix,
122+
result::SMC.StarSetColoringResult{<:CuSparseMatrixCSC},
123+
uplo::Symbol=:F,
124+
)
125+
if uplo != :F
126+
throw(
127+
SMC.UnsupportedDecompressionError(
128+
"Single-triangle decompression is not supported on GPU matrices"
129+
),
130+
)
131+
end
132+
compressed_indices = result.additional_info.compressed_indices_gpu_csc
133+
copyto!(A.nzVal, view(B, compressed_indices))
134+
return A
135+
end
136+
137+
function SMC.decompress!(
138+
A::CuSparseMatrixCSR,
139+
B::CuMatrix,
140+
result::SMC.StarSetColoringResult{<:CuSparseMatrixCSR},
141+
uplo::Symbol=:F,
142+
)
143+
if uplo != :F
144+
throw(
145+
SMC.UnsupportedDecompressionError(
146+
"Single-triangle decompression is not supported on GPU matrices"
147+
),
148+
)
149+
end
150+
compressed_indices = result.additional_info.compressed_indices_gpu_csr
151+
copyto!(A.nzVal, view(B, compressed_indices))
152+
return A
153+
end
154+
155+
end

‎src/decompression.jl‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,14 @@ function compress(
106106
return Br, Bc
107107
end
108108

109+
struct UnsupportedDecompressionError
110+
msg::String
111+
end
112+
113+
function Base.showerror(io::IO, err::UnsupportedDecompressionError)
114+
return print(io, "UnsupportedDecompressionError: $(err.msg)")
115+
end
116+
109117
"""
110118
decompress(B::AbstractMatrix, result::AbstractColoringResult{_,:column/:row})
111119
decompress(Br::AbstractMatrix, Bc::AbstractMatrix, result::AbstractColoringResult{_,:bidirectional})

‎src/graph.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ end
100100
Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
101101
"""
102102
function bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
103-
bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
103+
return bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
104104
end
105105

106106
function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T}

‎src/matrices.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ Return various versions of the same matrix:
1010
1111
Used for internal testing.
1212
"""
13-
function matrix_versions(A)
13+
function matrix_versions(A::AbstractMatrix)
1414
A_dense = Matrix(A)
1515
A_sparse = sparse(A)
1616
versions = [

0 commit comments

Comments
 (0)