diff --git a/.github/workflows/test_GNNLux_Reactant.yml b/.github/workflows/test_GNNLux_Reactant.yml new file mode 100644 index 000000000..b1a1b3608 --- /dev/null +++ b/.github/workflows/test_GNNLux_Reactant.yml @@ -0,0 +1,38 @@ +name: GNNLux Reactant + +on: + pull_request: + branches: + - master + push: + branches: + - master + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/master' }} + +permissions: + contents: read + +jobs: + test: + name: Julia 1.11 - Reactant CPU + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: '1.11' + - uses: julia-actions/cache@v3 + - name: Install Reactant test dependencies + run: julia --project=GNNLux/test/reactant -e 'using Pkg; Pkg.instantiate()' + - name: Run optional Reactant integration tests + run: julia --code-coverage=all --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: GNNLux/src,GNNlib/src,GNNlib/ext,GNNGraphs/src,GNNGraphs/ext + - uses: codecov/codecov-action@v7 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: lcov.info diff --git a/CHANGELOG.md b/CHANGELOG.md index 7611c27fe..5c1704b9b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,11 @@ 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.7.0) + +**Fixed** +- Reactant traced scalar types are now converted to their underlying numeric type before graph-side degree and COO adjacency allocation. COO graph indices can also remain traced device inputs during compiled GNNLux model calls. Reactant remains an optional GNNGraphs integration ([#710]). + ## GNNGraphs.jl 1.6.0 - 2026-11-04 **Added** @@ -27,6 +32,11 @@ Entries link to the pull request that introduced them. - `GNNChain` now returns its state with the same structure it receives, `(; layers = (; layer_1, ...))`, as the Lux layer interface requires. It previously returned the per-layer states without the `layers` wrapper, so feeding the returned state back into the chain (as `Lux.Training.single_train_step!` does) failed with a `FieldError`. This is breaking for code that worked around the bug by re-wrapping the state as `(; layers = st)`, as the docs examples did: drop the wrapper. - `NNConv` now returns its state as `(; nn = ...)`, the same structure it receives. It previously returned the inner network's state without the `nn` wrapper, so passing the returned state back into the layer (e.g. on the second step of `Lux.Training.single_train_step!`) failed with a `FieldError` ([#728]). +## GNNlib.jl — Unreleased (towards 1.5.0) + +**Fixed** +- Reactant compiled message passing now uses gather/scatter with traced COO indices instead of building a sparse adjacency matrix, allowing one compiled GCNConv to run on different graphs with the same shapes ([#710]). + ## GNNlib.jl 1.4.1 - 2026-11-04 **Fixed** @@ -251,4 +261,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722 [#718]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/718 [#728]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/728 +[#710]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/710 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662 diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml index 29c2d2272..3347d8833 100644 --- a/GNNGraphs/Project.toml +++ b/GNNGraphs/Project.toml @@ -1,6 +1,6 @@ name = "GNNGraphs" uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c" -version = "1.6.0" +version = "1.7.0-DEV" authors = ["Carlo Lucibello and contributors"] [workspace] @@ -25,12 +25,14 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622" [extensions] GNNGraphsCUDAExt = "CUDA" GNNGraphsEnzymeCoreExt = "EnzymeCore" GNNGraphsMooncakeExt = "Mooncake" +GNNGraphsReactantExt = "Reactant" GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs" [compat] @@ -46,6 +48,7 @@ MLUtils = "0.4" Mooncake = "0.5.24" NNlib = "0.9" NearestNeighbors = "0.4" +Reactant = "0.2" Random = "1" SimpleWeightedGraphs = "1.4.0" SparseArrays = "1" diff --git a/GNNGraphs/ext/GNNGraphsReactantExt.jl b/GNNGraphs/ext/GNNGraphsReactantExt.jl new file mode 100644 index 000000000..b74ced782 --- /dev/null +++ b/GNNGraphs/ext/GNNGraphsReactantExt.jl @@ -0,0 +1,10 @@ +module GNNGraphsReactantExt + +using GNNGraphs +using Reactant: TracedRNumber + +GNNGraphs._check_index_eltype(::Type{<:TracedRNumber{T}}) where {T <: Integer} = true + +GNNGraphs._graph_scalar_eltype(::Type{TracedRNumber{T}}) where {T} = T + +end diff --git a/GNNGraphs/src/abstracttypes.jl b/GNNGraphs/src/abstracttypes.jl index 73146160f..5153063a3 100644 --- a/GNNGraphs/src/abstracttypes.jl +++ b/GNNGraphs/src/abstracttypes.jl @@ -1,10 +1,16 @@ -const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V <: Union{Nothing, AbstractVector}} +# COO index arrays may have non-integer wrapper eltypes while traced by a compiler +# such as Reactant. `_check_index_eltype` validates that the wrapped type is integral. +const COO_T = Tuple{T, T, V} where {T <: AbstractVector, V <: Union{Nothing, AbstractVector}} const ADJLIST_T = AbstractVector{T} where {T <: AbstractVector{<:Integer}} const ADJMAT_T = AbstractMatrix const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T -const AVecI = AbstractVector{<:Integer} +# Graph indicator arrays may likewise be traced arrays whose eltype wraps integers. +const AVecI = AbstractVector + +_check_index_eltype(::Type{T}) where {T} = + T <: Integer || throw(ArgumentError("graph indices must be integers, got eltype $T")) # All concrete graph types should be subtypes of AbstractGNNGraph{T}. # GNNGraph and GNNHeteroGraph are the two concrete types. diff --git a/GNNGraphs/src/convert.jl b/GNNGraphs/src/convert.jl index f5675ffff..e78fec966 100644 --- a/GNNGraphs/src/convert.jl +++ b/GNNGraphs/src/convert.jl @@ -30,6 +30,9 @@ function to_coo(coo::COO_T; dir = :out, num_nodes = nothing, weighted = true, hetero = false) s, t, val = coo + _check_index_eltype(eltype(s)) + _check_index_eltype(eltype(t)) + if isnothing(num_nodes) ns = maximum(s) nt = maximum(t) @@ -146,6 +149,7 @@ 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 + T = _graph_scalar_eltype(T) num_nodes = size(A, 1) @assert num_nodes == size(A, 2) # @assert all(x -> (x == 1) || (x == 0), A) @@ -190,6 +194,7 @@ function to_dense(coo::COO_T, T = nothing; dir = :out, num_nodes = nothing, weig if T === nothing T = isnothing(val) ? eltype(s) : eltype(val) end + T = _graph_scalar_eltype(T) if val === nothing || !weighted val = ones_like(s, T) end @@ -216,6 +221,7 @@ function to_sparse(A::ADJMAT_T, T = nothing; dir = :out, num_nodes = nothing, num_nodes = size(A, 1) @assert num_nodes == size(A, 2) T = T === nothing ? eltype(A) : T + T = _graph_scalar_eltype(T) num_edges = A isa AbstractSparseMatrix ? nnz(A) : count(!=(0), A) if dir == :in A = A' @@ -242,6 +248,7 @@ 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 + T = _graph_scalar_eltype(T) if eweight === nothing || !weighted eweight = fill!(similar(s, T), 1) diff --git a/GNNGraphs/src/gatherscatter.jl b/GNNGraphs/src/gatherscatter.jl index e897399ed..05aff8295 100644 --- a/GNNGraphs/src/gatherscatter.jl +++ b/GNNGraphs/src/gatherscatter.jl @@ -11,7 +11,7 @@ _scatter(aggr, src::Dict, idx, n) = Dict([k => _scatter(aggr, v, idx, n) for (k, function _scatter(aggr, src::AbstractArray, - idx::AbstractVector{<:Integer}, + idx::AbstractVector, n::Integer) dstsize = (size(src)[1:(end - 1)]..., n) return NNlib.scatter(aggr, src, idx; dstsize) diff --git a/GNNGraphs/src/gnngraph.jl b/GNNGraphs/src/gnngraph.jl index dccaf8d67..b24322e14 100644 --- a/GNNGraphs/src/gnngraph.jl +++ b/GNNGraphs/src/gnngraph.jl @@ -125,6 +125,9 @@ function GNNGraph(graph::T, ndata::DataStore, edata::DataStore, gdata::DataStore) where {T <: Union{COO_T, ADJMAT_T}} + if !isnothing(graph_indicator) + _check_index_eltype(eltype(graph_indicator)) + end return GNNGraph{T}(graph, num_nodes, num_edges, num_graphs, graph_indicator, ndata, edata, gdata, false) end @@ -154,6 +157,9 @@ function GNNGraph(data::D; graph, num_nodes, num_edges = to_sparse(data; num_nodes, dir) end + if !isnothing(graph_indicator) + _check_index_eltype(eltype(graph_indicator)) + end num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1 ndata = normalize_graphdata(ndata, default_name = :x, n = num_nodes) diff --git a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl b/GNNGraphs/src/gnnheterograph/gnnheterograph.jl index 7c2b9ede5..9c00e0856 100644 --- a/GNNGraphs/src/gnnheterograph/gnnheterograph.jl +++ b/GNNGraphs/src/gnnheterograph/gnnheterograph.jl @@ -117,6 +117,10 @@ function GNNHeteroGraph(data::EDict; @assert dir ∈ [:in, :out] @assert graph_type==:coo "only :coo graph_type is supported for now" + if !isnothing(graph_indicator) + foreach(gi -> _check_index_eltype(eltype(gi)), values(graph_indicator)) + end + if num_nodes !== nothing num_nodes = Dict(num_nodes) end diff --git a/GNNGraphs/src/query.jl b/GNNGraphs/src/query.jl index aaad2d634..db8b4ac98 100644 --- a/GNNGraphs/src/query.jl +++ b/GNNGraphs/src/query.jl @@ -350,6 +350,7 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T::TT = nothing; dir = :out, else T end + T = _graph_scalar_eltype(T) return _degree((s, t), T, dir, ew, g.num_nodes) end @@ -371,10 +372,15 @@ function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT = nothing; dir = :out, T = Nt end end + T = _graph_scalar_eltype(T) A = adjacency_matrix(g) return _degree(A, T, dir, edge_weight, g.num_nodes) end +# Backends can specialize this hook when their scalar wrapper type should not be +# used as the element type of ordinary graph-side allocations. +_graph_scalar_eltype(::Type{T}) where {T} = T + function _degree((s, t)::Tuple, T::Type, dir::Symbol, edge_weight::Nothing, num_nodes::Int) _degree((s, t), T, dir, ones_like(s, T), num_nodes) end diff --git a/GNNGraphs/test/gnngraph.jl b/GNNGraphs/test/gnngraph.jl index ca9c8af02..534d5f938 100644 --- a/GNNGraphs/test/gnngraph.jl +++ b/GNNGraphs/test/gnngraph.jl @@ -1,6 +1,20 @@ # TODO test that the graph type is preserved # when constructing a GNNGraph from another +@testitem "Constructor: COO indices must be integers" setup=[GraphsTestModule] begin + @test GNNGraph(([1, 2], [2, 1]); num_nodes = 2).num_edges == 2 + @test_throws ArgumentError GNNGraph(([1.0, 2.0], [2.0, 1.0]); num_nodes = 2) + @test_throws ArgumentError GNNGraph(([1.5, 2.0], [2.0, 1.0]); num_nodes = 2) + @test_throws ArgumentError GNNGraph((["1", "2"], ["2", "1"]); num_nodes = 2) + + @test_throws ArgumentError GNNGraph( + ([1, 2], [2, 1]); num_nodes = 2, graph_indicator = [1.0, 1.0] + ) + @test_throws ArgumentError GNNGraph( + ([1, 2], [2, 1]); num_nodes = 2, graph_indicator = ["1", "1"] + ) +end + @testitem "Constructor: adjacency matrix" setup=[GraphsTestModule] begin using .GraphsTestModule for GRAPH_T in GRAPH_TYPES diff --git a/GNNLux/test/reactant/Project.toml b/GNNLux/test/reactant/Project.toml new file mode 100644 index 000000000..433c9c510 --- /dev/null +++ b/GNNLux/test/reactant/Project.toml @@ -0,0 +1,17 @@ +[deps] +GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c" +GNNLux = "e8545f4d-a905-48ac-a8c4-ca114b98986d" +GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +GNNGraphs = {path = "../../../GNNGraphs"} +GNNLux = {path = "../.."} +GNNlib = {path = "../../../GNNlib"} + +[compat] +Reactant = "0.2" +julia = "1.11" diff --git a/GNNLux/test/reactant/runtests.jl b/GNNLux/test/reactant/runtests.jl new file mode 100644 index 000000000..c595a053d --- /dev/null +++ b/GNNLux/test/reactant/runtests.jl @@ -0,0 +1,76 @@ +# Optional integration tests; the regular GNNLux test project does not depend on +# Reactant. From the repository root, instantiate and run with: +# julia --project=GNNLux/test/reactant -e 'using Pkg; Pkg.instantiate()' +# julia --project=GNNLux/test/reactant GNNLux/test/reactant/runtests.jl + +using GNNGraphs +using GNNLux +using Lux +using Random +using Reactant +using Test + +degree_for_features(g, x) = GNNGraphs.degree(g, eltype(x); dir = :in, edge_weight = false) + +function test_graph() + return GNNGraph( + ([1, 2, 3, 4, 5, 6], [2, 3, 1, 5, 6, 4]); + num_nodes = 6, + ) +end + +@testset "Reactant degree uses feature numeric type" begin + g = test_graph() + x = randn(Xoshiro(42), Float32, 3, 6) + x_d = x |> reactant_device() + + compiled_degree = @compile degree_for_features(g, x_d) + degree_d = compiled_degree(g, x_d) + + @test Array(degree_d) == ones(Float32, 6) +end + +@testset "Reactant GCNConv compiles, executes, and matches CPU" begin + rng = Xoshiro(42) + model = GCNConv(3 => 2, relu) + ps, st = Lux.setup(rng, model) + g = test_graph() + x = randn(rng, Float32, 3, 6) + y_cpu, _ = model(g, x, ps, st) + + dev = reactant_device() + ps_d, st_d = (ps, st) |> dev + x_d = x |> dev + compiled = @compile model(g, x_d, ps_d, st_d) + y_d, _ = compiled(g, x_d, ps_d, st_d) + + @test size(y_d) == (2, 6) + @test Array(y_d) ≈ y_cpu rtol = 1.0f-4 atol = 1.0f-5 +end + +@testset "Reactant GCNConv keeps device graph indices dynamic" begin + rng = Xoshiro(42) + model = GCNConv(3 => 2, relu) + ps, st = Lux.setup(rng, model) + g1 = test_graph() + g2 = GNNGraph( + ([1, 1, 1, 2, 2, 3], [4, 5, 6, 5, 6, 6]); + num_nodes = 6, + ) + x = randn(rng, Float32, 3, 6) + f(g, x, ps, st) = first(model(g, x, ps, st)) + + dev = reactant_device() + ps_d, st_d, x_d, g1_d, g2_d = (ps, st, x, g1, g2) |> dev + compiled = @compile f(g1_d, x_d, ps_d, st_d) + + y1_cpu = f(g1, x, ps, st) + y2_cpu = f(g2, x, ps, st) + y1_compiled = Array(compiled(g1_d, x_d, ps_d, st_d)) + y2_compiled = Array(compiled(g2_d, x_d, ps_d, st_d)) + + @test size(y1_compiled) == (2, 6) + @test y1_compiled ≈ y1_cpu rtol = 1.0f-4 atol = 1.0f-5 + @test size(y2_compiled) == (2, 6) + @test y2_compiled ≈ y2_cpu rtol = 1.0f-4 atol = 1.0f-5 +end diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml index a0dd0597e..d2a979fba 100644 --- a/GNNlib/Project.toml +++ b/GNNlib/Project.toml @@ -1,7 +1,7 @@ name = "GNNlib" uuid = "a6a84749-d869-43f8-aacc-be26a1996e48" authors = ["Carlo Lucibello and contributors"] -version = "1.4.1" +version = "1.5.0-DEV" [workspace] projects = ["test", "docs", "../GNNGraphs"] @@ -21,24 +21,27 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" [extensions] GNNlibAMDGPUExt = "AMDGPU" GNNlibCUDAExt = "CUDA" GNNlibMetalExt = "Metal" GNNlibMooncakeExt = "Mooncake" +GNNlibReactantExt = "Reactant" [compat] AMDGPU = "1, 2.7" CUDA = "5, 6" ChainRulesCore = "1.24" DataStructures = "0.18, 0.19" -GNNGraphs = "1.6" +GNNGraphs = "1.7" LinearAlgebra = "1" Metal = "1.0" MLUtils = "0.4" Mooncake = "0.5" NNlib = "0.9" Random = "1" +Reactant = "0.2" Statistics = "1" julia = "1.10" diff --git a/GNNlib/ext/GNNlibReactantExt.jl b/GNNlib/ext/GNNlibReactantExt.jl new file mode 100644 index 000000000..db27c3433 --- /dev/null +++ b/GNNlib/ext/GNNlibReactantExt.jl @@ -0,0 +1,25 @@ +module GNNlibReactantExt + +using Reactant: AnyTracedRMatrix +using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj +using GNNGraphs: GNNGraph, COO_T, get_edge_weight + +# The generic fast paths construct a sparse adjacency matrix, which Reactant +# cannot trace when its indices are dynamic. Gather/scatter keeps graph indices +# as inputs to the compiled program. +function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e) + propagate((xi, xj, e) -> copy_xj(xi, xj, e), g, +, xi, xj, e) +end + +function GNNlib.propagate(::typeof(e_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e::AbstractVector) + propagate((xi, xj, e) -> e_mul_xj(xi, xj, e), g, +, xi, xj, e) +end + +function GNNlib.propagate(::typeof(w_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+), + xi, xj::AnyTracedRMatrix, e::Nothing) + propagate((xi, xj, e) -> w_mul_xj(xi, xj, e), g, +, xi, xj, get_edge_weight(g)) +end + +end # module diff --git a/GNNlib/src/GNNlib.jl b/GNNlib/src/GNNlib.jl index 8fe51c9a8..0ae7892db 100644 --- a/GNNlib/src/GNNlib.jl +++ b/GNNlib/src/GNNlib.jl @@ -77,4 +77,3 @@ export global_pool, # include("layers/heteroconv.jl") # no functional part at the moment end #module - \ No newline at end of file