Re: [PR] [Unity][Disco] Use default stream in disco [tvm]

["junrushao (via GitHub)" <gi...@apache.org>]

junrushao commented on code in PR #15827:
URL: https://github.com/apache/tvm/pull/15827#discussion_r1345245896


##########
src/runtime/disco/nccl/nccl.cc:
##########
@@ -137,20 +137,12 @@ inline ncclRedOp_t AsNCCLRedOp(ReduceKind kind) {
 struct CCLThreadLocalContext {
   DiscoWorker* worker;
   int device_id;
-  deviceStream_t default_stream;
+  deviceStream_t default_stream = nullptr;
   ncclComm_t comm;
 
-  void Clear() {
-    NCCL_CALL(ncclCommDestroy(comm));
-    StreamDestroy(default_stream);
-  }
+  void Clear() { NCCL_CALL(ncclCommDestroy(comm)); }
 
-  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;
-  }

Review Comment:
   it seems to break cudagraph integration because cudagraph provides a default stream



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org