From 3ed702bf3696b22e1dd513a890e9581d7cf72303 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 27 Nov 2024 20:46:43 +0100
Subject: [PATCH] buildkite AMDGPU (#529)

---
 .buildkite/pipeline.yml                       | 26 +++++++++++++++++++
 .../test/examples/node_classification_cora.jl |  2 +-
 GraphNeuralNetworks/test/test_module.jl       | 25 +++++++++---------
 3 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index dfa94fc4e..4fd28dd07 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -23,3 +23,29 @@ steps:
       GNN_TEST_CUDA: "true"
       GNN_TEST_CPU: "false"
     timeout_in_minutes: 60
+
+  - label: "GNN AMDGPU"
+    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(["AMDGPU"])
+        Pkg.test("GraphNeuralNetworks")'
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    env:
+      GNN_TEST_AMDGPU: "true"
+      GNN_TEST_CPU: "false"
+    timeout_in_minutes: 60
diff --git a/GraphNeuralNetworks/test/examples/node_classification_cora.jl b/GraphNeuralNetworks/test/examples/node_classification_cora.jl
index 54241d50f..839c7ff80 100644
--- a/GraphNeuralNetworks/test/examples/node_classification_cora.jl
+++ b/GraphNeuralNetworks/test/examples/node_classification_cora.jl
@@ -40,7 +40,7 @@
         classes = dataset.metadata["classes"]
         g = mldataset2gnngraph(dataset) |> device
         X = g.ndata.features
-        y = onehotbatch(g.ndata.targets, classes)
+        y = onehotbatch(g.ndata.targets |> cpu, classes) |> device # https://github.com/FluxML/OneHotArrays.jl/issues/16
         train_mask = g.ndata.train_mask
         test_mask = g.ndata.test_mask
         ytrain = y[:, train_mask]
diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl
index a2cca1952..e7f9a0031 100644
--- a/GraphNeuralNetworks/test/test_module.jl
+++ b/GraphNeuralNetworks/test/test_module.jl
@@ -1,14 +1,5 @@
 @testmodule TestModule begin
 
-using GraphNeuralNetworks
-using Test
-using Statistics, Random
-using Flux
-using Functors: fmapstructure_with_path
-using Graphs
-using ChainRulesTestUtils, FiniteDifferences
-using Zygote
-using SparseArrays
 using Pkg
 
 ## Uncomment below to change the default test settings
@@ -18,21 +9,31 @@ 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
 
+using GraphNeuralNetworks
+using Test
+using Statistics, Random
+using Flux
+using Functors: fmapstructure_with_path
+using Graphs
+using ChainRulesTestUtils, FiniteDifferences
+using Zygote
+using SparseArrays
+
 
 # from Base
 export mean, randn, SparseArrays, AbstractSparseMatrix