use RAII wrapper around enabling and disabling the shared allocation …

…tracker to make it exception-safe
nmm0 · Jan 24, 2024 · d778994 · d778994
1 parent 862b545
commit d778994
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 34 deletions.
diff --git 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
@@ -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
@@ -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 */