[CUSPARSE] Interface color reordering (#1794)

Author: amontoison

lib/cusparse/CUSPARSE.jl

@@ -37,6 +37,7 @@ include("level2.jl")
 include("level3.jl")
 include("extra.jl")
 include("preconditioners.jl")
+include("reorderings.jl")
 include("conversions.jl")
 include("generic.jl")
 

lib/cusparse/reorderings.jl

@@ -0,0 +1,46 @@
+# routines that implement reorderings
+
+export color
+
+"""
+    color(A::CuSparseMatrixCSC, index::SparseChar; percentage::Number=1.0)
+    color(A::CuSparseMatrixCSR, index::SparseChar; percentage::Number=1.0)
+
+This function performs the coloring of the adjacency graph associated with the matrix A.
+The coloring is an assignment of colors (integer numbers) to nodes, such that neighboring nodes have distinct colors.
+An approximate coloring algorithm is used in this routine, and is stopped when a certain percentage of nodes has been colored.
+The rest of the nodes are assigned distinct colors (an increasing sequence of integers numbers, starting from the last integer used previously).
+The reordering is such that nodes that have been assigned the same color are reordered to be next to each other.
+
+The matrix A passed to this routine, must be stored as a general matrix and have a symmetric sparsity pattern.
+If the matrix is non-symmetric the user should pass A + Aᵀ as a parameter to this routine.
+"""
+function color end
+
+for (fname, subty, elty) in ((:cusparseScsrcolor, :Float32, :Float32),
+                             (:cusparseDcsrcolor, :Float64, :Float64),
+                             (:cusparseCcsrcolor, :Float32, :ComplexF32),
+                             (:cusparseZcsrcolor, :Float64, :ComplexF64))
+    @eval begin
+        function color(A::Union{CuSparseMatrixCSR{$elty},CuSparseMatrixCSC{$elty}}, index::SparseChar; percentage::Number=1.0)
+            desc = CuMatrixDescriptor('G', 'L', 'N', index)
+            m, n = size(A)
+            (m != n) && throw(DimensionMismatch("A must be square, but has dimensions ($m,$n)!"))
+
+            info = cusparseColorInfo_t[0]
+            cusparseCreateColorInfo(info)
+
+            ncolors = Ref{Cint}(-1)
+            coloring = CuVector{Cint}(undef, m)
+            reordering = CuVector{Cint}(undef, m)
+
+            if isa(A, CuSparseMatrixCSR)
+                $fname(handle(), m, nnz(A), desc, nonzeros(A), A.rowPtr, A.colVal, Ref{$subty}(percentage), ncolors, coloring, reordering, info[1])
+            else
+                $fname(handle(), m, nnz(A), desc, nonzeros(A), A.colPtr, A.rowVal, Ref{$subty}(percentage), ncolors, coloring, reordering, info[1])
+            end
+            cusparseDestroyColorInfo(info[1])
+            ncolors[], coloring, reordering
+        end
+    end
+end

test/cusparse.jl

@@ -930,3 +930,45 @@ end
         end
     end
 end
+
+for SparseMatrixType in [CuSparseMatrixCSC, CuSparseMatrixCSR]
+    @testset "$SparseMatrixType -- color" begin
+        @testset "color $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+            A = sprand(T, 50, 50, 0.03)
+            A = A + A'
+            dA = SparseMatrixType(A)
+            ncolors, coloring, reordering = color(dA, 'O')
+            @test 1 ≤ ncolors ≤ 50
+            @test maximum(coloring) == ncolors
+            @test minimum(reordering) == 1
+            @test maximum(reordering) == 50
+            @test CUDA.@allowscalar isperm(reordering)
+        end
+
+        # The routine color returns the wrong result with small sparse matrices.
+        # NVIDIA investigates the issue.
+        if false
+            A = [ 1 1 0 0 0 ;
+                  1 1 1 1 0 ;
+                  0 1 1 0 1 ;
+                  0 1 0 1 0 ;
+                  0 0 1 0 1 ]
+            A = sparse(A)
+            # The adjacency graph of A has at least two colors, one color for
+            # {1, 3, 4} and another one for {2, 5}.
+            @testset "5x5 example -- color $T" for T in [Float32, Float64, ComplexF32, ComplexF64]
+                dA = SparseMatrixType{T}(A)
+                ncolors, coloring, reordering = color(dA, 'O')
+                @test ncolors == 2
+                @test minimum(reordering) == 1
+                @test maximum(reordering) == 5
+                CUDA.allowscalar() do
+                    @test isperm(reordering)
+                    @test coloring[1] == coloring[3]  == coloring[4]
+                    @test coloring[2] == coloring[5]
+                    @test coloring[1] != coloring[2]
+                end
+            end
+        end
+    end
+end