Adding support for kokkos-tools

CMakeLists.txt

@@ -43,6 +43,7 @@ message(STATUS "Installation prefix: ${CMAKE_INSTALL_PREFIX}")
 # TODO: this is never used in EKAT src folder, so it should be removed.
 #       Downstream projects should handle bfb logic autonomously
 option (EKAT_DEFAULT_BFB "Whether EKAT should default to BFB behavior whenever possible/appropriate." ${EKAT_IS_DEBUG_BUILD})
+option (EKAT_ENABLE_KOKKOS_TOOLS "Whether to build kokkos-tools and use it to instrument tests at runtime" OFF)
 option (EKAT_ENABLE_VALGRIND "Whether to run tests with valgrind" OFF)
 option (EKAT_ENABLE_CUDA_MEMCHECK "Whether to run tests with cuda-memcheck" OFF)
 option (EKAT_ENABLE_COMPUTE_SANITIZER "Whether to run tests with nvidia's compute-sanitizer" OFF)

README.md

@@ -102,6 +102,7 @@ ctest --output-on-failure
 | `EKAT_ENABLE_ALGORITHM` | `OFF` | Enable the Algorithm sub-package. |
 | `EKAT_ENABLE_EXPRESSION` | `OFF` | Enable the Expression sub-package. |
 | `EKAT_ENABLE_KOKKOS` | `OFF` | Enable the KokkosUtils sub-package. |
+| `EKAT_ENABLE_KOKKOS_TOOLS` | `OFF` | Enable support for kokkos-tools. |
 | `EKAT_ENABLE_LOGGING` | `OFF` | Enable the Logging sub-package. |
 | `EKAT_ENABLE_PACK` | `OFF` | Enable the Pack sub-package. |
 | `EKAT_ENABLE_YAML_PARSER` | `OFF` | Enable the Parser sub-package. |

cmake/tpls/EkatBuildKokkosTools.cmake

@@ -0,0 +1,19 @@
+# Fetch kokkos-tools via FetchContent
+include (FetchContent)
+include (EkatUtils)
+
+message (STATUS "  Fetching kokkos-tools via FetchContent")
+
+set(KOKKOS_TOOLS_GIT_TAG e222e7b6cc7faae5d096e3520d95ca832fb9f889)
+
+# Fixed options
+set(KokkosTools_ENABLE_MPI      ON  CACHE BOOL "Enable MPI for kokkos-tools"   FORCE)
+set(KokkosTools_ENABLE_EXAMPLES OFF CACHE BOOL "Disable kokkos-tools examples" FORCE)
+set(KokkosTools_ENABLE_TESTS    OFF CACHE BOOL "Disable kokkos-tools tests"    FORCE)
+
+# If TPL is already present with correct sha simply adds subdir, otherwise uses FetchContent first
+ekat_fetch_content(KokkosTools
+  GIT_REPOSITORY https://github.com/kokkos/kokkos-tools.git
+  GIT_TAG        ${KOKKOS_TOOLS_GIT_TAG}
+  SOURCE_DIR     ${EKAT_SOURCE_DIR}/extern/kokkos-tools
+)

src/testing-support/CMakeLists.txt

@@ -1,3 +1,21 @@
+### Kokkos Tools for kernel profiling ###
+
+if (EKAT_ENABLE_KOKKOS_TOOLS)
+  set (KokkosTools_FOUND FALSE)
+  if (TARGET kokkostools)
+    set(KokkosTools_FOUND TRUE)
+    message(STATUS "EKAT: Using kokkos-tools provided by parent project")
+  else() # not found -- we fetch it ourselves
+    include (EkatBuildKokkosTools)
+  endif()
+
+  # Build a list of kokkos-tools libraries
+  set(profiling_libs simple-kernel-timer)
+  foreach (lib_path ${profiling_libs})
+    list(APPEND KOKKOS_TOOLS_LIBS ${kokkostools_BINARY_DIR}/profiling/${lib_path})
+  endforeach()
+endif()
+
 ### CATCH MAIN ###
 
 set (Catch2_FOUND FALSE)
@@ -84,6 +102,18 @@ set_target_properties(ekat_testsession PROPERTIES
   EXPORT_NAME DefaultTestSession
 )
 
+# If we're using kokkos-tools, define some symbols for ekat_catchmain and ekat_testsession
+if (EKAT_ENABLE_KOKKOS_TOOLS)
+  target_compile_definitions(ekat_catchmain PRIVATE
+    EKAT_ENABLE_KOKKOS_TOOLS
+    KOKKOS_TOOLS_LIBS="${KOKKOS_TOOLS_LIBS}"
+  )
+  target_compile_definitions(ekat_testsession PRIVATE
+    EKAT_ENABLE_KOKKOS_TOOLS
+    KOKKOS_TOOLS_LIBS="${KOKKOS_TOOLS_LIBS}"
+  )
+endif()
+
 # Install the packages
 install (TARGETS ekat_testsession
          EXPORT EkatTargets

src/testing-support/ekat_catch_main.cpp

@@ -49,6 +49,10 @@ Args split_args (int argc, char** argv)
 
 int main (int argc, char **argv) {
 
+#ifdef EKAT_ENABLE_KOKKOS_TOOLS
+  setenv("KOKKOS_TOOLS_LIBS", KOKKOS_TOOLS_LIBS, 1);
+#endif
+
 #ifdef EKAT_ENABLE_MPI
   // Initialize MPI
   MPI_Init(&argc,&argv);

src/testing-support/ekat_test_session.cpp

@@ -6,6 +6,10 @@
 #include "ekat_arch.hpp"
 
 void ekat_initialize_test_session (int argc, char** argv, const bool print_config) {
+#ifdef EKAT_ENABLE_KOKKOS_TOOLS
+  setenv("KOKKOS_TOOLS_LIBS", KOKKOS_TOOLS_LIBS, 1);
+#endif
+
 #ifdef EKAT_HAS_KOKKOS
   ekat::initialize_kokkos_session (argc,argv,print_config);
 #endif