test with sparse cuda arrays

Author: CarloLucibello

src/GNNGraphs/GNNGraphs.jl

@@ -3,6 +3,7 @@ module GNNGraphs
 using SparseArrays
 using Functors: @functor
 using CUDA 
+using CUDA.CUSPARSE
 import Graphs
 using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree
 import Flux

src/GNNGraphs/convert.jl

@@ -137,11 +137,39 @@ function to_sparse(coo::COO_T, T::DataType=Int; dir=:out, num_nodes=nothing)
     s, t, eweight  = coo
     eweight = isnothing(eweight) ? fill!(similar(s, T), 1) : eweight
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
-    A = sparse(s, t, eweight, num_nodes, num_nodes)
+    A = _sparse(s, t, eweight, num_nodes, num_nodes)
     num_edges = length(s)
     return A, num_nodes, num_edges
 end
 
+_sparse(s, t, eweight, n, m) = sparse(s, t, eweight, n, m)
+
+function _sparse(I::CuVector, J::CuVector, V::CuVector, m, n)
+    spcoo = CuSparseMatrixCOO{Float32, Int32}(Int32.(I), Int32.(J), Float32.(V), (m, n))
+    return CuSparseMatrixCSR(spcoo)
+end
+
+# function _sparse(I::CuVector, J::CuVector, V::CuVector, m, n; fmt=:csr)
+#     # Tv = Int32
+#     spcoo = CuSparseMatrixCOO{Float32, Int32}(Int32.(I), Int32.(J), Float32.(V), (m, n))
+#     if fmt == :csc
+#         return CuSparseMatrixCSC(spcoo)
+#     elseif fmt == :csr
+#         return CuSparseMatrixCSR(spcoo)
+#     elseif fmt == :coo
+#         return spcoo
+#     else
+#         error("Format :$fmt not available, use :csc, :csr, or :coo.")
+#     end
+# end
+
+
+# Workaround for https://github.com/JuliaGPU/CUDA.jl/issues/1113#issuecomment-955759875
+function Base.:*(A::CuMatrix, B::CuSparseMatrixCSR)
+    @assert size(A, 2) == size(B, 1)
+    return CuMatrix((B' * A')')
+end
+
 
 @non_differentiable to_coo(x...)
 @non_differentiable to_dense(x...)

src/GNNGraphs/gnngraph.jl

@@ -192,8 +192,8 @@ function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata, grap
             ndata, edata, gdata) 
 end
 
-function Base.show(io::IO, g::GNNGraph)
-    println(io, "GNNGraph:
+function Base.show(io::IO, g::GNNGraph{T}) where T
+    println(io, "GNNGraph{$T}:
     num_nodes = $(g.num_nodes)
     num_edges = $(g.num_edges)
     num_graphs = $(g.num_graphs)")

src/GNNGraphs/query.jl

@@ -74,12 +74,7 @@ function adjacency_list(g::GNNGraph; dir=:out)
 end
 
 function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=Int; dir=:out)
-    if g.graph[1] isa CuVector
-        # TODO revisit after https://github.com/JuliaGPU/CUDA.jl/pull/1152
-        A, n, m = to_dense(g.graph, T, num_nodes=g.num_nodes)
-    else
-        A, n, m = to_sparse(g.graph, T, num_nodes=g.num_nodes)
-    end
+    A, n, m = to_sparse(g.graph, T, num_nodes=g.num_nodes)
     @assert size(A) == (n, n)
     return dir == :out ? A : A'
 end

src/msgpass.jl

@@ -152,10 +152,10 @@ function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(+), xi, xj::AbstractM
     return xj * A
 end
 
-## avoid the fast path on gpu until we have better cuda support
-function propagate(::typeof(copyxj), g::GNNGraph{<:Union{COO_T,SPARSE_T}}, ::typeof(+), xi, xj::AnyCuMatrix, e)
-    propagate((xi,xj,e)->copyxj(xi,xj,e), g, +, xi, xj, e)
-end
+# ## avoid the fast path on gpu until we have better cuda support
+# function propagate(::typeof(copyxj), g::GNNGraph{<:Union{COO_T,SPARSE_T}}, ::typeof(+), xi, xj::AnyCuMatrix, e)
+#     propagate((xi,xj,e) -> copyxj(xi,xj,e), g, +, xi, xj, e)
+# end
 
 # function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
 #     A = adjacency_matrix(g)