44 #ifndef KOKKOS_SCRATCHSPACE_HPP 45 #define KOKKOS_SCRATCHSPACE_HPP 48 #include <Kokkos_Core_fwd.hpp> 49 #include <impl/Kokkos_Tags.hpp> 58 template<
class ExecSpace >
60 static_assert (Impl::is_execution_space<ExecSpace>::value,
"Instantiating ScratchMemorySpace on non-execution-space type.");
69 mutable char * m_iter_L0 ;
71 mutable char * m_iter_L1 ;
75 mutable int m_multiplier;
77 mutable int m_default_level;
82 enum { MASK = ALIGN - 1 };
88 typedef ExecSpace execution_space ;
92 typedef typename ExecSpace::array_layout array_layout ;
93 typedef typename ExecSpace::size_type size_type ;
95 template<
typename IntType >
96 KOKKOS_INLINE_FUNCTION
static 97 IntType align(
const IntType & size )
98 {
return ( size + MASK ) & ~MASK ; }
100 template<
typename IntType >
101 KOKKOS_INLINE_FUNCTION
102 void* get_shmem (
const IntType& size,
int level = -1)
const {
104 level = m_default_level;
106 void* tmp = m_iter_L0 + m_offset * align (size);
107 if (m_end_L0 < (m_iter_L0 += align (size) * m_multiplier)) {
108 m_iter_L0 -= align (size) * m_multiplier;
109 #ifdef KOKKOS_HAVE_DEBUG 113 printf (
"ScratchMemorySpace<...>::get_shmem: Failed to allocate " 114 "%ld byte(s); remaining capacity is %ld byte(s)\n",
long(size),
115 long(m_end_L0-m_iter_L0));
116 #endif // KOKKOS_HAVE_DEBUG 121 void* tmp = m_iter_L1 + m_offset * align (size);
122 if (m_end_L1 < (m_iter_L1 += align (size) * m_multiplier)) {
123 m_iter_L1 -= align (size) * m_multiplier;
124 #ifdef KOKKOS_HAVE_DEBUG 128 printf (
"ScratchMemorySpace<...>::get_shmem: Failed to allocate " 129 "%ld byte(s); remaining capacity is %ld byte(s)\n",
long(size),
130 long(m_end_L1-m_iter_L1));
131 #endif // KOKKOS_HAVE_DEBUG 139 template<
typename IntType >
140 KOKKOS_INLINE_FUNCTION
141 ScratchMemorySpace(
void * ptr_L0 ,
const IntType & size_L0 ,
void * ptr_L1 = NULL ,
const IntType & size_L1 = 0)
142 : m_iter_L0( (char *) ptr_L0 )
143 , m_end_L0( m_iter_L0 + size_L0 )
144 , m_iter_L1( (char *) ptr_L1 )
145 , m_end_L1( m_iter_L1 + size_L1 )
148 , m_default_level( 0 )
151 KOKKOS_INLINE_FUNCTION
152 const ScratchMemorySpace& set_team_thread_mode(
const int& level,
const int& multiplier,
const int& offset)
const {
153 m_default_level = level;
154 m_multiplier = multiplier;
Scratch memory space associated with an execution space.
Memory space for main process and CPU execution spaces.
ScratchMemorySpace memory_space
Tag this class as a memory space.
Kokkos::Device< execution_space, memory_space > device_type
This execution space preferred device_type.