Skip to content

Commit dac2926

Browse files
amontoisonamontoison
committed
Support sparse matrices from rocSPARSE
1 parent 108c7c1 commit dac2926

4 files changed

Lines changed: 232 additions & 4 deletions

File tree

‎Project.toml‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,23 +12,26 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
1212
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1313

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

1920
[extensions]
21+
SparseMatrixColoringsAMDGPUExt = "AMDGPU"
2022
SparseMatrixColoringsCUDAExt = "CUDA"
2123
SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
2224
SparseMatrixColoringsColorsExt = "Colors"
2325

2426
[compat]
2527
ADTypes = "1.2.1"
28+
AMDGPU = "1.3.3"
2629
CUDA = "5.8.2"
2730
CliqueTrees = "1"
2831
Colors = "0.12.11, 0.13"
2932
DocStringExtensions = "0.8,0.9"
30-
LinearAlgebra = "<0.0.1, 1"
33+
LinearAlgebra = "1.10"
3134
PrecompileTools = "1.2.1"
32-
Random = "<0.0.1, 1"
33-
SparseArrays = "<0.0.1, 1"
35+
Random = "1.10"
36+
SparseArrays = "1.10"
3437
julia = "1.10"

‎ext/SparseMatrixColoringsAMDGPUExt.jl‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
module SparseMatrixColoringsAMDGPUExt
2+
3+
import SparseMatrixColorings as SMC
4+
using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
5+
using AMDGPU: ROCVector, ROCMatrix
6+
using AMDGPU.rocSPARSE: AbstractROCSparseMatrix, ROCSparseMatrixCSC, ROCSparseMatrixCSR
7+
8+
SMC.matrix_versions(A::AbstractROCSparseMatrix) = (A,)
9+
10+
## Compression (slow, through CPU)
11+
12+
function SMC.compress(
13+
A::AbstractROCSparseMatrix, result::SMC.AbstractColoringResult{structure,:column}
14+
) where {structure}
15+
return ROCMatrix(SMC.compress(SparseMatrixCSC(A), result))
16+
end
17+
18+
function SMC.compress(
19+
A::AbstractROCSparseMatrix, result::SMC.AbstractColoringResult{structure,:row}
20+
) where {structure}
21+
return ROCMatrix(SMC.compress(SparseMatrixCSC(A), result))
22+
end
23+
24+
## CSC Result
25+
26+
function SMC.ColumnColoringResult(
27+
A::ROCSparseMatrixCSC, 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=ROCVector(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::ROCSparseMatrixCSC, 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=ROCVector(compressed_indices))
43+
return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
44+
end
45+
46+
function SMC.StarSetColoringResult(
47+
A::ROCSparseMatrixCSC,
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=ROCVector(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::ROCSparseMatrixCSR, 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=ROCVector(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::ROCSparseMatrixCSR, 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=ROCVector(compressed_indices_csr))
81+
return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info)
82+
end
83+
84+
function SMC.StarSetColoringResult(
85+
A::ROCSparseMatrixCSR,
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=ROCVector(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::ROCSparseMatrixCSC, B::ROCMatrix, result::SMC.$R{<:ROCSparseMatrixCSC}
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::ROCSparseMatrixCSR, B::ROCMatrix, result::SMC.$R{<:ROCSparseMatrixCSR}
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::ROCSparseMatrixCSC,
121+
B::ROCMatrix,
122+
result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSC},
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::ROCSparseMatrixCSR,
139+
B::ROCMatrix,
140+
result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSR},
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

‎test/rocm.jl‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
using AMDGPU.rocSPARSE: ROCSparseMatrixCSC, ROCSparseMatrixCSR
2+
using LinearAlgebra
3+
using SparseArrays
4+
using SparseMatrixColorings
5+
import SparseMatrixColorings as SMC
6+
using StableRNGs
7+
using Test
8+
9+
include("utils.jl")
10+
11+
rng = StableRNG(63)
12+
13+
asymmetric_params = vcat(
14+
[(10, 20, p) for p in (0.0:0.2:0.5)],
15+
[(20, 10, p) for p in (0.0:0.2:0.5)],
16+
[(100, 200, p) for p in (0.01:0.02:0.05)],
17+
[(200, 100, p) for p in (0.01:0.02:0.05)],
18+
)
19+
20+
symmetric_params = vcat(
21+
[(10, p) for p in (0.0:0.2:0.5)], #
22+
[(100, p) for p in (0.01:0.02:0.05)],
23+
)
24+
25+
@testset verbose = true "Column coloring & decompression" begin
26+
problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
27+
algo = GreedyColoringAlgorithm(; decompression=:direct)
28+
@testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR)
29+
@testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
30+
A0 = T(sprand(rng, m, n, p))
31+
test_coloring_decompression(A0, problem, algo; gpu=true)
32+
end
33+
end
34+
end;
35+
36+
@testset verbose = true "Row coloring & decompression" begin
37+
problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
38+
algo = GreedyColoringAlgorithm(; decompression=:direct)
39+
@testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR)
40+
@testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
41+
A0 = T(sprand(rng, m, n, p))
42+
test_coloring_decompression(A0, problem, algo; gpu=true)
43+
end
44+
end
45+
end;
46+
47+
@testset verbose = true "Symmetric coloring & direct decompression" begin
48+
problem = ColoringProblem(; structure=:symmetric, partition=:column)
49+
algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:direct)
50+
@testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR)
51+
@testset "$((; n, p))" for (n, p) in symmetric_params
52+
A0 = T(sparse(Symmetric(sprand(rng, n, n, p))))
53+
test_coloring_decompression(A0, problem, algo; gpu=true)
54+
end
55+
A0 = T(sparse(Diagonal(ones(10))))
56+
result = coloring(A0, problem, algo)
57+
B = compress(A0, result)
58+
@test_throws SMC.UnsupportedDecompressionError decompress!(
59+
similar(A0), B, result, :U
60+
)
61+
end
62+
end;

‎test/runtests.jl‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,15 @@ include("utils.jl")
1414
if get(ENV, "JULIA_SMC_TEST_GROUP", nothing) == "GPU"
1515
@testset "CUDA" begin
1616
using CUDA
17-
include("cuda.jl")
17+
if CUDA.functional()
18+
include("cuda.jl")
19+
end
20+
end
21+
@testset "ROCm" begin
22+
using AMDGPU
23+
if AMDGPU.functional()
24+
include("rocm.jl")
25+
end
1826
end
1927
else
2028
@testset verbose = true "Code quality" begin

0 commit comments

Comments
 (0)