Change submitHost parameter from exReq to cpuFlag

Summary:
The GPE thread only needs the `std::atomic_flag* cpuFlag` for signaling completion, so passing the whole `CtranExRequestImpl* exReq` object is unnecessary. This simplifies using `broadcastBinomialTree` with just a flag instead of requiring the full `exReq` object.

Changes include:
- Updated function signatures in CtranGpe.h, CtranGpe.cc, CtranGpeImpl.h, and CtranGpeImpl.cc
- Changed CtranGpeCmd struct to use `cpuFlag` instead of `exReq`
- Updated GPE thread implementation to call `cpuFlag->test_and_set()` instead of `exReq->complete()`
- Updated `CtranAlgo::broadcastBinomialTree` signature and implementation
- Updated call sites in ncclx v2_27 and v2_28 wrappers to pass `&reqImpl->bcast.complete`

This change maintains the same functionality while providing a cleaner API that only exposes what the GPE thread actually needs.

Reviewed By: Regina8023

Differential Revision: D86544452

fbshipit-source-id: 7a71623db10e762cf27d852cabddd0a3f2db5420

Author: Saif Hasan

comms/ctran/algos/Broadcast/BroadcastBinomialTree.cc

@@ -440,7 +440,7 @@ commResult_t CtranAlgo::broadcastBinomialTree(
     size_t count,
     commDataType_t datatype,
     int root,
-    CtranExRequestImpl* exReq) {
+    std::atomic_flag* cpuFlag) {
   auto opCount = ctran_->getOpCount();
   CTRAN_HOST_COLL_INFO(
       broadcastAlgoName(myAlgo).c_str(),
@@ -451,7 +451,7 @@ commResult_t CtranAlgo::broadcastBinomialTree(
       root,
       comm_,
       ctran_,
-      exReq);
+      cpuFlag);
   const auto statex = comm_->statex_.get();
 
   if (sendbuff != recvbuff && statex->rank() == root) {
@@ -489,8 +489,8 @@ commResult_t CtranAlgo::broadcastBinomialTree(
   config.args.collective.broadcast.datatype = datatype;
   config.args.collective.broadcast.count = count;
 
-  FB_COMMCHECK(
-      comm_->ctran_->gpe->submitHost(std::move(opGroup), impl, config, exReq));
+  FB_COMMCHECK(comm_->ctran_->gpe->submitHost(
+      std::move(opGroup), impl, config, cpuFlag));
 
   return commSuccess;
 }

comms/ctran/algos/CtranAlgo.h

@@ -57,7 +57,7 @@ class CtranAlgo {
       size_t count,
       commDataType_t datatype,
       int root,
-      ::ctran::CtranExRequestImpl* exReq);
+      std::atomic_flag* cpuFlag);
 
   commResult_t initTmpBufs();
   commResult_t initAllReduceDirectResource(int nBlocks, cudaStream_t stream);

comms/ctran/gpe/CtranGpe.cc

@@ -389,13 +389,13 @@ commResult_t CtranGpe::submitHost(
     std::vector<std::unique_ptr<struct OpElem>> opGroup,
     opFunc func,
     KernelConfig& kernelConfig,
-    CtranExRequestImpl* exReq) {
+    std::atomic_flag* cpuFlag) {
   return this->pimpl->submitHost(
       CtranGpeCmd::TypeEnum::GRAPH_ENQUEUE,
       std::move(opGroup),
       func,
       kernelConfig,
-      exReq);
+      cpuFlag);
 }
 
 commResult_t CtranGpe::allocKernelElems(

comms/ctran/gpe/CtranGpe.h

@@ -385,14 +385,14 @@ class CtranGpe {
   // Submit host mem communication. No kernel is launched, and only the host
   // side func will be submitted to the GPE thread. Also the op won't be
   // captured by cudagraph.
-  // Completion of the operation is tracked by exReq. exReq can be nullptr,
+  // Completion of the operation is tracked by cpuFlag. cpuFlag can be nullptr,
   // indicating that the caller doesn't care about the completion of the
   // operation.
   commResult_t submitHost(
       std::vector<std::unique_ptr<struct OpElem>> opGroup,
       opFunc func,
       KernelConfig& kernelConfig,
-      ::ctran::CtranExRequestImpl* exReq);
+      std::atomic_flag* cpuFlag);
 
   // Allocate numElems number of p2pElem objects from internal pool.
   // When free objects are not enough, it will be in blocking wait and reclaim

comms/ctran/gpe/CtranGpeImpl.cc

@@ -394,13 +394,13 @@ commResult_t CtranGpe::Impl::submitHost(
     std::vector<std::unique_ptr<struct OpElem>> opGroup,
     opFunc func,
     KernelConfig& kernelConfig,
-    CtranExRequestImpl* exReq) {
+    std::atomic_flag* cpuFlag) {
   // Enqueue op to gpeThread if any op is appended
   if (!opGroup.empty()) {
     class CtranGpeCmd* cmd = new class CtranGpeCmd;
     cmd->type = type;
     cmd->kernelFlag = nullptr;
-    cmd->exReq = exReq;
+    cmd->cpuFlag = cpuFlag;
 
     if (type == CtranGpeCmd::TypeEnum::GRAPH_ENQUEUE) {
       cmd->coll.opGroup = std::move(opGroup);
@@ -531,8 +531,8 @@ void CtranGpe::Impl::gpeThreadFn() {
         // Ensure the host communication request completes, irrespective of
         // outcome of collective function - success, failure, or an exception
         SCOPE_EXIT {
-          if (cmd->exReq) {
-            cmd->exReq->complete();
+          if (cmd->cpuFlag) {
+            cmd->cpuFlag->test_and_set();
           }
         };
 

comms/ctran/gpe/CtranGpeImpl.h

@@ -123,8 +123,8 @@ class CtranGpeCmd {
 
   // kernelFlag to assist device mem communication
   KernelFlagItem* kernelFlag{nullptr};
-  // request to track completion of host mem communication
-  ::ctran::CtranExRequestImpl* exReq{nullptr};
+  // cpuFlag to track completion of host mem communication
+  std::atomic_flag* cpuFlag{nullptr};
 
   bool persistent{false};
 
@@ -202,7 +202,7 @@ class CtranGpe::Impl {
       std::vector<std::unique_ptr<struct OpElem>> opGroup,
       opFunc func,
       KernelConfig& kernelConfig,
-      ::ctran::CtranExRequestImpl* exReq);
+      std::atomic_flag* cpuFlag);
 
   // start the GPE thread.
   void start();

comms/ncclx/v2_27/meta/wrapper/CtranExComm.cc

@@ -139,7 +139,12 @@ ncclResult_t CtranExComm::broadcast(
 
   // TODO: only tree supports host mem for now. Add support for direct.
   NCCLCHECK(metaCommToNccl(ctranComm->ctran_->algo->broadcastBinomialTree(
-      sendbuff, recvbuff, count, ncclToMetaComm(datatype), root, reqImpl)));
+      sendbuff,
+      recvbuff,
+      count,
+      ncclToMetaComm(datatype),
+      root,
+      &reqImpl->bcast.complete)));
 
   *req = reqPtr;
   return ncclSuccess;