Support sparse matrices from rocSPARSE

.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'
+      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'

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"

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

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"

test/cuda.jl

@@ -6,6 +6,8 @@ import SparseMatrixColorings as SMC
 using StableRNGs
 using Test
 
+include("utils.jl")
+
 rng = StableRNG(63)
 
 asymmetric_params = vcat(

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;

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)