Add Mooncake rule for propagate e_mul_xj fast path (#678)

Author: Parvm1102

GNNlib/ext/GNNlibMooncakeExt.jl

@@ -1,17 +1,14 @@
 module GNNlibMooncakeExt
 
-using GNNlib: GNNlib, propagate, copy_xj
-using GNNGraphs: GNNGraph, adjacency_matrix
+using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj
+using GNNGraphs: GNNGraph, adjacency_matrix, edge_index, set_edge_weight
 using LinearAlgebra: adjoint
 using Base: IEEEFloat
 import Mooncake
 using Mooncake: CoDual, DefaultCtx, NoRData, @is_primitive
 
-# Mooncake reverse rule for the sparse message-passing fast path
-#     propagate(copy_xj, g, +, xi, xj, e)  ==  xj * adjacency_matrix(g)
-# `A` is constant w.r.t. the inputs, so the pullback is just `dxj = dy * A'`.
-# Without it Mooncake differentiates the generic sparse matmul, which is far
-# slower than Zygote.
+# Reverse rule for the fast path `propagate(copy_xj, g, +, xj) == xj * A`.
+# `A` is constant w.r.t. the inputs, so `dxj = dy * A'`.
 
 @is_primitive DefaultCtx Tuple{
     typeof(propagate),
@@ -47,4 +44,45 @@ function Mooncake.rrule!!(
     return res, propagate_copy_xj_add_pullback!!
 end
 
+# Reverse rule for the weighted fast path `propagate(e_mul_xj, g, +, xj, e) == xj * A(e)`.
+# `e` enters `A`, so we also return `de_k = Σ_f xj[f, s_k] * dy[f, t_k]`.
+
+@is_primitive DefaultCtx Tuple{
+    typeof(propagate),
+    typeof(e_mul_xj),
+    GNNGraph,
+    typeof(+),
+    Nothing,
+    AbstractMatrix{P},
+    AbstractVector{P},
+} where {P <: IEEEFloat}
+
+function Mooncake.rrule!!(
+    ::CoDual{typeof(propagate)},
+    ::CoDual{typeof(e_mul_xj)},
+    g::CoDual{<:GNNGraph},
+    ::CoDual{typeof(+)},
+    ::CoDual{Nothing},
+    xj::CoDual{<:AbstractMatrix{P}},
+    e::CoDual{<:AbstractVector{P}},
+) where {P <: IEEEFloat}
+    pg = Mooncake.primal(g)
+    pxj = Mooncake.primal(xj)
+    pe = Mooncake.primal(e)
+    s, t = edge_index(pg)
+    A = adjacency_matrix(set_edge_weight(pg, pe), P; weighted = true)
+    y = pxj * A
+    res = Mooncake.zero_fcodual(y)
+    function propagate_e_mul_xj_add_pullback!!(::NoRData)
+        dy = Mooncake.tangent(res)
+        dxj = Mooncake.tangent(xj)
+        de = Mooncake.tangent(e)
+        dxj .+= dy * adjoint(A)
+        de .+= vec(sum(view(pxj, :, s) .* view(dy, :, t); dims = 1))
+        return NoRData(), NoRData(), Mooncake.zero_rdata(pg),
+               NoRData(), NoRData(), NoRData(), NoRData()
+    end
+    return res, propagate_e_mul_xj_add_pullback!!
+end
+
 end # module

GNNlib/test/msgpass.jl

@@ -151,6 +151,16 @@ end
         end
     end
 
+    @testset "e_mul_xj + (scalar edge weight)" begin
+        # Scalar per-edge weights hit the spmm fast path that GNNlibMooncakeExt
+        # has a dedicated Mooncake rule for (matrix `e` above stays generic).
+        for g in TEST_GRAPHS
+            e = rand(Float32, g.num_edges)
+            f(g, x, e) = propagate(e_mul_xj, g, +; xj = x, e)
+            test_gradients(f, g, g.x, e; test_grad_f=false, test_mooncake=TEST_MOONCAKE)
+        end
+    end
+
     @testset "w_mul_xj +" begin
         for g in TEST_GRAPHS
             w = rand(Float32, g.num_edges)