Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_LogicalSpaces.hpp
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16
17#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
18#include <Kokkos_Macros.hpp>
19static_assert(false,
20 "Including non-public Kokkos header files is not allowed.");
21#endif
22#ifndef KOKKOS_LOGICALSPACES_HPP
23#define KOKKOS_LOGICALSPACES_HPP
24
25#include <Kokkos_Macros.hpp>
26#include <Kokkos_Core_fwd.hpp>
27#include <Kokkos_ScratchSpace.hpp>
28#include <impl/Kokkos_MemorySpace.hpp>
29#include <impl/Kokkos_Error.hpp>
30#include <impl/Kokkos_SharedAlloc.hpp>
31#include <impl/Kokkos_Profiling.hpp>
32#include <cstring>
33namespace Kokkos {
34namespace Experimental {
35struct DefaultMemorySpaceNamer {
36 static constexpr const char* get_name() {
37 return "DefaultLogicalMemorySpaceName";
38 }
39};
40
41struct LogicalSpaceSharesAccess {
42 struct shared_access {};
43 struct no_shared_access {};
44};
45
51template <class BaseSpace, class DefaultBaseExecutionSpace = void,
52 class Namer = DefaultMemorySpaceNamer,
53 class SharesAccessWithBase = LogicalSpaceSharesAccess::shared_access>
55#ifdef KOKKOS_ENABLE_OPENMPTARGET
56 // [DZP] For some reason I don't yet know, using LogicalMemorySpaces
57 // inside an OpenMPTarget build causes errors in the
58 // SharedAllocationRecords of other types. This is my way of erroring
59 // a build if we instantiate a LogicalMemSpace in an OMPTarget build
60 static_assert(!std::is_same<BaseSpace, BaseSpace>::value,
61 "Can't use LogicalMemorySpaces in an OpenMPTarget build, we're "
62 "debugging memory issues");
63#endif
64 public:
68 using size_type = typename BaseSpace::size_type;
69
76
78 std::conditional_t<std::is_void<DefaultBaseExecutionSpace>::value,
79 typename BaseSpace::execution_space,
81
83
84 LogicalMemorySpace() = default;
85
86 template <typename... Args>
87 LogicalMemorySpace(Args&&... args) : underlying_space((Args &&) args...) {}
88
90 void* allocate(const size_t arg_alloc_size) const {
91 return allocate("[unlabeled]", arg_alloc_size);
92 }
93 void* allocate(const char* arg_label, const size_t arg_alloc_size,
94 const size_t arg_logical_size = 0) const {
95 return impl_allocate(arg_label, arg_alloc_size, arg_logical_size);
96 }
97
99 void deallocate(void* const arg_alloc_ptr,
100 const size_t arg_alloc_size) const {
101 deallocate("[unlabeled]", arg_alloc_ptr, arg_alloc_size);
102 }
103 void deallocate(const char* arg_label, void* const arg_alloc_ptr,
104 const size_t arg_alloc_size,
105 const size_t arg_logical_size = 0) const {
107 }
108
110 constexpr static const char* name() { return Namer::get_name(); }
111
112 private:
113 BaseSpace underlying_space;
114 template <class, class, class, class>
115 friend class LogicalMemorySpace;
116 friend class Kokkos::Impl::SharedAllocationRecord<memory_space, void>;
117
118 void* impl_allocate(const char* arg_label, const size_t arg_alloc_size,
119 const size_t arg_logical_size = 0,
120 Kokkos::Tools::SpaceHandle arg_handle =
121 Kokkos::Tools::make_space_handle(name())) const {
122 return underlying_space.impl_allocate(arg_label, arg_alloc_size,
124 }
125 void impl_deallocate(const char* arg_label, void* const arg_alloc_ptr,
126 const size_t arg_alloc_size,
127 const size_t arg_logical_size = 0,
128 const Kokkos::Tools::SpaceHandle arg_handle =
129 Kokkos::Tools::make_space_handle(name())) const {
130 underlying_space.impl_deallocate(arg_label, arg_alloc_ptr, arg_alloc_size,
132 }
133};
134} // namespace Experimental
135} // namespace Kokkos
136
137//----------------------------------------------------------------------------
138
139namespace Kokkos {
140
141namespace Impl {
142
143template <typename BaseSpace, typename DefaultBaseExecutionSpace, class Namer,
144 typename OtherSpace>
145struct MemorySpaceAccess<
146 Kokkos::Experimental::LogicalMemorySpace<
147 BaseSpace, DefaultBaseExecutionSpace, Namer,
148 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>,
149 OtherSpace> {
150 enum { assignable = MemorySpaceAccess<BaseSpace, OtherSpace>::assignable };
151 enum { accessible = MemorySpaceAccess<BaseSpace, OtherSpace>::accessible };
152 enum { deepcopy = MemorySpaceAccess<BaseSpace, OtherSpace>::deepcopy };
153};
154
155template <typename BaseSpace, typename DefaultBaseExecutionSpace, class Namer,
156 typename OtherSpace>
157struct MemorySpaceAccess<
158 OtherSpace,
159 Kokkos::Experimental::LogicalMemorySpace<
160 BaseSpace, DefaultBaseExecutionSpace, Namer,
161 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>> {
162 enum { assignable = MemorySpaceAccess<OtherSpace, BaseSpace>::assignable };
163 enum { accessible = MemorySpaceAccess<OtherSpace, BaseSpace>::accessible };
164 enum { deepcopy = MemorySpaceAccess<OtherSpace, BaseSpace>::deepcopy };
165};
166
167template <typename BaseSpace, typename DefaultBaseExecutionSpace, class Namer>
168struct MemorySpaceAccess<
169 Kokkos::Experimental::LogicalMemorySpace<
170 BaseSpace, DefaultBaseExecutionSpace, Namer,
171 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>,
173 BaseSpace, DefaultBaseExecutionSpace, Namer,
174 Kokkos::Experimental::LogicalSpaceSharesAccess::shared_access>> {
175 enum { assignable = true };
176 enum { accessible = true };
177 enum { deepcopy = true };
178};
179
180} // namespace Impl
181
182} // namespace Kokkos
183
184//----------------------------------------------------------------------------
185
186namespace Kokkos {
187
188namespace Impl {
189template <class BaseSpace, class DefaultBaseExecutionSpace, class Namer,
190 class SharesAccessSemanticsWithBase>
191class SharedAllocationRecord<Kokkos::Experimental::LogicalMemorySpace<
192 BaseSpace, DefaultBaseExecutionSpace, Namer,
193 SharesAccessSemanticsWithBase>,
194 void> : public SharedAllocationRecord<void, void> {
195 private:
196 using SpaceType =
198 DefaultBaseExecutionSpace, Namer,
199 SharesAccessSemanticsWithBase>;
200 using RecordBase = SharedAllocationRecord<void, void>;
201
202 SharedAllocationRecord(const SharedAllocationRecord&) = delete;
203 SharedAllocationRecord& operator=(const SharedAllocationRecord&) = delete;
204
205 static void deallocate(RecordBase* arg_rec) {
206 delete static_cast<SharedAllocationRecord*>(arg_rec);
207 }
208
209#ifdef KOKKOS_ENABLE_DEBUG
212 static RecordBase s_root_record;
213#endif
214
215 const SpaceType m_space;
216
217 protected:
218 ~SharedAllocationRecord() {
219 m_space.deallocate(RecordBase::m_alloc_ptr->m_label,
220 SharedAllocationRecord<void, void>::m_alloc_ptr,
221 SharedAllocationRecord<void, void>::m_alloc_size,
222 (SharedAllocationRecord<void, void>::m_alloc_size -
223 sizeof(SharedAllocationHeader)));
224 }
225 SharedAllocationRecord() = default;
226
227 template <typename ExecutionSpace>
228 SharedAllocationRecord(
229 const ExecutionSpace& /*exec_space*/, const SpaceType& arg_space,
230 const std::string& arg_label, const size_t arg_alloc_size,
231 const RecordBase::function_type arg_dealloc = &deallocate)
232 : SharedAllocationRecord(arg_space, arg_label, arg_alloc_size,
233 arg_dealloc) {}
234
235 SharedAllocationRecord(
236 const SpaceType& arg_space, const std::string& arg_label,
237 const size_t arg_alloc_size,
238 const RecordBase::function_type arg_dealloc = &deallocate)
239 : SharedAllocationRecord<void, void>(
240#ifdef KOKKOS_ENABLE_DEBUG
241 &SharedAllocationRecord<SpaceType, void>::s_root_record,
242#endif
243 Impl::checked_allocation_with_header(arg_space, arg_label,
244 arg_alloc_size),
245 sizeof(SharedAllocationHeader) + arg_alloc_size, arg_dealloc,
246 arg_label),
247 m_space(arg_space) {
248 // Fill in the Header information
249 RecordBase::m_alloc_ptr->m_record =
250 static_cast<SharedAllocationRecord<void, void>*>(this);
251
252 strncpy(RecordBase::m_alloc_ptr->m_label, arg_label.c_str(),
253 SharedAllocationHeader::maximum_label_length - 1);
254 // Set last element zero, in case c_str is too long
255 RecordBase::m_alloc_ptr
256 ->m_label[SharedAllocationHeader::maximum_label_length - 1] = '\0';
257 }
258
259 public:
260 inline std::string get_label() const {
261 return std::string(RecordBase::head()->m_label);
262 }
263 KOKKOS_INLINE_FUNCTION static SharedAllocationRecord* allocate(
264 const SpaceType& arg_space, const std::string& arg_label,
265 const size_t arg_alloc_size) {
266 KOKKOS_IF_ON_HOST((return new SharedAllocationRecord(arg_space, arg_label,
267 arg_alloc_size);))
268 KOKKOS_IF_ON_DEVICE(((void)arg_space; (void)arg_label; (void)arg_alloc_size;
269 return nullptr;))
270 }
271
273 static void* allocate_tracked(const SpaceType& arg_space,
274 const std::string& arg_label,
275 const size_t arg_alloc_size) {
276 if (!arg_alloc_size) return (void*)nullptr;
277
278 SharedAllocationRecord* const r =
279 allocate(arg_space, arg_label, arg_alloc_size);
280
281 RecordBase::increment(r);
282
283 return r->data();
284 }
285
287 static void* reallocate_tracked(void* const arg_alloc_ptr,
288 const size_t arg_alloc_size) {
289 SharedAllocationRecord* const r_old = get_record(arg_alloc_ptr);
290 SharedAllocationRecord* const r_new =
291 allocate(r_old->m_space, r_old->get_label(), arg_alloc_size);
292
293 Kokkos::Impl::DeepCopy<SpaceType, SpaceType>(
294 r_new->data(), r_old->data(), std::min(r_old->size(), r_new->size()));
295 Kokkos::fence(
296 "SharedAllocationRecord<Kokkos::Experimental::LogicalMemorySpace, "
297 "void>::reallocate_tracked: fence after copying data");
298
299 RecordBase::increment(r_new);
300 RecordBase::decrement(r_old);
301
302 return r_new->data();
303 }
305 static void deallocate_tracked(void* const arg_alloc_ptr) {
306 if (arg_alloc_ptr != nullptr) {
307 SharedAllocationRecord* const r = get_record(arg_alloc_ptr);
308
309 RecordBase::decrement(r);
310 }
311 }
312
313 static SharedAllocationRecord* get_record(void* alloc_ptr) {
314 using Header = SharedAllocationHeader;
315 using RecordHost = SharedAllocationRecord<SpaceType, void>;
316
317 SharedAllocationHeader const* const head =
318 alloc_ptr ? Header::get_header(alloc_ptr)
319 : (SharedAllocationHeader*)nullptr;
320 RecordHost* const record =
321 head ? static_cast<RecordHost*>(head->m_record) : (RecordHost*)nullptr;
322
323 if (!alloc_ptr || record->m_alloc_ptr != head) {
324 Kokkos::Impl::throw_runtime_exception(std::string(
325 "Kokkos::Impl::SharedAllocationRecord< LogicalMemorySpace<> , "
326 "void >::get_record ERROR"));
327 }
328
329 return record;
330 }
331#ifdef KOKKOS_ENABLE_DEBUG
332 static void print_records(std::ostream& s, const SpaceType&,
333 bool detail = false) {
334 SharedAllocationRecord<void, void>::print_host_accessible_records(
335 s, "HostSpace", &s_root_record, detail);
336 }
337#else
338 static void print_records(std::ostream&, const SpaceType&,
339 bool detail = false) {
340 (void)detail;
341 throw_runtime_exception(
342 "SharedAllocationRecord<HostSpace>::print_records only works "
343 "with KOKKOS_ENABLE_DEBUG enabled");
344 }
345#endif
346};
347#ifdef KOKKOS_ENABLE_DEBUG
350template <class BaseSpace, class DefaultBaseExecutionSpace, class Namer,
351 class SharesAccessSemanticsWithBase>
352SharedAllocationRecord<void, void>
353 SharedAllocationRecord<Kokkos::Experimental::LogicalMemorySpace<
354 BaseSpace, DefaultBaseExecutionSpace, Namer,
355 SharesAccessSemanticsWithBase>,
356 void>::s_root_record;
357#endif
358
359} // namespace Impl
360
361} // namespace Kokkos
362
363//----------------------------------------------------------------------------
364
365namespace Kokkos {
366
367namespace Impl {
368
369template <class Namer, class BaseSpace, class DefaultBaseExecutionSpace,
370 class SharesAccess, class ExecutionSpace>
371struct DeepCopy<Kokkos::Experimental::LogicalMemorySpace<
372 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
374 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
375 ExecutionSpace> {
376 DeepCopy(void* dst, void* src, size_t n) {
377 DeepCopy<BaseSpace, BaseSpace, ExecutionSpace>(dst, src, n);
378 }
379 DeepCopy(const ExecutionSpace& exec, void* dst, void* src, size_t n) {
380 DeepCopy<BaseSpace, BaseSpace, ExecutionSpace>(exec, dst, src, n);
381 }
382};
383
384template <class Namer, class BaseSpace, class DefaultBaseExecutionSpace,
385 class SharesAccess, class ExecutionSpace, class SourceSpace>
386struct DeepCopy<SourceSpace,
387 Kokkos::Experimental::LogicalMemorySpace<
388 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
389 ExecutionSpace> {
390 DeepCopy(void* dst, void* src, size_t n) {
391 DeepCopy<SourceSpace, BaseSpace, ExecutionSpace>(dst, src, n);
392 }
393 DeepCopy(const ExecutionSpace& exec, void* dst, void* src, size_t n) {
394 DeepCopy<SourceSpace, BaseSpace, ExecutionSpace>(exec, dst, src, n);
395 }
396};
397
398template <class Namer, class BaseSpace, class DefaultBaseExecutionSpace,
399 class SharesAccess, class ExecutionSpace, class DestinationSpace>
400struct DeepCopy<Kokkos::Experimental::LogicalMemorySpace<
401 BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccess>,
402 DestinationSpace, ExecutionSpace> {
403 DeepCopy(void* dst, void* src, size_t n) {
404 DeepCopy<BaseSpace, DestinationSpace, ExecutionSpace>(dst, src, n);
405 }
406 DeepCopy(const ExecutionSpace& exec, void* dst, void* src, size_t n) {
407 DeepCopy<BaseSpace, DestinationSpace, ExecutionSpace>(exec, dst, src, n);
408 }
409};
410} // namespace Impl
411
412} // namespace Kokkos
413#endif // KOKKOS_LOGICALSPACES_HPP
A thread safe view to a bitset.
LogicalMemorySpace is a space that is identical to another space, but differentiable by name and temp...
LogicalMemorySpace< BaseSpace, DefaultBaseExecutionSpace, Namer, SharesAccessWithBase > memory_space
Tag this class as a kokkos memory space.
void * allocate(const size_t arg_alloc_size) const
Allocate untracked memory in the space.
static constexpr const char * name()
Return Name of the MemorySpace.
std::conditional_t< std::is_void< DefaultBaseExecutionSpace >::value, typename BaseSpace::execution_space, DefaultBaseExecutionSpace > 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.