45 #ifndef KOKKOS_LOGICALSPACES_HPP 46 #define KOKKOS_LOGICALSPACES_HPP 48 #include <Kokkos_Macros.hpp> 49 #include <Kokkos_Core_fwd.hpp> 50 #include <Kokkos_ScratchSpace.hpp> 51 #include <impl/Kokkos_MemorySpace.hpp> 52 #include <impl/Kokkos_Error.hpp> 53 #include <impl/Kokkos_SharedAlloc.hpp> 54 #include <impl/Kokkos_Profiling.hpp> 58 struct DefaultMemorySpaceNamer {
59 static constexpr
const char* get_name() {
60 return "DefaultLogicalMemorySpaceName";
64 struct LogicalSpaceSharesAccess {
65 struct shared_access {};
66 struct no_shared_access {};
74 template <
class BaseSpace,
class DefaultBaseExecutionSpace = void,
75 class Namer = DefaultMemorySpaceNamer,
76 class SharesAccessWithBase = LogicalSpaceSharesAccess::shared_access>
77 class LogicalMemorySpace {
78 #ifdef KOKKOS_ENABLE_OPENMPTARGET 83 static_assert(!std::is_same<BaseSpace, BaseSpace>::value,
84 "Can't use LogicalMemorySpaces in an OpenMPTarget build, we're " 85 "debugging memory issues");
89 using memory_space = LogicalMemorySpace<BaseSpace, DefaultBaseExecutionSpace,
90 Namer, SharesAccessWithBase>;
91 using size_type =
typename BaseSpace::size_type;
101 typename std::conditional<std::is_void<DefaultBaseExecutionSpace>::value,
102 typename BaseSpace::execution_space,
103 DefaultBaseExecutionSpace>::type;
105 using device_type = Kokkos::Device<execution_space, memory_space>;
109 template <
typename... Args>
113 void*
allocate(
const size_t arg_alloc_size)
const {
114 return allocate(
"[unlabeled]", arg_alloc_size);
116 void*
allocate(
const char* arg_label,
const size_t arg_alloc_size,
117 const size_t arg_logical_size = 0)
const {
118 return impl_allocate(arg_label, arg_alloc_size, arg_logical_size);
123 const size_t arg_alloc_size)
const {
124 deallocate(
"[unlabeled]", arg_alloc_ptr, arg_alloc_size);
126 void deallocate(
const char* arg_label,
void*
const arg_alloc_ptr,
127 const size_t arg_alloc_size,
128 const size_t arg_logical_size = 0)
const {
129 impl_deallocate(arg_label, arg_alloc_ptr, arg_alloc_size, arg_logical_size);
133 constexpr
static const char*
name() {
return Namer::get_name(); }
136 BaseSpace underlying_space;
137 template <
class,
class,
class,
class>
139 friend class Kokkos::Impl::SharedAllocationRecord<
memory_space, void>;
141 void* impl_allocate(
const char* arg_label,
const size_t arg_alloc_size,
142 const size_t arg_logical_size = 0,
143 Kokkos::Tools::SpaceHandle arg_handle =
144 Kokkos::Tools::make_space_handle(
name()))
const {
145 return underlying_space.impl_allocate(arg_label, arg_alloc_size,
146 arg_logical_size, arg_handle);
148 void impl_deallocate(
const char* arg_label,
void*
const arg_alloc_ptr,
149 const size_t arg_alloc_size,
150 const size_t arg_logical_size = 0,
151 const Kokkos::Tools::SpaceHandle arg_handle =
152 Kokkos::Tools::make_space_handle(
name()))
const {
153 underlying_space.impl_deallocate(arg_label, arg_alloc_ptr, arg_alloc_size,
154 arg_logical_size, arg_handle);
166 template <
typename BaseSpace,
typename DefaultBaseExecutionSpace,
class Namer,
168 struct MemorySpaceAccess<
169 Kokkos::Experimental::LogicalMemorySpace<
170 BaseSpace, DefaultBaseExecutionSpace, Namer,
171 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>,
173 enum { assignable = MemorySpaceAccess<BaseSpace, OtherSpace>::assignable };
174 enum { accessible = MemorySpaceAccess<BaseSpace, OtherSpace>::accessible };
175 enum { deepcopy = MemorySpaceAccess<BaseSpace, OtherSpace>::deepcopy };
178 template <
typename BaseSpace,
typename DefaultBaseExecutionSpace,
class Namer,
180 struct MemorySpaceAccess<
182 Kokkos::Experimental::LogicalMemorySpace<
183 BaseSpace, DefaultBaseExecutionSpace, Namer,
184 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>> {
185 enum { assignable = MemorySpaceAccess<OtherSpace, BaseSpace>::assignable };
186 enum { accessible = MemorySpaceAccess<OtherSpace, BaseSpace>::accessible };
187 enum { deepcopy = MemorySpaceAccess<OtherSpace, BaseSpace>::deepcopy };
190 template <
typename BaseSpace,
typename DefaultBaseExecutionSpace,
class Namer>
191 struct MemorySpaceAccess<
192 Kokkos::Experimental::LogicalMemorySpace<
193 BaseSpace, DefaultBaseExecutionSpace, Namer,
194 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>,
196 BaseSpace, DefaultBaseExecutionSpace, Namer,
197 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>> {
198 enum { assignable =
true };
199 enum { accessible =
true };
200 enum { deepcopy =
true };
212 template <
class BaseSpace,
class DefaultBaseExecutionSpace,
class Namer,
213 class SharesAccessSemanticsWithBase>
214 class SharedAllocationRecord<
Kokkos::Experimental::LogicalMemorySpace<
215 BaseSpace, DefaultBaseExecutionSpace, Namer,
216 SharesAccessSemanticsWithBase>,
217 void> :
public SharedAllocationRecord<void, void> {
221 DefaultBaseExecutionSpace, Namer,
222 SharesAccessSemanticsWithBase>;
223 using RecordBase = SharedAllocationRecord<void, void>;
225 SharedAllocationRecord(
const SharedAllocationRecord&) =
delete;
226 SharedAllocationRecord& operator=(
const SharedAllocationRecord&) =
delete;
228 static void deallocate(RecordBase* arg_rec) {
229 delete static_cast<SharedAllocationRecord*
>(arg_rec);
232 #ifdef KOKKOS_ENABLE_DEBUG 235 static RecordBase s_root_record;
238 const SpaceType m_space;
241 ~SharedAllocationRecord() {
242 m_space.deallocate(RecordBase::m_alloc_ptr->m_label,
243 SharedAllocationRecord<void, void>::m_alloc_ptr,
244 SharedAllocationRecord<void, void>::m_alloc_size,
245 (SharedAllocationRecord<void, void>::m_alloc_size -
246 sizeof(SharedAllocationHeader)));
248 SharedAllocationRecord() =
default;
250 SharedAllocationRecord(
251 const SpaceType& arg_space,
const std::string& arg_label,
252 const size_t arg_alloc_size,
253 const RecordBase::function_type arg_dealloc = &deallocate)
254 : SharedAllocationRecord<void, void>(
255 #ifdef KOKKOS_ENABLE_DEBUG
256 &SharedAllocationRecord<SpaceType, void>::s_root_record,
258 Impl::checked_allocation_with_header(arg_space, arg_label,
260 sizeof(SharedAllocationHeader) + arg_alloc_size, arg_dealloc),
263 RecordBase::m_alloc_ptr->m_record =
264 static_cast<SharedAllocationRecord<void, void>*
>(
this);
266 strncpy(RecordBase::m_alloc_ptr->m_label, arg_label.c_str(),
267 SharedAllocationHeader::maximum_label_length - 1);
269 RecordBase::m_alloc_ptr
270 ->m_label[SharedAllocationHeader::maximum_label_length - 1] =
'\0';
274 inline std::string get_label()
const {
275 return std::string(RecordBase::head()->m_label);
277 KOKKOS_INLINE_FUNCTION
static SharedAllocationRecord* allocate(
278 const SpaceType& arg_space,
const std::string& arg_label,
279 const size_t arg_alloc_size) {
280 #if defined(KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST) 281 return new SharedAllocationRecord(arg_space, arg_label, arg_alloc_size);
285 (void)arg_alloc_size;
286 return (SharedAllocationRecord*)
nullptr;
291 static void* allocate_tracked(
const SpaceType& arg_space,
292 const std::string& arg_label,
293 const size_t arg_alloc_size) {
294 if (!arg_alloc_size)
return (
void*)
nullptr;
296 SharedAllocationRecord*
const r =
297 allocate(arg_space, arg_label, arg_alloc_size);
299 RecordBase::increment(r);
305 static void* reallocate_tracked(
void*
const arg_alloc_ptr,
306 const size_t arg_alloc_size) {
307 SharedAllocationRecord*
const r_old = get_record(arg_alloc_ptr);
308 SharedAllocationRecord*
const r_new =
309 allocate(r_old->m_space, r_old->get_label(), arg_alloc_size);
311 Kokkos::Impl::DeepCopy<SpaceType, SpaceType>(
312 r_new->data(), r_old->data(), std::min(r_old->size(), r_new->size()));
314 RecordBase::increment(r_new);
315 RecordBase::decrement(r_old);
317 return r_new->data();
320 static void deallocate_tracked(
void*
const arg_alloc_ptr) {
321 if (arg_alloc_ptr !=
nullptr) {
322 SharedAllocationRecord*
const r = get_record(arg_alloc_ptr);
324 RecordBase::decrement(r);
328 static SharedAllocationRecord* get_record(
void* alloc_ptr) {
329 using Header = SharedAllocationHeader;
330 using RecordHost = SharedAllocationRecord<SpaceType, void>;
332 SharedAllocationHeader
const*
const head =
333 alloc_ptr ? Header::get_header(alloc_ptr)
334 : (SharedAllocationHeader*)nullptr;
335 RecordHost*
const record =
336 head ?
static_cast<RecordHost*
>(head->m_record) : (RecordHost*)
nullptr;
338 if (!alloc_ptr || record->m_alloc_ptr != head) {
339 Kokkos::Impl::throw_runtime_exception(std::string(
340 "Kokkos::Impl::SharedAllocationRecord< LogicalMemorySpace<> , " 341 "void >::get_record ERROR"));
346 #ifdef KOKKOS_ENABLE_DEBUG 347 static void print_records(std::ostream& s,
const SpaceType&,
348 bool detail =
false) {
349 SharedAllocationRecord<void, void>::print_host_accessible_records(
350 s,
"HostSpace", &s_root_record, detail);
353 static void print_records(std::ostream&,
const SpaceType&,
354 bool detail =
false) {
356 throw_runtime_exception(
357 "SharedAllocationRecord<HostSpace>::print_records only works " 358 "with KOKKOS_ENABLE_DEBUG enabled");
362 #ifdef KOKKOS_ENABLE_DEBUG 365 template <
class BaseSpace,
class DefaultBaseExecutionSpace,
class Namer,
366 class SharesAccessSemanticsWithBase>
367 SharedAllocationRecord<void, void>
369 BaseSpace, DefaultBaseExecutionSpace, Namer,
370 SharesAccessSemanticsWithBase>,
371 void>::s_root_record;
384 template <
class Namer,
class BaseSpace,
class DefaultBaseExecutionSpace,
385 class SharesAccess,
class ExecutionSpace>
386 struct DeepCopy<
Kokkos::Experimental::LogicalMemorySpace<
387 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
389 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
391 DeepCopy(
void* dst,
void* src,
size_t n) {
392 DeepCopy<BaseSpace, BaseSpace, ExecutionSpace>(dst, src, n);
394 DeepCopy(
const ExecutionSpace& exec,
void* dst,
void* src,
size_t n) {
395 DeepCopy<BaseSpace, BaseSpace, ExecutionSpace>(exec, dst, src, n);
399 template <
class Namer,
class BaseSpace,
class DefaultBaseExecutionSpace,
400 class SharesAccess,
class ExecutionSpace,
class SourceSpace>
401 struct DeepCopy<SourceSpace,
402 Kokkos::Experimental::LogicalMemorySpace<
403 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
405 DeepCopy(
void* dst,
void* src,
size_t n) {
406 DeepCopy<SourceSpace, BaseSpace, ExecutionSpace>(dst, src, n);
408 DeepCopy(
const ExecutionSpace& exec,
void* dst,
void* src,
size_t n) {
409 DeepCopy<SourceSpace, BaseSpace, ExecutionSpace>(exec, dst, src, n);
413 template <
class Namer,
class BaseSpace,
class DefaultBaseExecutionSpace,
414 class SharesAccess,
class ExecutionSpace,
class DestinationSpace>
415 struct DeepCopy<
Kokkos::Experimental::LogicalMemorySpace<
416 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
417 DestinationSpace, ExecutionSpace> {
418 DeepCopy(
void* dst,
void* src,
size_t n) {
419 DeepCopy<BaseSpace, DestinationSpace, ExecutionSpace>(dst, src, n);
421 DeepCopy(
const ExecutionSpace& exec,
void* dst,
void* src,
size_t n) {
422 DeepCopy<BaseSpace, DestinationSpace, ExecutionSpace>(exec, dst, src, n);
428 #endif // KOKKOS_LOGICALSPACES_HPP
LogicalMemorySpace< BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccessWithBase > memory_space
Tag this class as a kokkos memory space.
static constexpr const char * name()
Return Name of the MemorySpace.
void * allocate(const size_t arg_alloc_size) const
Allocate untracked memory in the space.
typename std::conditional< std::is_void< DefaultBaseExecutionSpace >::value, typename BaseSpace::execution_space, DefaultBaseExecutionSpace >::type execution_space
Default execution space for this memory space.
void deallocate(void *const arg_alloc_ptr, const size_t arg_alloc_size) const
Deallocate untracked memory in the space.
LogicalMemorySpace is a space that is identical to another space, but differentiable by name and temp...