Make CtranExRequestImpl.bcast.complete a shared_ptr

Summary:
Change `std::atomic_flag complete` in `CtranExRequestImpl.bcast` to be a `std::shared_ptr<std::atomic_flag>` instead of a raw `std::atomic_flag`.

This simplifies the ownership model and removes the need for the no-op deleter trick when creating shared_ptrs in the ncclx wrappers. The `CtranExRequestImpl` now owns the flag directly as a `shared_ptr`, and this can be passed to `submitHost` without any special handling.

This change maintains the same functionality while providing cleaner code and a more intuitive ownership model where the `CtranExRequestImpl` owns the completion flag.

Reviewed By: Regina8023

Differential Revision: D86548364

fbshipit-source-id: 43e1edf1e0640075da8da0f5a19098bfa6995ed0

Author: Saif Hasan

comms/ctran/CtranExImpl.h

@@ -86,7 +86,7 @@ class CtranExRequestImpl {
     } recvSyncCtrl;
     struct {
       // completion is set by GPE thread and checked by calling thread.
-      std::atomic_flag complete;
+      std::shared_ptr<std::atomic_flag> complete;
     } bcast;
   };
 

comms/ctran/CtranExRequest.cc

@@ -58,7 +58,8 @@ void CtranExRequestImpl::initialize(Type type, CtranIb* ctranIb) {
 void CtranExRequestImpl::initialize(Type type, CtranComm* ctranComm) {
   switch (type) {
     case BCAST:
-      bcast.complete.clear();
+      bcast.complete = std::make_shared<std::atomic_flag>();
+      bcast.complete->clear();
       break;
     default:
       FB_CHECKABORT(
@@ -75,7 +76,7 @@ void CtranExRequestImpl::initialize(Type type, CtranComm* ctranComm) {
 void CtranExRequestImpl::complete() {
   switch (type) {
     case BCAST:
-      bcast.complete.test_and_set();
+      bcast.complete->test_and_set();
       break;
     // no op for other types
     default:
@@ -113,7 +114,7 @@ void CtranExRequestImpl::atComplete(CtranExRequest* req) {
 bool CtranExRequest::isComplete() const {
   auto reqImpl = reinterpret_cast<CtranExRequestImpl*>(impl_);
   if (reqImpl->type == CtranExRequestImpl::BCAST) {
-    return reqImpl->bcast.complete.test();
+    return reqImpl->bcast.complete->test();
   }
   return reqImpl->ibReq.isComplete();
 }
@@ -125,7 +126,7 @@ commResult_t CtranExRequest::test(bool& complete) {
   // marks completion. Thus, no need to polling backend progress by the calling
   // thread.
   if (reqImpl->type == CtranExRequestImpl::BCAST) {
-    complete = reqImpl->bcast.complete.test();
+    complete = reqImpl->bcast.complete->test();
 
     // Check if there is any error reported by the GPE thread;
     // if so, return the error code.
@@ -162,7 +163,7 @@ commResult_t CtranExRequest::wait() {
 
   if (reqImpl->type == CtranExRequestImpl::BCAST) {
     // GPE thread is handling the communcation, wait for it to complete
-    reqImpl->bcast.complete.wait(true);
+    reqImpl->bcast.complete->wait(true);
     // Check if there is any error reported by the GPE thread;
     // if so, return the error code.
     if (reqImpl->asyncErr) {

comms/ctran/algos/Broadcast/BroadcastBinomialTree.cc

@@ -440,7 +440,7 @@ commResult_t CtranAlgo::broadcastBinomialTree(
     size_t count,
     commDataType_t datatype,
     int root,
-    std::atomic_flag* cpuFlag) {
+    std::shared_ptr<std::atomic_flag> cpuFlag) {
   auto opCount = ctran_->getOpCount();
   CTRAN_HOST_COLL_INFO(
       broadcastAlgoName(myAlgo).c_str(),

comms/ctran/algos/CtranAlgo.h

@@ -57,7 +57,7 @@ class CtranAlgo {
       size_t count,
       commDataType_t datatype,
       int root,
-      std::atomic_flag* cpuFlag);
+      std::shared_ptr<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,
-    std::atomic_flag* cpuFlag) {
+    std::shared_ptr<std::atomic_flag> cpuFlag) {
   return this->pimpl->submitHost(
       CtranGpeCmd::TypeEnum::GRAPH_ENQUEUE,
       std::move(opGroup),
       func,
       kernelConfig,
-      cpuFlag);
+      std::move(cpuFlag));
 }
 
 commResult_t CtranGpe::allocKernelElems(

comms/ctran/gpe/CtranGpe.h

@@ -392,7 +392,7 @@ class CtranGpe {
       std::vector<std::unique_ptr<struct OpElem>> opGroup,
       opFunc func,
       KernelConfig& kernelConfig,
-      std::atomic_flag* cpuFlag);
+      std::shared_ptr<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,
-    std::atomic_flag* cpuFlag) {
+    std::shared_ptr<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->cpuFlag = cpuFlag;
+    cmd->cpuFlag = std::move(cpuFlag);
 
     if (type == CtranGpeCmd::TypeEnum::GRAPH_ENQUEUE) {
       cmd->coll.opGroup = std::move(opGroup);

comms/ctran/gpe/CtranGpeImpl.h

@@ -124,7 +124,7 @@ class CtranGpeCmd {
   // kernelFlag to assist device mem communication
   KernelFlagItem* kernelFlag{nullptr};
   // cpuFlag to track completion of host mem communication
-  std::atomic_flag* cpuFlag{nullptr};
+  std::shared_ptr<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,
-      std::atomic_flag* cpuFlag);
+      std::shared_ptr<std::atomic_flag> cpuFlag);
 
   // start the GPE thread.
   void start();

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

@@ -144,7 +144,7 @@ ncclResult_t CtranExComm::broadcast(
       count,
       ncclToMetaComm(datatype),
       root,
-      &reqImpl->bcast.complete)));
+      reqImpl->bcast.complete)));
 
   *req = reqPtr;
   return ncclSuccess;