Avoid scalar indexing in remove_edges - #672
Merged
Merged
Conversation
Contributor
Author
|
Failures seem unrelated |
Member
|
Can you add a gpu test in GNNGraphs/test/transform.jl? |
Contributor
Author
|
Of course. Is it enough to add them to that file for the GPU code to be tested, or should I modify the testing configuration/flags/functions? |
CarloLucibello
added a commit
that referenced
this pull request
Jul 22, 2026
* Fix GPU scalar indexing in remove_edges and getgraph remove_edges (#668): the follow-up to #672 built the keep-mask with `fill_like(edges_to_remove, true)`, which produced a mask of the wrong length (that of `edges_to_remove`, not the number of edges) and wrong eltype (Int instead of Bool) — breaking even the existing CPU tests. Build the mask from `s` with `fill_like(s, true, Bool, length(s))` so it lives on the graph's device and does logical indexing correctly. getgraph (#161): the subgraph extraction uses scalar Dict lookups and comprehensions that can't run on the GPU. Detect a non-CPU device, run the extraction on the CPU, and move the result back to the original device. Also restores the indentation of the commented add_snapshot!/ remove_snapshot! exports mangled by #672, and adds :gpu-tagged regression tests for both functions. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Reference PR #691 in the changelog entries Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Make getgraph GPU-native for COO graphs instead of CPU round-trip The first cut fixed getgraph on the GPU by moving the graph to the CPU and back. That works but forces a full device<->host sync and its cost grows with the graph size. Instead, express the COO path (the common case, which has a dense graph_indicator) entirely with scatter/gather/cumsum, which run natively on the GPU: build the node keep-mask from a per-graph boolean lookup, relabel nodes via cumsum and graphs via a scatter, and select edges with `node_mask[s]`. Indexing goes through the integer vectors `kept_nodes`/`kept_edges` (one `findall` each) rather than repeated boolean masks, to keep the number of GPU synchronizations low. Adjacency-matrix graphs still fall back to the CPU: they need submatrix indexing and store `graph_indicator` as a `SparseVector`, neither of which is GPU-friendly. The result stays on the original device and, on real hardware, avoids the host transfer entirely; on large batches it is dramatically faster than the round-trip. The GPU test now covers all graph types (COO native + sparse/dense fallback), node/edge features, weights, and `nmap`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(AGENTS): document running and reproducing GPU tests Add a local CUDA test command and note the cuDNN / allowscalar(false) setup needed to reproduce a :gpu test item in the REPL — the gotcha being that gpu_device() silently falls back to the CPU when cuDNN isn't loaded. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #668 by creating a boolean mask on GPU when the inputs are, without modifying the effects in normal CPU mode