Mooncake on CUDA issue tracker #702

I ran Mooncake on GPU where the test sites for CPU were enabled (Mooncake 0.5.41, CUDA.jl 6.2.1, Flux 0.16.10, Julia 1.12.5).
CPU is fully green, but many fail on GPU. Here are my findings:

Already works on GPU (COO storage): GraphConv, SAGEConv, GINConv, EdgeConv, NNConv, ResGatedGraphConv, CGConv, MEGNetConv, GlobalPool, GlobalAttentionPool, GCNConv with add_self_loops=false, and all the propagate fast paths (copy_xj with +/mean, e_mul_xj, w_mul_xj). The rules in GNNlibMooncakeExt.

Zygote uses ChainRules for non-differentiable functions, Mooncake doesn't, so it the hits GPU code it can't differentiate:

• add_self_loops for COO does nodes = convert(typeof(s), [1:n;]), i.e. builds a CPU vector and copies it to the GPU. That lands on a CPU→GPU unsafe_copyto! containing try/catch, which Mooncake can't trace (its cross-device copy rule only covers float arrays, and here the destination is an integer index array). It affects all layers with self loops: GCNConv (default), GATConv, GATv2Conv, TransformerConv, AGNNConv, SGConv, TAGConv, TGCN.
• On :dense graphs, edge_index/to_coo call _findnz_idx, i.e. findall on a CuArray{Bool} (no Mooncake rule), and right after that v = A[nz] indexes with a vector of CartesianIndex (also no rule). This kills every layer on dense storage.

We can add GNNGraphs/ext/GNNGraphsMooncakeExt.jl (Mooncake as weakdep) giving Mooncake the same semantics the ChainRules markers give Zygote:

Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(add_self_loops), GNNGraph}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(GNNGraphs._findnz_idx), Any}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(to_coo), GNNGraphs.ADJMAT_T}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), NamedTuple, typeof(to_coo), GNNGraphs.ADJMAT_T}

The Core.kwcall rule is required because all in-repo call sites call to_coo with keyword arguments. I validated all of this by defining the rules via type piracy in a test script: with them in place, GCNConv (coo and dense), EdgeConv/dense, ResGatedGraphConv/dense, SGConv, AGNNConv, GMMConv and GatedGraphConv pass on GPU with gradients matching Zygote to ~1e-7.

Two caveats worth recording in the ext:

• add_self_loops: same caveat as the existing @non_differentiable in transform.jl (the graph carries feature arrays; fine in practice since layers take x explicitly).
• to_coo: Zygote only skips _findnz_idx and still differentiates A[nz], so it delivers dA for weighted dense adjacencies. A zero-derivative to_coo drops that under Mooncake. If Mooncake gains a CartesianIndex getindex rule upstream (item 4 below), we can drop the to_coo rules and keep full parity.

Upstream Mooncake.jl gaps found (I think I can open some issues on Mooncake regarding this):

1. Correctness bug: broadcasts over non-contiguous SubArrays of CuArrays silently drop gradients. This is why GatedGraphConv returns wrong gradients with no error on GPU. Flux's GRUCell slices its gates with chunk(Wi*x, 3, dims=1).
2. No kwcall rule for sum(x::CuArray; dims). It breaks GMMConv and other attention score computation.
3. Cross-device unsafe_copyto! rule covers float destinations only, not integer.
4. No getindex(::CuArray, ::Vector{CartesianIndex}) rule.
5. CuMatrix + CuMatrix → cuBLAS.geam! has no rule (CuRef tangent type missing).
6. No repeat(::CuArray, ...) rule (hits TGCNCell's repeat(h, 1, g.num_nodes)).
7. task_local_storage() IdDict fdata mismatch, reached through CUDA's default RNG by dropout.

I will work on these once it gets green flag.