45 #ifndef KOKKOS_SCRATCHSPACE_HPP
46 #define KOKKOS_SCRATCHSPACE_HPP
49 #include <Kokkos_Core_fwd.hpp>
50 #include <Kokkos_Concepts.hpp>
59 template <
class ExecSpace>
62 is_execution_space<ExecSpace>::value,
63 "Instantiating ScratchMemorySpace on non-execution-space type.");
71 mutable char* m_iter_L0;
73 mutable char* m_iter_L1;
76 mutable int m_multiplier;
78 mutable int m_default_level;
83 enum { MASK = ALIGN - 1 };
88 typedef ExecSpace execution_space;
90 typedef Kokkos::Device<execution_space, memory_space>
device_type;
92 typedef typename ExecSpace::array_layout array_layout;
93 typedef typename ExecSpace::size_type size_type;
95 static constexpr
const char* name() {
return "ScratchMemorySpace"; }
97 template <
typename IntType>
98 KOKKOS_INLINE_FUNCTION
static IntType align(
const IntType& size) {
99 return (size + MASK) & ~MASK;
102 template <
typename IntType>
103 KOKKOS_INLINE_FUNCTION
void* get_shmem(
const IntType& size,
104 int level = -1)
const {
105 if (level == -1) level = m_default_level;
107 void* tmp = m_iter_L0 + m_offset * align(size);
108 if (m_end_L0 < (m_iter_L0 += align(size) * m_multiplier)) {
109 m_iter_L0 -= align(size) * m_multiplier;
115 "ScratchMemorySpace<...>::get_shmem: Failed to allocate "
116 "%ld byte(s); remaining capacity is %ld byte(s)\n",
117 long(size),
long(m_end_L0 - m_iter_L0));
118 #endif // KOKKOS_DEBUG
123 void* tmp = m_iter_L1 + m_offset * align(size);
124 if (m_end_L1 < (m_iter_L1 += align(size) * m_multiplier)) {
125 m_iter_L1 -= align(size) * m_multiplier;
131 "ScratchMemorySpace<...>::get_shmem: Failed to allocate "
132 "%ld byte(s); remaining capacity is %ld byte(s)\n",
133 long(size),
long(m_end_L1 - m_iter_L1));
134 #endif // KOKKOS_DEBUG
141 KOKKOS_INLINE_FUNCTION
142 void* get_shmem_aligned(
const ptrdiff_t size,
const ptrdiff_t alignment,
143 int level = -1)
const {
144 if (level == -1) level = m_default_level;
146 char* previous = m_iter_L0;
147 const ptrdiff_t missalign = size_t(m_iter_L0) % alignment;
148 if (missalign) m_iter_L0 += alignment - missalign;
150 void* tmp = m_iter_L0 + m_offset * size;
151 if (m_end_L0 < (m_iter_L0 += size * m_multiplier)) {
152 m_iter_L0 = previous;
158 "ScratchMemorySpace<...>::get_shmem: Failed to allocate "
159 "%ld byte(s); remaining capacity is %ld byte(s)\n",
160 long(size),
long(m_end_L0 - m_iter_L0));
161 #endif // KOKKOS_DEBUG
166 char* previous = m_iter_L1;
167 const ptrdiff_t missalign = size_t(m_iter_L1) % alignment;
168 if (missalign) m_iter_L1 += alignment - missalign;
170 void* tmp = m_iter_L1 + m_offset * size;
171 if (m_end_L1 < (m_iter_L1 += size * m_multiplier)) {
172 m_iter_L1 = previous;
178 "ScratchMemorySpace<...>::get_shmem: Failed to allocate "
179 "%ld byte(s); remaining capacity is %ld byte(s)\n",
180 long(size),
long(m_end_L1 - m_iter_L1));
181 #endif // KOKKOS_DEBUG
188 template <
typename IntType>
189 KOKKOS_INLINE_FUNCTION ScratchMemorySpace(
void* ptr_L0,
190 const IntType& size_L0,
191 void* ptr_L1 =
nullptr,
192 const IntType& size_L1 = 0)
193 : m_iter_L0((char*)ptr_L0),
194 m_end_L0(m_iter_L0 + size_L0),
195 m_iter_L1((char*)ptr_L1),
196 m_end_L1(m_iter_L1 + size_L1),
199 m_default_level(0) {}
201 KOKKOS_INLINE_FUNCTION
202 const ScratchMemorySpace& set_team_thread_mode(
const int& level,
203 const int& multiplier,
204 const int& offset)
const {
205 m_default_level = level;
206 m_multiplier = multiplier;
Kokkos::Device< execution_space, memory_space > device_type
This execution space preferred device_type.
Scratch memory space associated with an execution space.
ScratchMemorySpace memory_space
Tag this class as a memory space.