45 #ifndef KOKKOS_SCRATCHSPACE_HPP 46 #define KOKKOS_SCRATCHSPACE_HPP 50 #include <Kokkos_Core_fwd.hpp> 51 #include <Kokkos_Concepts.hpp> 60 template <
class ExecSpace>
63 is_execution_space<ExecSpace>::value,
64 "Instantiating ScratchMemorySpace on non-execution-space type.");
72 mutable char* m_iter_L0 =
nullptr;
73 mutable char* m_iter_L1 =
nullptr;
74 char* m_end_L0 =
nullptr;
75 char* m_end_L1 =
nullptr;
77 mutable int m_multiplier = 0;
78 mutable int m_offset = 0;
79 mutable int m_default_level = 0;
81 enum { MASK = ALIGN - 1 };
86 using execution_space = ExecSpace;
88 using device_type = Kokkos::Device<execution_space, memory_space>;
90 using array_layout =
typename ExecSpace::array_layout;
91 using size_type =
typename ExecSpace::size_type;
93 static constexpr
const char* name() {
return "ScratchMemorySpace"; }
95 template <
typename IntType>
96 KOKKOS_INLINE_FUNCTION
static IntType align(
const IntType& size) {
97 return (size + MASK) & ~MASK;
100 template <
typename IntType>
101 KOKKOS_INLINE_FUNCTION
void* get_shmem(
const IntType& size,
102 int level = -1)
const {
103 return get_shmem_common<
false>(size, 1, level);
106 template <
typename IntType>
107 KOKKOS_INLINE_FUNCTION
void* get_shmem_aligned(
const IntType& size,
108 const ptrdiff_t alignment,
109 int level = -1)
const {
110 return get_shmem_common<
true>(size, alignment, level);
114 template <
bool aligned,
typename IntType>
115 KOKKOS_INLINE_FUNCTION
void* get_shmem_common(
const IntType& size,
116 const ptrdiff_t alignment,
117 int level = -1)
const {
118 if (level == -1) level = m_default_level;
119 auto& m_iter = (level == 0) ? m_iter_L0 : m_iter_L1;
120 auto& m_end = (level == 0) ? m_end_L0 : m_end_L1;
121 char* previous = m_iter;
122 const ptrdiff_t missalign = size_t(m_iter) % alignment;
123 if (missalign) m_iter += alignment - missalign;
125 void* tmp = m_iter + m_offset * (aligned ? size : align(size));
126 if (m_end < (m_iter += (aligned ? size : align(size)) * m_multiplier)) {
128 #ifdef KOKKOS_ENABLE_DEBUG 132 KOKKOS_IMPL_DO_NOT_USE_PRINTF(
133 "ScratchMemorySpace<...>::get_shmem: Failed to allocate " 134 "%ld byte(s); remaining capacity is %ld byte(s)\n",
135 long(size),
long(m_end - m_iter));
136 #endif // KOKKOS_ENABLE_DEBUG 143 KOKKOS_DEFAULTED_FUNCTION
144 ScratchMemorySpace() =
default;
146 template <
typename IntType>
147 KOKKOS_INLINE_FUNCTION ScratchMemorySpace(
void* ptr_L0,
148 const IntType& size_L0,
149 void* ptr_L1 =
nullptr,
150 const IntType& size_L1 = 0)
151 : m_iter_L0((char*)ptr_L0),
152 m_iter_L1((char*)ptr_L1),
153 m_end_L0((char*)ptr_L0 + size_L0),
154 m_end_L1((char*)ptr_L1 + size_L1),
157 m_default_level(0) {}
159 KOKKOS_INLINE_FUNCTION
160 const ScratchMemorySpace& set_team_thread_mode(
const int& level,
161 const int& multiplier,
162 const int& offset)
const {
163 m_default_level = level;
164 m_multiplier = multiplier;
Scratch memory space associated with an execution space.
Kokkos::Device< execution_space, memory_space > device_type
This execution space preferred device_type.