Kokkos Core Kernels Package  Version of the Day
Kokkos_GraphNode.hpp
1 /*
2 //@HEADER
3 // ************************************************************************
4 //
5 // Kokkos v. 3.0
6 // Copyright (2020) National Technology & Engineering
7 // Solutions of Sandia, LLC (NTESS).
8 //
9 // Under the terms of Contract DE-NA0003525 with NTESS,
10 // the U.S. Government retains certain rights in this software.
11 //
12 // Redistribution and use in source and binary forms, with or without
13 // modification, are permitted provided that the following conditions are
14 // met:
15 //
16 // 1. Redistributions of source code must retain the above copyright
17 // notice, this list of conditions and the following disclaimer.
18 //
19 // 2. Redistributions in binary form must reproduce the above copyright
20 // notice, this list of conditions and the following disclaimer in the
21 // documentation and/or other materials provided with the distribution.
22 //
23 // 3. Neither the name of the Corporation nor the names of the
24 // contributors may be used to endorse or promote products derived from
25 // this software without specific prior written permission.
26 //
27 // THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
28 // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
29 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
30 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
31 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
32 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
33 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
34 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
35 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
36 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
37 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
38 //
39 // Questions? Contact Christian R. Trott (crtrott@sandia.gov)
40 //
41 // ************************************************************************
42 //@HEADER
43 */
44 
45 #ifndef KOKKOS_KOKKOS_GRAPHNODE_HPP
46 #define KOKKOS_KOKKOS_GRAPHNODE_HPP
47 
48 #include <Kokkos_Macros.hpp>
49 
50 #include <impl/Kokkos_Error.hpp> // contract macros
51 
52 #include <Kokkos_Core_fwd.hpp>
53 #include <Kokkos_Graph_fwd.hpp>
54 #include <impl/Kokkos_GraphImpl_fwd.hpp>
55 #include <Kokkos_Parallel_Reduce.hpp>
56 #include <impl/Kokkos_GraphImpl_Utilities.hpp>
57 #include <impl/Kokkos_GraphImpl.hpp> // GraphAccess
58 
59 #include <memory> // std::shared_ptr
60 
61 namespace Kokkos {
62 namespace Experimental {
63 
64 template <class ExecutionSpace, class Kernel /*= TypeErasedTag*/,
65  class Predecessor /*= TypeErasedTag*/>
66 class GraphNodeRef {
67  //----------------------------------------------------------------------------
68  // <editor-fold desc="template parameter constraints"> {{{2
69 
70  // Note: because of these assertions, instantiating this class template is not
71  // intended to be SFINAE-safe, so do validation before you instantiate.
72 
73 // WORKAROUND Could not get it to compile with IBM XL V16.1.1
74 #ifndef KOKKOS_COMPILER_IBM
75  static_assert(
76  std::is_same<Predecessor, TypeErasedTag>::value ||
77  Kokkos::Impl::is_specialization_of<Predecessor, GraphNodeRef>::value,
78  "Invalid predecessor template parameter given to GraphNodeRef");
79 #endif
80 
81  static_assert(
82  Kokkos::is_execution_space<ExecutionSpace>::value,
83  "Invalid execution space template parameter given to GraphNodeRef");
84 
85  static_assert(std::is_same<Predecessor, TypeErasedTag>::value ||
86  Kokkos::Impl::is_graph_kernel<Kernel>::value,
87  "Invalid kernel template parameter given to GraphNodeRef");
88 
89  static_assert(!Kokkos::Impl::is_more_type_erased<Kernel, Predecessor>::value,
90  "The kernel of a graph node can't be more type-erased than the "
91  "predecessor");
92 
93  // </editor-fold> end template parameter constraints }}}2
94  //----------------------------------------------------------------------------
95 
96  public:
97  //----------------------------------------------------------------------------
98  // <editor-fold desc="public member types"> {{{2
99 
100  using execution_space = ExecutionSpace;
101  using graph_kernel = Kernel;
102  using graph_predecessor = Predecessor;
103 
104  // </editor-fold> end public member types }}}2
105  //----------------------------------------------------------------------------
106 
107  private:
108  //----------------------------------------------------------------------------
109  // <editor-fold desc="Friends"> {{{2
110 
111  template <class, class, class>
112  friend class GraphNodeRef;
113  friend struct Kokkos::Impl::GraphAccess;
114 
115  // </editor-fold> end Friends }}}2
116  //----------------------------------------------------------------------------
117 
118  //----------------------------------------------------------------------------
119  // <editor-fold desc="Private Data Members"> {{{2
120 
121  using graph_impl_t = Kokkos::Impl::GraphImpl<ExecutionSpace>;
122  std::weak_ptr<graph_impl_t> m_graph_impl;
123 
124  // TODO @graphs figure out if we can get away with a weak reference here?
125  // GraphNodeRef instances shouldn't be stored by users outside
126  // of the create_graph closure, and so if we restructure things
127  // slightly, we could make it so that the graph owns the
128  // node_impl_t instance and this only holds a std::weak_ptr to
129  // it.
130  using node_impl_t =
131  Kokkos::Impl::GraphNodeImpl<ExecutionSpace, Kernel, Predecessor>;
132  std::shared_ptr<node_impl_t> m_node_impl;
133 
134  // </editor-fold> end Private Data Members }}}2
135  //----------------------------------------------------------------------------
136 
137  //----------------------------------------------------------------------------
138  // <editor-fold desc="Implementation detail accessors"> {{{2
139 
140  // Internally, use shallow constness
141  node_impl_t& get_node_impl() const { return *m_node_impl.get(); }
142  std::shared_ptr<node_impl_t> const& get_node_ptr() const& {
143  return m_node_impl;
144  }
145  std::shared_ptr<node_impl_t> get_node_ptr() && {
146  return std::move(m_node_impl);
147  }
148  std::weak_ptr<graph_impl_t> get_graph_weak_ptr() const {
149  return m_graph_impl;
150  }
151 
152  // </editor-fold> end Implementation detail accessors }}}2
153  //----------------------------------------------------------------------------
154 
155  // TODO kernel name propagation and exposure
156 
157  template <class NextKernelDeduced>
158  auto _then_kernel(NextKernelDeduced&& arg_kernel) const {
159  // readability note:
160  // std::remove_cvref_t<NextKernelDeduced> is a specialization of
161  // Kokkos::Impl::GraphNodeKernelImpl:
162  static_assert(Kokkos::Impl::is_specialization_of<
163  Kokkos::Impl::remove_cvref_t<NextKernelDeduced>,
164  Kokkos::Impl::GraphNodeKernelImpl>::value,
165  "Kokkos internal error");
166 
167  auto graph_ptr = m_graph_impl.lock();
168  KOKKOS_EXPECTS(bool(graph_ptr))
169 
170  using next_kernel_t = Kokkos::Impl::remove_cvref_t<NextKernelDeduced>;
171 
172  using return_t = GraphNodeRef<ExecutionSpace, next_kernel_t, GraphNodeRef>;
173 
174  auto rv = Kokkos::Impl::GraphAccess::make_graph_node_ref(
175  m_graph_impl,
176  Kokkos::Impl::GraphAccess::make_node_shared_ptr<
177  typename return_t::node_impl_t>(
178  m_node_impl->execution_space_instance(),
179  Kokkos::Impl::_graph_node_kernel_ctor_tag{},
180  (NextKernelDeduced &&) arg_kernel,
181  // *this is the predecessor
182  Kokkos::Impl::_graph_node_predecessor_ctor_tag{}, *this));
183 
184  // Add the node itself to the backend's graph data structure, now that
185  // everything is set up.
186  graph_ptr->add_node(rv.m_node_impl);
187  // Add the predecessaor we stored in the constructor above in the backend's
188  // data structure, now that everything is set up.
189  graph_ptr->add_predecessor(rv.m_node_impl, *this);
190  KOKKOS_ENSURES(bool(rv.m_node_impl))
191  return rv;
192  }
193 
194  //----------------------------------------------------------------------------
195  // <editor-fold desc="Private constructors"> {{{2
196 
197  GraphNodeRef(std::weak_ptr<graph_impl_t> arg_graph_impl,
198  std::shared_ptr<node_impl_t> arg_node_impl)
199  : m_graph_impl(std::move(arg_graph_impl)),
200  m_node_impl(std::move(arg_node_impl)) {}
201 
202  // </editor-fold> end Private constructors }}}2
203  //----------------------------------------------------------------------------
204 
205  public:
206  //----------------------------------------------------------------------------
207  // <editor-fold desc="Constructors, destructors, and assignment"> {{{2
208 
209  //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
210  // <editor-fold desc="rule of 6 ctors"> {{{3
211 
212  // Copyable and movable (basically just shared_ptr semantics
213  GraphNodeRef() noexcept = default;
214  GraphNodeRef(GraphNodeRef const&) = default;
215  GraphNodeRef(GraphNodeRef&&) noexcept = default;
216  GraphNodeRef& operator=(GraphNodeRef const&) = default;
217  GraphNodeRef& operator=(GraphNodeRef&&) noexcept = default;
218  ~GraphNodeRef() = default;
219 
220  // </editor-fold> end rule of 6 ctors }}}3
221  //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
222 
223  //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
224  // <editor-fold desc="Type-erasing converting ctor and assignment"> {{{3
225 
226  template <
227  class OtherKernel, class OtherPredecessor,
228  typename std::enable_if_t<
229  // Not a copy/move constructor
230  !std::is_same<GraphNodeRef, GraphNodeRef<execution_space, OtherKernel,
231  OtherPredecessor>>::value &&
232  // must be an allowed type erasure of the kernel
233  Kokkos::Impl::is_compatible_type_erasure<OtherKernel,
234  graph_kernel>::value &&
235  // must be an allowed type erasure of the predecessor
236  Kokkos::Impl::is_compatible_type_erasure<
237  OtherPredecessor, graph_predecessor>::value,
238  int> = 0>
239  /* implicit */
240  GraphNodeRef(
241  GraphNodeRef<execution_space, OtherKernel, OtherPredecessor> const& other)
242  : m_graph_impl(other.m_graph_impl), m_node_impl(other.m_node_impl) {}
243 
244  // Note: because this is an implicit conversion (as is supposed to be the
245  // case with most type-erasing wrappers like this), we don't also need
246  // a converting assignment operator.
247 
248  // </editor-fold> end Type-erasing converting ctor and assignment }}}3
249  //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
250 
251  // </editor-fold> end Constructors, destructors, and assignment }}}2
252  //----------------------------------------------------------------------------
253 
254  //----------------------------------------------------------------------------
255  // <editor-fold desc="then_parallel_for"> {{{2
256 
257  template <
258  class Policy, class Functor,
259  typename std::enable_if<
260  // equivalent to:
261  // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
262  is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
263  // --------------------
264  int>::type = 0>
265  auto then_parallel_for(std::string arg_name, Policy&& arg_policy,
266  Functor&& functor) const {
267  //----------------------------------------
268  KOKKOS_EXPECTS(!m_graph_impl.expired())
269  KOKKOS_EXPECTS(bool(m_node_impl))
270  // TODO @graph restore this expectation once we add comparability to space
271  // instances
272  // KOKKOS_EXPECTS(
273  // arg_policy.space() == m_graph_impl->get_execution_space());
274 
275  // needs to static assert constraint: DataParallelFunctor<Functor>
276 
277  using policy_t = Kokkos::Impl::remove_cvref_t<Policy>;
278  // constraint check: same execution space type (or defaulted, maybe?)
279  static_assert(
280  std::is_same<typename policy_t::execution_space,
281  execution_space>::value,
282  // TODO @graph make defaulted execution space work
283  //|| policy_t::execution_space_is_defaulted,
284  "Execution Space mismatch between execution policy and graph");
285 
286  auto policy = Experimental::require((Policy &&) arg_policy,
287  Kokkos::Impl::KernelInGraphProperty{});
288 
289  using next_policy_t = decltype(policy);
290  using next_kernel_t =
291  Kokkos::Impl::GraphNodeKernelImpl<ExecutionSpace, next_policy_t,
292  std::decay_t<Functor>,
293  Kokkos::ParallelForTag>;
294  return this->_then_kernel(next_kernel_t{std::move(arg_name), policy.space(),
295  (Functor &&) functor,
296  (Policy &&) policy});
297  }
298 
299  template <
300  class Policy, class Functor,
301  typename std::enable_if<
302  // equivalent to:
303  // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
304  is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
305  // --------------------
306  int>::type = 0>
307  auto then_parallel_for(Policy&& policy, Functor&& functor) const {
308  // needs to static assert constraint: DataParallelFunctor<Functor>
309  return this->then_parallel_for("", (Policy &&) policy,
310  (Functor &&) functor);
311  }
312 
313  template <class Functor>
314  auto then_parallel_for(std::string name, std::size_t n,
315  Functor&& functor) const {
316  // needs to static assert constraint: DataParallelFunctor<Functor>
317  return this->then_parallel_for(std::move(name),
319  (Functor &&) functor);
320  }
321 
322  template <class Functor>
323  auto then_parallel_for(std::size_t n, Functor&& functor) const {
324  // needs to static assert constraint: DataParallelFunctor<Functor>
325  return this->then_parallel_for("", n, (Functor &&) functor);
326  }
327 
328  // </editor-fold> end then_parallel_for }}}2
329  //----------------------------------------------------------------------------
330 
331  //----------------------------------------------------------------------------
332  // <editor-fold desc="then_parallel_reduce"> {{{2
333 
334  template <
335  class Policy, class Functor, class ReturnType,
336  typename std::enable_if<
337  // equivalent to:
338  // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
339  is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
340  // --------------------
341  int>::type = 0>
342  auto then_parallel_reduce(std::string arg_name, Policy&& arg_policy,
343  Functor&& functor,
344  ReturnType&& return_value) const {
345  auto graph_impl_ptr = m_graph_impl.lock();
346  KOKKOS_EXPECTS(bool(graph_impl_ptr))
347  KOKKOS_EXPECTS(bool(m_node_impl))
348  // TODO @graph restore this expectation once we add comparability to space
349  // instances
350  // KOKKOS_EXPECTS(
351  // arg_policy.space() == m_graph_impl->get_execution_space());
352 
353  // needs static assertion of constraint:
354  // DataParallelReductionFunctor<Functor, ReturnType>
355 
356  using policy_t = typename std::remove_cv<
357  typename std::remove_reference<Policy>::type>::type;
358  static_assert(
359  std::is_same<typename policy_t::execution_space,
360  execution_space>::value,
361  // TODO @graph make defaulted execution space work
362  // || policy_t::execution_space_is_defaulted,
363  "Execution Space mismatch between execution policy and graph");
364 
365  // This is also just an expectation, but it's one that we expect the user
366  // to interact with (even in release mode), so we should throw an exception
367  // with an explanation rather than just doing a contract assertion.
368  // We can't static_assert this because of the way that Reducers store
369  // whether or not they point to a View as a runtime boolean rather than part
370  // of the type.
371  if (Kokkos::Impl::parallel_reduce_needs_fence(
372  graph_impl_ptr->get_execution_space(), return_value)) {
373  Kokkos::Impl::throw_runtime_exception(
374  "Parallel reductions in graphs can't operate on Reducers that "
375  "reference a scalar because they can't complete synchronously. Use a "
376  "Kokkos::View instead and keep in mind the result will only be "
377  "available once the graph is submitted (or in tasks that depend on "
378  "this one).");
379  }
380 
381  //----------------------------------------
382  // This is a disaster, but I guess it's not a my disaster to fix right now
383  using return_type_remove_cvref = typename std::remove_cv<
384  typename std::remove_reference<ReturnType>::type>::type;
385  static_assert(Kokkos::is_view<return_type_remove_cvref>::value ||
386  Kokkos::is_reducer<return_type_remove_cvref>::value,
387  "Output argument to parallel reduce in a graph must be a "
388  "View or a Reducer");
389  using return_type =
390  // Yes, you do really have to do this...
391  std::conditional_t<Kokkos::is_reducer<return_type_remove_cvref>::value,
392  return_type_remove_cvref,
393  const return_type_remove_cvref>;
394  using functor_type = Kokkos::Impl::remove_cvref_t<Functor>;
395  // see Kokkos_Parallel_Reduce.hpp for how these details are used there;
396  // we're just doing the same thing here
397  using return_value_adapter =
398  Kokkos::Impl::ParallelReduceReturnValue<void, return_type,
399  functor_type>;
400  using functor_adaptor = Kokkos::Impl::ParallelReduceFunctorType<
401  functor_type, Policy, typename return_value_adapter::value_type,
402  execution_space>;
403  // End of Kokkos reducer disaster
404  //----------------------------------------
405 
406  auto policy = Experimental::require((Policy &&) arg_policy,
407  Kokkos::Impl::KernelInGraphProperty{});
408 
409  using next_policy_t = decltype(policy);
410  using next_kernel_t = Kokkos::Impl::GraphNodeKernelImpl<
411  ExecutionSpace, next_policy_t, typename functor_adaptor::functor_type,
412  Kokkos::ParallelReduceTag, typename return_value_adapter::reducer_type>;
413 
414  return this->_then_kernel(next_kernel_t{
415  std::move(arg_name), graph_impl_ptr->get_execution_space(),
416  (Functor &&) functor, (Policy &&) policy,
417  return_value_adapter::return_value(return_value, functor)});
418  }
419 
420  template <
421  class Policy, class Functor, class ReturnType,
422  typename std::enable_if<
423  // equivalent to:
424  // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
425  is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
426  // --------------------
427  int>::type = 0>
428  auto then_parallel_reduce(Policy&& arg_policy, Functor&& functor,
429  ReturnType&& return_value) const {
430  return this->then_parallel_reduce("", (Policy &&) arg_policy,
431  (Functor &&) functor,
432  (ReturnType &&) return_value);
433  }
434 
435  template <class Functor, class ReturnType>
436  auto then_parallel_reduce(std::string label,
437  typename execution_space::size_type idx_end,
438  Functor&& functor,
439  ReturnType&& return_value) const {
440  return this->then_parallel_reduce(
441  std::move(label), Kokkos::RangePolicy<execution_space>{0, idx_end},
442  (Functor &&) functor, (ReturnType &&) return_value);
443  }
444 
445  template <class Functor, class ReturnType>
446  auto then_parallel_reduce(typename execution_space::size_type idx_end,
447  Functor&& functor,
448  ReturnType&& return_value) const {
449  return this->then_parallel_reduce("", idx_end, (Functor &&) functor,
450  (ReturnType &&) return_value);
451  }
452 
453  // </editor-fold> end then_parallel_reduce }}}2
454  //----------------------------------------------------------------------------
455 
456  // TODO @graph parallel scan, deep copy, etc.
457 };
458 
459 } // end namespace Experimental
460 } // end namespace Kokkos
461 
462 #endif // KOKKOS_KOKKOS_GRAPHNODE_HPP
ReturnType
Execution policy for work over a range of an integral type.
Definition: dummy.cpp:3