Refactor CUDA message-passing

Author: yuehhua

src/cuda/msgpass.jl

@@ -9,16 +9,13 @@ end
 end
 
 @inline function update_batch_edge(mp::T, adj, E::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
+    n = size(adj, 1)
     edge_idx = edge_index_table(adj)
-    E_ = Vector[]
-    for (i, js) = enumerate(adj)
-        for j = js
-            k = edge_idx[(i,j)]
-            m = message(mp, get_feature(X, i), get_feature(X, j), get_feature(E, k))
-            push!(E_, m)
-        end
-    end
-    hcat(E_...)
+    hcat([apply_batch_message(mp, i, adj[i], edge_idx, E, X) for i in 1:n]...)
+end
+
+@inline function apply_batch_message(mp::T, i, js, edge_idx, E::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
+    hcat([message(mp, get_feature(X, i), get_feature(X, j), get_feature(E, edge_idx[(i,j)])) for j = js]...)
 end
 
 @inline function update_batch_vertex(mp::T, M::AbstractMatrix, X::CuMatrix) where {T<:MessagePassing}
@@ -32,10 +29,6 @@ end
 end
 
 @inline function update_batch_vertex(mp::T, M::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
-    X_ = Vector[]
-    for i = 1:size(X,2)
-        x = update(mp, get_feature(M, i), get_feature(X, i))
-        push!(X_, x)
-    end
+    X_ = [update(mp, get_feature(M, i), get_feature(X, i)) for i = 1:size(X,2)]
     hcat(X_...)
 end

test/cuda/msgpass.jl

@@ -1,6 +1,7 @@
 in_channel = 10
 out_channel = 5
 N = 6
+T = Float32
 adj = [0. 1. 0. 0. 0. 0.;
        1. 0. 0. 1. 1. 1.;
        0. 0. 0. 0. 0. 1.;
@@ -18,8 +19,9 @@ NewCudaLayer(m, n) = NewCudaLayer(randn(m,n))
 GeometricFlux.message(n::NewCudaLayer, x_i, x_j, e_ij) = n.weight * x_j
 GeometricFlux.update(::NewCudaLayer, m, x) = m
 
-X = rand(Float32, in_channel, N) |> gpu
+X = rand(T, in_channel, N) |> gpu
 fg = FeaturedGraph(adj, X)
+fg.ef = Fill(zero(T), 0, 2num_E)
 l = NewCudaLayer(out_channel, in_channel) |> gpu
 
 @testset "cuda/msgpass" begin