Add binarize function for CuSparseMatrixCSC and update its usage

Author: dferre97

GNNGraphs/ext/GNNGraphsCUDAExt.jl

@@ -4,6 +4,7 @@ using CUDA
 using Random, Statistics, LinearAlgebra
 using GNNGraphs
 using GNNGraphs: COO_T, ADJMAT_T, SPARSE_T 
+using SparseArrays
 
 const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}
 
@@ -19,6 +20,18 @@ GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz
 # Utils
 
 GNNGraphs.iscuarray(x::AnyCuArray) = true
+function GNNGraphs.binarize(Mat::CUSPARSE.CuSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+    @debug "Binarizing CuSparseMatrixCSC with type $(Tv) and index type $(Ti)"
+    bin_vals = fill!(similar(nonzeros(Mat)), one(Tv))
+    # Binarize a CuSparseMatrixCSC by setting all nonzero values to one(Tv)
+    return CUSPARSE.CuSparseMatrixCSC(Mat.colPtr, rowvals(Mat), bin_vals, size(Mat))
+end
+function GNNGraphs.binarize(Mat::CUSPARSE.CuSparseMatrixCSC, T::DataType)
+    @debug "Binarizing CuSparseMatrixCSC with type $(T)"
+    bin_vals = fill!(similar(nonzeros(Mat)), one(T))
+    # Binarize a CuSparseMatrixCSC by setting all nonzero values to one(T)
+    return CUSPARSE.CuSparseMatrixCSC(Mat.colPtr, rowvals(Mat), bin_vals, size(Mat))
+end
 
 
 function sort_edge_index(u::AnyCuArray, v::AnyCuArray)

GNNGraphs/src/query.jl

@@ -235,7 +235,7 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType = eltype(g
     @assert dir ∈ [:in, :out]
     A = g.graph
     if !weighted
-        A = binarize(A)
+        A = binarize(A, T)
     end
     A = T != eltype(A) ? T.(A) : A
     return dir == :out ? A : A'
@@ -377,7 +377,7 @@ end
 
 function _degree(A::AbstractMatrix, T::Type, dir::Symbol, edge_weight::Bool, num_nodes::Int)
     if edge_weight === false
-        A = binarize(A)
+        A = binarize(A, T)
     end
     A = eltype(A) != T ? T.(A) : A
     return dir == :out ? vec(sum(A, dims = 2)) :

GNNGraphs/src/utils.jl

@@ -296,6 +296,8 @@ function _rand_edges(rng, (n1, n2), m)
 end
 
 binarize(x) = map(>(0), x)
+# here just to allow CUDA extension to overload this function with correct type casting
+binarize(x, T::DataType) = binarize(x)
 
 CRC.@non_differentiable binarize(x...)
 CRC.@non_differentiable edge_encoding(x...)

GNNlib/ext/GNNlibCUDAExt.jl

@@ -7,14 +7,6 @@ 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::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

GNNlib/src/msgpass.jl

@@ -213,7 +213,8 @@ end
 ## COPY_XJ 
 
 function propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
-    A = adjacency_matrix(g, weighted = false)
+    @debug "copy_xj: propagating with type $(typeof(xj))"
+    A = adjacency_matrix(g, eltype(xj); weighted = false)
     return xj * A
 end