Thread loop optimizations RAJA launch

include/RAJA/pattern/launch/launch_context_policy.hpp

@@ -0,0 +1,124 @@
+/*!
+ ******************************************************************************
+ *
+ * \file
+ *
+ * \brief   RAJA header file containing template types of RAJA::LaunchContextT
+ *
+ ******************************************************************************
+ */
+
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+// Copyright (c) 2016-25, Lawrence Livermore National Security, LLC
+// and RAJA project contributors. See the RAJA/LICENSE file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
+
+#ifndef RAJA_pattern_context_policy_HPP
+#define RAJA_pattern_context_policy_HPP
+
+namespace RAJA
+{
+
+template<typename LaunchContextPolicy>
+class LaunchContextT;
+
+class LaunchContextDefaultPolicy;
+
+#if defined(RAJA_CUDA_ACTIVE) || defined(RAJA_HIP_ACTIVE)
+class LaunchContextDim3Policy;
+#endif
+
+namespace detail
+{
+
+
+template<typename T, typename = void>
+struct has_single_call_operator : std::false_type
+{};
+
+template<typename T>
+struct has_single_call_operator<
+    T,
+    std::enable_if_t<
+        !std::is_same_v<decltype(&std::decay_t<T>::operator()), void>>>
+    : std::true_type
+{};
+
+template<typename T>
+struct function_traits
+{};
+
+template<typename R, typename... Args>
+struct function_traits<R(Args...)>
+{
+  using result_type                  = R;
+  static constexpr std::size_t arity = sizeof...(Args);
+
+  template<std::size_t N>
+  struct arg
+  {
+    static_assert(N < arity, "argument index out of range");
+    using type = typename std::tuple_element<N, std::tuple<Args...>>::type;
+  };
+};
+
+template<typename R, typename... Args>
+struct function_traits<R (*)(Args...)> : function_traits<R(Args...)>
+{};
+
+template<typename R, typename... Args>
+struct function_traits<R (&)(Args...)> : function_traits<R(Args...)>
+{};
+
+template<typename C, typename R, typename... Args>
+struct function_traits<R (C::*)(Args...) const> : function_traits<R(Args...)>
+{
+  using functional_type = C;
+};
+
+template<typename C, typename R, typename... Args>
+struct function_traits<R (C::*)(Args...)> : function_traits<R(Args...)>
+{
+  using functional_type = C;
+};
+
+template<typename T,
+         bool hasCallOp = has_single_call_operator<std::decay_t<T>>::value>
+struct functional_traits : function_traits<std::decay_t<T>>
+{};
+
+template<typename T>
+struct functional_traits<T, true>
+    : function_traits<decltype(&std::decay_t<T>::operator())>
+{};
+
+template<typename T, typename = void>
+struct has_arg0 : std::false_type
+{};
+
+template<typename T>
+struct has_arg0<T,
+                typename std::enable_if_t<!std::is_same_v<
+                    typename functional_traits<T>::template arg<0>::type,
+                    void>>> : std::true_type
+{};
+
+template<typename T, bool HasArg0 = has_arg0<T>::value>
+struct launch_context_type
+{
+  using type = LaunchContextT<LaunchContextDefaultPolicy>;
+};
+
+template<typename T>
+struct launch_context_type<T, true>
+{
+  using type = typename functional_traits<T>::template arg<0>::type;
+};
+
+
+}  // namespace detail
+
+}  // namespace RAJA
+#endif

include/RAJA/pattern/launch/launch_core.hpp

@@ -20,6 +20,7 @@
 
 #include "RAJA/config.hpp"
 #include "RAJA/internal/get_platform.hpp"
+#include "RAJA/pattern/launch/launch_context_policy.hpp"
 #include "RAJA/util/StaticLayout.hpp"
 #include "RAJA/util/macros.hpp"
 #include "RAJA/util/plugins.hpp"
@@ -176,21 +177,20 @@ struct LaunchParams
   Threads apply(Threads const& a) { return (threads = a); }
 };
 
-class LaunchContext
+class LaunchContextBase
 {
 public:
   // Bump style allocator used to
   // get memory from the pool
   size_t shared_mem_offset;
-
   void* shared_mem_ptr;
 
 #if defined(RAJA_SYCL_ACTIVE)
   // SGS ODR issue
   mutable ::sycl::nd_item<3>* itm;
 #endif
 
-  RAJA_HOST_DEVICE LaunchContext()
+  RAJA_HOST_DEVICE LaunchContextBase()
       : shared_mem_offset(0),
         shared_mem_ptr(nullptr)
   {}
@@ -209,20 +209,6 @@ class LaunchContext
     return static_cast<T*>(mem_ptr);
   }
 
-  /*
-  //Odd dependecy with atomics is breaking CI builds
-  template<typename T, size_t DIM, typename IDX_T=RAJA::Index_type, ptrdiff_t
-  z_stride=DIM-1, typename arg, typename... args> RAJA_HOST_DEVICE auto
-  getSharedMemoryView(size_t bytes, arg idx, args... idxs)
-  {
-    T * mem_ptr = &((T*) shared_mem_ptr)[shared_mem_offset];
-
-    shared_mem_offset += bytes*sizeof(T);
-    return RAJA::View<T, RAJA::Layout<DIM, IDX_T, z_stride>>(mem_ptr, idx,
-  idxs...);
-  }
-  */
-
   RAJA_HOST_DEVICE void releaseSharedMemory()
   {
     // On the cpu/gpu we want to restart the count
@@ -243,6 +229,39 @@ class LaunchContext
   }
 };
 
+template<>
+class LaunchContextT<LaunchContextDefaultPolicy> : public LaunchContextBase
+{
+public:
+  static constexpr bool hasDim3 = false;
+
+  using LaunchContextBase::LaunchContextBase;
+};
+
+// Preserve backwards compatibility
+using LaunchContext = LaunchContextT<LaunchContextDefaultPolicy>;
+
+#if defined(RAJA_CUDA_ACTIVE) || defined(RAJA_HIP_ACTIVE)
+template<>
+class LaunchContextT<LaunchContextDim3Policy> : public LaunchContextBase
+{
+public:
+  static constexpr bool hasDim3 = true;
+
+  dim3 thread_id;
+  dim3 block_dim;
+
+  LaunchContextT() : LaunchContextBase(), thread_id(), block_dim() {}
+
+  RAJA_DEVICE
+  LaunchContextT(dim3 thread, dim3 block)
+      : LaunchContextBase(),
+        thread_id(thread),
+        block_dim(block)
+  {}
+};
+#endif
+
 template<typename LAUNCH_POLICY>
 struct LaunchExecute;