add buildkite workflow (#526)

Author: CarloLucibello

.buildkite/pipeline.yml

@@ -0,0 +1,25 @@
+steps:
+  - label: "GNN CUDA"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-coverage#v1:
+          dirs:
+            - GraphNeuralNetworks/src
+    command: |
+      julia --color=yes --depwarn=yes --project=GraphNeuralNetworks/test -e '
+        import Pkg
+        dev_pkgs = Pkg.PackageSpec[]
+        for pkg in ("GNNGraphs", "GNNlib", "GraphNeuralNetworks")
+          push!(dev_pkgs, Pkg.PackageSpec(path=pkg));
+        end
+        Pkg.develop(dev_pkgs)
+        Pkg.add(["CUDA", "cuDNN"])
+        Pkg.test("GraphNeuralNetworks")'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      GNN_TEST_CUDA: "true"
+      GNN_TEST_CPU: "false"
+    timeout_in_minutes: 60

GNNGraphs/test/test_utils.jl

@@ -1,6 +1,12 @@
 using ChainRulesTestUtils, FiniteDifferences, Zygote, Adapt, CUDA
 CUDA.allowscalar(false)
 
+# Using this until https://github.com/JuliaDiff/FiniteDifferences.jl/issues/188 is fixed
+function FiniteDifferences.to_vec(x::Integer)
+    Integer_from_vec(v) = x
+    return Int[x], Integer_from_vec
+end
+
 function ngradient(f, x...)
     fdm = central_fdm(5, 1)
     return FiniteDifferences.grad(fdm, f, x...)

GraphNeuralNetworks/Project.toml

@@ -6,10 +6,8 @@ version = "0.6.22"
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
 GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
-Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -18,41 +16,16 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
-[weakdeps]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-
 [compat]
-CUDA = "4, 5"
 ChainRulesCore = "1"
 Flux = "0.14"
-Functors = "0.4.1"
 GNNGraphs = "1.0"
 GNNlib = "0.2"
-Graphs = "1.12"
 LinearAlgebra = "1"
 MLUtils = "0.4"
 MacroTools = "0.5"
 NNlib = "0.9"
-Pkg = "1"
 Random = "1"
 Reexport = "1"
 Statistics = "1"
-TestItemRunner = "1.0.5"
 julia = "1.10"
-
-[extras]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
-InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
-MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["Test", "TestItemRunner", "Pkg", "MLDatasets", "Adapt", "DataFrames", "InlineStrings", "SparseArrays", "Graphs", "Zygote", "FiniteDifferences", "ChainRulesTestUtils"]

GraphNeuralNetworks/docs/src/dev.md

@@ -1,30 +1,62 @@
 # Developer Notes
 
-## Develop and Managing the Monorepo
 
-### Development Enviroment
+## Development Enviroment
 GraphNeuralNetworks.jl is package hosted in a monorepo that contains multiple packages. 
-The GraphNeuralNetworks.jl package depends on GNNGraphs.jl, also hosted in the same monorepo.
+The GraphNeuralNetworks.jl package depends on GNNGraphs.jl and GNNlib.jl, also hosted in the same monorepo.
+In order 
 
 ```julia
 pkg> activate .
 
 pkg> dev ./GNNGraphs
 ```
 
-### Add a New Layer 
+## Add a New Layer 
 
 To add a new graph convolutional layer and make it available in both the Flux-based frontend (GraphNeuralNetworks.jl) and the Lux-based frontend (GNNLux), you need to:
+
 1. Add the functional version to GNNlib
 2. Add the stateful version to GraphNeuralNetworks
 3. Add the stateless version to GNNLux
 4. Add the layer to the table in docs/api/conv.md
 
-### Versions and Tagging
+We suggest to start with implementing a self-contained Flux layer in GraphNeuralNetworks.jl, add the corresponding tests, and then when everything is working, move the implementation of the forward pass to GNNlib.jl. At this point, you can add the stateless version to GNNLux.jl.
+
+It could also be convenient to use the `@structdef` macro from [Autostruct.jl](https://github.com/CarloLucibello/AutoStructs.jl) to simultaneously generate the struct and the constructor for the layer.
+For example, the Flux implementation of [`MEGNetConv`](@ref) layer can be written as follows:
+
+```julia
+using Flux, GraphNeuralNetworks, AutoStructs
+
+@structdef function MEGNetConv(ch::Pair{Int, Int}; aggr = mean)
+    nin, nout = ch
+    ϕe = Chain(Dense(3nin, nout, relu),
+               Dense(nout, nout))
+
+    ϕv = Chain(Dense(nin + nout, nout, relu),
+               Dense(nout, nout))
+
+    return MEGNetConv(ϕe, ϕv, aggr)
+end
+
+Flux.@layer MEGNetConv
+
+function (l::MEGNetConv)(g::AbstractGraph, x::AbstractMatrix, e::AbstractMatrix)
+    ē = apply_edges(g, xi = x, xj = x, e = e) do xi, xj, e
+        l.ϕe(vcat(xi, xj, e))
+    end
+    xᵉ = aggregate_neighbors(g, l.aggr, ē)
+    x̄ = l.ϕv(vcat(x, xᵉ))
+    return x̄, ē
+end
+```
+
+## Versions and Tagging
 Each PR should update the version number in the Porject.toml file of each involved package if needed by semnatic versioning. For instance, when adding new features GNNGraphs could move from "1.17.5" to "1.18.0-DEV". The "DEV" will be removed when the package is tagged and released. Pay also attention to updating
 the compat bounds, e.g. GraphNeuralNetworks might require a newer version of GNNGraphs.
 
-### Generate Documentation Locally
+## Generate Documentation Locally
 For generating the documentation locally
 ```
 cd docs

GraphNeuralNetworks/src/GraphNeuralNetworks.jl

@@ -6,11 +6,9 @@ using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, batch
 using MacroTools: @forward
 using NNlib
-using NNlib: scatter, gather
 using ChainRulesCore
-using Reexport
+using Reexport: @reexport
 using MLUtils: zeros_like
-using Graphs: Graphs
 
 using GNNGraphs:  COO_T, ADJMAT_T, SPARSE_T,
                   check_num_nodes, check_num_edges,

GraphNeuralNetworks/test/Project.toml

@@ -0,0 +1,15 @@
+[deps]
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+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"

GraphNeuralNetworks/test/examples/node_classification_cora.jl

@@ -108,7 +108,8 @@ end # module
     TrainingExampleModule.train_many()
 end
 
-@testitem "training example GPU" setup=[TrainingExampleModule] tags=[:gpu] begin
+@testitem "training example GPU" setup=[TestModule, TrainingExampleModule] tags=[:gpu] begin
+    using .TestModule # for loading gpu packages
     using .TrainingExampleModule
     TrainingExampleModule.train_many(use_gpu = true)
 end

GraphNeuralNetworks/test/layers/conv.jl

@@ -101,6 +101,7 @@ end
     k = 2
     l = ChebConv(D_IN => D_OUT, k)
     for g in TEST_GRAPHS
+        has_isolated_nodes(g) && continue
         g.graph isa AbstractSparseMatrix && continue
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
         test_gradients(l, g, g.x, rtol = RTOL_LOW, test_gpu = true, compare_finite_diff = false)
@@ -377,6 +378,7 @@ end
     l = CGConv((D_IN, edim) => D_OUT, tanh, residual = false, bias = true)
     for g in TEST_GRAPHS
         g.graph isa AbstractSparseMatrix && continue
+        g = GNNGraph(g, edata = rand(Float32, edim, g.num_edges))
         @test size(l(g, g.x, g.e)) == (D_OUT, g.num_nodes)
         test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
     end   
@@ -432,6 +434,7 @@ end
     l = MEGNetConv(D_IN => D_OUT, aggr = +)
     for g in TEST_GRAPHS
         g.graph isa AbstractSparseMatrix && continue
+        g = GNNGraph(g, edata = rand(Float32, D_IN, g.num_edges))
         y = l(g, g.x, g.e)
         @test size(y[1]) == (D_OUT, g.num_nodes)
         @test size(y[2]) == (D_OUT, g.num_edges)
@@ -462,6 +465,7 @@ end
     l = GMMConv((D_IN, ein_channel) => D_OUT, K = K)
     for g in TEST_GRAPHS
         g.graph isa AbstractSparseMatrix && continue
+        g = GNNGraph(g, edata = rand(Float32, ein_channel, g.num_edges))
         y = l(g, g.x, g.e)
         test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
     end   
@@ -585,6 +589,7 @@ end
                         bias_qkv = true)
     for g in TEST_GRAPHS
         g.graph isa AbstractSparseMatrix && continue
+        g = GNNGraph(g, edata = rand(Float32, ein, g.num_edges))
         @test size(l(g, g.x, g.e)) == (D_IN * heads, g.num_nodes)
         test_gradients(l, g, g.x, g.e, rtol = RTOL_LOW, test_gpu = true, compare_finite_diff = false)
     end

GraphNeuralNetworks/test/runtests.jl

@@ -4,12 +4,16 @@ using TestItemRunner
 ## for how to run the tests within VS Code.
 ## See test_module.jl for the test infrastructure.
 
-## Uncomment below to change the default test settings
+## 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

GraphNeuralNetworks/test/test_module.jl

@@ -3,9 +3,11 @@
 using GraphNeuralNetworks
 using Test
 using Statistics, Random
-using Flux, Functors
+using Flux
+using Functors: fmapstructure_with_path
 using Graphs
-using ChainRulesTestUtils, FiniteDifferences, Zygote, Adapt
+using ChainRulesTestUtils, FiniteDifferences
+using Zygote
 using SparseArrays
 using Pkg
 
@@ -16,22 +18,23 @@ using Pkg
 # ENV["GNN_TEST_Metal"] = "true"
 
 if get(ENV, "GNN_TEST_CUDA", "false") == "true"
-    Pkg.add(["CUDA", "cuDNN"])
+    # Pkg.add(["CUDA", "cuDNN"])
     using CUDA
     CUDA.allowscalar(false)
 end
 if get(ENV, "GNN_TEST_AMDGPU", "false") == "true"
-    Pkg.add("AMDGPU")
+    # Pkg.add("AMDGPU")
     using AMDGPU
     AMDGPU.allowscalar(false)
 end
 if get(ENV, "GNN_TEST_Metal", "false") == "true"
-    Pkg.add("Metal")
+    # Pkg.add("Metal")
     using Metal
     Metal.allowscalar(false)
 end
 
-# from Bse
+
+# from Base
 export mean, randn, SparseArrays, AbstractSparseMatrix
 
 # from other packages