[Bugfix] Fix Disco-CUDAGraph Integration

This PR fixed a bug introduced in #15827 since which the cudagraph's stream is discarded.
apache · Oct 4, 2023 · 88536e1 · 88536e1
1 parent 9f0ac49
commit 88536e1
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/src/runtime/disco/nccl/nccl.cc b/src/runtime/disco/nccl/nccl.cc
@@ -140,9 +140,19 @@ struct CCLThreadLocalContext {
   deviceStream_t default_stream = nullptr;
   ncclComm_t comm;
 
-  void Clear() { NCCL_CALL(ncclCommDestroy(comm)); }
+  void Clear() {
+    NCCL_CALL(ncclCommDestroy(comm));
+    if (default_stream != nullptr) {
+      StreamDestroy(default_stream);
+    }
+  }
 
-  deviceStream_t GetDefaultStream() { return nullptr; }
+  deviceStream_t GetDefaultStream() {
+    const auto* func = tvm::runtime::Registry::Get("runtime.get_" TVM_DISCO_DEVICE_NAME "_stream");
+    ICHECK(func != nullptr);
+    deviceStream_t stream = static_cast<deviceStream_t>((*func)().operator void*());
+    return stream == nullptr ? default_stream : stream;
+  }
 
   static CCLThreadLocalContext* Get() {
     thread_local static CCLThreadLocalContext ctx;