WIP: PROCESSES proof of concept

CMakeLists.txt

@@ -69,6 +69,9 @@ INCLUDE(${CMAKE_SOURCE_DIR}/ProjectName.cmake)
 # not in an include file :-(
 PROJECT(${PROJECT_NAME} NONE)
 
+SET(TRIBITS_TEST_EXTRA_ENVIRONMENT CTEST_KOKKOS_DEVICE_TYPE=gpus)
+SET(TRIBITS_RESOURCES_PER_PROCESS gpus:1)
+
 # Set default C++ standard to C++11
 SET(Trilinos_ENABLE_CXX11_DEFAULT ON)
 

cmake/tribits/core/package_arch/TribitsAddAdvancedTest.cmake

@@ -1394,6 +1394,18 @@ FUNCTION(TRIBITS_ADD_ADVANCED_TEST TEST_NAME_IN)
         PROCESSORS "${MAX_NUM_PROCESSORS_USED}")
     ENDIF()
 
+    IF(TRIBITS_TEST_EXTRA_ENVIRONMENT)
+      TRIBITS_SET_TEST_PROPERTY(${TEST_NAME} APPEND PROPERTY ENVIRONMENT "${TRIBITS_TEST_EXTRA_ENVIRONMENT}")
+    ENDIF()
+
+    IF(TRIBITS_RESOURCES_PER_PROCESS)
+      SET(NUM_PROCESSES ${MAX_NUM_PROCESSORS_USED})
+      IF(NOT NUM_PROCESSES OR NUM_PROCESSES LESS 1)
+        SET(NUM_PROCESSES 1)
+      ENDIF()
+      TRIBITS_SET_TESTS_PROPERTIES(${TEST_NAME} PROPERTIES PROCESSES "${NUM_PROCESSES},${TRIBITS_RESOURCES_PER_PROCESS}")
+    ENDIF()
+
     IF (SET_DISABLED_AND_MSG)
       TRIBITS_SET_TESTS_PROPERTIES(${TEST_NAME} PROPERTIES DISABLED ON)
     ENDIF()

cmake/tribits/core/package_arch/TribitsAddTestHelpers.cmake

