Add custom _adjacency_matrix for propagate CUDA COO graphs

- Leave public adjacency_matrix interface uniform, always returning a sparse adjacency_matrix
- Implement custom _adjacency_matrix for propagate copy_xj for CUDA COO graphs, converting to dense when more efficient

Author: dferre97

GNNGraphs/ext/GNNGraphsCUDAExt.jl

@@ -14,20 +14,6 @@ const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{
 
 GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1))
 
-function Graphs.adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out,
-                                 weighted = true)
-    @debug "Using CUDA adjacency_matrix for GNNGraph"
-    if !g.is_coalesced
-        # Revisit after 
-        # https://github.com/JuliaGPU/CUDA.jl/issues/1113
-        A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesce, construction of sparse matrix is slow
-    else
-        A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted)
-    end
-    @assert size(A) == (n, n)
-    return dir == :out ? A : A'
-end
-
 # Transform
 
 GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz)
@@ -54,18 +40,24 @@ end
 # Convert
 
 function GNNGraphs.to_sparse(coo::CUDA_COO_T, T = nothing; dir = :out, num_nodes = nothing,
-                   weighted = true)
+                   weighted = true, is_coalesced = false)
     s, t, eweight = coo
-    @debug "Using CUDA to_sparse for COO"
+    @debug "Using CUDA to_sparse for COO with is_coalesced=$is_coalesced"
     T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T
 
     if eweight === nothing || !weighted
         eweight = fill!(similar(s, T), 1)
     end
 
     num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes
-    A = CUDA.CUSPARSE.CuSparseMatrixCOO{T,eltype(s)}(s, t, eweight, (num_nodes, num_nodes)) # create sparse matrix in COO format
 
+    # if coalesced build directly sparse coo matrix
+    if is_coalesced
+        A = CUDA.CUSPARSE.CuSparseMatrixCOO{T,eltype(s)}(s, t, eweight, (num_nodes, num_nodes)) 
+    else
+        A = sparse(s, t, eweight, num_nodes, num_nodes)
+    end
+
     num_edges::Int = nnz(A)
     if eltype(A) != T
         A = T.(A)

GNNlib/ext/GNNlibCUDAExt.jl

@@ -3,7 +3,9 @@ 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, adjacency_matrix
+using GNNGraphs: GNNGraph, COO_T, SPARSE_T
+
+const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{Nothing, AnyCuArray}}
 
 ###### PROPAGATE SPECIALIZATIONS ####################
 
@@ -13,7 +15,7 @@ using GNNGraphs: GNNGraph, COO_T, SPARSE_T, adjacency_matrix
 function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+),
         xi, xj::AnyCuMatrix, e)
     @debug "Using CUDA propagate for copy_xj"
-    A = adjacency_matrix(g, eltype(xj); weighted = false)
+    A = _adjacency_matrix(g, eltype(xj); weighted = false)
 
     return xj * A
 end
@@ -45,4 +47,20 @@ end
 
 # Flux.Zygote.@nograd compute_degree
 
+## CUSTOM ADJACENCY_MATRIX IMPLEMENTATION FOR CUDA COO GRAPHS, returning dense matrix when not coalesced, more efficient 
+
+function _adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out,
+                                 weighted = true)
+    @debug "Using CUDA _adjacency_matrix for COO GNNGraph"
+    if !g.is_coalesced
+        # Revisit after 
+        # https://github.com/JuliaGPU/CUDA.jl/issues/1113
+        A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesced, construction of sparse matrix is slow
+    else
+        A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted, is_coalesced = true)
+    end
+    @assert size(A) == (n, n)
+    return dir == :out ? A : A'
+end
+
 end #module