reorg tests

Author: CarloLucibello

test/cuda/gnngraph.jl

@@ -11,6 +11,10 @@
 
         # core functionality
         g = GNNGraph(s, t; graph_type=GRAPH_T)
+        if TEST_GPU
+            g_gpu = g |> gpu
+        end
+
         @test g.num_edges == 8
         @test g.num_nodes == 4
         @test collect(edges(g)) |> sort == collect(zip(s, t)) |> sort
@@ -21,23 +25,69 @@
         @test s1 == s
         @test t1 == t
 
-        # adjacency
-        @test adjacency_matrix(g) == adj_mat
-        @test adjacency_matrix(g; dir=:in) == adj_mat
-        @test adjacency_matrix(g; dir=:out) == adj_mat
         @test sort.(adjacency_list(g; dir=:in)) == adj_list_in
         @test sort.(adjacency_list(g; dir=:out)) == adj_list_out
 
+        @testset "adjacency_matrix" begin
+            @test adjacency_matrix(g) == adj_mat
+            @test adjacency_matrix(g; dir=:in) == adj_mat
+            @test adjacency_matrix(g; dir=:out) == adj_mat
+            
+            if TEST_GPU
+                # See https://github.com/JuliaGPU/CUDA.jl/pull/1093
+                mat_gpu = adjacency_matrix(g_gpu)
+                @test mat_gpu isa ACUMatrix{Int}
+                @test Array(mat_gpu) == adj_mat 
+            end
+        end
+        
+        @testset "normalized_laplacian" begin
+            mat = normalized_laplacian(g)
+            if TEST_GPU
+                mat_gpu = normalized_laplacian(g_gpu)
+                @test mat_gpu isa ACUMatrix{Float32}
+                @test Array(mat_gpu) == mat 
+            end
+        end
+
+
+        @testset "scaled_laplacian" begin
+            if TEST_GPU
+                @test_broken begin 
+                    mat = scaled_laplacian(g)
+                    mat_gpu = scaled_laplacian(g_gpu)
+                    @test mat_gpu isa ACUMatrix{Float32}
+                    @test Array(mat_gpu) == mat
+                end
+            end
+        end
+
         @testset "constructors" begin
-            g = GNNGraph(adj_mat; graph_type=GRAPH_T)
             adjacency_matrix(g; dir=:out) == adj_mat
             adjacency_matrix(g; dir=:in) == adj_mat
         end 
 
         @testset "degree" begin
-            g = GNNGraph(adj_mat; graph_type=GRAPH_T)
             @test degree(g, dir=:out) == vec(sum(adj_mat, dims=2))
             @test degree(g, dir=:in) == vec(sum(adj_mat, dims=1))
+
+            if TEST_GPU
+                d = degree(g)
+                d_gpu = degree(g_gpu)
+                @test d_gpu isa CuVector
+                @test Array(d_gpu) == d
+            end
+        end
+
+        if TEST_GPU
+            @testset "functor" begin                
+                s_cpu, t_cpu = edge_index(g)
+                s_gpu, t_gpu = edge_index(g_gpu)
+                @test s_gpu isa CuVector{Int}
+                @test Array(s_gpu) == s_cpu
+                @test t_gpu isa CuVector{Int}
+                @test Array(t_gpu) == t_cpu
+            end
         end
     end
 
@@ -59,6 +109,10 @@
 
         # core functionality
         g = GNNGraph(s, t; graph_type=GRAPH_T)
+        if TEST_GPU
+            g_gpu = g |> gpu
+        end
+
         @test g.num_edges == 4
         @test g.num_nodes == 4
         @test collect(edges(g)) |> sort == collect(zip(s, t)) |> sort
@@ -78,7 +132,6 @@
         @test adjacency_list(g, dir=:in) ==  adj_list_in
 
         @testset "degree" begin
-            g = GNNGraph(adj_mat_out; graph_type=GRAPH_T)
             @test degree(g, dir=:out) == vec(sum(adj_mat_out, dims=2))
             @test degree(g, dir=:in) == vec(sum(adj_mat_out, dims=1))
         end
@@ -221,3 +274,5 @@
         @test first(d) == getgraph(g, 1:2)
     end
 end
+
+

test/gnngraph.jl

@@ -11,6 +11,8 @@ using Test
 using MLDatasets
 CUDA.allowscalar(false)
 
+const ACUMatrix{T} = Union{CuMatrix{T}, CUDA.CUSPARSE.CuSparseMatrix{T}}
+
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
 
 include("test_utils.jl")
@@ -26,17 +28,12 @@ tests = [
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
-# Testing all graph types. :sparse is a bit broken at the moment
 @testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo,:sparse,:dense)
 
     global GRAPH_T = graph_type
     global TEST_GPU = CUDA.functional()# && GRAPH_T != :sparse
 
     for t in tests
         include("$t.jl")
-
-        if TEST_GPU && isfile("cuda/$t.jl")
-            include("cuda/$t.jl")
-        end
     end
 end