@@ -852,6 +852,18 @@ FUNCTION(TRIBITS_PRIVATE_ADD_TEST_POST_PROCESS_ADDED_TEST  TEST_NAME_IN
     TRIBITS_SET_TESTS_PROPERTIES(${TEST_NAME_IN} PROPERTIES DISABLED ON)
   ENDIF()
 
+  IF(TRIBITS_TEST_EXTRA_ENVIRONMENT)
+    TRIBITS_SET_TEST_PROPERTY(${TEST_NAME_IN} APPEND PROPERTY ENVIRONMENT "${TRIBITS_TEST_EXTRA_ENVIRONMENT}")
+  ENDIF()
+
+  IF(TRIBITS_RESOURCES_PER_PROCESS)
+    SET(NUM_PROCESSES ${NUM_PROCS_USED_IN})
+    IF(NUM_PROCESSES LESS 1)
+      SET(NUM_PROCESSES 1)
+    ENDIF()
+    TRIBITS_SET_TESTS_PROPERTIES(${TEST_NAME_IN} PROPERTIES PROCESSES "${NUM_PROCESSES},${TRIBITS_RESOURCES_PER_PROCESS}")
+  ENDIF()
+
   TRIBITS_PRIVATE_ADD_TEST_ADD_LABEL_AND_KEYWORDS(${TEST_NAME_IN})
 
   TRIBITS_PRIVATE_ADD_TEST_PRINT_ADDED(${TEST_NAME_IN}

packages/kokkos/core/src/Kokkos_Core.hpp

@@ -166,6 +166,9 @@ void fence();
 /** \brief Print "Bill of Materials" */
 void print_configuration( std::ostream & , const bool detail = false );
 
+/** \brief Get GPU assigned by CTest */
+int get_ctest_gpu(const char* local_rank_str);
+
 } // namespace Kokkos
 
 //----------------------------------------------------------------------------

packages/kokkos/core/src/impl/Kokkos_Core.cpp

@@ -50,6 +50,7 @@
 #include <cstdlib>
 #include <stack>
 #include <cerrno>
+#include <string>
 #include <unistd.h>
 
 //----------------------------------------------------------------------------
@@ -98,17 +99,24 @@ setenv("MEMKIND_HBW_NODES", "1", 0);
   const int ndevices = args.ndevices;
   const int skip_device = args.skip_device;
   // if the exact device is not set, but ndevices was given, assign round-robin using on-node MPI rank
-  if (use_gpu < 0 && ndevices >= 0) {
-    auto local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"); //OpenMPI
+  if (use_gpu < 0) {
+    auto const* local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"); //OpenMPI
     if (!local_rank_str) local_rank_str = std::getenv("MV2_COMM_WORLD_LOCAL_RANK"); //MVAPICH2
     if (!local_rank_str) local_rank_str = std::getenv("SLURM_LOCALID"); //SLURM
-    if (local_rank_str) {
-      auto local_rank = std::atoi(local_rank_str);
-      use_gpu = local_rank % ndevices;
-    } else {
-      // user only gave us ndevices, but the MPI environment variable wasn't set.
-      // start with GPU 0 at this point
-      use_gpu = 0;
+
+    auto const* ctest_kokkos_device_type = std::getenv("CTEST_KOKKOS_DEVICE_TYPE"); //CTest
+    auto const* ctest_process_count_str = std::getenv("CTEST_PROCESS_COUNT"); //CTest
+    if (ctest_kokkos_device_type && ctest_process_count_str && local_rank_str) {
+      use_gpu = get_ctest_gpu(local_rank_str);
+    } else if (ndevices >= 0) {
+      if (local_rank_str) {
+        auto local_rank = std::atoi(local_rank_str);
+        use_gpu = local_rank % ndevices;
+      } else {
+        // user only gave us ndevices, but the MPI environment variable wasn't set.
+        // start with GPU 0 at this point
+        use_gpu = 0;
+      }
     }
     // shift assignments over by one so no one is assigned to "skip_device"
     if (use_gpu >= skip_device) ++use_gpu;
@@ -1037,6 +1045,74 @@ void print_configuration( std::ostream & out , const bool detail )
   out << msg.str() << std::endl;
 }
 
+int get_ctest_gpu(const char* local_rank_str)
+{
+  auto const* ctest_kokkos_device_type = std::getenv("CTEST_KOKKOS_DEVICE_TYPE");
+  if (!ctest_kokkos_device_type) {
+    return 0;
+  }
+
+  auto const* ctest_process_count_str = std::getenv("CTEST_PROCESS_COUNT");
+  if (!ctest_process_count_str) {
+    return 0;
+  }
+
+  // Make sure rank is within bounds of processes specified by CTest
+  auto process_count = std::atoi(ctest_process_count_str);
+  auto local_rank = std::atoi(local_rank_str);
+  if (local_rank >= process_count) {
+    return 0;
+  }
+
+  // Get the resource types allocated to this process
+  std::ostringstream ctest_process;
+  ctest_process << "CTEST_PROCESS_" << local_rank;
+  auto const* ctest_process_str = std::getenv(ctest_process.str().c_str());
+  if (!ctest_process_str) {
+    return 0;
+  }
+
+  // Look for the device type specified in CTEST_KOKKOS_DEVICE_TYPE
+  bool found_device = false;
+  std::string ctest_process_cxx_str = ctest_process_str;
+  std::istringstream instream(ctest_process_cxx_str);
+  while (true) {
+    std::string devName;
+    std::getline(instream, devName, ',');
+    if (devName == ctest_kokkos_device_type) {
+      found_device = true;
+      break;
+    }
+    if (instream.eof() || devName.length() == 0) {
+      break;
+    }
+  }
+
+  if (!found_device) {
+    return 0;
+  }
+
+  // Get the device ID
+  std::string ctest_device_type_upper = ctest_kokkos_device_type;
+  for (auto& c : ctest_device_type_upper) {
+    c = std::toupper(c);
+  }
+  ctest_process << "_" << ctest_device_type_upper;
+
+  auto resource_str = std::getenv(ctest_process.str().c_str());
+  if (!resource_str) {
+    return 0;
+  }
+
+  auto const* comma = std::strchr(resource_str, ',');
+  if (!comma || strncmp(resource_str, "id:", 3)) {
+    return 0;
+  }
+
+  std::string id(resource_str + 3, comma - resource_str - 3);
+  return std::atoi(id.c_str());
+}
+
 bool is_initialized() noexcept { return g_is_initialized; }
 
 bool show_warnings() noexcept { return g_show_warnings; }

packages/kokkos/core/unit_test/CMakeLists.txt

@@ -856,6 +856,15 @@ TRIBITS_ADD_EXECUTABLE_AND_TEST(
     TESTONLYLIBS kokkos_gtest ${TEST_LINK_TARGETS}
 )
 
+TRIBITS_ADD_EXECUTABLE_AND_TEST(
+  UnitTest_CTestDevice
+  SOURCES UnitTestMain.cpp TestCTestDevice.cpp
+  COMM serial mpi
+  NUM_MPI_PROCS 1
+  FAIL_REGULAR_EXPRESSION "  FAILED  "
+    TESTONLYLIBS kokkos_gtest ${TEST_LINK_TARGETS}
+)
+
 #
 # Compile-only tests
 #

packages/kokkos/core/unit_test/TestCTestDevice.cpp

@@ -0,0 +1,92 @@
+#include <Kokkos_Core.hpp>
+
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <iostream>
+
+namespace Test {
+
+class ctest_environment : public ::testing::Test {
+protected:
+  void SetUp();
+};
+
+void ctest_environment::SetUp()
+{
+  setenv("CTEST_KOKKOS_DEVICE_TYPE", "gpus", 1);
+  setenv("CTEST_PROCESS_COUNT", "10", 1);
+  unsetenv("CTEST_PROCESS_0");
+  setenv("CTEST_PROCESS_1", "threads", 1);
+  setenv("CTEST_PROCESS_2", "threads,cores", 1);
+
+  setenv("CTEST_PROCESS_3", "gpus", 1);
+  unsetenv("CTEST_PROCESS_3_GPUS");
+
+  setenv("CTEST_PROCESS_4", "gpus", 1);
+  setenv("CTEST_PROCESS_4_GPUS", "id:2", 1);
+
+  setenv("CTEST_PROCESS_5", "gpus", 1);
+  setenv("CTEST_PROCESS_5_GPUS", "slots:1,id:2", 1);
+
+  setenv("CTEST_PROCESS_6", "gpus", 1);
+  setenv("CTEST_PROCESS_6_GPUS", "id:2,slots:1", 1);
+
+  setenv("CTEST_PROCESS_7", "threads,gpus", 1);
+  setenv("CTEST_PROCESS_7_GPUS", "id:3,slots:1", 1);
+
+  setenv("CTEST_PROCESS_8", "gpus,threads", 1);
+  setenv("CTEST_PROCESS_8_GPUS", "id:1,slots:1", 1);
+
+  setenv("CTEST_PROCESS_9", "cores,gpus,threads", 1);
+  setenv("CTEST_PROCESS_9_GPUS", "id:4,slots:1", 1);
+}
+
+TEST_F(ctest_environment, no_device_type)
+{
+  unsetenv("CTEST_KOKKOS_DEVICE_TYPE");
+  EXPECT_EQ(Kokkos::get_ctest_gpu("0"), 0);
+}
+
+TEST_F(ctest_environment, no_process_count)
+{
+  unsetenv("CTEST_PROCESS_COUNT");
+  EXPECT_EQ(Kokkos::get_ctest_gpu("0"), 0);
+}
+
+TEST_F(ctest_environment, invalid_rank)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("10"), 0);
+}
+
+TEST_F(ctest_environment, no_type_str)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("0"), 0);
+}
+
+TEST_F(ctest_environment, missing_type)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("1"), 0);
+  EXPECT_EQ(Kokkos::get_ctest_gpu("2"), 0);
+}
+
+TEST_F(ctest_environment, no_id_str)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("3"), 0);
+}
+
+TEST_F(ctest_environment, invalid_id_str)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("4"), 0);
+  EXPECT_EQ(Kokkos::get_ctest_gpu("5"), 0);
+}
+
+TEST_F(ctest_environment, good)
+{
+  EXPECT_EQ(Kokkos::get_ctest_gpu("6"), 2);
+  EXPECT_EQ(Kokkos::get_ctest_gpu("7"), 3);
+  EXPECT_EQ(Kokkos::get_ctest_gpu("8"), 1);
+  EXPECT_EQ(Kokkos::get_ctest_gpu("9"), 4);
+}
+
+} // namespace Test