make Gpe host side wait cancellable

Summary:
As title.

abortCtrl is not held by KernelElem since it is a struct shared by both host & device side. Arguably we could create a wrapper class for KernelElem on the host side, but I feel like that might be overkill, so only a new function is added.

Reviewed By: saifhhasan

Differential Revision: D85371956

fbshipit-source-id: cdad607103d81735b87328e36e660167789f4caf

Author: arttianezhu

comms/ctran/gpe/CtranGpeDev.h

@@ -8,6 +8,7 @@
 
 #include "comms/ctran/algos/CtranAlgoArgDev.h"
 #include "comms/ctran/algos/CtranAlgoDev.h"
+#include "comms/ctran/utils/Abort.h"
 #include "comms/utils/commSpecs.h"
 
 #ifdef CTRAN_DISABLE_TCPDM
@@ -139,6 +140,7 @@ struct alignas(16) KernelElem {
   // need make progress. It can be safely called only when algorithm ensures no
   // network progress is needed.
   void wait(int groupId = -1);
+  void wait(std::shared_ptr<ctran::utils::Abort> abort, int groupId = -1);
 };
 
 template <>

comms/ctran/gpe/CtranGpeImpl.cc

@@ -779,6 +779,14 @@ void KernelElem::wait(int groupId) {
   }
 }
 
+void KernelElem::wait(std::shared_ptr<ctran::utils::Abort> abort, int groupId) {
+  // wait for all thread blocks to complete
+  while (!this->isComplete(groupId) && !abort->Test()) {
+    // friendly spin so we don't hog CPU
+    std::this_thread::yield();
+  }
+}
+
 KernelElemPool::KernelElemPool(size_t capacity) : capacity_(capacity) {
   FB_CUDACHECKTHROW(cudaHostAlloc(
       &this->memPtr_,

comms/ctran/gpe/tests/KernelElemPoolUT.cc

@@ -1,15 +1,19 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 
+#include <csignal>
+#include <thread>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include <csignal>
+
 #include "comms/ctran/Ctran.h"
 #include "comms/ctran/gpe/CtranGpeDev.h"
 #include "comms/ctran/gpe/CtranGpeImpl.h"
 #include "comms/ctran/utils/CudaWrap.h"
 // FIXME [REBASE]: update the path once moved to fbcode/comms
 #include "comms/ctran/gpe/tests/KernelElemPoolUTKernels.h"
 #include "comms/ctran/tests/CtranXPlatUtUtils.h"
+#include "comms/ctran/utils/Abort.h"
 
 class KernelElemPoolTest : public ::testing::Test {
  public:
@@ -395,3 +399,106 @@ TEST_F(KernelElemPoolAbortTest, CheckFreeWithInvalidNGroups) {
   ASSERT_EXIT(elem->isFree(), testing::KilledBySignal(SIGABRT), "")
       << "Expect abort when revoking with invalid ngroups";
 }
+
+TEST_F(KernelElemPoolTest, WaitWithoutAbortCtrl) {
+  constexpr int poolSize = 10;
+  auto elemPool = std::make_unique<KernelElemPool>(poolSize);
+  ASSERT_NE(elemPool, nullptr);
+
+  constexpr int ngroups = 5;
+  KernelElem* elem = elemPool->pop(ngroups);
+  ASSERT_NE(elem, nullptr);
+
+  // Set elem status to simulate completion
+  for (int i = 0; i < ngroups; i++) {
+    elem->status[i] = KernelElem::ElemStatus::INUSE;
+  }
+
+  std::thread completer([elem, ngroups]() {
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    for (int i = 0; i < ngroups; i++) {
+      elem->status[i] = KernelElem::ElemStatus::DONE;
+    }
+  });
+  elem->wait();
+  EXPECT_TRUE(elem->isComplete());
+
+  completer.join();
+
+  // Clean up
+  elem->free();
+  elemPool->reclaim();
+}
+
+TEST_F(KernelElemPoolTest, WaitWithAbortCtrlWithoutSet) {
+  constexpr int poolSize = 10;
+  auto elemPool = std::make_unique<KernelElemPool>(poolSize);
+  ASSERT_NE(elemPool, nullptr);
+
+  constexpr int ngroups = 5;
+  KernelElem* elem = elemPool->pop(ngroups);
+  ASSERT_NE(elem, nullptr);
+
+  // Enable abortCtrl
+  auto abortCtrl = ctran::utils::createAbort(/*enabled=*/true);
+  EXPECT_TRUE(abortCtrl->Enabled());
+  EXPECT_FALSE(abortCtrl->Test());
+
+  // Set elem status to simulate ongoing work
+  for (int i = 0; i < ngroups; i++) {
+    elem->status[i] = KernelElem::ElemStatus::INUSE;
+  }
+
+  std::thread completer([elem, ngroups]() {
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    for (int i = 0; i < ngroups; i++) {
+      elem->status[i] = KernelElem::ElemStatus::DONE;
+    }
+  });
+  elem->wait(abortCtrl);
+  EXPECT_TRUE(elem->isComplete());
+  EXPECT_FALSE(abortCtrl->Test());
+
+  completer.join();
+
+  // Clean up
+  elem->free();
+  elemPool->reclaim();
+}
+
+TEST_F(KernelElemPoolTest, WaitWithAbortCtrlUnblockOnSet) {
+  constexpr int poolSize = 10;
+  auto elemPool = std::make_unique<KernelElemPool>(poolSize);
+  ASSERT_NE(elemPool, nullptr);
+
+  constexpr int ngroups = 5;
+  KernelElem* elem = elemPool->pop(ngroups);
+  ASSERT_NE(elem, nullptr);
+
+  // Enable abortCtrl
+  auto abortCtrl = ctran::utils::createAbort(/*enabled=*/true);
+  EXPECT_TRUE(abortCtrl->Enabled());
+  EXPECT_FALSE(abortCtrl->Test());
+
+  // Set elem status to simulate ongoing work that won't complete
+  for (int i = 0; i < ngroups; i++) {
+    elem->status[i] = KernelElem::ElemStatus::INUSE;
+  }
+
+  std::thread aborter([abortCtrl]() {
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    abortCtrl->Set();
+  });
+  elem->wait(abortCtrl);
+  EXPECT_FALSE(elem->isComplete());
+  EXPECT_TRUE(abortCtrl->Test());
+
+  aborter.join();
+
+  // Clean up - since abort is set and element is not complete,
+  // we need to manually set status to allow free
+  for (int i = 0; i < ngroups; i++) {
+    elem->status[i] = KernelElem::ElemStatus::RESET;
+  }
+  elemPool->reclaim();
+}