Update OpenACC FunctorAdapter (#6077)

* Update OpenACC FunctorAdapter to be able to handle functors containing parallel loops. * Update OpenACC FunctorAdapter implementation to use enum types for template specialization. * Minor revision as suggested by the code review.
sandialabs · Apr 28, 2023 · de5c017 · de5c017
1 parent 55bbd9f
commit de5c017
Show file tree

Hide file tree

Showing 8 changed files with 59 additions and 25 deletions.
diff --git a/core/src/OpenACC/Kokkos_OpenACC_FunctorAdapter.hpp b/core/src/OpenACC/Kokkos_OpenACC_FunctorAdapter.hpp
@@ -17,27 +17,40 @@
 #ifndef KOKKOS_OPENACC_FUNCTOR_ADAPTER_HPP
 #define KOKKOS_OPENACC_FUNCTOR_ADAPTER_HPP
 
+#include <OpenACC/Kokkos_OpenACC_Macros.hpp>
 #include <type_traits>
 
 namespace Kokkos::Experimental::Impl {
 
-template <class Functor, class Policy>
-class FunctorAdapter {
-  Functor m_functor;
-  using WorkTag = typename Policy::work_tag;
-
- public:
-  FunctorAdapter(Functor const &functor) : m_functor(functor) {}
-
-  template <class... Args>
-  KOKKOS_FUNCTION void operator()(Args &&... args) const {
-    if constexpr (std::is_void_v<WorkTag>) {
-      m_functor(static_cast<Args &&>(args)...);
-    } else {
-      m_functor(WorkTag(), static_cast<Args &&>(args)...);
-    }
+enum class RoutineClause { worker, seq };
+
+template <class Functor, class Policy, RoutineClause>
+class FunctorAdapter;
+
+#define KOKKOS_IMPL_ACC_FUNCTOR_ADAPTER(CLAUSE)                    \
+  template <class Functor, class Policy>                           \
+  class FunctorAdapter<Functor, Policy, RoutineClause::CLAUSE> {   \
+    Functor m_functor;                                             \
+    using WorkTag = typename Policy::work_tag;                     \
+                                                                   \
+   public:                                                         \
+    FunctorAdapter(Functor const &functor) : m_functor(functor) {} \
+                                                                   \
+    KOKKOS_IMPL_ACC_PRAGMA_HELPER(routine CLAUSE)                  \
+    template <class... Args>                                       \
+    KOKKOS_FUNCTION void operator()(Args &&... args) const {       \
+      if constexpr (std::is_void_v<WorkTag>) {                     \
+        m_functor(static_cast<Args &&>(args)...);                  \
+      } else {                                                     \
+        m_functor(WorkTag(), static_cast<Args &&>(args)...);       \
+      }                                                            \
+    }                                                              \
   }
-};
+
+KOKKOS_IMPL_ACC_FUNCTOR_ADAPTER(worker);
+KOKKOS_IMPL_ACC_FUNCTOR_ADAPTER(seq);
+
+#undef KOKKOS_IMPL_ACC_FUNCTOR_ADAPTER
 
 }  // namespace Kokkos::Experimental::Impl
 

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_MDRange.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_MDRange.hpp
@@ -640,7 +640,9 @@ template <class Functor, class... Traits>
 class Kokkos::Impl::ParallelFor<Functor, Kokkos::MDRangePolicy<Traits...>,
                                 Kokkos::Experimental::OpenACC> {
   using Policy = MDRangePolicy<Traits...>;
-  Kokkos::Experimental::Impl::FunctorAdapter<Functor, Policy> m_functor;
+  Kokkos::Experimental::Impl::FunctorAdapter<
+      Functor, Policy, Kokkos::Experimental::Impl::RoutineClause::seq>
+      m_functor;
   Policy m_policy;
 
  public:

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_Range.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_Range.hpp
@@ -78,7 +78,9 @@ template <class Functor, class... Traits>
 class Kokkos::Impl::ParallelFor<Functor, Kokkos::RangePolicy<Traits...>,
                                 Kokkos::Experimental::OpenACC> {
   using Policy = Kokkos::RangePolicy<Traits...>;
-  Kokkos::Experimental::Impl::FunctorAdapter<Functor, Policy> m_functor;
+  Kokkos::Experimental::Impl::FunctorAdapter<
+      Functor, Policy, Kokkos::Experimental::Impl::RoutineClause::seq>
+      m_functor;
   Policy m_policy;
   using ScheduleType = Kokkos::Experimental::Impl::OpenACCScheduleType<Policy>;
 

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_Team.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelFor_Team.hpp
@@ -31,7 +31,9 @@ class Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>,
  private:
   using Policy = Kokkos::Impl::TeamPolicyInternal<Kokkos::Experimental::OpenACC,
                                                   Properties...>;
-  Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy> m_functor;
+  Kokkos::Experimental::Impl::FunctorAdapter<
+      FunctorType, Policy, Kokkos::Experimental::Impl::RoutineClause::seq>
+      m_functor;
   using Member = typename Policy::member_type;
 
   const Policy m_policy;
@@ -130,7 +132,8 @@ class Kokkos::Impl::ParallelFor<FunctorType, Kokkos::TeamPolicy<Properties...>,
  private:
   using Policy = Kokkos::Impl::TeamPolicyInternal<Kokkos::Experimental::OpenACC,
                                                   Properties...>;
-  Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy> m_functor;
+  Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy, worker>
+      m_functor;
   using Member = typename Policy::member_type;
 
   const Policy m_policy;

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_MDRange.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_MDRange.hpp
@@ -76,7 +76,9 @@ class Kokkos::Impl::ParallelReduce<CombinedFunctorReducerType,
     reducer.init(&val);
 
     Kokkos::Experimental::Impl::OpenACCParallelReduceMDRangeHelper(
-        Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy>(
+        Kokkos::Experimental::Impl::FunctorAdapter<
+            FunctorType, Policy,
+            Kokkos::Experimental::Impl::RoutineClause::seq>(
             m_functor_reducer.get_functor()),
         std::conditional_t<
             std::is_same_v<FunctorType, typename ReducerType::functor_type>,

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_Range.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_Range.hpp
@@ -74,7 +74,9 @@ class Kokkos::Impl::ParallelReduce<CombinedFunctorReducerType,
     reducer.init(&val);
 
     Kokkos::Experimental::Impl::OpenACCParallelReduceHelper(
-        Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy>(
+        Kokkos::Experimental::Impl::FunctorAdapter<
+            FunctorType, Policy,
+            Kokkos::Experimental::Impl::RoutineClause::seq>(
             m_functor_reducer.get_functor()),
         std::conditional_t<
             std::is_same_v<FunctorType, typename ReducerType::functor_type>,

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_Team.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelReduce_Team.hpp
@@ -21,6 +21,14 @@
 #include <OpenACC/Kokkos_OpenACC_FunctorAdapter.hpp>
 #include <OpenACC/Kokkos_OpenACC_Macros.hpp>
 
+#ifdef KOKKOS_ENABLE_OPENACC_COLLAPSE_HIERARCHICAL_CONSTRUCTS
+#define KOKKOS_IMPL_OPENACC_LOOP_CLAUSE \
+  Kokkos::Experimental::Impl::RoutineClause::seq
+#else
+#define KOKKOS_IMPL_OPENACC_LOOP_CLAUSE \
+  Kokkos::Experimental::Impl::RoutineClause::worker
+#endif
+
 //----------------------------------------------------------------------------
 //----------------------------------------------------------------------------
 // Hierarchical Parallelism -> Team level implementation
@@ -67,7 +75,8 @@ class Kokkos::Impl::ParallelReduce<CombinedFunctorReducerType,
     reducer.init(&tmp);
 
     Kokkos::Experimental::Impl::OpenACCParallelReduceTeamHelper(
-        Kokkos::Experimental::Impl::FunctorAdapter<FunctorType, Policy>(
+        Kokkos::Experimental::Impl::FunctorAdapter<
+            FunctorType, Policy, KOKKOS_IMPL_OPENACC_LOOP_CLAUSE>(
             m_functor_reducer.get_functor()),
         std::conditional_t<
             std::is_same_v<FunctorType, typename ReducerType::functor_type>,

diff --git a/core/src/OpenACC/Kokkos_OpenACC_ParallelScan_Range.hpp b/core/src/OpenACC/Kokkos_OpenACC_ParallelScan_Range.hpp
@@ -63,8 +63,9 @@ class Kokkos::Impl::ParallelScan<Functor, Kokkos::RangePolicy<Traits...>,
     } else {
       chunk_size = default_scan_chunk_size;
     }
-    const Kokkos::Experimental::Impl::FunctorAdapter<Functor, Policy> functor(
-        m_functor);
+    const Kokkos::Experimental::Impl::FunctorAdapter<
+        Functor, Policy, Kokkos::Experimental::Impl::RoutineClause::seq>
+        functor(m_functor);
     const IndexType N        = end - begin;
     const IndexType n_chunks = (N + chunk_size - 1) / chunk_size;
     Kokkos::View<ValueType*, Kokkos::Experimental::OpenACCSpace> chunk_values(