From f359027e7936c9c6aedf48fde3c83cf058daf54e Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Sat, 31 Oct 2026 20:37:16 +0100 Subject: [PATCH 1/2] Create the eigensolver start vector on the device of the matrix scaled_laplacian estimates the largest eigenvalue with KrylovKit, starting from a random vector that was always created on the CPU. For a GPU graph, Symmetric(A) * x0 then went to CPU BLAS with a GPU matrix. CUDA avoided this with its own method in the extension, but on AMDGPU OpenBLAS read GPU memory and the process crashed with a segmentation fault, which is why the GNNGraphs AMDGPU job has failed in every build since at least August. The start vector is now allocated with similar, so it lives where the matrix does on any backend. Sparse matrices keep a dense CPU vector, since similar would give a sparse one. --- CHANGELOG.md | 1 + GNNGraphs/src/query.jl | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ae90818b..9bdcfd301 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ Entries link to the pull request that introduced them. - 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** +- Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph. - Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]). - 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]). diff --git a/GNNGraphs/src/query.jl b/GNNGraphs/src/query.jl index 2c2f4332a..aaad2d634 100644 --- a/GNNGraphs/src/query.jl +++ b/GNNGraphs/src/query.jl @@ -508,7 +508,10 @@ function _eigmax(A) KrylovKit.eigsolve(Symmetric(A), x0, 1, :LR)[1][1] # also eigs(A, x0, nev, mode) available end -_rand_dense_vector(A::AbstractMatrix{T}) where {T} = randn(float(T), size(A, 1)) +# On the same device as `A`, otherwise a GPU `A` would be passed to CPU BLAS. +_rand_dense_vector(A::AbstractMatrix{T}) where {T} = randn!(similar(A, float(T), size(A, 1))) +# `similar` would give a sparse vector here. +_rand_dense_vector(A::AbstractSparseMatrix{T}) where {T} = randn(float(T), size(A, 1)) # Eigenvalues for cuarray don't seem to be well supported. # https://github.com/JuliaGPU/CUDA.jl/issues/154 From abf2c46151b60bea3ca99ccedfe1ed3f7c6555b4 Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Sat, 31 Oct 2026 20:52:26 +0100 Subject: [PATCH 2/2] Link the changelog entry to #722 --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bdcfd301..19081e2a8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ Entries link to the pull request that introduced them. - 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** -- Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph. +- Fixed `scaled_laplacian` (and so `ChebConv`) crashing with a segmentation fault on AMDGPU graphs: the eigensolver's start vector was created on the CPU, so CPU BLAS was handed GPU memory. It is now created on the same device as the graph ([#722]). - Added an `EnzymeCore` package extension marking `scaled_laplacian` inactive for Enzyme, matching its existing `@non_differentiable` ChainRules declaration. Enzyme previously differentiated the Krylov eigensolve inside it and failed; this unblocks `ChebConv`, `GConvGRUCell` and `GConvLSTMCell` ([#706]). - 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]). @@ -239,4 +239,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704 [#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707 [#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706 +[#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662