SparseMatrixColorings.jl/test/rocm.jl at 03514b381ac1a9d950041470b93eceb23ad1e44d · JuliaDiff/SparseMatrixColorings.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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;