Kokkos Core Kernels Package  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Groups Pages
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>
19 static_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 
44 namespace 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
49 template <typename ValueType, typename ExecutionSpace, typename QueueType>
50 class 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 
246 template <typename ValueType, typename Scheduler>
247 class 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
415 template <typename, typename ExecSpace = void>
416 struct is_future : public std::false_type {};
417 
418 template <typename ValueType, typename Scheduler, typename ExecSpace>
419 struct 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 
428 namespace Impl {
429 
430 template <class Arg1, class Arg2>
431 class 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 
465 template <class Arg1 = void, class Arg2 = void>
466 using 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 */
bool is_null(const boost::shared_ptr< T > &p)
Atomic functions.