amdgpu extension

Author: CarloLucibello

GNNlib/Project.toml

@@ -14,12 +14,15 @@ 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]
 GNNlibCUDAExt = "CUDA"
+GNNlibAMDGPUExt = "AMDGPU"
 
 [compat]
+AMDGPU = "1"
 CUDA = "4, 5"
 ChainRulesCore = "1.24"
 DataStructures = "0.18"
@@ -30,12 +33,3 @@ 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"]

GNNlib/ext/GNNlibAMDGPUExt.jl

@@ -0,0 +1,34 @@
+module GNNlibAMDGPUExt
+
+using AMDGPU
+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