Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Future.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_FUTURE_HPP
23#define KOKKOS_FUTURE_HPP
24
25//----------------------------------------------------------------------------
26
27#include <Kokkos_Macros.hpp>
28#if defined(KOKKOS_ENABLE_TASKDAG)
29
30#include <Kokkos_Core_fwd.hpp>
31#include <Kokkos_TaskScheduler_fwd.hpp>
32//----------------------------------------------------------------------------
33
34#include <impl/Kokkos_TaskQueue.hpp>
35#include <impl/Kokkos_TaskResult.hpp>
36#include <impl/Kokkos_TaskBase.hpp>
37#include <Kokkos_Atomic.hpp>
38
39#include <Kokkos_Concepts.hpp> // is_space
40
41//----------------------------------------------------------------------------
42//----------------------------------------------------------------------------
43
44namespace Kokkos {
45
46// For now, hack this in as a partial specialization
47// TODO @tasking @cleanup Make this the "normal" class template and make the old
48// code the specialization
49template <typename ValueType, typename ExecutionSpace, typename QueueType>
50class BasicFuture<ValueType, SimpleTaskScheduler<ExecutionSpace, QueueType>> {
51 public:
52 using value_type = ValueType;
53 using execution_space = ExecutionSpace;
54 using scheduler_type = SimpleTaskScheduler<ExecutionSpace, QueueType>;
55 using queue_type = typename scheduler_type::task_queue_type;
56
57 private:
58 template <class, class>
59 friend class SimpleTaskScheduler;
60 template <class, class>
61 friend class BasicFuture;
62
63 using task_base_type = typename scheduler_type::task_base_type;
64 using task_queue_type = typename scheduler_type::task_queue_type;
65
66 using task_queue_traits = typename scheduler_type::task_queue_traits;
67 using task_scheduling_info_type =
68 typename scheduler_type::task_scheduling_info_type;
69
70 using result_storage_type = Impl::TaskResultStorage<
71 ValueType,
72 Impl::SchedulingInfoStorage<Impl::RunnableTaskBase<task_queue_traits>,
73 task_scheduling_info_type>>;
74
75 OwningRawPtr<task_base_type> m_task = nullptr;
76
77 KOKKOS_INLINE_FUNCTION
78 explicit BasicFuture(task_base_type* task) : m_task(task) {
79 // Note: reference count starts at 2 to account for initial increment
80 // TODO @tasking @minor DSH verify reference count here and/or encapsulate
81 // starting reference count closer to here
82 }
83
84 public:
85 KOKKOS_INLINE_FUNCTION
86 BasicFuture() noexcept : m_task(nullptr) {}
87
88 KOKKOS_INLINE_FUNCTION
89 BasicFuture(BasicFuture&& rhs) noexcept : m_task(std::move(rhs.m_task)) {
90 rhs.m_task = nullptr;
91 }
92
93 KOKKOS_INLINE_FUNCTION
94 BasicFuture(BasicFuture const& rhs)
95 // : m_task(rhs.m_task)
96 : m_task(nullptr) {
97 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
98 if (m_task) m_task->increment_reference_count();
99 }
100
101 KOKKOS_INLINE_FUNCTION
102 BasicFuture& operator=(BasicFuture&& rhs) noexcept {
103 if (m_task != rhs.m_task) {
104 clear();
105 // m_task = std::move(rhs.m_task);
106 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
107 // rhs.m_task reference count is unchanged, since this is a move
108 } else {
109 // They're the same, but this is a move, so 1 fewer references now
110 rhs.clear();
111 }
112 rhs.m_task = nullptr;
113 return *this;
114 }
115
116 KOKKOS_INLINE_FUNCTION
117 BasicFuture& operator=(BasicFuture const& rhs) {
118 if (m_task != rhs.m_task) {
119 clear();
120 // m_task = rhs.m_task;
121 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
122 }
123 if (m_task != nullptr) {
124 m_task->increment_reference_count();
125 }
126 return *this;
127 }
128
129 //----------------------------------------
130
131 template <class T, class S>
132 KOKKOS_INLINE_FUNCTION BasicFuture(
133 BasicFuture<T, S>&& rhs) noexcept // NOLINT(google-explicit-constructor)
134 : m_task(std::move(rhs.m_task)) {
135 static_assert(std::is_void<scheduler_type>::value ||
136 std::is_same<scheduler_type, S>::value,
137 "Moved Futures must have the same scheduler");
138
139 static_assert(
140 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
141 "Moved Futures must have the same value_type");
142
143 // reference counts are unchanged, since this is a move
144 rhs.m_task = nullptr;
145 }
146
147 template <class T, class S>
148 KOKKOS_INLINE_FUNCTION BasicFuture(
149 BasicFuture<T, S> const& rhs) // NOLINT(google-explicit-constructor)
150 //: m_task(rhs.m_task)
151 : m_task(nullptr) {
152 static_assert(std::is_void<scheduler_type>::value ||
153 std::is_same<scheduler_type, S>::value,
154 "Copied Futures must have the same scheduler");
155
156 static_assert(
157 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
158 "Copied Futures must have the same value_type");
159
160 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
161 if (m_task) m_task->increment_reference_count();
162 }
163
164 template <class T, class S>
165 KOKKOS_INLINE_FUNCTION BasicFuture& operator=(BasicFuture<T, S> const& rhs) {
166 static_assert(std::is_void<scheduler_type>::value ||
167 std::is_same<scheduler_type, S>::value,
168 "Assigned Futures must have the same scheduler");
169
170 static_assert(
171 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
172 "Assigned Futures must have the same value_type");
173
174 if (m_task != rhs.m_task) {
175 clear();
176 // m_task = rhs.m_task;
177 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
178 if (m_task != nullptr) {
179 m_task->increment_reference_count();
180 }
181 }
182 return *this;
183 }
184
185 template <class T, class S>
186 KOKKOS_INLINE_FUNCTION BasicFuture& operator=(BasicFuture<T, S>&& rhs) {
187 static_assert(std::is_void<scheduler_type>::value ||
188 std::is_same<scheduler_type, S>::value,
189 "Assigned Futures must have the same scheduler");
190
191 static_assert(
192 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
193 "Assigned Futures must have the same value_type");
194
195 if (m_task != rhs.m_task) {
196 clear();
197 // m_task = std::move(rhs.m_task);
198 *static_cast<task_base_type* volatile*>(&m_task) = rhs.m_task;
199 // rhs.m_task reference count is unchanged, since this is a move
200 } else {
201 // They're the same, but this is a move, so 1 fewer references now
202 rhs.clear();
203 }
204 rhs.m_task = nullptr;
205 return *this;
206 }
207
208 KOKKOS_INLINE_FUNCTION
209 ~BasicFuture() noexcept { clear(); }
210
211 //----------------------------------------
212
213 KOKKOS_INLINE_FUNCTION
214 void clear() noexcept {
215 if (m_task) {
216 bool should_delete = m_task->decrement_and_check_reference_count();
217 if (should_delete) {
218 static_cast<task_queue_type*>(m_task->ready_queue_base_ptr())
219 ->deallocate(std::move(*m_task));
220 }
221 }
222 // m_task = nullptr;
223 *static_cast<task_base_type* volatile*>(&m_task) = nullptr;
224 }
225
226 KOKKOS_INLINE_FUNCTION
227 bool is_null() const noexcept { return m_task == nullptr; }
228
229 KOKKOS_INLINE_FUNCTION
230 bool is_ready() const noexcept {
231 return (m_task == nullptr) || m_task->wait_queue_is_consumed();
232 }
233
234 KOKKOS_INLINE_FUNCTION
235 const typename Impl::TaskResult<ValueType>::reference_type get() const {
236 KOKKOS_EXPECTS(is_ready());
237 return static_cast<result_storage_type*>(m_task)->value_reference();
238 // return Impl::TaskResult<ValueType>::get(m_task);
239 }
240};
241
243// OLD CODE
245
246template <typename ValueType, typename Scheduler>
247class BasicFuture {
248 private:
249 template <typename, typename>
250 friend class BasicTaskScheduler;
251 template <typename, typename>
252 friend class BasicFuture;
253 friend class Impl::TaskBase;
254 template <typename, typename, typename>
255 friend class Impl::Task;
256
257 //----------------------------------------
258
259 public:
260 //----------------------------------------
261
262 using scheduler_type = Scheduler;
263 using queue_type = typename scheduler_type::queue_type;
264 using execution_space = typename scheduler_type::execution_space;
265 using value_type = ValueType;
266
267 //----------------------------------------
268
269 private:
270 //----------------------------------------
271
272 using task_base = Impl::TaskBase;
273
274 task_base* m_task;
275
276 KOKKOS_INLINE_FUNCTION explicit BasicFuture(task_base* task)
277 : m_task(nullptr) {
278 if (task) queue_type::assign(&m_task, task);
279 }
280
281 //----------------------------------------
282
283 public:
284 //----------------------------------------
285
286 KOKKOS_INLINE_FUNCTION
287 bool is_null() const { return nullptr == m_task; }
288
289 KOKKOS_INLINE_FUNCTION
290 int reference_count() const {
291 return nullptr != m_task ? m_task->reference_count() : 0;
292 }
293
294 //----------------------------------------
295
296 KOKKOS_INLINE_FUNCTION
297 void clear() {
298 if (m_task) queue_type::assign(&m_task, nullptr);
299 }
300
301 //----------------------------------------
302
303 KOKKOS_INLINE_FUNCTION
304 ~BasicFuture() { clear(); }
305
306 //----------------------------------------
307
308 KOKKOS_INLINE_FUNCTION
309 BasicFuture() noexcept : m_task(nullptr) {}
310
311 KOKKOS_INLINE_FUNCTION
312 BasicFuture(BasicFuture&& rhs) noexcept : m_task(rhs.m_task) {
313 rhs.m_task = nullptr;
314 }
315
316 KOKKOS_INLINE_FUNCTION
317 BasicFuture(const BasicFuture& rhs) : m_task(nullptr) {
318 if (rhs.m_task) queue_type::assign(&m_task, rhs.m_task);
319 }
320
321 KOKKOS_INLINE_FUNCTION
322 BasicFuture& operator=(BasicFuture&& rhs) noexcept {
323 clear();
324 m_task = rhs.m_task;
325 rhs.m_task = nullptr;
326 return *this;
327 }
328
329 KOKKOS_INLINE_FUNCTION
330 BasicFuture& operator=(BasicFuture const& rhs) {
331 if (m_task || rhs.m_task) queue_type::assign(&m_task, rhs.m_task);
332 return *this;
333 }
334
335 //----------------------------------------
336
337 template <class T, class S>
338 KOKKOS_INLINE_FUNCTION BasicFuture(
339 BasicFuture<T, S>&& rhs) noexcept // NOLINT(google-explicit-constructor)
340 : m_task(rhs.m_task) {
341 static_assert(std::is_void<scheduler_type>::value ||
342 std::is_same<scheduler_type, S>::value,
343 "Assigned Futures must have the same scheduler");
344
345 static_assert(
346 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
347 "Assigned Futures must have the same value_type");
348
349 rhs.m_task = 0;
350 }
351
352 template <class T, class S>
353 KOKKOS_INLINE_FUNCTION BasicFuture(
354 BasicFuture<T, S> const& rhs) // NOLINT(google-explicit-constructor)
355 : m_task(nullptr) {
356 static_assert(std::is_void<scheduler_type>::value ||
357 std::is_same<scheduler_type, S>::value,
358 "Assigned Futures must have the same scheduler");
359
360 static_assert(
361 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
362 "Assigned Futures must have the same value_type");
363
364 if (rhs.m_task) queue_type::assign(&m_task, rhs.m_task);
365 }
366
367 template <class T, class S>
368 KOKKOS_INLINE_FUNCTION BasicFuture& operator=(BasicFuture<T, S> const& rhs) {
369 static_assert(std::is_void<scheduler_type>::value ||
370 std::is_same<scheduler_type, S>::value,
371 "Assigned Futures must have the same scheduler");
372
373 static_assert(
374 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
375 "Assigned Futures must have the same value_type");
376
377 if (m_task || rhs.m_task) queue_type::assign(&m_task, rhs.m_task);
378 return *this;
379 }
380
381 template <class T, class S>
382 KOKKOS_INLINE_FUNCTION BasicFuture& operator=(BasicFuture<T, S>&& rhs) {
383 static_assert(std::is_void<scheduler_type>::value ||
384 std::is_same<scheduler_type, S>::value,
385 "Assigned Futures must have the same scheduler");
386
387 static_assert(
388 std::is_void<value_type>::value || std::is_same<value_type, T>::value,
389 "Assigned Futures must have the same value_type");
390
391 clear();
392 m_task = rhs.m_task;
393 rhs.m_task = 0;
394 return *this;
395 }
396
397 //----------------------------------------
398
399 KOKKOS_INLINE_FUNCTION
400 int is_ready() const noexcept {
401 return (nullptr == m_task) ||
402 (reinterpret_cast<task_base*>(task_base::LockTag) == m_task->m_wait);
403 }
404
405 KOKKOS_INLINE_FUNCTION
406 const typename Impl::TaskResult<ValueType>::reference_type get() const {
407 if (nullptr == m_task) {
408 Kokkos::abort("Kokkos:::Future::get ERROR: is_null()");
409 }
410 return Impl::TaskResult<ValueType>::get(m_task);
411 }
412};
413
414// Is a Future with the given execution space
415template <typename, typename ExecSpace = void>
416struct is_future : public std::false_type {};
417
418template <typename ValueType, typename Scheduler, typename ExecSpace>
419struct is_future<BasicFuture<ValueType, Scheduler>, ExecSpace>
420 : std::bool_constant<
421 std::is_same<ExecSpace, typename Scheduler::execution_space>::value ||
422 std::is_void<ExecSpace>::value> {};
423
425// END OLD CODE
427
428namespace Impl {
429
430template <class Arg1, class Arg2>
431class ResolveFutureArgOrder {
432 private:
433 enum { Arg1_is_space = Kokkos::is_space<Arg1>::value };
434 enum { Arg2_is_space = Kokkos::is_space<Arg2>::value };
435 enum { Arg1_is_value = !Arg1_is_space && !std::is_void<Arg1>::value };
436 enum { Arg2_is_value = !Arg2_is_space && !std::is_void<Arg2>::value };
437
438 static_assert(!(Arg1_is_space && Arg2_is_space),
439 "Future cannot be given two spaces");
440
441 static_assert(!(Arg1_is_value && Arg2_is_value),
442 "Future cannot be given two value types");
443
444 using value_type =
445 std::conditional_t<Arg1_is_value, Arg1,
446 std::conditional_t<Arg2_is_value, Arg2, void>>;
447
448 using execution_space = typename std::conditional_t<
449 Arg1_is_space, Arg1,
450 std::conditional_t<Arg2_is_space, Arg2, void>>::execution_space;
451
452 public:
453 using type = BasicFuture<value_type, TaskScheduler<execution_space>>;
454};
455
456} // end namespace Impl
457
465template <class Arg1 = void, class Arg2 = void>
466using Future = typename Impl::ResolveFutureArgOrder<Arg1, Arg2>::type;
467
468} // namespace Kokkos
469
470//----------------------------------------------------------------------------
471//----------------------------------------------------------------------------
472
473#endif /* #if defined( KOKKOS_ENABLE_TASKDAG ) */
474#endif /* #ifndef KOKKOS_FUTURE */
Atomic functions.
A thread safe view to a bitset.