From 83b582d597f7807bc59f7d23037ecbe0d19af7ba Mon Sep 17 00:00:00 2001 From: xuejiakn Date: Tue, 21 Jul 2026 06:06:12 +0000 Subject: [PATCH] fix(npu): sync PyTorch NPU stream for graph-structure queries and segment_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 --- python/dgl/_ffi/streams.py | 14 +++++++++---- python/dgl/_sparse_ops.py | 6 ++++++ python/dgl/core.py | 2 +- python/dgl/heterograph.py | 42 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 5 deletions(-) diff --git a/python/dgl/_ffi/streams.py b/python/dgl/_ffi/streams.py index dcc44a7e63ee..77eac8f7c987 100644 --- a/python/dgl/_ffi/streams.py +++ b/python/dgl/_ffi/streams.py @@ -11,18 +11,24 @@ def to_dgl_stream_handle(cuda_stream): - """Convert torch.cuda.Stream to DGL stream handle + """Convert torch Stream to DGL stream handle + + Supports both torch.cuda.Stream (``.cuda_stream``) and + torch.npu.Stream (``.npu_stream``). Parameters ---------- - cuda_stream : torch.cuda.Stream. + cuda_stream : torch.cuda.Stream or torch.npu.Stream. Returns ------- DGLStreamHandle - DGLStreamHandle of the input ``cuda_stream``. + DGLStreamHandle of the input stream. """ - return ctypes.c_void_p(cuda_stream.cuda_stream) + raw = getattr(cuda_stream, "npu_stream", None) + if raw is None: + raw = cuda_stream.cuda_stream + return ctypes.c_void_p(raw) def _dgl_get_stream(ctx): diff --git a/python/dgl/_sparse_ops.py b/python/dgl/_sparse_ops.py index b3a57d740256..5bb2e5c79943 100644 --- a/python/dgl/_sparse_ops.py +++ b/python/dgl/_sparse_ops.py @@ -680,6 +680,12 @@ def _segment_reduce(op, feat, offsets): if op in ["min", "max"]: arg = F.zeros(out_shp, idtype, ctx) arg_nd = to_dgl_nd_for_write(arg) + # Sync PyTorch NPU stream before DGL kernel to prevent data races + try: + import torch + torch.npu.synchronize() + except Exception: + pass _CAPI_DGLKernelSegmentReduce( op, to_dgl_nd(feat), diff --git a/python/dgl/core.py b/python/dgl/core.py index 0be71fc20f30..16b06be74670 100644 --- a/python/dgl/core.py +++ b/python/dgl/core.py @@ -142,7 +142,7 @@ def invoke_udf_reduce(graph, func, msgdata, *, orig_nid=None): # order the incoming edges per node by edge ID eid_bkt = F.zerocopy_to_numpy(graph.in_edges(node_bkt, form="eid")) - assert len(eid_bkt) == deg * len(node_bkt) + assert len(eid_bkt) == int(deg) * len(node_bkt) eid_bkt = np.sort(eid_bkt.reshape((len(node_bkt), deg)), 1) eid_bkt = F.zerocopy_from_numpy(eid_bkt.flatten()) diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py index f7a9a2aba1d5..681b3eb44257 100644 --- a/python/dgl/heterograph.py +++ b/python/dgl/heterograph.py @@ -3415,7 +3415,9 @@ def in_edges(self, v, form="uv", etype=None): out_edges """ v = utils.prepare_tensor(self, v, "v") + self._sync_npu() src, dst, eid = self._graph.in_edges(self.get_etype_id(etype), v) + self._sync_npu() if form == "all": return src, dst, eid elif form == "uv": @@ -3588,7 +3590,9 @@ def all_edges(self, form="uv", order="eid", etype=None): in_edges out_edges """ + self._sync_npu() src, dst, eid = self._graph.edges(self.get_etype_id(etype), order) + self._sync_npu() if form == "all": return src, dst, eid elif form == "uv": @@ -3670,7 +3674,9 @@ def in_degrees(self, v=ALL, etype=None): if is_all(v): v = self.dstnodes(dsttype) v_tensor = utils.prepare_tensor(self, v, "v") + self._sync_npu() deg = self._graph.in_degrees(etid, v_tensor) + self._sync_npu() if isinstance(v, numbers.Integral): return F.as_scalar(deg) else: @@ -3750,7 +3756,9 @@ def out_degrees(self, u=ALL, etype=None): F.sum(self.has_nodes(u_tensor, ntype=srctype), dim=0) ) != len(u_tensor): raise DGLError("u contains invalid node IDs") + self._sync_npu() deg = self._graph.out_degrees(etid, utils.prepare_tensor(self, u, "u")) + self._sync_npu() if isinstance(u, numbers.Integral): return F.as_scalar(deg) else: @@ -5739,8 +5747,42 @@ def to(self, device, **kwargs): # pylint: disable=invalid-name } ret._batch_num_edges = new_bne + # 3. Record the PyTorch NPU stream on the graph so DGL's memory + # management can track it. DGL-Ascend runs its kernels on the default + # ACL stream (nullptr) which is decoupled from PyTorch's NPU stream; + # without explicit synchronization, DGL graph-structure queries + # (in_degrees, in_edges, etc.) can return corrupted values because + # DGL's aclrtMalloc-allocated buffers race with PyTorch ops on a + # different stream. Recording the stream is necessary (but not + # sufficient); structure-query methods also sync via _sync_npu(). + if F.device_type(ret.device) == "npu": + try: + import torch + cur_stream = torch.npu.current_stream() + ret.record_stream(cur_stream) + except Exception: + pass + return ret + def _sync_npu(self): + """Synchronize PyTorch NPU stream before/after DGL graph-structure queries. + + DGL-Ascend runs kernels on the default ACL stream (nullptr), decoupled + from PyTorch's NPU stream. DGL also allocates NPU device memory via + aclrtMalloc, independent of PyTorch's caching allocator. Without + synchronization, DGL structure queries (in_degrees, edges, etc.) can + read stale or corrupted data because the PyTorch stream may still have + pending ops that overwrite the same physical memory. This is a + no-op on non-NPU devices. + """ + if F.device_type(self.device) == "npu": + try: + import torch + torch.npu.synchronize() + except Exception: + pass + def cpu(self): """Return a new copy of this graph on CPU.