refine

Oneflow-Inc · Jan 25, 2025 · b014608 · b014608
1 parent c56c527
commit b014608
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/oneflow/user/kernels/collective_communication/cpu/cpu_recv.cpp b/oneflow/user/kernels/collective_communication/cpu/cpu_recv.cpp
@@ -41,7 +41,7 @@ class CpuRecvImpl final : public Recv {
   }
 
   void Launch(ep::Stream* stream, void* out, size_t elem_cnt, int64_t src,
-              const CclComm& ccl_comm) const override {
+              CclComm ccl_comm) const override {
     Launch(stream, out, elem_cnt, src);
   }
 

diff --git a/oneflow/user/kernels/collective_communication/cuda/cuda_send.cpp b/oneflow/user/kernels/collective_communication/cuda/cuda_send.cpp
@@ -34,6 +34,7 @@ class CudaSend final : public Send {
   void Launch(ep::Stream* stream, const void* in, size_t elem_cnt, int64_t dst) const override {
 #if HAS_NCCL_SEND_RECV
     const auto& comm_and_peer_rank = GetNcclCommAndPeerNcclRank(dst);
+    printf("\n CudaSend >>> Launch >>> communication_ctx");
     OF_NCCL_CHECK(ncclSend(in, elem_cnt, nccl_datatype_, comm_and_peer_rank.second,
                            comm_and_peer_rank.first, stream->As<ep::CudaStream>()->cuda_stream()));
 #else
@@ -42,14 +43,15 @@ class CudaSend final : public Send {
   }
 
   void Launch(ep::Stream* stream, const void* in, size_t elem_cnt, int64_t dst,
-              CclComm ccl_comm) const override {
+              ccl::CclComm ccl_comm) const override {
 #if HAS_NCCL_SEND_RECV
     ncclComm_t* comm = reinterpret_cast<ncclComm_t*>(ccl_comm.getComm());
     OF_NCCL_CHECK(ncclSend(in, elem_cnt, nccl_datatype_, dst, *comm,
                            stream->As<ep::CudaStream>()->cuda_stream()));
 #else
     UNIMPLEMENTED() << "GPU send is only supported when nccl version >= 2.7"
 #endif  // HAS_NCCL_SEND_RECV
+    printf("\n CudaSend >>> Launch >>> ccl::CclComm");
   }
 
  private:

diff --git a/oneflow/user/kernels/collective_communication/include/recv.h b/oneflow/user/kernels/collective_communication/include/recv.h
@@ -33,7 +33,7 @@ class Recv : public CollectiveCommunication {
   virtual void Launch(ep::Stream* stream, void* out, size_t elem_cnt, int64_t src) const = 0;
 
   virtual void Launch(ep::Stream* stream, void* out, size_t elem_cnt, int64_t src,
-                      const CclComm& ccl_comm) const = 0;
+                      CclComm ccl_comm) const = 0;
 };
 
 inline bool IsRecvRegistered(DeviceType device_type) {

diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp
@@ -42,7 +42,7 @@ class NcclLogical2DSameDim0KernelCommState : public user_op::OpKernelState {
   }
   ~NcclLogical2DSameDim0KernelCommState() override = default;
 
-  const ccl::CclComm& ccl_comm() const {
+  ccl::CclComm ccl_comm() {
     if (!is_init_) { Init(); }
     return ccl_comm_;
   }