diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 0ff0cc02..f0593aef 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -2,7 +2,7 @@ env: SECRET_CODECOV_TOKEN: "l7PxpMgxHEiz2l3UDDj5122Xirv5du1QFQiTf09YISC8VOz/nXeQC2yYY6iTMbccQG3xrumfWke9vtLbAc9LPH75kw3zQODuWMbz0eGSf9wtw7foqO6KAo5neMTeZJ0ZFgFMa2Q89T5SNIQsi4b5Zs7BVto6qB9Z/3QEs/BpsR25cYkY4Y6JBU6XuqZ6GRAc6BtlB4OmBEp3BBsXatx6y64zF0qbp4rocmcBQEoeQkcXtxf0dfA0KNAHpweWPrzIAMZ0aUYp6iEikxwLY5TjVFhOIcpVXUlxIdhl2qaaDF6b6WlgGXiGLpBjsLQfvlOgIlXH59Ddg0IVxboF2a37OA==;U2FsdGVkX19MveWMe1aasYZoJwaeiKO5XqwMZ/utOFj7h1CE9UX4nduBMXTyJ77tpNmCBYsQNJ/PadbX2224hQ==" steps: - - label: "Julia v1" + - label: "NVIDIA GPUs -- Julia v1" plugins: - JuliaCI/julia#v1: version: "1" @@ -14,4 +14,19 @@ steps: if: build.message !~ /\[skip tests\]/ timeout_in_minutes: 60 env: - JULIA_SMC_TEST_GROUP: 'GPU' \ No newline at end of file + JULIA_SMC_TEST_GROUP: 'GPU' + + - label: "AMD GPUs -- Julia v1" + plugins: + - JuliaCI/julia#v1: + version: "1" + - JuliaCI/julia-test#v1: ~ + - JuliaCI/julia-coverage#v1: + codecov: true + agents: + queue: "rocm" + rocmgpu: "*" + if: build.message !~ /\[skip tests\]/ + timeout_in_minutes: 60 + env: + JULIA_SMC_TEST_GROUP: 'GPU' diff --git a/Project.toml b/Project.toml index 5e649540..e72441d7 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [weakdeps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" @@ -21,6 +22,7 @@ MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060" [extensions] +SparseMatrixColoringsAMDGPUExt = "AMDGPU" SparseMatrixColoringsCUDAExt = ["CUDA", "cuSPARSE"] SparseMatrixColoringsCliqueTreesExt = "CliqueTrees" SparseMatrixColoringsColorsExt = "Colors" @@ -29,16 +31,17 @@ SparseMatrixColoringsJuMPExt = ["JuMP", "MathOptInterface"] [compat] ADTypes = "1.2.1" +AMDGPU = "2.2.0" CUDA = "6.0.0" CliqueTrees = "1" Colors = "0.12.11, 0.13" DocStringExtensions = "0.8,0.9" GPUArrays = "11.5.0" JuMP = "1.29.1" -LinearAlgebra = "1" +LinearAlgebra = "1.10" MathOptInterface = "1.45.0" PrecompileTools = "1.2.1" -Random = "1" -SparseArrays = "1" +Random = "1.10" +SparseArrays = "1.10" cuSPARSE = "6.0.0" julia = "1.10" diff --git a/ext/SparseMatrixColoringsAMDGPUExt.jl b/ext/SparseMatrixColoringsAMDGPUExt.jl new file mode 100644 index 00000000..58714ad2 --- /dev/null +++ b/ext/SparseMatrixColoringsAMDGPUExt.jl @@ -0,0 +1,139 @@ +module SparseMatrixColoringsAMDGPUExt + +import SparseMatrixColorings as SMC +using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange +using AMDGPU: ROCVector, ROCMatrix +using AMDGPU.rocSPARSE: AbstractROCSparseMatrix, ROCSparseMatrixCSC, ROCSparseMatrixCSR + +## CSC Result + +function SMC.ColumnColoringResult( + A::ROCSparseMatrixCSC, 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=ROCVector(compressed_indices)) + return SMC.ColumnColoringResult( + A, bg, color, group, compressed_indices, additional_info + ) +end + +function SMC.RowColoringResult( + A::ROCSparseMatrixCSC, 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=ROCVector(compressed_indices)) + return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info) +end + +function SMC.StarSetColoringResult( + A::ROCSparseMatrixCSC, + 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=ROCVector(compressed_indices)) + return SMC.StarSetColoringResult( + A, ag, color, group, compressed_indices, additional_info + ) +end + +## CSR Result + +function SMC.ColumnColoringResult( + A::ROCSparseMatrixCSR, 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=ROCVector(compressed_indices_csr)) + return SMC.ColumnColoringResult( + A, bg, color, group, compressed_indices, additional_info + ) +end + +function SMC.RowColoringResult( + A::ROCSparseMatrixCSR, 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=ROCVector(compressed_indices_csr)) + return SMC.RowColoringResult(A, bg, color, group, compressed_indices, additional_info) +end + +function SMC.StarSetColoringResult( + A::ROCSparseMatrixCSR, + 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=ROCVector(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::ROCSparseMatrixCSC, B::ROCMatrix, result::SMC.$R{<:ROCSparseMatrixCSC} + ) + compressed_indices = result.additional_info.compressed_indices_gpu_csc + copyto!(A.nzVal, view(B, compressed_indices)) + return A + end + + @eval function SMC.decompress!( + A::ROCSparseMatrixCSR, B::ROCMatrix, result::SMC.$R{<:ROCSparseMatrixCSR} + ) + compressed_indices = result.additional_info.compressed_indices_gpu_csr + copyto!(A.nzVal, view(B, compressed_indices)) + return A + end +end + +function SMC.decompress!( + A::ROCSparseMatrixCSC, + B::ROCMatrix, + result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSC}, + 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::ROCSparseMatrixCSR, + B::ROCMatrix, + result::SMC.StarSetColoringResult{<:ROCSparseMatrixCSR}, + 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 diff --git a/test/Project.toml b/test/Project.toml index 5e41ceb2..b335f90c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" BandedMatrices = "aae01518-5342-5314-be14-df237901396f" diff --git a/test/cuda.jl b/test/cuda.jl index 8b27e07b..6ae73060 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -6,6 +6,8 @@ import SparseMatrixColorings as SMC using StableRNGs using Test +include("utils.jl") + rng = StableRNG(63) asymmetric_params = vcat( diff --git a/test/rocm.jl b/test/rocm.jl new file mode 100644 index 00000000..e2aa7840 --- /dev/null +++ b/test/rocm.jl @@ -0,0 +1,62 @@ +using AMDGPU.rocSPARSE: ROCSparseMatrixCSC, ROCSparseMatrixCSR +using LinearAlgebra +using SparseArrays +using SparseMatrixColorings +import SparseMatrixColorings as SMC +using StableRNGs +using Test + +include("utils.jl") + +rng = StableRNG(63) + +asymmetric_params = vcat( + [(10, 20, p) for p in (0.0:0.2:0.5)], + [(20, 10, p) for p in (0.0:0.2:0.5)], + [(100, 200, p) for p in (0.01:0.02:0.05)], + [(200, 100, p) for p in (0.01:0.02:0.05)], +) + +symmetric_params = vcat( + [(10, p) for p in (0.0:0.2:0.5)], # + [(100, p) for p in (0.01:0.02:0.05)], +) + +@testset verbose = true "Column coloring & decompression" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:column) + algo = GreedyColoringAlgorithm(; decompression=:direct) + @testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR) + @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 verbose = true "Row coloring & decompression" begin + problem = ColoringProblem(; structure=:nonsymmetric, partition=:row) + algo = GreedyColoringAlgorithm(; decompression=:direct) + @testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR) + @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 verbose = true "Symmetric coloring & direct decompression" begin + problem = ColoringProblem(; structure=:symmetric, partition=:column) + algo = GreedyColoringAlgorithm(; postprocessing=false, decompression=:direct) + @testset for T in (ROCSparseMatrixCSC, ROCSparseMatrixCSR) + @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) + end + A0 = T(sparse(Diagonal(ones(10)))) + result = coloring(A0, problem, algo) + B = compress(A0, result) + @test_throws SMC.UnsupportedDecompressionError decompress!( + similar(A0), B, result, :U + ) + end +end; diff --git a/test/runtests.jl b/test/runtests.jl index ec679d4e..877ebbbe 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,15 +7,22 @@ using Test # Load package extensions to test them with JET using Colors: Colors -include("utils.jl") - @testset verbose = true "SparseMatrixColorings" begin if get(ENV, "JULIA_SMC_TEST_GROUP", nothing) == "GPU" @testset "CUDA" begin using CUDA, cuSPARSE - include("cuda.jl") + if CUDA.functional() + include("cuda.jl") + end + end + @testset "ROCm" begin + using AMDGPU + if AMDGPU.functional() + include("rocm.jl") + end end else + include("utils.jl") @testset verbose = true "Code quality" begin @testset "Aqua" begin Aqua.test_all(SparseMatrixColorings; undocumented_names=true)