diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fa253314..5e412e143 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the packages adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). Entries link to the pull request that introduced them. -## GNNGraphs.jl — Unreleased (towards 1.5.2) +## GNNGraphs.jl — Unreleased (towards 1.6.0) + +**Added** +- Added a `Mooncake` package extension, so that Mooncake can differentiate on CUDA through `add_self_loops` and through adjacency-matrix graphs. Float edge weights of adjacency-matrix graphs stay differentiable ([#704]). **Fixed** - Enzyme can now differentiate through the adjacency-matrix → COO graph conversion: a new internal keyword-free helper `_to_coo_graph` avoids the union-typed keyword handling of the `GNNGraph(g; graph_type)` constructor and of `to_coo`, which Enzyme's type analysis cannot compile. Together with upstream fixes in Enzyme ≥ 0.13.197 this removes the `EnzymeInternalError` crash on `:dense`/`:sparse` graphs ([#703]). @@ -230,4 +233,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#695]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/695 [#696]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/696 [#703]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/703 +[#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662 diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml index 9304c71d1..015e9a68c 100644 --- a/GNNGraphs/Project.toml +++ b/GNNGraphs/Project.toml @@ -1,6 +1,6 @@ name = "GNNGraphs" uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c" -version = "1.5.2-DEV" +version = "1.6.0-DEV" authors = ["Carlo Lucibello and contributors"] [workspace] @@ -23,10 +23,12 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622" [extensions] GNNGraphsCUDAExt = "CUDA" +GNNGraphsMooncakeExt = "Mooncake" GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs" [compat] @@ -38,6 +40,7 @@ KrylovKit = "0.8, 0.9, 0.10" LinearAlgebra = "1" MLDataDevices = "1.0" MLUtils = "0.4" +Mooncake = "0.5.24" NNlib = "0.9" NearestNeighbors = "0.4" Random = "1" diff --git a/GNNGraphs/ext/GNNGraphsMooncakeExt.jl b/GNNGraphs/ext/GNNGraphsMooncakeExt.jl new file mode 100644 index 000000000..73ae817a8 --- /dev/null +++ b/GNNGraphs/ext/GNNGraphsMooncakeExt.jl @@ -0,0 +1,16 @@ +module GNNGraphsMooncakeExt + +using GNNGraphs +import Mooncake + +# Without these rules Mooncake traces into graph-structure code and fails on CUDA: +# host→device index copies in `add_self_loops`, `findall` on `CuArray` in +# `_findnz_idx`, bounds checks in `_edge_values` (Mooncake's CUDA `getindex` rule +# covers only float/complex arrays). No gradient is lost: all three yield integers, +# whose Mooncake tangent is `NoTangent` anyway. Float edge values stay on the AD path. +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(GNNGraphs._edge_values), AbstractMatrix{<:Integer}, Any, Any} + +end # module diff --git a/GNNGraphs/src/convert.jl b/GNNGraphs/src/convert.jl index 577617e82..f5675ffff 100644 --- a/GNNGraphs/src/convert.jl +++ b/GNNGraphs/src/convert.jl @@ -80,9 +80,13 @@ end CRC.@non_differentiable _findnz_idx(A) +# Values of `A` at the extracted edge positions. Linear indexing keeps float edge +# weights differentiable on GPU, unlike `A[nz]` with CartesianIndex. +_edge_values(A::AbstractMatrix, s, t) = vec(A)[s .+ (t .- 1) .* size(A, 1)] + function to_coo(A::ADJMAT_T; dir = :out, num_nodes = nothing, weighted = true) - s, t, nz = _findnz_idx(A) - v = A[nz] + s, t, _ = _findnz_idx(A) + v = _edge_values(A, s, t) if dir == :in s, t = t, s end @@ -126,8 +130,8 @@ end function _to_coo_graph(g::GNNGraph{<:ADJMAT_T}) A = g.graph - s, t, nz = _findnz_idx(A) - v = A[nz] + s, t, _ = _findnz_idx(A) + v = _edge_values(A, s, t) return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs, g.graph_indicator, g.ndata, g.edata, g.gdata) end diff --git a/GNNGraphs/test/Project.toml b/GNNGraphs/test/Project.toml index 5c0294d13..ab94cbce6 100644 --- a/GNNGraphs/test/Project.toml +++ b/GNNGraphs/test/Project.toml @@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -26,5 +27,6 @@ GNNGraphs = {path = ".."} [compat] Enzyme = "0.13.197" +Mooncake = "0.5.24" NNlib = "0.9.38" StableRNGs = "1.0.4" diff --git a/GNNGraphs/test/ext/Mooncake.jl b/GNNGraphs/test/ext/Mooncake.jl new file mode 100644 index 000000000..ec4d06a73 --- /dev/null +++ b/GNNGraphs/test/ext/Mooncake.jl @@ -0,0 +1,79 @@ +@testitem "GNNGraphsMooncakeExt zero-derivative rules" setup=[GraphsTestModule] begin + using .GraphsTestModule + import Mooncake + import NNlib + + # Loading Mooncake triggers the extension. + @test Base.get_extension(GNNGraphs, :GNNGraphsMooncakeExt) !== nothing + + # Mooncake is only exercised on Julia >= 1.12, as in the other test suites. + if VERSION >= v"1.12" + function mooncake_gradient(f, x) + cache = Mooncake.prepare_gradient_cache(f, x) + _, grads = Mooncake.value_and_gradient!!(cache, f, x) + return grads[2] + end + + rng = MersenneTwister(17) + + # Gradients through `add_self_loops` (zero-derivative rule). + for graph_type in GRAPH_TYPES + g = rand_graph(rng, 6, 10; graph_type) + x = randn(rng, Float32, 3, g.num_nodes) + loss(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g))[1])) + gm = mooncake_gradient(loss, x) + @test gm≈gradient(loss, x)[1] rtol=1e-4 + @test gm≈ngradient(loss, x)[1] rtol=1e-4 + end + + # Gradients through `edge_index` on adjacency-matrix graphs. + for graph_type in (:dense, :sparse) + g = rand_graph(rng, 6, 10; graph_type) + x = randn(rng, Float32, 3, g.num_nodes) + loss(x) = sum(abs2, NNlib.gather(x, edge_index(g)[1])) + gm = mooncake_gradient(loss, x) + @test gm≈gradient(loss, x)[1] rtol=1e-4 + @test gm≈ngradient(loss, x)[1] rtol=1e-4 + end + + # Float adjacency: edge weights stay on the AD path. + A = Float32[0 1 0 2; 1 0 3 0; 0 3 0 1; 2 0 1 0] + loss_w(A) = sum(abs2, get_edge_weight(GNNGraph(A, graph_type = :dense))) + gw = mooncake_gradient(loss_w, copy(A)) + @test gw≈2 .* A rtol=1e-4 + @test gw≈gradient(loss_w, A)[1] rtol=1e-4 + end +end + +@testitem "GNNGraphsMooncakeExt on GPU" setup=[GraphsTestModule] tags=[:gpu] begin + using .GraphsTestModule + import Mooncake + import NNlib + + dev = gpu_device(force = true) + # Mooncake's GPU rules are CUDA-only. + if VERSION >= v"1.12" && dev isa CUDADevice + function mooncake_gradient(f, x) + cache = Mooncake.prepare_gradient_cache(f, x) + _, grads = Mooncake.value_and_gradient!!(cache, f, x) + return grads[2] + end + + rng = MersenneTwister(17) + + # Without the rules Mooncake errors here instead of returning a gradient. + for graph_type in (:coo, :dense) + g = rand_graph(rng, 6, 10; graph_type) + x = randn(rng, Float32, 3, g.num_nodes) + g_gpu, x_gpu = dev(g), dev(x) + loss(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g))[1])) + loss_gpu(x) = sum(abs2, NNlib.gather(x, edge_index(add_self_loops(g_gpu))[1])) + @test Array(mooncake_gradient(loss_gpu, x_gpu))≈mooncake_gradient(loss, x) rtol=1e-4 + end + + # Float adjacency: edge weights stay on the AD path. + A = Float32[0 1 0 2; 1 0 3 0; 0 3 0 1; 2 0 1 0] + loss_w(A) = sum(abs2, get_edge_weight(GNNGraph(A, graph_type = :dense))) + @test Array(mooncake_gradient(loss_w, dev(A)))≈2 .* A rtol=1e-4 + end +end