changing fifo poll function

include/mscclpp/fifo.hpp

@@ -25,8 +25,8 @@ class Fifo {
 
   /// Polls the FIFO for a trigger.
   ///
-  /// @param trigger A pointer to the trigger to be filled.
-  void poll(ProxyTrigger* trigger);
+  /// Returns @ref ProxyTrigger which is the trigger at the head of fifo.
+  ProxyTrigger poll();
 
   /// Pops a trigger from the FIFO.
   void pop();

python/fifo_py.cpp

@@ -21,7 +21,7 @@ void register_fifo(nb::module_& m) {
 
   nb::class_<Fifo>(m, "Fifo")
       .def(nb::init<>())
-      .def("poll", &Fifo::poll, nb::arg("trigger"))
+      .def("poll", &Fifo::poll)
       .def("pop", &Fifo::pop)
       .def("flush_tail", &Fifo::flushTail, nb::arg("sync") = false)
       .def("device_handle", &Fifo::deviceHandle);

src/fifo.cc

@@ -39,9 +39,11 @@ struct Fifo::Impl {
 MSCCLPP_API_CPP Fifo::Fifo() : pimpl(std::make_unique<Impl>()) {}
 MSCCLPP_API_CPP Fifo::~Fifo() = default;
 
-MSCCLPP_API_CPP void Fifo::poll(ProxyTrigger* trigger) {
+MSCCLPP_API_CPP ProxyTrigger Fifo::poll() {
+  ProxyTrigger trigger;
   __m128i xmm0 = _mm_load_si128((__m128i*)&pimpl->triggers.get()[pimpl->hostTail % MSCCLPP_PROXY_FIFO_SIZE]);
-  _mm_store_si128((__m128i*)trigger, xmm0);
+  _mm_store_si128((__m128i*)&trigger, xmm0);
+  return trigger;
 }
 
 MSCCLPP_API_CPP void Fifo::pop() {

src/proxy.cc

@@ -67,7 +67,7 @@ MSCCLPP_API_CPP void Proxy::start() {
         }
       }
       // Poll to see if we are ready to send anything
-      fifo.poll(&trigger);
+      trigger = fifo.poll();
       if (trigger.fst == 0) {  // TODO: this check is a potential pitfall for custom triggers
         continue;              // there is one in progress
       }

test/unit/fifo_tests.cu

@@ -52,7 +52,7 @@ TEST(FifoTest, Fifo) {
   mscclpp::Timer timer(3);
   for (uint64_t i = 0; i < ITER; ++i) {
     while (trigger.fst == 0) {
-      hostFifo.poll(&trigger);
+      trigger = hostFifo.poll();
 
       if (spin++ > 1000000) {
         FAIL() << "Polling is stuck.";