[GNNlib] fix cuda ext (#465)

JuliaGraphs · Jul 28, 2024 · a49677e · a49677e · CarloLucibello · Jul 28, 2024
1 parent 80c672a
commit a49677e
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 19 deletions.
diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml
@@ -1,7 +1,7 @@
 name = "GNNlib"
 uuid = "a6a84749-d869-43f8-aacc-be26a1996e48"
 authors = ["Carlo Lucibello and contributors"]
-version = "0.2.0-DEV"
+version = "0.2.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

diff --git a/GNNlib/ext/GNNlibCUDAExt/msgpass.jl → GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt/msgpass.jl → GNNlib/ext/GNNlibCUDAExt.jl
@@ -1,31 +1,37 @@
+module GNNlibCUDAExt
+
+using CUDA
+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 propagate(::typeof(copy_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
-                   xi, xj::AnyCuMatrix, e)
+function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
+        xi, xj::AnyCuMatrix, 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 propagate(::typeof(e_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
-                   xi, xj::AnyCuMatrix, e::AbstractVector)
+function GNNlib.propagate(::typeof(e_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
+        xi, xj::AnyCuMatrix, 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 cuda support
-function propagate(::typeof(w_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
-                   xi, xj::AnyCuMatrix, e::Nothing)
+function GNNlib.propagate(::typeof(w_mul_xj), g::GNNGraph{<:Union{COO_T, SPARSE_T}}, ::typeof(+),
+        xi, xj::AnyCuMatrix, e::Nothing)
     propagate((xi, xj, e) -> w_mul_xj(xi, xj, e), g, +, xi, xj, e)
 end
 
-# function propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
+# function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
 #     A = adjacency_matrix(g, weighted=false)
 #     D = compute_degree(A)
 #     return xj * A * D
@@ -35,3 +41,5 @@ end
 # compute_degree(A) = Diagonal(1f0 ./ vec(sum(A; dims=2)))
 
 # Flux.Zygote.@nograd compute_degree
+
+end #module
diff --git a/GNNlib/ext/GNNlibCUDAExt/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt/GNNlibCUDAExt.jl