Refactor cudagraph-aware AllToAllvDynamic: move cudagraph-aware out from GPE submit path

Summary:
Similar to AllToAll refactor, moved cudagraph prepare function from gpe to algo.

This diff also moves `prepareCudagraphAware*` implementation to `CudaGraphUtilsImpl.cc` since `PersistentObj` is a variant type: if implement the function in corresponding Pimpl.cc file, they need to include the header files of all possible types part of the variant (alltoallPImpl.cc needs to include alltoallvdynamicP.h, and vice versa). So moved all impl to 1 file to only include all headers once.

Reviewed By: minsii

Differential Revision: D85492033

fbshipit-source-id: d53c833b78e60ee5671355f5c9b18269453d3331

Author: Regina8023

comms/ctran/algos/AllToAll/AllToAllPImpl.cc

@@ -391,32 +391,4 @@ commResult_t AlgoImpl::updatePersistentFuncAndOp(
       (void*)op);
   return commSuccess;
 }
-
-commResult_t prepareCudagraphAwareAllToAll(
-    opFunc& opFunc,
-    struct OpElem* op,
-    PersistentObj& pObj) {
-  pObj = std::make_unique<AlgoImpl>(op->comm_, op->stream);
-  auto algoImplPtr = std::get<std::unique_ptr<AlgoImpl>>(pObj).get();
-  if (!algoImplPtr) {
-    return commSystemError;
-  }
-
-  FB_COMMCHECK(algoImplPtr->setPArgs(
-      op->alltoall.recvbuff,
-      op->alltoall.count * op->comm_->statex_->nRanks(),
-      true /* skipCtrlMsg */,
-      op->alltoall.datatype));
-
-  // Exchange mem handles and record in pArgs. This will not be captured
-  // by cudagraph.
-  FB_COMMCHECK(algoImplPtr->init());
-
-  // Replace gpe func by the persistent version (skip exchanging mem
-  // handle); and OpGroup by the persistent op which has the remote
-  // handles recorded.
-
-  FB_COMMCHECK(algoImplPtr->updatePersistentFuncAndOp(opFunc, op));
-  return commSuccess;
-}
 } // namespace ctran::alltoallp

comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.cc

@@ -1,6 +1,7 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 
 #include "comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.h"
+#include "Types.h"
 #include "comms/ctran/CtranComm.h"
 #include "comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.h"
 #include "comms/ctran/algos/CtranAlgo.h"
@@ -224,22 +225,16 @@ commResult_t AlgoImpl::init() {
 
 commResult_t AlgoImpl::updatePersistFuncAndOp(
     opFunc& opFunc,
-    std::vector<std::unique_ptr<struct OpElem>>& opGroup,
     struct OpElem* op) {
   opFunc = gpeFn;
-  auto new_op = std::make_unique<OpElem>(op);
-  // The original op is not needed and will/may be destroyed. So set kElem to
-  // nullptr to avoid it be freed.
-  op->alltoallv_dynamic.kElem = nullptr;
   // FIXME: only support split_non_contig for now
-  new_op->type = OpElem::opType::ALLTOALLV_DYNAMIC_SPLIT_NON_CONTIG_P;
-  new_op->alltoallv_dynamic.pArgs = &pArgs;
-  opGroup.push_back(std::move(new_op));
+  op->type = OpElem::opType::ALLTOALLV_DYNAMIC_SPLIT_NON_CONTIG_P;
+  op->alltoallv_dynamic.pArgs = &pArgs;
   CLOGF_TRACE(
       COLL,
       "AllToAllvDynamicP: rank {} updated op to {} and gpeFn to persistent version.",
       comm_->statex_->rank(),
-      (void*)opGroup.front().get());
+      (void*)op);
   return commSuccess;
 }
 } // namespace ctran::alltoallvdynamicp

comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.h

@@ -3,22 +3,13 @@
 #pragma once
 
 #include <folly/synchronization/CallOnce.h>
