fused e_mul_xj and weighted option for adjacency_matrix (#107)

* weighted option for adjacency_matrix

* fused e_mul_xj

Author: CarloLucibello

Project.toml

@@ -1,7 +1,7 @@
 name = "GraphNeuralNetworks"
 uuid = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 authors = ["Carlo Lucibello and contributors"]
-version = "0.3.9"
+version = "0.3.10"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

perf/bench_gnn.jl

@@ -0,0 +1,73 @@
+using SparseArrays
+using GraphNeuralNetworks
+using BenchmarkTools
+import Random: seed!
+using LinearAlgebra
+
+n = 1024
+seed!(0)
+A = sprand(n, n, 0.01)
+b = rand(1, n)
+B = rand(100, n)
+
+g = GNNGraph(
+    A,
+    ndata=(; b=b, B=B),
+    edata=(; A=reshape(A.nzval, 1, :)),
+    graph_type=:coo  # changing to :sparse has little effect on performance
+)
+
+function spmv(g)
+    propagate(
+        (xi, xj, e) -> e .* xj ,  # same as e_mul_xj
+        g, +; xj=g.ndata.b, e=g.edata.A
+        )
+end
+
+function spmm1(g)
+    propagate(
+        (xi, xj, e) -> e .* xj ,  # same as e_mul_xj
+        g, +; xj=g.ndata.B, e=g.edata.A
+        )
+end
+function spmm2(g)
+    propagate(
+        e_mul_xj,
+        g, +; xj=g.ndata.B, e=vec(g.edata.A)
+        )
+end
+
+# @assert isequal(spmv(g),  b * A)  # true
+# @btime spmv(g)  # ~5 ms
+# @btime b * A  # ~32 us
+
+@assert isequal(spmm1(g), B * A)  # true
+@assert isequal(spmm2(g), B * A)  # true
+@btime spmm1(g)  # ~9 ms
+@btime spmm2(g)  # ~9 ms
+@btime B * A  # ~400 us
+
+
+function spmm_copyxj_fused(g)
+    propagate(
+        copy_xj,
+        g, +; xj=g.ndata.B
+        )
+end
+
+function spmm_copyxj_unfused(g)
+    propagate(
+        (xi, xj, e) -> xj,
+        g, +; xj=g.ndata.B
+        )
+end
+
+Adj = map(x -> x > 0 ? 1 : 0, A)
+@assert spmm_copyxj_unfused(g) ≈ B * Adj
+@assert spmm_copyxj_fused(g) ≈ B * Adj # bug fixed in #107
+
+@btime spmm_copyxj_fused(g)  # 268.614 μs (22 allocations: 1.13 MiB)
+@btime spmm_copyxj_unfused(g)  # 4.263 ms (52855 allocations: 12.23 MiB)
+@btime B * Adj  # 196.135 μs (2 allocations: 800.05 KiB)
+
+println()

src/GNNGraphs/convert.jl

@@ -1,6 +1,6 @@
 ### CONVERT_TO_COO REPRESENTATION ########
 
-function to_coo(coo::COO_T; dir=:out, num_nodes=nothing)
+function to_coo(coo::COO_T; dir=:out, num_nodes=nothing, weighted=true)
     s, t, val = coo   
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     @assert isnothing(val) || length(val) == length(s)
@@ -10,21 +10,26 @@ function to_coo(coo::COO_T; dir=:out, num_nodes=nothing)
         @assert max(maximum(s), maximum(t)) <= num_nodes 
     end
     num_edges = length(s)
+    if !weighted
+        coo = (s, t, nothing)
+    end
     return coo, num_nodes, num_edges
 end
 
-function to_coo(A::SPARSE_T; dir=:out, num_nodes=nothing)
+function to_coo(A::SPARSE_T; dir=:out, num_nodes=nothing, weighted=true)
     s, t, v = findnz(A)
     if dir == :in
         s, t = t, s
     end
     num_nodes = isnothing(num_nodes) ? size(A, 1) : num_nodes 
     num_edges = length(s)
-
+    if !weighted
+        v = nothing
+    end
     return (s, t, v), num_nodes, num_edges
 end
 
-function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing)
+function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing, weighted=true)
     nz = findall(!=(0), A) # vec of cartesian indexes
     s, t = ntuple(i -> map(t->t[i], nz), 2)
     v = A[nz] 
