From 49b6398039aef0fb6cf9e5e3ba62aaa633367dd8 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Sat, 30 Nov 2024 13:00:57 +0100 Subject: [PATCH] buildkite for GNNlib + get_graph_type (#530) --- .buildkite/pipeline.yml | 51 ++++++ GNNGraphs/src/GNNGraphs.jl | 1 + GNNGraphs/src/operators.jl | 4 +- GNNGraphs/src/query.jl | 58 +++++- GNNGraphs/src/transform.jl | 1 + GNNGraphs/test/gnngraph.jl | 3 + GNNGraphs/test/query.jl | 17 ++ GNNlib/Project.toml | 18 +- GNNlib/ext/GNNlibAMDGPUExt.jl | 34 ++++ GNNlib/src/msgpass.jl | 2 +- GNNlib/test/Project.toml | 22 +++ GNNlib/test/msgpass.jl | 219 +++++++++++++++++++++++ GNNlib/test/msgpass_tests.jl | 140 --------------- GNNlib/test/runtests.jl | 32 +++- GNNlib/test/shared_testsetup.jl | 12 -- GNNlib/test/test_module.jl | 180 +++++++++++++++++++ GNNlib/test/{utils_tests.jl => utils.jl} | 3 +- GraphNeuralNetworks/test/Project.toml | 4 + GraphNeuralNetworks/test/test_module.jl | 57 +++--- 19 files changed, 658 insertions(+), 200 deletions(-) create mode 100644 GNNlib/ext/GNNlibAMDGPUExt.jl create mode 100644 GNNlib/test/Project.toml create mode 100644 GNNlib/test/msgpass.jl delete mode 100644 GNNlib/test/msgpass_tests.jl delete mode 100644 GNNlib/test/shared_testsetup.jl create mode 100644 GNNlib/test/test_module.jl rename GNNlib/test/{utils_tests.jl => utils.jl} (96%) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 4fd28dd07..49127a403 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -49,3 +49,54 @@ steps: GNN_TEST_AMDGPU: "true" GNN_TEST_CPU: "false" timeout_in_minutes: 60 + + - label: "GNNlib CUDA" + plugins: + - JuliaCI/julia#v1: + version: "1" + - JuliaCI/julia-coverage#v1: + dirs: + - GNNlib/src + command: | + julia --color=yes --depwarn=yes --project=GNNlib/test -e ' + import Pkg + dev_pkgs = Pkg.PackageSpec[] + for pkg in ("GNNGraphs", "GNNlib") + push!(dev_pkgs, Pkg.PackageSpec(path=pkg)); + end + Pkg.develop(dev_pkgs) + Pkg.add(["CUDA", "cuDNN"]) + Pkg.test("GNNlib")' + agents: + queue: "juliagpu" + cuda: "*" + env: + GNN_TEST_CUDA: "true" + GNN_TEST_CPU: "false" + timeout_in_minutes: 60 + + - label: "GNNlib AMDGPU" + plugins: + - JuliaCI/julia#v1: + version: "1" + - JuliaCI/julia-coverage#v1: + dirs: + - GNNlib/src + command: | + julia --color=yes --depwarn=yes --project=GNNlib/test -e ' + import Pkg + dev_pkgs = Pkg.PackageSpec[] + for pkg in ("GNNGraphs", "GNNlib") + push!(dev_pkgs, Pkg.PackageSpec(path=pkg)); + end + Pkg.develop(dev_pkgs) + Pkg.add(["AMDGPU"]) + Pkg.test("GNNlib")' + agents: + queue: "juliagpu" + rocm: "*" + rocmgpu: "*" + env: + GNN_TEST_AMDGPU: "true" + GNN_TEST_CPU: "false" + timeout_in_minutes: 60 diff --git a/GNNGraphs/src/GNNGraphs.jl b/GNNGraphs/src/GNNGraphs.jl index 5a5b5fe66..3054a9ab8 100644 --- a/GNNGraphs/src/GNNGraphs.jl +++ b/GNNGraphs/src/GNNGraphs.jl @@ -47,6 +47,7 @@ include("query.jl") export adjacency_list, edge_index, get_edge_weight, + get_graph_type, graph_indicator, has_multi_edges, is_directed, diff --git a/GNNGraphs/src/operators.jl b/GNNGraphs/src/operators.jl index 4fdd6ac87..1faa4adcb 100644 --- a/GNNGraphs/src/operators.jl +++ b/GNNGraphs/src/operators.jl @@ -6,8 +6,8 @@ Intersect two graphs by keeping only the common edges. """ function Base.intersect(g1::GNNGraph, g2::GNNGraph) @assert g1.num_nodes == g2.num_nodes - @assert graph_type_symbol(g1) == graph_type_symbol(g2) - graph_type = graph_type_symbol(g1) + @assert get_graph_type(g1) == get_graph_type(g2) + graph_type = get_graph_type(g1) num_nodes = g1.num_nodes idx1, _ = edge_encoding(edge_index(g1)..., num_nodes) diff --git a/GNNGraphs/src/query.jl b/GNNGraphs/src/query.jl index 18622af21..a11eb564c 100644 --- a/GNNGraphs/src/query.jl +++ b/GNNGraphs/src/query.jl @@ -80,9 +80,61 @@ function Graphs.has_edge(g::GNNHeteroGraph, edge_t::EType, i::Integer, j::Intege return any((s .== i) .& (t .== j)) end -graph_type_symbol(::GNNGraph{<:COO_T}) = :coo -graph_type_symbol(::GNNGraph{<:SPARSE_T}) = :sparse -graph_type_symbol(::GNNGraph{<:ADJMAT_T}) = :dense +""" + get_graph_type(g::GNNGraph) + +Return the underlying representation for the graph `g` as a symbol. + +Possible values are: +- `:coo`: Coordinate list representation. The graph is stored as a tuple of vectors `(s, t, w)`, + where `s` and `t` are the source and target nodes of the edges, and `w` is the edge weights. +- `:sparse`: Sparse matrix representation. The graph is stored as a sparse matrix representing the weighted adjacency matrix. +- `:dense`: Dense matrix representation. The graph is stored as a dense matrix representing the weighted adjacency matrix. + +The default representation for graph constructors GNNGraphs.jl is `:coo`. +The underlying representation can be accessed through the `g.graph` field. + +See also [`GNNGraph`](@ref). + +# Examples + +The default representation for graph constructors GNNGraphs.jl is `:coo`. +```jldoctest +julia> g = rand_graph(5, 10) +GNNGraph: + num_nodes: 5 + num_edges: 10 + +julia> get_graph_type(g) +:coo +``` +The `GNNGraph` constructor can also be used to create graphs with different representations. +```jldoctest +julia> g = GNNGraph([2,3,5], [1,2,4], graph_type=:sparse) +GNNGraph: + num_nodes: 5 + num_edges: 3 + +julia> g.graph +5×5 SparseArrays.SparseMatrixCSC{Int64, Int64} with 3 stored entries: + ⋅ ⋅ ⋅ ⋅ ⋅ + 1 ⋅ ⋅ ⋅ ⋅ + ⋅ 1 ⋅ ⋅ ⋅ + ⋅ ⋅ ⋅ ⋅ ⋅ + ⋅ ⋅ ⋅ 1 ⋅ + +julia> get_graph_type(g) +:sparse + +julia> gcoo = GNNGraph(g, graph_type=:coo); + +julia> gcoo.graph +([2, 3, 5], [1, 2, 4], [1, 1, 1]) +``` +""" +get_graph_type(::GNNGraph{<:COO_T}) = :coo +get_graph_type(::GNNGraph{<:SPARSE_T}) = :sparse +get_graph_type(::GNNGraph{<:ADJMAT_T}) = :dense Graphs.nv(g::GNNGraph) = g.num_nodes Graphs.ne(g::GNNGraph) = g.num_edges diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl index 325a20f5c..6e0b564cd 100644 --- a/GNNGraphs/src/transform.jl +++ b/GNNGraphs/src/transform.jl @@ -731,6 +731,7 @@ end Set `w` as edge weights in the returned graph. """ function set_edge_weight(g::GNNGraph, w::AbstractVector) + # TODO preserve the representation instead of converting to COO s, t = edge_index(g) @assert length(w) == length(s) diff --git a/GNNGraphs/test/gnngraph.jl b/GNNGraphs/test/gnngraph.jl index d65f7adda..0c15e6cff 100644 --- a/GNNGraphs/test/gnngraph.jl +++ b/GNNGraphs/test/gnngraph.jl @@ -1,3 +1,6 @@ +# TODO test that the graph type is preserved +# when constructing a GNNGraph from another + @testset "Constructor: adjacency matrix" begin A = sprand(10, 10, 0.5) sA, tA, vA = findnz(A) diff --git a/GNNGraphs/test/query.jl b/GNNGraphs/test/query.jl index e7f55e76a..a345ae779 100644 --- a/GNNGraphs/test/query.jl +++ b/GNNGraphs/test/query.jl @@ -257,3 +257,20 @@ if GRAPH_T == :coo end end +@testset "get_graph_type" begin + g = rand_graph(10, 20, graph_type = GRAPH_T) + @test get_graph_type(g) == GRAPH_T + + gsparse = GNNGraph(g, graph_type=:sparse) + @test get_graph_type(gsparse) == :sparse + @test gsparse.graph isa SparseMatrixCSC + + gcoo = GNNGraph(g, graph_type=:coo) + @test get_graph_type(gcoo) == :coo + @test gcoo.graph[1:2] isa Tuple{Vector{Int}, Vector{Int}} + + + gdense = GNNGraph(g, graph_type=:dense) + @test get_graph_type(gdense) == :dense + @test gdense.graph isa Matrix{Int} +end diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml index 20cf840c3..4acfceaaa 100644 --- a/GNNlib/Project.toml +++ b/GNNlib/Project.toml @@ -7,6 +7,7 @@ version = "0.2.2" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" @@ -14,28 +15,27 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" [extensions] +GNNlibAMDGPUExt = "AMDGPU" GNNlibCUDAExt = "CUDA" +# GPUArraysCore is not needed as a direct dependency +# but pinning it to 0.1 avoids problems when we do Pkg.add("CUDA") in testing +# See https://github.com/JuliaGPU/CUDA.jl/issues/2564 + [compat] +AMDGPU = "1" CUDA = "4, 5" ChainRulesCore = "1.24" DataStructures = "0.18" GNNGraphs = "1.0" +GPUArraysCore = "0.1" LinearAlgebra = "1" MLUtils = "0.4" NNlib = "0.9" Random = "1" Statistics = "1" julia = "1.10" - -[extras] -ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" -Reexport = "189a3867-3050-52da-a836-e630ba90ab69" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[targets] -test = ["Test", "ReTestItems", "Reexport", "SparseArrays"] diff --git a/GNNlib/ext/GNNlibAMDGPUExt.jl b/GNNlib/ext/GNNlibAMDGPUExt.jl new file mode 100644 index 000000000..83a4591c4 --- /dev/null +++ b/GNNlib/ext/GNNlibAMDGPUExt.jl @@ -0,0 +1,34 @@ +module GNNlibAMDGPUExt + +using AMDGPU: AnyROCMatrix +using Random, Statistics, LinearAlgebra +using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj +using GNNGraphs: GNNGraph, COO_T, SPARSE_T + +###### PROPAGATE SPECIALIZATIONS #################### + +## COPY_XJ + +## avoid the fast path on gpu until we have better cuda support +function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+), + xi, xj::AnyROCMatrix, e) + propagate((xi, xj, e) -> copy_xj(xi, xj, e), g, +, xi, xj, e) +end + +## E_MUL_XJ + +## avoid the fast path on gpu until we have better cuda support +function GNNlib.propagate(::typeof(e_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+), + xi, xj::AnyROCMatrix, e::AbstractVector) + propagate((xi, xj, e) -> e_mul_xj(xi, xj, e), g, +, xi, xj, e) +end + +## W_MUL_XJ + +## avoid the fast path on gpu until we have better support +function GNNlib.propagate(::typeof(w_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+), + xi, xj::AnyROCMatrix, e::Nothing) + propagate((xi, xj, e) -> w_mul_xj(xi, xj, e), g, +, xi, xj, e) +end + +end #module diff --git a/GNNlib/src/msgpass.jl b/GNNlib/src/msgpass.jl index 1aa17437a..f6a1cf659 100644 --- a/GNNlib/src/msgpass.jl +++ b/GNNlib/src/msgpass.jl @@ -184,7 +184,7 @@ xj_sub_xi(xi, xj, e) = xj .- xi """ e_mul_xj(xi, xj, e) = reshape(e, (...)) .* xj -Reshape `e` into broadcast compatible shape with `xj` +Reshape `e` into a broadcast compatible shape with `xj` (by prepending singleton dimensions) then perform broadcasted multiplication. """ diff --git a/GNNlib/test/Project.toml b/GNNlib/test/Project.toml new file mode 100644 index 000000000..36fcae23b --- /dev/null +++ b/GNNlib/test/Project.toml @@ -0,0 +1,22 @@ +[deps] +FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c" +GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +GPUArraysCore = "0.1" diff --git a/GNNlib/test/msgpass.jl b/GNNlib/test/msgpass.jl new file mode 100644 index 000000000..5741cdb5a --- /dev/null +++ b/GNNlib/test/msgpass.jl @@ -0,0 +1,219 @@ +@testitem "msgpass" setup=[TestModuleGNNlib] begin + using .TestModuleGNNlib + #TODO test all graph types + g = TEST_GRAPHS[1] + out_channel = size(g.x, 1) + num_V = g.num_nodes + num_E = g.num_edges + g = GNNGraph(g, edata = rand(Float32, size(g.x, 1), g.num_edges)) + + @testset "propagate" begin + function message(xi, xj, e) + @test xi === nothing + @test e === nothing + ones(Float32, out_channel, size(xj, 2)) + end + + m = propagate(message, g, +, xj = g.x) + + @test size(m) == (out_channel, num_V) + + @testset "isolated nodes" begin + x1 = rand(1, 6) + g1 = GNNGraph(collect(1:5), collect(1:5), num_nodes = 6) + y1 = propagate((xi, xj, e) -> xj, g1, +, xj = x1) + @test size(y1) == (1, 6) + end + end + + @testset "apply_edges" begin + m = apply_edges(g, e = g.e) do xi, xj, e + @test xi === nothing + @test xj === nothing + ones(out_channel, size(e, 2)) + end + + @test m == ones(out_channel, num_E) + + # With NamedTuple input + m = apply_edges(g, xj = (; a = g.x, b = 2g.x), e = g.e) do xi, xj, e + @test xi === nothing + @test xj.b == 2 * xj.a + @test size(xj.a, 2) == size(xj.b, 2) == size(e, 2) + ones(out_channel, size(e, 2)) + end + + # NamedTuple output + m = apply_edges(g, e = g.e) do xi, xj, e + @test xi === nothing + @test xj === nothing + (; a = ones(out_channel, size(e, 2))) + end + + @test m.a == ones(out_channel, num_E) + + @testset "sizecheck" begin + x = rand(3, g.num_nodes - 1) + @test_throws AssertionError apply_edges(copy_xj, g, xj = x) + @test_throws AssertionError apply_edges(copy_xj, g, xi = x) + + x = (a = rand(3, g.num_nodes), b = rand(3, g.num_nodes + 1)) + @test_throws AssertionError apply_edges(copy_xj, g, xj = x) + @test_throws AssertionError apply_edges(copy_xj, g, xi = x) + + e = rand(3, g.num_edges - 1) + @test_throws AssertionError apply_edges(copy_xj, g, e = e) + end + end + + @testset "copy_xj" begin + n = 128 + A = sprand(n, n, 0.1) + Adj = map(x -> x > 0 ? 1 : 0, A) + X = rand(10, n) + + g = GNNGraph(A, ndata = X, graph_type = :coo) + + function spmm_copyxj_fused(g) + propagate(copy_xj, + g, +; xj = g.ndata.x) + end + + function spmm_copyxj_unfused(g) + propagate((xi, xj, e) -> xj, + g, +; xj = g.ndata.x) + end + + @test spmm_copyxj_unfused(g) ≈ X * Adj + @test spmm_copyxj_fused(g) ≈ X * Adj + end + + @testset "e_mul_xj and w_mul_xj for weighted conv" begin + n = 128 + A = sprand(n, n, 0.1) + Adj = map(x -> x > 0 ? 1 : 0, A) + X = rand(10, n) + + g = GNNGraph(A, ndata = X, edata = A.nzval, graph_type = :coo) + + function spmm_unfused(g) + propagate((xi, xj, e) -> reshape(e, 1, :) .* xj, + g, +; xj = g.ndata.x, e = g.edata.e) + end + function spmm_fused(g) + propagate(e_mul_xj, + g, +; xj = g.ndata.x, e = g.edata.e) + end + + function spmm_fused2(g) + propagate(w_mul_xj, + g, +; xj = g.ndata.x) + end + + @test spmm_unfused(g) ≈ X * A + @test spmm_fused(g) ≈ X * A + @test spmm_fused2(g) ≈ X * A + end + + @testset "aggregate_neighbors" begin + @testset "sizecheck" begin + m = rand(2, g.num_edges - 1) + @test_throws AssertionError aggregate_neighbors(g, +, m) + + m = (a = rand(2, g.num_edges + 1), b = nothing) + @test_throws AssertionError aggregate_neighbors(g, +, m) + end + end +end + +@testitem "propagate" setup=[TestModuleGNNlib] begin + using .TestModuleGNNlib + + @testset "copy_xj +" begin + for g in TEST_GRAPHS + f(g, x) = propagate(copy_xj, g, +, xj = x) + test_gradients(f, g, g.x; test_grad_f=false) + end + end + + @testset "copy_xj mean" begin + for g in TEST_GRAPHS + f(g, x) = propagate(copy_xj, g, mean, xj = x) + test_gradients(f, g, g.x; test_grad_f=false) + end + end + + @testset "e_mul_xj +" begin + for g in TEST_GRAPHS + e = rand(Float32, size(g.x, 1), g.num_edges) + f(g, x, e) = propagate(e_mul_xj, g, +; xj = x, e) + test_gradients(f, g, g.x, e; test_grad_f=false) + end + end + + @testset "w_mul_xj +" begin + for g in TEST_GRAPHS + w = rand(Float32, g.num_edges) + function f(g, x, w) + g = set_edge_weight(g, w) + return propagate(w_mul_xj, g, +, xj = x) + end + test_gradients(f, g, g.x, w; test_grad_f=false) + end + end +end + +@testitem "propagate GPU" setup=[TestModuleGNNlib] tags=[:gpu] begin + using .TestModuleGNNlib + + @testset "copy_xj +" begin + for g in TEST_GRAPHS + dev = gpu_device(force=true) + broken = get_graph_type(g) == :sparse && dev isa AMDGPUDevice + f(g, x) = propagate(copy_xj, g, +, xj = x) + @test test_gradients( + f, g, g.x; test_gpu=true, test_grad_f=false, compare_finite_diff=false + ) broken=broken + end + end + + @testset "copy_xj mean" begin + for g in TEST_GRAPHS + dev = gpu_device(force=true) + broken = get_graph_type(g) == :sparse && dev isa AMDGPUDevice + f(g, x) = propagate(copy_xj, g, mean, xj = x) + @test test_gradients( + f, g, g.x; test_gpu=true, test_grad_f=false, compare_finite_diff=false + ) broken=broken + end + end + + @testset "e_mul_xj +" begin + for g in TEST_GRAPHS + dev = gpu_device(force=true) + broken = get_graph_type(g) == :sparse && dev isa AMDGPUDevice + e = rand(Float32, size(g.x, 1), g.num_edges) + f(g, x, e) = propagate(e_mul_xj, g, +; xj = x, e) + @test test_gradients( + f, g, g.x, e; test_gpu=true, test_grad_f=false, compare_finite_diff=false + ) broken=broken + end + end + + @testset "w_mul_xj +" begin + for g in TEST_GRAPHS + w = rand(Float32, g.num_edges) + function f(g, x, w) + g = set_edge_weight(g, w) + return propagate(w_mul_xj, g, +, xj = x) + end + dev = gpu_device(force=true) + # @show get_graph_type(g) has_isolated_nodes(g) + # broken = get_graph_type(g) == :sparse + broken = true + @test test_gradients( + f, g, g.x, w; test_gpu=true, test_grad_f=false, compare_finite_diff=false + ) broken=broken + end + end +end diff --git a/GNNlib/test/msgpass_tests.jl b/GNNlib/test/msgpass_tests.jl deleted file mode 100644 index 60d13fcb0..000000000 --- a/GNNlib/test/msgpass_tests.jl +++ /dev/null @@ -1,140 +0,0 @@ -@testitem "msgpass" setup=[SharedTestSetup] begin - #TODO test all graph types - GRAPH_T = :coo - in_channel = 10 - out_channel = 5 - num_V = 6 - num_E = 14 - T = Float32 - - adj = [0 1 0 0 0 0 - 1 0 0 1 1 1 - 0 0 0 0 0 1 - 0 1 0 0 1 0 - 0 1 0 1 0 1 - 0 1 1 0 1 0] - - X = rand(T, in_channel, num_V) - E = rand(T, in_channel, num_E) - - g = GNNGraph(adj, graph_type = GRAPH_T) - - @testset "propagate" begin - function message(xi, xj, e) - @test xi === nothing - @test e === nothing - ones(T, out_channel, size(xj, 2)) - end - - m = propagate(message, g, +, xj = X) - - @test size(m) == (out_channel, num_V) - - @testset "isolated nodes" begin - x1 = rand(1, 6) - g1 = GNNGraph(collect(1:5), collect(1:5), num_nodes = 6) - y1 = propagate((xi, xj, e) -> xj, g, +, xj = x1) - @test size(y1) == (1, 6) - end - end - - @testset "apply_edges" begin - m = apply_edges(g, e = E) do xi, xj, e - @test xi === nothing - @test xj === nothing - ones(out_channel, size(e, 2)) - end - - @test m == ones(out_channel, num_E) - - # With NamedTuple input - m = apply_edges(g, xj = (; a = X, b = 2X), e = E) do xi, xj, e - @test xi === nothing - @test xj.b == 2 * xj.a - @test size(xj.a, 2) == size(xj.b, 2) == size(e, 2) - ones(out_channel, size(e, 2)) - end - - # NamedTuple output - m = apply_edges(g, e = E) do xi, xj, e - @test xi === nothing - @test xj === nothing - (; a = ones(out_channel, size(e, 2))) - end - - @test m.a == ones(out_channel, num_E) - - @testset "sizecheck" begin - x = rand(3, g.num_nodes - 1) - @test_throws AssertionError apply_edges(copy_xj, g, xj = x) - @test_throws AssertionError apply_edges(copy_xj, g, xi = x) - - x = (a = rand(3, g.num_nodes), b = rand(3, g.num_nodes + 1)) - @test_throws AssertionError apply_edges(copy_xj, g, xj = x) - @test_throws AssertionError apply_edges(copy_xj, g, xi = x) - - e = rand(3, g.num_edges - 1) - @test_throws AssertionError apply_edges(copy_xj, g, e = e) - end - end - - @testset "copy_xj" begin - n = 128 - A = sprand(n, n, 0.1) - Adj = map(x -> x > 0 ? 1 : 0, A) - X = rand(10, n) - - g = GNNGraph(A, ndata = X, graph_type = GRAPH_T) - - function spmm_copyxj_fused(g) - propagate(copy_xj, - g, +; xj = g.ndata.x) - end - - function spmm_copyxj_unfused(g) - propagate((xi, xj, e) -> xj, - g, +; xj = g.ndata.x) - end - - @test spmm_copyxj_unfused(g) ≈ X * Adj - @test spmm_copyxj_fused(g) ≈ X * Adj - end - - @testset "e_mul_xj and w_mul_xj for weighted conv" begin - n = 128 - A = sprand(n, n, 0.1) - Adj = map(x -> x > 0 ? 1 : 0, A) - X = rand(10, n) - - g = GNNGraph(A, ndata = X, edata = A.nzval, graph_type = GRAPH_T) - - function spmm_unfused(g) - propagate((xi, xj, e) -> reshape(e, 1, :) .* xj, - g, +; xj = g.ndata.x, e = g.edata.e) - end - function spmm_fused(g) - propagate(e_mul_xj, - g, +; xj = g.ndata.x, e = g.edata.e) - end - - function spmm_fused2(g) - propagate(w_mul_xj, - g, +; xj = g.ndata.x) - end - - @test spmm_unfused(g) ≈ X * A - @test spmm_fused(g) ≈ X * A - @test spmm_fused2(g) ≈ X * A - end - - @testset "aggregate_neighbors" begin - @testset "sizecheck" begin - m = rand(2, g.num_edges - 1) - @test_throws AssertionError aggregate_neighbors(g, +, m) - - m = (a = rand(2, g.num_edges + 1), b = nothing) - @test_throws AssertionError aggregate_neighbors(g, +, m) - end - end - -end \ No newline at end of file diff --git a/GNNlib/test/runtests.jl b/GNNlib/test/runtests.jl index e4c4512b4..d420f5a6c 100644 --- a/GNNlib/test/runtests.jl +++ b/GNNlib/test/runtests.jl @@ -1,6 +1,28 @@ -using GNNlib -using Test -using ReTestItems -using Random, Statistics +using TestItemRunner -runtests(GNNlib) +## See https://www.julia-vscode.org/docs/stable/userguide/testitems/ +## for how to run the tests within VS Code. +## See test_module.jl for the test infrastructure. + +## Uncomment below and in test_module.jl to change the default test settings +# ENV["GNN_TEST_CPU"] = "false" +# ENV["GNN_TEST_CUDA"] = "true" +# ENV["GNN_TEST_AMDGPU"] = "true" +# ENV["GNN_TEST_Metal"] = "true" + +# The only available tag at the moment is :gpu +# Tests not tagged with :gpu are considered to be CPU tests +# Tests tagged with :gpu should run on all GPU backends + +if get(ENV, "GNN_TEST_CPU", "true") == "true" + @run_package_tests filter = ti -> :gpu ∉ ti.tags +end +if get(ENV, "GNN_TEST_CUDA", "false") == "true" + @run_package_tests filter = ti -> :gpu ∈ ti.tags +end +if get(ENV, "GNN_TEST_AMDGPU", "false") == "true" + @run_package_tests filter = ti -> :gpu ∈ ti.tags +end +if get(ENV, "GNN_TEST_Metal", "false") == "true" + @run_package_tests filter = ti -> :gpu ∈ ti.tags +end diff --git a/GNNlib/test/shared_testsetup.jl b/GNNlib/test/shared_testsetup.jl deleted file mode 100644 index 106db5159..000000000 --- a/GNNlib/test/shared_testsetup.jl +++ /dev/null @@ -1,12 +0,0 @@ -@testsetup module SharedTestSetup - -import Reexport: @reexport - -@reexport using GNNlib -@reexport using GNNGraphs -@reexport using NNlib -@reexport using MLUtils -@reexport using SparseArrays -@reexport using Test, Random, Statistics - -end \ No newline at end of file diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl new file mode 100644 index 000000000..27a83154c --- /dev/null +++ b/GNNlib/test/test_module.jl @@ -0,0 +1,180 @@ +@testmodule TestModuleGNNlib begin + +using Pkg + +### GPU backends settings ############ +# tried to put this in __init__ but is not executed for some reason + +## Uncomment below to change the default test settings +# ENV["GNN_TEST_CUDA"] = "true" +# ENV["GNN_TEST_AMDGPU"] = "true" +# ENV["GNN_TEST_Metal"] = "true" + +to_test(backend) = get(ENV, "GNN_TEST_$(backend)", "false") == "true" +has_dependecies(pkgs) = all(pkg -> haskey(Pkg.project().dependencies, pkg), pkgs) +deps_dict = Dict(:CUDA => ["CUDA", "cuDNN"], :AMDGPU => ["AMDGPU"], :Metal => ["Metal"]) + +for (backend, deps) in deps_dict + if to_test(backend) + if !has_dependecies(deps) + Pkg.add(deps) + end + @eval using $backend + if backend == :CUDA + @eval using cuDNN + end + @eval $backend.allowscalar(false) + end +end +###################################### + +import Reexport: @reexport + +@reexport using GNNlib +@reexport using GNNGraphs +@reexport using NNlib +@reexport using MLUtils +@reexport using SparseArrays +@reexport using Test, Random, Statistics +@reexport using MLDataDevices +using Functors: fmapstructure_with_path +using FiniteDifferences: FiniteDifferences +using Zygote: Zygote +using Flux: Flux + +# from this module +export D_IN, D_OUT, GRAPH_TYPES, TEST_GRAPHS, + test_gradients, finitediff_withgradient, + check_equal_leaves + + +const D_IN = 3 +const D_OUT = 5 + +function finitediff_withgradient(f, x...) + y = f(x...) + # We set a range to avoid domain errors + fdm = FiniteDifferences.central_fdm(5, 1, max_range=1e-2) + return y, FiniteDifferences.grad(fdm, f, x...) +end + +function check_equal_leaves(a, b; rtol=1e-4, atol=1e-4) + equal = true + fmapstructure_with_path(a, b) do kp, x, y + if x isa AbstractArray + # @show kp + # @assert isapprox(x, y; rtol, atol) + if !isapprox(x, y; rtol, atol) + equal = false + end + end + end + @assert equal +end + +function test_gradients( + f, + graph::GNNGraph, + xs...; + rtol=1e-5, atol=1e-5, + test_gpu = false, + test_grad_f = true, + test_grad_x = true, + compare_finite_diff = true, + loss = (f, g, xs...) -> mean(f(g, xs...)), + ) + + if !test_gpu && !compare_finite_diff + error("You should either compare finite diff vs CPU AD \ + or CPU AD vs GPU AD.") + end + + ## Let's make sure first that the forward pass works. + l = loss(f, graph, xs...) + @assert l isa Number + if test_gpu + gpu_dev = gpu_device(force=true) + cpu_dev = cpu_device() + graph_gpu = graph |> gpu_dev + xs_gpu = xs |> gpu_dev + f_gpu = f |> gpu_dev + l_gpu = loss(f_gpu, graph_gpu, xs_gpu...) + @assert l_gpu isa Number + end + + if test_grad_x + # Zygote gradient with respect to input. + y, g = Zygote.withgradient((xs...) -> loss(f, graph, xs...), xs...) + + if compare_finite_diff + # Cast to Float64 to avoid precision issues. + f64 = f |> Flux.f64 + xs64 = xs .|> Flux.f64 + y_fd, g_fd = finitediff_withgradient((xs...) -> loss(f64, graph, xs...), xs64...) + @assert isapprox(y, y_fd; rtol, atol) + check_equal_leaves(g, g_fd; rtol, atol) + end + + if test_gpu + # Zygote gradient with respect to input on GPU. + y_gpu, g_gpu = Zygote.withgradient((xs...) -> loss(f_gpu, graph_gpu, xs...), xs_gpu...) + @assert get_device(g_gpu) == get_device(xs_gpu) + @assert isapprox(y_gpu, y; rtol, atol) + check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol) + end + end + + if test_grad_f + # Zygote gradient with respect to f. + y, g = Zygote.withgradient(f -> loss(f, graph, xs...), f) + + if compare_finite_diff + # Cast to Float64 to avoid precision issues. + f64 = f |> Flux.f64 + ps, re = Flux.destructure(f64) + y_fd, g_fd = finitediff_withgradient(ps -> loss(re(ps),graph, xs...), ps) + g_fd = (re(g_fd[1]),) + @assert isapprox(y, y_fd; rtol, atol) + check_equal_leaves(g, g_fd; rtol, atol) + end + + if test_gpu + # Zygote gradient with respect to f on GPU. + y_gpu, g_gpu = Zygote.withgradient(f -> loss(f,graph_gpu, xs_gpu...), f_gpu) + # @assert get_device(g_gpu) == get_device(xs_gpu) + @assert isapprox(y_gpu, y; rtol, atol) + check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol) + end + end + @test true # if we reach here, the test passed + return true +end + +function generate_test_graphs(graph_type) + adj1 = [0 1 0 1 + 1 0 1 0 + 0 1 0 1 + 1 0 1 0] + + g1 = GNNGraph(adj1, + ndata = rand(Float32, D_IN, 4); + graph_type) + + adj_single_vertex = [0 0 0 1 + 0 0 0 0 + 0 0 0 1 + 1 0 1 0] + + g_single_vertex = GNNGraph(adj_single_vertex, + ndata = rand(Float32, D_IN, 4); + graph_type) + + return (g1, g_single_vertex) +end + +GRAPH_TYPES = [:coo, :dense, :sparse] +TEST_GRAPHS = [generate_test_graphs(:coo)..., + generate_test_graphs(:dense)..., + generate_test_graphs(:sparse)...] + +end # module \ No newline at end of file diff --git a/GNNlib/test/utils_tests.jl b/GNNlib/test/utils.jl similarity index 96% rename from GNNlib/test/utils_tests.jl rename to GNNlib/test/utils.jl index 762ba58b9..bf06f86fd 100644 --- a/GNNlib/test/utils_tests.jl +++ b/GNNlib/test/utils.jl @@ -1,4 +1,5 @@ -@testitem "utils" setup=[SharedTestSetup] begin +@testitem "utils" setup=[TestModuleGNNlib] begin + using .TestModuleGNNlib # TODO test all graph types GRAPH_T = :coo De, Dx = 3, 2 diff --git a/GraphNeuralNetworks/test/Project.toml b/GraphNeuralNetworks/test/Project.toml index 5e554715c..d6f77b391 100644 --- a/GraphNeuralNetworks/test/Project.toml +++ b/GraphNeuralNetworks/test/Project.toml @@ -3,6 +3,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" @@ -13,3 +14,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +GPUArraysCore = "0.1" diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl index e7f9a0031..fa50e5821 100644 --- a/GraphNeuralNetworks/test/test_module.jl +++ b/GraphNeuralNetworks/test/test_module.jl @@ -3,25 +3,25 @@ using Pkg ## Uncomment below to change the default test settings -# ENV["GNN_TEST_CPU"] = "false" # ENV["GNN_TEST_CUDA"] = "true" # ENV["GNN_TEST_AMDGPU"] = "true" # ENV["GNN_TEST_Metal"] = "true" -if get(ENV, "GNN_TEST_CUDA", "false") == "true" - Pkg.add(["CUDA", "cuDNN"]) - using CUDA - CUDA.allowscalar(false) -end -if get(ENV, "GNN_TEST_AMDGPU", "false") == "true" - Pkg.add("AMDGPU") - using AMDGPU - AMDGPU.allowscalar(false) -end -if get(ENV, "GNN_TEST_Metal", "false") == "true" - Pkg.add("Metal") - using Metal - Metal.allowscalar(false) +to_test(backend) = get(ENV, "GNN_TEST_$(backend)", "false") == "true" +has_dependecies(pkgs) = all(pkg -> haskey(Pkg.project().dependencies, pkg), pkgs) +deps_dict = Dict(:CUDA => ["CUDA", "cuDNN"], :AMDGPU => ["AMDGPU"], :Metal => ["Metal"]) + +for (backend, deps) in deps_dict + if to_test(backend) + if !has_dependecies(deps) + Pkg.add(deps) + end + @eval using $backend + if backend == :CUDA + @eval using cuDNN + end + @eval $backend.allowscalar(false) + end end using GraphNeuralNetworks @@ -59,15 +59,17 @@ function finitediff_withgradient(f, x...) end function check_equal_leaves(a, b; rtol=1e-4, atol=1e-4) + equal = true fmapstructure_with_path(a, b) do kp, x, y if x isa AbstractArray # @show kp - @test x ≈ y rtol=rtol atol=atol - # elseif x isa Number - # @show kp - # @test x ≈ y rtol=rtol atol=atol + # @assert isapprox(x, y; rtol, atol) + if !isapprox(x, y; rtol, atol) + equal = false + end end end + @assert equal end function test_gradients( @@ -89,7 +91,7 @@ function test_gradients( ## Let's make sure first that the forward pass works. l = loss(f, graph, xs...) - @test l isa Number + @assert l isa Number if test_gpu gpu_dev = gpu_device(force=true) cpu_dev = cpu_device() @@ -97,7 +99,7 @@ function test_gradients( xs_gpu = xs |> gpu_dev f_gpu = f |> gpu_dev l_gpu = loss(f_gpu, graph_gpu, xs_gpu...) - @test l_gpu isa Number + @assert l_gpu isa Number end if test_grad_x @@ -109,15 +111,15 @@ function test_gradients( f64 = f |> Flux.f64 xs64 = xs .|> Flux.f64 y_fd, g_fd = finitediff_withgradient((xs...) -> loss(f64, graph, xs...), xs64...) - @test y ≈ y_fd rtol=rtol atol=atol + @assert isapprox(y, y_fd; rtol, atol) check_equal_leaves(g, g_fd; rtol, atol) end if test_gpu # Zygote gradient with respect to input on GPU. y_gpu, g_gpu = Zygote.withgradient((xs...) -> loss(f_gpu, graph_gpu, xs...), xs_gpu...) - @test get_device(g_gpu) == get_device(xs_gpu) - @test y_gpu ≈ y rtol=rtol atol=atol + @assert get_device(g_gpu) == get_device(xs_gpu) + @assert isapprox(y_gpu, y; rtol, atol) check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol) end end @@ -132,18 +134,19 @@ function test_gradients( ps, re = Flux.destructure(f64) y_fd, g_fd = finitediff_withgradient(ps -> loss(re(ps),graph, xs...), ps) g_fd = (re(g_fd[1]),) - @test y ≈ y_fd rtol=rtol atol=atol + @assert isapprox(y, y_fd; rtol, atol) check_equal_leaves(g, g_fd; rtol, atol) end if test_gpu # Zygote gradient with respect to f on GPU. y_gpu, g_gpu = Zygote.withgradient(f -> loss(f,graph_gpu, xs_gpu...), f_gpu) - # @test get_device(g_gpu) == get_device(xs_gpu) - @test y_gpu ≈ y rtol=rtol atol=atol + # @assert get_device(g_gpu) == get_device(xs_gpu) + @assert isapprox(y_gpu, y; rtol, atol) check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol) end end + @test true # if we reach here, the test passed return true end