fix(npu): sync PyTorch NPU stream for graph-structure queries - #13
Open
xuejiakn wants to merge 1 commit into
Open
fix(npu): sync PyTorch NPU stream for graph-structure queries#13xuejiakn wants to merge 1 commit into
xuejiakn wants to merge 1 commit into
Conversation
…ment_reduce
- streams.py: to_dgl_stream_handle() now supports torch.npu.Stream (.npu_stream)
in addition to torch.cuda.Stream (.cuda_stream)
- heterograph.py:
- DGLGraph.to() calls record_stream() with current PyTorch NPU stream
- New _sync_npu() helper calls torch.npu.synchronize() before/after every
structure query (in_degrees, out_degrees, edges, in_edges) on NPU graphs
- _sparse_ops.py: add torch.npu.synchronize() before _CAPI_DGLKernelSegmentReduce
to prevent data races between PyTorch NPU stream and DGL default ACL stream
- core.py: fix assert to use int(deg) for numpy type compatibility
xuejiakn
force-pushed
the
fix/npu-stream-sync
branch
from
July 28, 2026 06:40
39744cb to
83b582d
Compare
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
1 participant
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.
Summary
Sync PyTorch NPU stream for graph-structure queries to fix data races.
DGL-Ascend runs C++ kernels on the default ACL stream, decoupled from PyTorch's NPU stream. This caused data races: graph-structure queries (in_degrees, out_degrees, edges, in_edges) returned corrupted values when called after PyTorch tensor ops on the same graph.
Changes
streams.py:to_dgl_stream_handle()now supportstorch.npu.Streamheterograph.py:DGLGraph.to()callsrecord_stream()with NPU stream; new_sync_npu()helper around structure queries_sparse_ops.py/core.py: stream sync for segment_reduceRestructure note
This PR has been restructured to consolidate the stream-sync commits that were duplicated across #14[3] and #15[6] (functionally identical supersets). The branch now contains a single clean commit representing the latest version.
Verified:
test_cgcnn.py::test_cgcnn_regressionpasses stably on Ascend 910B3 (CANN 8.3.RC1, torch_npu 2.7.1).