@@ -33,10 +38,13 @@ function to_coo(A::ADJMAT_T; dir=:out, num_nodes=nothing)
     end
     num_nodes = isnothing(num_nodes) ? size(A, 1) : num_nodes 
     num_edges = length(s)
+    if !weighted
+        v = nothing
+    end
     return (s, t, v), num_nodes, num_edges
 end
 
-function to_coo(adj_list::ADJLIST_T; dir=:out, num_nodes=nothing)
+function to_coo(adj_list::ADJLIST_T; dir=:out, num_nodes=nothing, weighted=true)
     @assert dir ∈ [:out, :in]
     num_nodes = length(adj_list)
     num_edges = sum(length.(adj_list))
@@ -64,7 +72,7 @@ end
 
 to_dense(A::AbstractSparseMatrix, x...; kws...) = to_dense(collect(A), x...; kws...)
 
-function to_dense(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_dense(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     @assert dir ∈ [:out, :in]
     T = T === nothing ? eltype(A) : T
     num_nodes = size(A, 1)
@@ -77,10 +85,13 @@ function to_dense(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     if T != eltype(A)
         A = T.(A)
     end
+    if !weighted
+        A = map(x -> x > 0 ? T(1) : T(0), A)
+    end
     return A, num_nodes, num_edges
 end
 
-function to_dense(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_dense(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     @assert dir ∈ [:out, :in]
     num_nodes = length(adj_list)
     num_edges = sum(length.(adj_list))
@@ -99,13 +110,16 @@ function to_dense(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
     A, num_nodes, num_edges
 end
 
-function to_dense(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_dense(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     # `dir` will be ignored since the input `coo` is always in source -> target format.
     # The output will always be a adjmat in :out format (e.g. A[i,j] denotes from i to j)
     s, t, val = coo
     n = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes
     val = isnothing(val) ? eltype(s)(1) : val
     T = T === nothing ? eltype(val) : T
+    if !weighted
+        val = T(1)
+    end
     A = fill!(similar(s, T, (n, n)), 0)
     v = vec(A)
     idxs = s .+ n .* (t .- 1) 
@@ -116,7 +130,7 @@ end
 
 ### SPARSE #############
 
-function to_sparse(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_sparse(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     @assert dir ∈ [:out, :in]
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
@@ -131,18 +145,25 @@ function to_sparse(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     if !(A isa AbstractSparseMatrix)
         A = sparse(A)
     end
+    if !weighted
+        A = map(x -> x > 0 ? T(1) : T(0), A)
+    end
     return A, num_nodes, num_edges
 end
 
-function to_sparse(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_sparse(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     coo, num_nodes, num_edges = to_coo(adj_list; dir, num_nodes)
     return to_sparse(coo; num_nodes)
 end
 
-function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)
+function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing, weighted=true)
     s, t, eweight  = coo
     T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T
-    eweight = eweight === nothing ? fill!(similar(s, T), 1) : eweight
+    
+    if eweight === nothing || !weighted
+        eweight = fill!(similar(s, T), 1)
+    end
+
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)
     num_edges = nnz(A)

src/GNNGraphs/query.jl

@@ -132,20 +132,35 @@ end
 adjacency_list(g::GNNGraph; dir=:out) = adjacency_list(g, 1:g.num_nodes; dir)
 
 
-function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=eltype(g); dir=:out)
+"""
+    adjacency_matrix(g::GNNGraph, T=eltype(g); dir=:out, weighted=true)
+
+Return the adjacency matrix `A` for the graph `g`. 
+
+If `dir=:out`, `A[i,j] > 0` denotes the presence of an edge from node `i` to node `j`.
+If `dir=:in` instead, `A[i,j] > 0` denotes the presence of an edge from node `j` to node `i`.
+
+User may specify the eltype `T` of the returned matrix. 
+
+If `weighted=true`, the `A` will contain the edge weigths if any, otherwise the elements of `A` will be either 0 or 1.
+"""
+function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=eltype(g); dir=:out, weighted=true)
     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)
+        A, n, m = to_dense(g.graph, T; num_nodes=g.num_nodes, weighted)
     else
-        A, n, m = to_sparse(g.graph, T, num_nodes=g.num_nodes)
+        A, n, m = to_sparse(g.graph, T; num_nodes=g.num_nodes, weighted)
     end
     @assert size(A) == (n, n)
     return dir == :out ? A : A'
 end
 
-function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g); dir=:out)
+function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g); dir=:out, weighted=true)
     @assert dir ∈ [:in, :out]
     A = g.graph
+    if !weighted
+        A = map(>(0), A)
+    end
     A = T != eltype(A) ? T.(A) : A
     return dir == :out ? A : A'
 end

src/layers/conv.jl

@@ -25,8 +25,8 @@ and optionally an edge weight vector.
 - `bias`: Add learnable bias. Default `true`.
 - `init`: Weights' initializer. Default `glorot_uniform`.
 - `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
-- `use_edge_weight`. If `true`, consider the edge weights in the input graph (if available).
-                 If `add_self_loops=true` the new weights will be set to 1. Default `false`.
+- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
+                     If `add_self_loops=true` the new weights will be set to 1. Default `false`.
 
 # Examples
 
@@ -81,17 +81,13 @@ function (l::GCNConv)(g::GNNGraph{<:COO_T}, x::AbstractMatrix)
     return l(g, x, edge_weight)
 end
 
-function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix)
-    edge_weight = nothing
-    return l(g, x, edge_weight)
-end
-
-function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}, edge_weight::EW) where 
+function (l::GCNConv)(g::GNNGraph{<:COO_T}, x::AbstractMatrix{T}, edge_weight::EW) where 
     {T, EW<:Union{Nothing,AbstractVector}}
 
     if l.add_self_loops
         g = add_self_loops(g)
         if edge_weight !== nothing
+            # TODO for ADJMAT_T the new edges are not generally at the end
             edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
             @assert length(edge_weight) == g.num_edges
         end
@@ -116,6 +112,33 @@ function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}, edge_weight::EW) where
     return l.σ.(x .+ l.bias)
 end
 
+# TODO merge the ADJMAT_T and COO_T methods
+# The main problem is handling the weighted case for both.
+function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix{T}) where T
+    if l.add_self_loops
+        g = add_self_loops(g)
+    end
+    Dout, Din = size(l.weight)
+    if Dout < Din
+        # multiply before convolution if it is more convenient, otherwise multiply after
+        x = l.weight * x
+    end
+    d = degree(g, T; dir=:in, edge_weight=l.use_edge_weight)
+    c = 1 ./ sqrt.(d)
+    x = x .* c'
+    A = adjacency_matrix(g, weighted=l.use_edge_weight)
+    x = x * A
+    x = x .* c'
+    if Dout >= Din
+        x = l.weight * x
+    end
+    return l.σ.(x .+ l.bias)
+end
+
+function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix, edge_weight::AbstractVector)
+    g = GNNGraph(edge_index(g)...; g.num_nodes)  # convert to COO
+    return l(g, x, edge_weight)
+end
 
 function Base.show(io::IO, l::GCNConv)
     out, in = size(l.weight)

src/msgpass.jl

@@ -179,17 +179,27 @@ function e_mul_xj(xi, xj::AbstractArray{Tj,Nj}, e::AbstractArray{Te,Ne}) where {
 end
 
 function propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
-    A = adjacency_matrix(g)
+    A = adjacency_matrix(g, weighted=false)
     return xj * A
 end
 
+# for weighted convolution
+function propagate(::typeof(e_mul_xj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e::AbstractVector)
+    s, t = edge_index(g)
+    g = GNNGraph((s, t, e); g.num_nodes)
+    A = adjacency_matrix(g, weighted=true)
+    return xj * A
+end
+
+
+
 ## 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)
     propagate((xi,xj,e)->copy_xj(xi,xj,e), g, +, xi, xj, e)
 end
 
 # function propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
-#     A = adjacency_matrix(g)
+#     A = adjacency_matrix(g, weigthed=false)
 #     D = compute_degree(A)
 #     return xj * A * D
 # end

test/GNNGraphs/query.jl

@@ -95,9 +95,15 @@
 
     @testset "adjacency_matrix" begin
         a = sprand(5, 5, 0.5)
+        abin = map(x -> x > 0 ? 1 : 0, a)
+        
         g = GNNGraph(a, graph_type=GRAPH_T)
         A = adjacency_matrix(g, Float32)
-        @test a ≈ A
+        @test A ≈ a
         @test eltype(A) == Float32
+    
+        Abin = adjacency_matrix(g, Float32, weighted=false)
+        @test Abin ≈ abin
+        @test eltype(Abin) == Float32    
     end
 end