Specify numa policy in examples

NVIDIA · Nov 2, 2023 · bb39d12 · bb39d12
1 parent b496f23
commit bb39d12
Show file tree

Hide file tree

Showing 9 changed files with 42 additions and 12 deletions.
diff --git a/examples/benchmark/common.hpp b/examples/benchmark/common.hpp
@@ -97,7 +97,7 @@ struct numa_deleter {
 };
 
 template <class Pool, class RunThread>
-void my_main(int argc, char** argv) {
+void my_main(int argc, char** argv, exec::numa_policy* policy = exec::get_numa_policy()) {
   int nthreads = std::thread::hardware_concurrency();
   if (argc > 1) {
     nthreads = std::atoi(argv[1]);
@@ -112,7 +112,6 @@ void my_main(int argc, char** argv) {
   std::atomic<bool> stop{false};
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
   std::size_t buffer_size = 1000 << 20;
-  exec::numa_policy* policy = exec::get_numa_policy();
   for (std::size_t i = 0; i < static_cast<std::size_t>(nthreads); ++i) {
     exec::numa_allocator<char> alloc(policy->thread_index_to_node(i));
     buffers.push_back(std::unique_ptr<char, numa_deleter>{alloc.allocate(buffer_size), numa_deleter{buffer_size, alloc}});
@@ -128,7 +127,8 @@ void my_main(int argc, char** argv) {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
       std::span<char>{buffers[i].get(), buffer_size},
 #endif
-      std::ref(stop));
+      std::ref(stop),
+      policy);
   }
   std::size_t nRuns = 100;
   std::size_t warmup = 1;

diff --git a/examples/benchmark/static_thread_pool.cpp b/examples/benchmark/static_thread_pool.cpp
@@ -25,7 +25,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;
@@ -80,7 +83,13 @@ struct RunThread {
   }
 };
 
+struct my_numa_distribution : public exec::default_numa_policy {
+  std::size_t thread_index_to_node(std::size_t index) override {
+    return default_numa_policy::thread_index_to_node(2 * index);
+  }
+};
 
 int main(int argc, char** argv) {
-  my_main<exec::static_thread_pool, RunThread>(argc, argv);
+  my_numa_distribution numa{};
+  my_main<exec::static_thread_pool, RunThread>(argc, argv, &numa);
 }
diff --git a/examples/benchmark/static_thread_pool_bulk_enqueue.cpp b/examples/benchmark/static_thread_pool_bulk_enqueue.cpp
@@ -30,7 +30,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     while (true) {
       barrier.arrive_and_wait();
       if (stop.load()) {

diff --git a/examples/benchmark/static_thread_pool_bulk_enqueue_nested.cpp b/examples/benchmark/static_thread_pool_bulk_enqueue_nested.cpp
@@ -30,7 +30,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     while (true) {
       barrier.arrive_and_wait();

diff --git a/examples/benchmark/static_thread_pool_nested.cpp b/examples/benchmark/static_thread_pool_nested.cpp
@@ -26,7 +26,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;

diff --git a/examples/benchmark/static_thread_pool_nested_old.cpp b/examples/benchmark/static_thread_pool_nested_old.cpp
@@ -26,7 +26,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;

diff --git a/examples/benchmark/static_thread_pool_old.cpp b/examples/benchmark/static_thread_pool_old.cpp
@@ -26,7 +26,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;

diff --git a/examples/benchmark/tbb_thread_pool.cpp b/examples/benchmark/tbb_thread_pool.cpp
@@ -26,7 +26,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;

diff --git a/examples/benchmark/tbb_thread_pool_nested.cpp b/examples/benchmark/tbb_thread_pool_nested.cpp
@@ -28,7 +28,10 @@ struct RunThread {
 #ifndef STDEXEC_NO_MONOTONIC_BUFFER_RESOURCE
     [[maybe_unused]] std::span<char> buffer,
 #endif
-    std::atomic<bool>& stop) {
+    std::atomic<bool>& stop,
+    exec::numa_policy* numa) {
+    std::size_t numa_node = numa->thread_index_to_node(tid);
+    numa->bind_to_node(numa_node);
     auto scheduler = pool.get_scheduler();
     std::mutex mut;
     std::condition_variable cv;