different work stealing interface for prep and DrvX

pando-rt/src/drvx/cores.cpp

@@ -108,4 +108,17 @@ Cores::TaskQueue* Cores::getTaskQueue(Place place) noexcept {
 void Cores::finalize() {
 }
 
+void Cores::workStealing(std::optional<pando::Task>& task) {
+  const auto thisPlace = getCurrentPlace();
+  const auto coreDims = getCoreDims();
+
+  // onyl steal from the neighbor core
+  for(std::int8_t i = thisPlace.core.x + 1; i < thisPlace.core.x + 2; i++) {
+    auto* otherQueue =  pando::Cores::getTaskQueue(pando::Place{thisPlace.node, thisPlace.pod, pando::CoreIndex(i % coreDims.x, 0)});
+    if(otherQueue->getApproxSize() > STEAL_THRESH_HOLD_SIZE) {
+      task = otherQueue->tryDequeue();
+    }
+  }
+}
+
 } // namespace pando

pando-rt/src/drvx/cores.hpp

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <tuple>
 
+#include "../start.hpp"
 #include "../queue.hpp"
 #include "pando-rt/execution/task.hpp"
 #include "pando-rt/index.hpp"
@@ -54,6 +55,11 @@ class Cores {
    * @brief Returns a flag to check if the core is active.
    */
   static CoreActiveFlag getCoreActiveFlag() noexcept;
+
+  /**
+   * @brief Steals work from other cores.
+   */
+  static void workStealing(std::optional<pando::Task>& task);
 };
 
 } // namespace pando

pando-rt/src/prep/cores.cpp

@@ -23,6 +23,8 @@
 #include "log.hpp"
 #include "status.hpp"
 
+#include <pando-rt/benchmark/counters.hpp>
+
 namespace pando {
 
 namespace {
@@ -451,4 +453,30 @@ Cores::CoreActiveFlag Cores::getCoreActiveFlag() noexcept {
   return CoreActiveFlag{hartContextGet()};
 }
 
+void Cores::workStealing(std::optional<pando::Task>& task, SchedulerFailState& failState, counter::Record<std::int64_t>& idleCount, counter::HighResolutionCount<IDLE_TIMER_ENABLE>& idleTimer) {
+  const auto thisPlace = getCurrentPlace();
+  const auto coreDims = getCoreDims();
+
+  switch(failState) {
+    case SchedulerFailState::YIELD:
+      counter::recordHighResolutionEvent(idleCount, idleTimer, false, thisPlace.core.x, coreDims.x);
+      pando::hartYield();
+      //In Drvx hart yielding is a 1000 cycle wait which is too much
+      idleTimer.start();
+      failState = SchedulerFailState::STEAL;
+      break;
+
+    case SchedulerFailState::STEAL:
+      // only steal work from the neighbor core
+      for(std::int8_t i = thisPlace.core.x + 1; i < thisPlace.core.x + 2; i++) {
+        auto* otherQueue =  pando::Cores::getTaskQueue(pando::Place{thisPlace.node, thisPlace.pod, pando::CoreIndex(i % coreDims.x, 0)});
+        if(otherQueue->getApproxSize() > STEAL_THRESH_HOLD_SIZE) {
+          task = otherQueue->tryDequeue();
+        }
+      }
+      failState = SchedulerFailState::YIELD;
+      break;
+  }
+}
+
 } // namespace pando

pando-rt/src/prep/cores.hpp

@@ -7,10 +7,12 @@
 #include <cstddef>
 #include <tuple>
 
+#include "../start.hpp"
 #include "../queue.hpp"
 #include "pando-rt/execution/task.hpp"
 #include "pando-rt/index.hpp"
 #include "pando-rt/status.hpp"
+#include <pando-rt/benchmark/counters.hpp>
 
 namespace pando {
 
@@ -120,6 +122,11 @@ class Cores {
    * @brief Returns a flag to check if the core is active.
    */
   static CoreActiveFlag getCoreActiveFlag() noexcept;
+
+  /**
+   * @brief Steals work from other cores.
+   */
+  static void workStealing(std::optional<pando::Task>& task, SchedulerFailState& failState, counter::Record<std::int64_t>& idleCount, counter::HighResolutionCount<IDLE_TIMER_ENABLE>& idleTimer);
 };
 
 } // namespace pando

pando-rt/src/start.cpp

@@ -22,16 +22,8 @@
 #include "drvx/drvx.hpp"
 #endif
 
-constexpr std::uint64_t STEAL_THRESH_HOLD_SIZE = 16;
-
-constexpr bool IDLE_TIMER_ENABLE = false;
 counter::Record<std::int64_t> idleCount = counter::Record<std::int64_t>();
 
-enum SchedulerFailState{
-  YIELD,
-  STEAL,
-};
-
 extern "C" int __start(int argc, char** argv) {
   const auto thisPlace = pando::getCurrentPlace();
   const auto coreDims = pando::getCoreDims();
@@ -54,38 +46,24 @@ extern "C" int __start(int argc, char** argv) {
     auto ctok = queue->makeConsumerToken();
 
     std::optional<pando::Task> task = std::nullopt;
-
+
+#ifdef PANDO_RT_USE_BACKEND_PREP
     SchedulerFailState failState = SchedulerFailState::YIELD;
+#endif
 
     do {
       idleTimer.start();
       task = queue->tryDequeue(ctok);
 
       if (!task.has_value()) {
-
-        switch(failState) {
-          case SchedulerFailState::YIELD:
 #ifdef PANDO_RT_USE_BACKEND_PREP
-            counter::recordHighResolutionEvent(idleCount, idleTimer, false, thisPlace.core.x, coreDims.x);
-            pando::hartYield();
-            //In Drvx hart yielding is a 1000 cycle wait which is too much
-            idleTimer.start();
+        pando::Cores::workStealing(task, failState, idleCount, idleTimer);
+#elif defined(PANDO_RT_USE_BACKEND_DRVX)
+        pando::Cores::workStealing(task);
 #endif
-            failState = SchedulerFailState::STEAL;
-            break;
-
-          case SchedulerFailState::STEAL:
-            for(std::int8_t i = thisPlace.core.x + 1; i < thisPlace.core.x + 2 && !task.has_value(); i++) {
-              auto* otherQueue =  pando::Cores::getTaskQueue(pando::Place{thisPlace.node, thisPlace.pod, pando::CoreIndex(i % coreDims.x, 0)});
-              if(!otherQueue || otherQueue == queue) {continue;}
-              if(otherQueue->getApproxSize() > STEAL_THRESH_HOLD_SIZE) {
-                task = otherQueue->tryDequeue();
-              }
-            }
-            failState = SchedulerFailState::YIELD;
-            break;
-        }
       }
+
+      // After work stealing
       if(task.has_value()) {
         (*task)();
         task = std::nullopt;

pando-rt/src/start.hpp

@@ -4,6 +4,17 @@
 #ifndef PANDO_RT_SRC_START_HPP_
 #define PANDO_RT_SRC_START_HPP_
 
+#include <cstdint>
+
+constexpr bool IDLE_TIMER_ENABLE = false;
+
+constexpr std::uint64_t STEAL_THRESH_HOLD_SIZE = 16;
+
+enum SchedulerFailState{
+  YIELD,
+  STEAL,
+};
+
 /**
  * @brief Start function for each hart.
  *