add branching

kokkos · Nov 19, 2023 · 374064a · 374064a
1 parent 61842b7
commit 374064a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/core/unit_test/TestDeviceAndThreads.py b/core/unit_test/TestDeviceAndThreads.py
@@ -43,7 +43,9 @@ def GetNumThreads(max_threads):
     result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     output = result.stdout.decode('utf-8')
     phys_cores_count = int(output)
-    looplist = [1] + [i*phys_cores_count for i in [1,2,3,4,5,6,7]]
+    looplist = [1] + [i*phys_cores_count for i in [1,2,3,4,5,6,7]] \
+        if GetFlag("hwloc_enabled") else [1,2,3,4,5]
+
     for x in looplist:
         if x >= max_threads:
             break

diff --git a/core/unit_test/UnitTest_DeviceAndThreads.cpp b/core/unit_test/UnitTest_DeviceAndThreads.cpp
@@ -68,6 +68,14 @@ int get_max_threads() {
 #endif
 }
 
+int get_hwloc_enabled() {
+#ifdef KOKKOS_ENABLE_HWLOC
+  return 1;
+#else
+  return 0;
+#endif
+}
+
 int get_num_threads() {
   int const num_threads = Kokkos::DefaultHostExecutionSpace().concurrency();
   assert(num_threads == Kokkos::num_threads());
@@ -93,6 +101,7 @@ int print_flag(std::string const& flag) {
   KOKKOS_TEST_PRINT_FLAG(device_count);
   KOKKOS_TEST_PRINT_FLAG(disable_warnings);
   KOKKOS_TEST_PRINT_FLAG(tune_internals);
+  KOKKOS_TEST_PRINT_FLAG(hwloc_enabled);
 
 #undef KOKKOS_TEST_PRINT_FLAG