Add free functions to initialize and finalize krepe

README.md

@@ -180,16 +180,17 @@ the `--kernel-replayer-dump` command line flag:
 
 ### Repeating a replay with the same inputs
 
-```cpp
-krepe::ScopeGuard replay_scope(argc, argv, /*enable_input_reset=*/true);
-Kokkos::ScopeGuard kokkos_scope(argc, argv);
+It is possible to reset the input views in the replayed program to their
+original value, by calling `krepe::reset_inputs()`. In order to enable this
+functionality, you should add the `--kernel-replayer-enable-input-reset` when
+launching the replayed program.
 
-replay_scope.reset_inputs();
+```sh
+./replay_prog --kernel-replayer-dump=krepe_scale_2.h5 --kernel-replayer-enable-input-reset
 ```
 
-This option stores a copy of the inputs in host memory. `reset_inputs()` waits
-for previous work and restores those bytes at the same addresses, completing
-all copies before returning.
+Note that calling `reset_inputs` will fence before and after the memory copies
+required to restore the data.
 
 ### Modifying the execution policy
 

src/krepe/replay/kernel_replayer.cpp

@@ -34,8 +34,8 @@ void remove_cli_args(int& argc, char* argv[], int pos, int n) {
   argc -= n;
 }
 
-std::string_view find_flag_argument(int& argc, char* argv[],
-                                    std::string_view flag) {
+std::string_view find_cli_option(int& argc, char* argv[],
+                                 std::string_view flag) {
   std::string_view argument;
   for (int i = 1; i < argc; i++) {
     std::string_view current_arg{argv[i]};
@@ -62,6 +62,17 @@ std::string_view find_flag_argument(int& argc, char* argv[],
   return argument;
 }
 
+bool find_cli_flag(int& argc, char* argv[], std::string_view flag) {
+  for (int i = 1; i < argc; i++) {
+    if (argv[i] == flag) {
+      remove_cli_args(argc, argv, i, 1);
+      return true;
+    }
+  }
+
+  return false;
+}
+
 void device_init() {
   int device_supports_virtual_address = 0;
 #if defined(KOKKOS_ENABLE_CUDA)
@@ -128,6 +139,13 @@ namespace impl {
 
 static const ReplayAllocations* replay_allocations = nullptr;
 static std::unordered_map<std::string, const std::string> metadata;
+static bool krepe_is_initialized = false;
+static bool input_reset_enabled  = false;
+// ScopeGuard used internally by krepe::initialize and krepe::finalize
+static std::unique_ptr<ScopeGuard> initialize_internal_scope_guard;
+// Pointer to the ScopeGuard that was either created by the user or by
+// krepe::initialize
+static ScopeGuard* scope_guard = nullptr;
 
 static void check_replay_allocations(
     [[maybe_unused]] MemorySpaceType memory_space) {
@@ -749,10 +767,17 @@ std::vector<std::pair<std::byte*, std::size_t>> compute_allocations(
 }
 }  // namespace impl
 
-ScopeGuard::ScopeGuard(int& argc, char* argv[], bool enable_input_reset)
-    : enable_input_reset_(enable_input_reset) {
+ScopeGuard::ScopeGuard(int& argc, char* argv[]) {
+  if (impl::krepe_is_initialized) {
+    throw std::runtime_error(
+        "Constructing a ScopeGuard but krepe has already been initialized");
+  }
+
+  impl::input_reset_enabled =
+      find_cli_flag(argc, argv, "--kernel-replayer-enable-input-reset");
+
   std::string_view hdf5_filename =
-      find_flag_argument(argc, argv, "--kernel-replayer-dump");
+      find_cli_option(argc, argv, "--kernel-replayer-dump");
   if (hdf5_filename.data() == nullptr) {
     throw std::runtime_error(
         "The kernel replayer expects the flag --kernel-replayer-dump");
@@ -814,7 +839,7 @@ ScopeGuard::ScopeGuard(int& argc, char* argv[], bool enable_input_reset)
         const auto space = impl::memory_space_type_from_string(entry.space);
         if (entry.size != 0) {
           impl::copy_data(space, entry.address, data, entry.size);
-          if (enable_input_reset_) {
+          if (impl::input_reset_enabled) {
             input_snapshots_.push_back(
                 {space, entry.address,
                  std::vector<std::byte>(data, data + entry.size)});
@@ -872,19 +897,26 @@ ScopeGuard::ScopeGuard(int& argc, char* argv[], bool enable_input_reset)
   file.close_checked();
 
   impl::replay_allocations = &replay_allocations_;
+
+  impl::krepe_is_initialized = true;
+  impl::scope_guard          = this;
 }
 
 ScopeGuard::~ScopeGuard() {
   if (impl::replay_allocations == &replay_allocations_) {
     impl::replay_allocations = nullptr;
   }
+  impl::scope_guard = nullptr;
 }
 
 void ScopeGuard::reset_inputs() {
-  if (!enable_input_reset_) {
+  if (!impl::input_reset_enabled) {
     throw std::runtime_error(
-        "Input reset must be enabled when constructing krepe::ScopeGuard");
+        "Input reset must be enabled via the "
+        "--kernel-replayer-enable-input-reset flag in order to be able to call "
+        "krepe::reset_inputs()");
   }
+  // FIXME: remove once the scopeguard requires kokkos to be initialized
   if (!Kokkos::is_initialized() || Kokkos::is_finalized()) {
     throw std::runtime_error(
         "reset_inputs can be called only inside Kokkos execution environment");
@@ -948,4 +980,46 @@ void ScopeGuard::allocate_output(impl::StoredAllocation& allocation,
   descriptor.reference_data = allocation.reference.get();
 }
 
+void initialize(int& argc, char* argv[]) {
+  if (impl::krepe_is_initialized) {
+    throw std::runtime_error(
+        "Trying to initialize krepe, but it has already been initialized");
+  }
+
+  impl::initialize_internal_scope_guard =
+      std::make_unique<ScopeGuard>(argc, argv);
+}
+
+void finalize() {
+  if (!impl::krepe_is_initialized) {
+    throw std::runtime_error(
+        "Trying to finalize krepe, but it has not been initialized");
+  }
+
+  // if the internal ScopeGuard is null, either krepe has been initialized by
+  // the user constructing a ScopeGuard (in which case impl::scope_guard is
+  // non-null) or krepe has already been finalized (impl::scope_guard is null)
+  if (!impl::initialize_internal_scope_guard) {
+    if (impl::scope_guard) {
+      throw std::runtime_error(
+          "Trying to call krepe::finalize, but krepe has been initialized by a "
+          "ScopeGuard");
+    } else {
+      throw std::runtime_error(
+          "Trying to call krepe::finalize, but krepe has already been "
+          "finalized");
+    }
+  }
+
+  impl::initialize_internal_scope_guard.reset();
+}
+
+void reset_inputs() {
+  if (impl::scope_guard == nullptr) {
+    throw std::runtime_error("krepe::reset_inputs(): krepe is not initialized");
+  }
+
+  impl::scope_guard->reset_inputs();
+}
+
 }  // namespace krepe

src/krepe/replay/kernel_replayer.hpp

@@ -424,7 +424,6 @@ class ScopeGuard {
     std::vector<std::byte> data;
   };
 
-  bool enable_input_reset_;
   std::vector<InputSnapshot> input_snapshots_;
   impl::ReplayAllocations replay_allocations_;
   std::vector<impl::Allocation> host_raw_allocations;
@@ -439,18 +438,26 @@ class ScopeGuard {
                        const impl::SnapshotAllocation& snapshot,
                        const std::byte* data);
 
+  void reset_inputs();
+
  public:
   // Enabling input reset keeps a host copy of all input allocation bytes.
-  ScopeGuard(int& argc, char* argv[], bool enable_input_reset = false);
+  ScopeGuard(int& argc, char* argv[]);
   ScopeGuard(const ScopeGuard&)            = delete;
   ScopeGuard& operator=(const ScopeGuard&) = delete;
   ~ScopeGuard();
 
-  // Restore inputs at their original addresses without reloading the dump.
-  // Requires input reset to be enabled and Kokkos to be initialized.
-  void reset_inputs();
+  friend void reset_inputs();
 };
 
+void initialize(int& argc, char* argv[]);
+
+void finalize();
+
+// Restore inputs at their original addresses without reloading the dump.
+// Requires input reset to be enabled and Kokkos to be initialized.
+void reset_inputs();
+
 std::optional<std::string> get_metadata(const std::string& key);
 
 template <class MemorySpace>

tests/CMakeLists.txt

@@ -9,18 +9,22 @@ set(TEST_LAUNCHER "${CMAKE_CURRENT_LIST_DIR}/launcher.sh")
 # the same replay tests inside a single Slurm allocation.
 add_custom_target(krepe_run_replay_tests)
 
-function(add_replay_test TEST_NAME COMPARE_OUTPUT)
+function(add_replay_test TEST_NAME)
+  cmake_parse_arguments(PARSE_ARGV 1 arg "COMPARE_OUTPUT" "" "TOOL_FLAGS;REPLAY_FLAGS")
   set(TEST_WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/work/${TEST_NAME}")
   file(MAKE_DIRECTORY "${TEST_WORKING_DIRECTORY}")
 
-  if(COMPARE_OUTPUT)
-    set(COMPARE_OUTPUT TRUE)
-  else()
-    set(COMPARE_OUTPUT FALSE)
+  if(NOT arg_TOOL_FLAGS)
+    set(arg_TOOL_FLAGS "\"\"")
   endif()
 
-  set(TEST_COMMAND "${TEST_LAUNCHER}" "$<TARGET_FILE:krepe_kokkos_hooks>" "$<TARGET_FILE:${TEST_NAME}>"
-                   "$<TARGET_FILE:${TEST_NAME}_replay>" "${COMPARE_OUTPUT}" ${ARGN}
+  if(NOT arg_REPLAY_FLAGS)
+    set(arg_REPLAY_FLAGS "\"\"")
+  endif()
+
+  set(TEST_COMMAND
+      "${TEST_LAUNCHER}" "$<TARGET_FILE:krepe_kokkos_hooks>" "$<TARGET_FILE:${TEST_NAME}>"
+      "$<TARGET_FILE:${TEST_NAME}_replay>" "${arg_COMPARE_OUTPUT}" "${arg_TOOL_FLAGS}" "${arg_REPLAY_FLAGS}"
   )
 
   add_test(NAME "${TEST_NAME}" COMMAND ${TEST_COMMAND} WORKING_DIRECTORY "${TEST_WORKING_DIRECTORY}")
@@ -60,6 +64,7 @@ add_subdirectory(team_policy)
 add_subdirectory(non_default_execution_space)
 add_subdirectory(index_type)
 add_subdirectory(empty_view)
+add_subdirectory(initialize_finalize)
 
 if(Kokkos_ENABLE_CUDA OR Kokkos_ENABLE_HIP)
   add_subdirectory(gpu_memory_spaces)

tests/empty_view/CMakeLists.txt

@@ -9,7 +9,7 @@ target_link_libraries(
   empty_view_replay PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_replayer_sanitized
 )
 
-add_replay_test(empty_view FALSE)
+add_replay_test(empty_view)
 
 add_executable(empty_view_without_extent without_extent_main.cpp)
 target_link_libraries(empty_view_without_extent PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_extractor)
@@ -19,4 +19,4 @@ target_link_libraries(
   empty_view_without_extent_replay PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_replayer_sanitized
 )
 
-add_replay_test(empty_view_without_extent FALSE)
+add_replay_test(empty_view_without_extent)

tests/force_policy/CMakeLists.txt

@@ -7,4 +7,4 @@ target_link_libraries(force_policy PRIVATE krepe_warnings krepe_sanitizers Kokko
 add_executable(force_policy_replay replay_main.cpp)
 target_link_libraries(force_policy_replay PRIVATE krepe_warnings Kokkos::kokkos kernel_replayer)
 
-add_replay_test(force_policy "")
+add_replay_test(force_policy)

tests/gpu_memory_spaces/CMakeLists.txt

@@ -7,4 +7,4 @@ target_link_libraries(gpu_memory_spaces PRIVATE krepe_warnings krepe_sanitizers
 add_executable(gpu_memory_spaces_replay replay_main.cpp)
 target_link_libraries(gpu_memory_spaces_replay PRIVATE krepe_warnings Kokkos::kokkos kernel_replayer)
 
-add_replay_test(gpu_memory_spaces FALSE)
+add_replay_test(gpu_memory_spaces REPLAY_FLAGS --kernel-replayer-enable-input-reset)

tests/gpu_memory_spaces/replay_main.cpp

@@ -106,15 +106,15 @@ void compare_values(const krepe::ReplayAllocation& allocation) {
 }  // namespace
 
 int main(int argc, char* argv[]) {
-  krepe::ScopeGuard replay_scope(argc, argv, true);
+  krepe::ScopeGuard replay_scope(argc, argv);
   Kokkos::ScopeGuard kokkos_scope(argc, argv);
 
   Kokkos::View<int*, DeviceSpace> device_values;
   Kokkos::View<int*, ManagedSpace> managed_values;
   Kokkos::View<int*, HostPinnedSpace> host_pinned_values;
 
   for (int repetition = 0; repetition < 3; ++repetition) {
-    replay_scope.reset_inputs();
+    krepe::reset_inputs();
     krepe::parallel_for(
         "test_kernel", 0, KOKKOS_LAMBDA(const int i) {
           device_values(i) *= 2;

tests/index_type/CMakeLists.txt

@@ -7,4 +7,4 @@ target_link_libraries(index_type PRIVATE krepe_warnings krepe_sanitizers Kokkos:
 add_executable(index_type_replay replay_main.cpp)
 target_link_libraries(index_type_replay PRIVATE krepe_warnings Kokkos::kokkos kernel_replayer)
 
-add_replay_test(index_type FALSE)
+add_replay_test(index_type)

tests/initialize_finalize/CMakeLists.txt

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: MIT
+# SPDX-FileCopyrightText: Copyright Contributors to the KREPE project
+
+add_executable(initialize_finalize main.cpp)
+target_link_libraries(initialize_finalize PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_extractor)
+
+add_executable(initialize_finalize_replay replay_main.cpp)
+target_link_libraries(initialize_finalize_replay PRIVATE krepe_warnings Kokkos::kokkos kernel_replayer)
+
+add_replay_test(initialize_finalize)