From d778994ab21fbffac0524937b4171d7101a5ac89 Mon Sep 17 00:00:00 2001
From: Nicolas Morales <nmmoral@sandia.gov>
Date: Mon, 17 Jul 2023 15:20:51 -0700
Subject: [PATCH] use RAII wrapper around enabling and disabling the shared
 allocation tracker to make it exception-safe

---
 core/src/Kokkos_Parallel.hpp        | 37 ++++++++++++++++-------------
 core/src/Kokkos_Parallel_Reduce.hpp | 28 ++++++++++++----------
 core/src/Kokkos_View.hpp            | 18 ++++++++++----
 3 files changed, 49 insertions(+), 34 deletions(-)

diff --git a/core/src/Kokkos_Parallel.hpp b/core/src/Kokkos_Parallel.hpp
index 484f6c0d5f4..fa7bfd6e833 100644
--- a/core/src/Kokkos_Parallel.hpp
+++ b/core/src/Kokkos_Parallel.hpp
@@ -137,9 +137,9 @@ inline void parallel_for(const std::string& str, const ExecPolicy& policy,
   ExecPolicy inner_policy = policy;
   Kokkos::Tools::Impl::begin_parallel_for(inner_policy, functor, str, kpID);
 
-  Kokkos::Impl::shared_allocation_tracking_disable();
-  Impl::ParallelFor<FunctorType, ExecPolicy> closure(functor, inner_policy);
-  Kokkos::Impl::shared_allocation_tracking_enable();
+  auto closure = Kokkos::Impl::with_shared_allocation_tracking_disabled([&]() {
+    return Impl::ParallelFor<FunctorType, ExecPolicy>(functor, inner_policy);
+  });
 
   closure.execute();
 
@@ -352,10 +352,10 @@ inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
   ExecutionPolicy inner_policy = policy;
   Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);
 
-  Kokkos::Impl::shared_allocation_tracking_disable();
-  Impl::ParallelScan<FunctorType, ExecutionPolicy> closure(functor,
-                                                           inner_policy);
-  Kokkos::Impl::shared_allocation_tracking_enable();
+  auto closure = Kokkos::Impl::with_shared_allocation_tracking_disabled([&]() {
+    return Impl::ParallelScan<FunctorType, ExecutionPolicy>(functor,
+                                                            inner_policy);
+  });
 
   closure.execute();
 
@@ -398,18 +398,21 @@ inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
   Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);
 
   if constexpr (Kokkos::is_view<ReturnType>::value) {
-    Kokkos::Impl::shared_allocation_tracking_disable();
-    Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy,
-                                typename ReturnType::value_type>
-        closure(functor, inner_policy, return_value);
-    Kokkos::Impl::shared_allocation_tracking_enable();
+    auto closure =
+        Kokkos::Impl::with_shared_allocation_tracking_disabled([&]() {
+          return Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy,
+                                             typename ReturnType::value_type>(
+              functor, inner_policy, return_value);
+        });
     closure.execute();
   } else {
-    Kokkos::Impl::shared_allocation_tracking_disable();
-    Kokkos::View<ReturnType, Kokkos::HostSpace> view(&return_value);
-    Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy, ReturnType>
-        closure(functor, inner_policy, view);
-    Kokkos::Impl::shared_allocation_tracking_enable();
+    auto closure =
+        Kokkos::Impl::with_shared_allocation_tracking_disabled([&]() {
+          Kokkos::View<ReturnType, Kokkos::HostSpace> view(&return_value);
+          return Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy,
+                                             ReturnType>(functor, inner_policy,
+                                                         view);
+        });
     closure.execute();
   }
 
diff --git a/core/src/Kokkos_Parallel_Reduce.hpp b/core/src/Kokkos_Parallel_Reduce.hpp
index d499eba6dcc..e4075d7377c 100644
--- a/core/src/Kokkos_Parallel_Reduce.hpp
+++ b/core/src/Kokkos_Parallel_Reduce.hpp
@@ -1502,18 +1502,22 @@ struct ParallelReduceAdaptor {
     using Analysis = FunctorAnalysis<FunctorPatternInterface::REDUCE,
                                      PolicyType, typename ReducerSelector::type,
                                      typename return_value_adapter::value_type>;
-    Kokkos::Impl::shared_allocation_tracking_disable();
-    CombinedFunctorReducer functor_reducer(
-        functor, typename Analysis::Reducer(
-                     ReducerSelector::select(functor, return_value)));
-
-    // FIXME Remove "Wrapper" once all backends implement the new interface
-    Impl::ParallelReduce<decltype(functor_reducer), PolicyType,
-                         typename Impl::FunctorPolicyExecutionSpace<
-                             FunctorType, PolicyType>::execution_space>
-        closure(functor_reducer, inner_policy,
-                return_value_adapter::return_value(return_value, functor));
-    Kokkos::Impl::shared_allocation_tracking_enable();
+
+    auto closure =
+        Kokkos::Impl::with_shared_allocation_tracking_disabled([&]() {
+          CombinedFunctorReducer functor_reducer(
+              functor, typename Analysis::Reducer(
+                           ReducerSelector::select(functor, return_value)));
+
+          // FIXME Remove "Wrapper" once all backends implement the new
+          // interface
+          return Impl::ParallelReduce<
+              decltype(functor_reducer), PolicyType,
+              typename Impl::FunctorPolicyExecutionSpace<
+                  FunctorType, PolicyType>::execution_space>(
+              functor_reducer, inner_policy,
+              return_value_adapter::return_value(return_value, functor));
+        });
     closure.execute();
 
     Kokkos::Tools::Impl::end_parallel_reduce<PassedReducerType>(
diff --git a/core/src/Kokkos_View.hpp b/core/src/Kokkos_View.hpp
index 484a0e6f62e..672b1bf9f23 100644
--- a/core/src/Kokkos_View.hpp
+++ b/core/src/Kokkos_View.hpp
@@ -26,6 +26,7 @@ static_assert(false,
 #include <string>
 #include <algorithm>
 #include <initializer_list>
+#include <functional>
 
 #include <Kokkos_Core_fwd.hpp>
 #include <Kokkos_HostSpace.hpp>
@@ -1877,13 +1878,20 @@ KOKKOS_INLINE_FUNCTION bool operator!=(const View<LT, LP...>& lhs,
 
 namespace Kokkos {
 namespace Impl {
+struct SharedAllocationDisableTrackingGuard {
+  SharedAllocationDisableTrackingGuard() {
+    Kokkos::Impl::SharedAllocationRecord<void, void>::tracking_disable();
+  }
 
-inline void shared_allocation_tracking_disable() {
-  Kokkos::Impl::SharedAllocationRecord<void, void>::tracking_disable();
-}
+  ~SharedAllocationDisableTrackingGuard() {
+    Kokkos::Impl::SharedAllocationRecord<void, void>::tracking_enable();
+  }
+};
 
-inline void shared_allocation_tracking_enable() {
-  Kokkos::Impl::SharedAllocationRecord<void, void>::tracking_enable();
+template <class F>
+inline decltype(auto) with_shared_allocation_tracking_disabled(F&& fun) {
+  [[maybe_unused]] auto guard = SharedAllocationDisableTrackingGuard{};
+  return std::invoke(std::forward<F>(fun));
 }
 
 } /* namespace Impl */