+#include "Types.h"
 #include "comms/ctran/CtranComm.h"
 #include "comms/ctran/gpe/CtranGpe.h"
 #include "comms/ctran/mapper/CtranMapperTypes.h"
 #include "comms/utils/cvars/nccl_cvars.h"
 
 namespace ctran::alltoallvdynamicp {
-struct PersistArgs {
-  std::vector<void*> recvbuffs;
-  std::vector<void*> recvHdls;
-  size_t maxSendCount;
-  size_t maxRecvCount;
-  commDataType_t datatype;
-  std::vector<void*> remoteRecvBuffs;
-  std::vector<struct CtranMapperRemoteAccessKey> remoteAccessKeys;
-};
-
 class AlgoImpl {
  public:
   PersistArgs pArgs;
@@ -29,10 +20,7 @@ class AlgoImpl {
 
   commResult_t init();
 
-  commResult_t updatePersistFuncAndOp(
-      opFunc& opFunc,
-      std::vector<std::unique_ptr<struct OpElem>>& opGroup,
-      struct OpElem* op);
+  commResult_t updatePersistFuncAndOp(opFunc& opFunc, struct OpElem* op);
 
   static inline const std::string algoName(enum NCCL_ALLTOALL_ALGO algo) {
     switch (algo) {
@@ -47,4 +35,9 @@ class AlgoImpl {
   CtranComm* comm_{nullptr};
   cudaStream_t stream_{nullptr};
 };
+
+commResult_t prepareCudagraphAwareAllToAllvDynamic(
+    opFunc& opFunc,
+    struct OpElem* op,
+    PersistentObj& pObj);
 } // namespace ctran::alltoallvdynamicp

comms/ctran/algos/AllToAll/AllToallvDynamicSplitNonContig.cc

@@ -3,8 +3,10 @@
 #include <cuda_fp16.h>
 #include <cstddef>
 
+#include "Types.h"
 #include "comms/ctran/CtranComm.h"
 #include "comms/ctran/algos/AllToAll/AllToAllvDynamicCommon.h"
+#include "comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.h"
 #include "comms/ctran/algos/CtranAlgo.h"
 #include "comms/ctran/gpe/CtranGpe.h"
 
@@ -88,11 +90,19 @@ commResult_t ctranAlltoallvDynamicSplitNonContig(
   XCHECK(alltoallvDynamicSplitNonContigKerns.contains(datatype))
       << "alltoallvDynamicSplitNonContigKerns does not contain datatype "
       << datatype;
+
+  ctran::PreLaunchGraphPrepareFn graphPrepareFn = nullptr;
+  if (NCCL_CTRAN_ALLTOALL_CUDAGRAPH_AWARE_ENABLE) {
+    graphPrepareFn =
+        ctran::alltoallvdynamicp::prepareCudagraphAwareAllToAllvDynamic;
+  }
   FB_COMMCHECK(comm->ctran_->gpe->submit(
       std::move(opGroup),
       opIbImpl,
       config,
-      alltoallvDynamicSplitNonContigKerns.at(datatype)));
+      alltoallvDynamicSplitNonContigKerns.at(datatype),
+      std::nullopt, /* timeout */
+      graphPrepareFn));
 
   return commSuccess;
 }

comms/ctran/algos/AllToAll/Types.h

@@ -6,7 +6,8 @@
 #include "comms/ctran/mapper/CtranMapperTypes.h"
 #include "comms/utils/commSpecs.h"
 
-namespace ctran::alltoallp {
+namespace ctran {
+namespace alltoallp {
 struct PersistArgs {
   void* recvbuff;
   void* recvHdl;
@@ -18,4 +19,19 @@ struct PersistArgs {
 };
 
 class AlgoImpl;
-} // namespace ctran::alltoallp
+} // namespace alltoallp
+
+namespace alltoallvdynamicp {
+struct PersistArgs {
+  std::vector<void*> recvbuffs;
+  std::vector<void*> recvHdls;
+  size_t maxSendCount;
+  size_t maxRecvCount;
+  commDataType_t datatype;
+  std::vector<void*> remoteRecvBuffs;
+  std::vector<struct CtranMapperRemoteAccessKey> remoteAccessKeys;
+};
+
+class AlgoImpl;
+} // namespace alltoallvdynamicp
+} // namespace ctran

comms/ctran/algos/CudaGraphUtilsImpl.cc

@@ -0,0 +1,74 @@
+#include "comms/ctran/CtranComm.h"
+#include "comms/ctran/algos/AllToAll/AllToAllPImpl.h"
+#include "comms/ctran/algos/AllToAll/AllToAllvDynamicPImpl.h"
+#include "comms/ctran/algos/AllToAll/Types.h"
+
+namespace ctran {
+namespace alltoallp {
+commResult_t prepareCudagraphAwareAllToAll(
+    opFunc& opFunc,
+    struct OpElem* op,
+    PersistentObj& pObj) {
+  pObj = std::make_unique<AlgoImpl>(op->comm_, op->stream);
+  auto algoImplPtr = std::get<std::unique_ptr<AlgoImpl>>(pObj).get();
+  if (!algoImplPtr) {
+    return commSystemError;
+  }
+
+  FB_COMMCHECK(algoImplPtr->setPArgs(
+      op->alltoall.recvbuff,
+      op->alltoall.count * op->comm_->statex_->nRanks(),
+      true /* skipCtrlMsg */,
+      op->alltoall.datatype));
+
+  // Exchange mem handles and record in pArgs. This will not be captured
+  // by cudagraph.
+  FB_COMMCHECK(algoImplPtr->init());
+
+  // Replace gpe func by the persistent version (skip exchanging mem
+  // handle); and OpGroup by the persistent op which has the remote
+  // handles recorded.
+
+  FB_COMMCHECK(algoImplPtr->updatePersistentFuncAndOp(opFunc, op));
+  return commSuccess;
+}
+
+} // namespace alltoallp
+namespace alltoallvdynamicp {
+commResult_t prepareCudagraphAwareAllToAllvDynamic(
+    opFunc& opFunc,
+    struct OpElem* op,
+    PersistentObj& pObj) {
+  pObj = std::make_unique<AlgoImpl>(op->comm_, op->stream);
+  auto algoImplPtr = std::get<std::unique_ptr<AlgoImpl>>(pObj).get();
+  if (!algoImplPtr) {
+    return commSystemError;
+  }
+
+  const int nRanks = op->comm_->statex_->nRanks();
+  std::vector<void*> recvbuffs(nRanks);
+  for (int i = 0; i < nRanks; i++) {
+    recvbuffs[i] = op->alltoallv_dynamic.recvbuffs[i];
+  }
+  // FIXME: confirm if sendbuffs are also persistent, so we don't need to
+  // search handle for sendbuffs every time
+  algoImplPtr->pArgs = {
+      .recvbuffs = recvbuffs,
+      .maxSendCount = op->alltoallv_dynamic.maxSendcount,
+      .maxRecvCount = op->alltoallv_dynamic.maxRecvcount,
+      .datatype = op->alltoallv_dynamic.datatype,
+  };
+
+  // Exchange mem handles and record in pArgs. This will not be captured
+  // by cudagraph.
+  FB_COMMCHECK(algoImplPtr->init());
+
+  // Replace gpe func by the persistent version (skip exchanging mem
+  // handle); and OpGroup by the persistent op which has the remote
+  // handles recorded.
+
+  FB_COMMCHECK(algoImplPtr->updatePersistFuncAndOp(opFunc, op));
+  return commSuccess;
+}
+} // namespace alltoallvdynamicp
+} // namespace ctran

comms/ctran/gpe/CtranGpe.h

@@ -22,8 +22,10 @@ typedef commResult_t (*opFunc)(
     const std::vector<std::unique_ptr<struct OpElem>>& opGroup);
 
 namespace ctran {
-using PersistentObj =
-    std::variant<std::monostate, std::unique_ptr<ctran::alltoallp::AlgoImpl>>;
+using PersistentObj = std::variant<
+    std::monostate,
+    std::unique_ptr<alltoallp::AlgoImpl>,
+    std::unique_ptr<alltoallvdynamicp::AlgoImpl>>;
 using PreLaunchGraphPrepareFn =
     commResult_t (*)(opFunc& opFunc, struct OpElem* op, PersistentObj& pObj);
 } // namespace ctran

comms/ctran/gpe/CtranGpeImpl.cc

@@ -46,41 +46,6 @@ static std::unordered_map<KernelConfig::KernelType, const std::string>
          "AllToAllvDynamicSplitNonContig"},
 };
 
-namespace {
-commResult_t updateGraphAwareAlltoAllvDynamicCmd(CtranGpeCmd* cmd) {
-  auto op = cmd->coll.opGroup.front().get();
-  ctran::alltoallvdynamicp::AlgoImpl* algo =
-      new ctran::alltoallvdynamicp::AlgoImpl(op->comm_, op->stream);
-  if (!algo) {
-    return commSystemError;
-  }
-  const int nRanks = op->comm_->statex_->nRanks();
-  std::vector<void*> recvbuffs(nRanks);
-  for (int i = 0; i < nRanks; i++) {
-    recvbuffs[i] = op->alltoallv_dynamic.recvbuffs[i];
-  }
-  // FIXME: confirm if sendbuffs are also persistent, so we don't need to
-  // search handle for sendbuffs every time
-  algo->pArgs = {
-      .recvbuffs = recvbuffs,
-      .maxSendCount = op->alltoallv_dynamic.maxSendcount,
-      .maxRecvCount = op->alltoallv_dynamic.maxRecvcount,
-      .datatype = op->alltoallv_dynamic.datatype,
-  };
-  // Exchange mem handles and record in pArgs. This will not be captured
-  // by cudagraph.
-  FB_COMMCHECK(algo->init());
-
-  // Replace gpe func by the persistent version (skip exchanging mem
-  // handle); and OpGroup by the persistent op which has the remote
-  // handles recorded.
-  std::vector<std::unique_ptr<struct OpElem>> newOpGroup;
-  FB_COMMCHECK(algo->updatePersistFuncAndOp(cmd->coll.func, newOpGroup, op));
-  cmd->coll.opGroup = std::move(newOpGroup);
-  return commSuccess;
-}
-} // namespace
-
 CtranGpe::Impl::Impl() {
   this->kernelFlagPool = std::unique_ptr<KernelFlagPool>(
       new KernelFlagPool(NCCL_CTRAN_NUM_KERNEL_FLAGS));
@@ -254,15 +219,6 @@ commResult_t CtranGpe::Impl::submit(
     if (streamCaptureInfo.status == cudaStreamCaptureStatusActive) {
       FB_COMMCHECK(preLaunchGraphPrepare(cmd, graphPrepareFn));
       struct cmdCbPlan* plan = new struct cmdCbPlan;
-      // cudagraph-aware alltoall: transfer alltoall to alltoallPersistent for
-      // perf optimization
-      auto op = cmd->coll.opGroup.front().get();
-      if (NCCL_CTRAN_ALLTOALL_CUDAGRAPH_AWARE_ENABLE &&
-          op->type == OpElem::opType::ALLTOALLV_DYNAMIC_SPLIT_NON_CONTIG) {
-        // FIXME: this should control by hints passed from user instead of CVAR
-        // so we can have per-collective control
-        updateGraphAwareAlltoAllvDynamicCmd(cmd);
-      }
       plan->cmd = cmd;
       plan->gpe = this->gpe;
       cmd->persistent = true;