fallback to dense adjmatrix on gpu

Author: CarloLucibello

src/GraphNeuralNetworks.jl

@@ -1,7 +1,7 @@
 module GraphNeuralNetworks
 
 using Statistics: mean
-using LinearAlgebra
+using LinearAlgebra, Random
 using SparseArrays
 import KrylovKit
 using Base: tail

src/gnngraph.jl

@@ -348,8 +348,10 @@ function scaled_laplacian(g::GNNGraph, T::DataType=Float32; dir=:out)
 end
 
 # _eigmax(A) = eigmax(Symmetric(A)) # Doesn't work on sparse arrays
-_eigmax(A) = KrylovKit.eigsolve(Symmetric(A), 1, :LR)[1][1] # also eigs(A, x0, nev, mode) available 
-
+function _eigmax(A)
+    x0 = randn!(similar(A, float(eltype(A)), size(A, 1)))
+    KrylovKit.eigsolve(Symmetric(A), x0, 1, :LR)[1][1] # also eigs(A, x0, nev, mode) available 
+end
 # Eigenvalues for cuarray don't seem to be well supported. 
 # https://github.com/JuliaGPU/CUDA.jl/issues/154
 # https://discourse.julialang.org/t/cuda-eigenvalues-of-a-sparse-matrix/46851/5

src/graph_conversions.jl

@@ -104,30 +104,38 @@ function to_sparse(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothi
     @assert dir ∈ [:out, :in]
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
-    num_edges = round(Int, sum(A))
+    num_edges = A isa AbstractSparseMatrix ? nnz(A) : count(!=(0), A)
     if dir == :in
         A = A'
     end
     if T != eltype(A)
         A = T.(A)
     end
-    return sparse(A), num_nodes, num_edges
+    if !(A isa AbstractSparseMatrix)
+        A = sparse(A)
+    end 
+    return A, num_nodes, num_edges
 end
 
 function to_sparse(adj_list::ADJLIST_T, T::DataType=Int; dir=:out, num_nodes=nothing)
     coo, num_nodes, num_edges = to_coo(adj_list; dir, num_nodes)
-    to_sparse(coo; dir, num_nodes)
+    return to_sparse(coo; dir, num_nodes)
 end
 
 function to_sparse(coo::COO_T, T::DataType=Int; dir=:out, num_nodes=nothing)
     s, t, eweight  = coo
+    s isa CuVector && return to_dense(coo, T; dir, num_nodes)
     eweight = isnothing(eweight) ? fill!(similar(s, T), 1) : eweight
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)
     num_edges = length(s)
-    A, num_nodes, num_edges
+    return A, num_nodes, num_edges
 end
 
+# _sparse(I, J, V, m, n) = sparse(I, J, V, m, n)
+
+# _sparse(I::CuVector, J::CuVector, V::CuVector, m, n)
+
 @non_differentiable to_coo(x...)
 @non_differentiable to_dense(x...)
 @non_differentiable to_sparse(x...)

test/cuda/gnngraph.jl

@@ -1,6 +1,8 @@
+const ACUMatrix{T} = Union{CuMatrix{T}, CUDA.CUSPARSE.CuSparseMatrix{T}}
+
 @testset "cuda/gnngraph" begin
-    s = [1,1,2,3,4,5,5,5]
-    t = [2,5,3,2,1,4,3,1]
+    s = [1,1,2,3,4]
+    t = [2,5,3,5,5]
     s, t = [s; t], [t; s]  #symmetrize
     g = GNNGraph(s, t, graph_type=GRAPH_T) 
     g_gpu = g |> gpu
@@ -15,41 +17,27 @@
     end
 
     @testset "adjacency_matrix" begin
-        function test_adj()
-            mat = adjacency_matrix(g)
-            mat_gpu = adjacency_matrix(g_gpu)
-            @test mat_gpu isa CuMatrix{Int}
-            true
-        end
+        # See https://github.com/JuliaGPU/CUDA.jl/pull/1093
 
-        if GRAPH_T == :coo
-            # See https://github.com/JuliaGPU/CUDA.jl/pull/1093
-            @test_broken test_adj()
-        else
-            test_adj()
-        end
+        mat = adjacency_matrix(g)
+        mat_gpu = adjacency_matrix(g_gpu)
+        @test mat_gpu isa ACUMatrix{Int}
+        @test Array(mat_gpu) == mat 
     end
 
     @testset "normalized_laplacian" begin
-        function test_normlapl()
-            mat = normalized_laplacian(g)
-            mat_gpu = normalized_laplacian(g_gpu)
-            @test mat_gpu isa CuMatrix{Float32}
-            true
-        end
-        if GRAPH_T == :coo
-            @test_broken test_normlapl()
-        else
-            test_normlapl()
-        end
+        mat = normalized_laplacian(g)
+        mat_gpu = normalized_laplacian(g_gpu)
+        @test mat_gpu isa ACUMatrix{Float32}
+        @test Array(mat_gpu) == mat 
     end
 
     @testset "scaled_laplacian" begin
-        @test_broken begin
+        @test_broken begin 
             mat = scaled_laplacian(g)
             mat_gpu = scaled_laplacian(g_gpu)
-            @test mat_gpu isa CuMatrix{Float32}
-            true
+            @test mat_gpu isa ACUMatrix{Float32}
+            @test Array(mat_gpu) == mat
         end
     end
-end
+end