fix(npu): sync PyTorch NPU stream for graph-structure queries

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):

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),

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())
 

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.