Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_GraphNode.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_KOKKOS_GRAPHNODE_HPP
23#define KOKKOS_KOKKOS_GRAPHNODE_HPP
24
25#include <Kokkos_Macros.hpp>
26
27#include <impl/Kokkos_Error.hpp> // contract macros
28
29#include <Kokkos_Core_fwd.hpp>
30#include <Kokkos_Graph_fwd.hpp>
31#include <impl/Kokkos_GraphImpl_fwd.hpp>
32#include <Kokkos_Parallel_Reduce.hpp>
33#include <impl/Kokkos_GraphImpl_Utilities.hpp>
34#include <impl/Kokkos_GraphImpl.hpp> // GraphAccess
35
36#include <memory> // std::shared_ptr
37
38namespace Kokkos {
39namespace Experimental {
40
41template <class ExecutionSpace, class Kernel /*= TypeErasedTag*/,
42 class Predecessor /*= TypeErasedTag*/>
43class GraphNodeRef {
44 //----------------------------------------------------------------------------
45 // <editor-fold desc="template parameter constraints"> {{{2
46
47 // Note: because of these assertions, instantiating this class template is not
48 // intended to be SFINAE-safe, so do validation before you instantiate.
49
50 static_assert(
51 std::is_same<Predecessor, TypeErasedTag>::value ||
52 Kokkos::Impl::is_specialization_of<Predecessor, GraphNodeRef>::value,
53 "Invalid predecessor template parameter given to GraphNodeRef");
54
55 static_assert(
57 "Invalid execution space template parameter given to GraphNodeRef");
58
59 static_assert(std::is_same<Predecessor, TypeErasedTag>::value ||
60 Kokkos::Impl::is_graph_kernel<Kernel>::value,
61 "Invalid kernel template parameter given to GraphNodeRef");
62
63 static_assert(!Kokkos::Impl::is_more_type_erased<Kernel, Predecessor>::value,
64 "The kernel of a graph node can't be more type-erased than the "
65 "predecessor");
66
67 // </editor-fold> end template parameter constraints }}}2
68 //----------------------------------------------------------------------------
69
70 public:
71 //----------------------------------------------------------------------------
72 // <editor-fold desc="public member types"> {{{2
73
74 using execution_space = ExecutionSpace;
75 using graph_kernel = Kernel;
76 using graph_predecessor = Predecessor;
77
78 // </editor-fold> end public member types }}}2
79 //----------------------------------------------------------------------------
80
81 private:
82 //----------------------------------------------------------------------------
83 // <editor-fold desc="Friends"> {{{2
84
85 template <class, class, class>
86 friend class GraphNodeRef;
87 friend struct Kokkos::Impl::GraphAccess;
88
89 // </editor-fold> end Friends }}}2
90 //----------------------------------------------------------------------------
91
92 //----------------------------------------------------------------------------
93 // <editor-fold desc="Private Data Members"> {{{2
94
95 using graph_impl_t = Kokkos::Impl::GraphImpl<ExecutionSpace>;
96 std::weak_ptr<graph_impl_t> m_graph_impl;
97
98 // TODO @graphs figure out if we can get away with a weak reference here?
99 // GraphNodeRef instances shouldn't be stored by users outside
100 // of the create_graph closure, and so if we restructure things
101 // slightly, we could make it so that the graph owns the
102 // node_impl_t instance and this only holds a std::weak_ptr to
103 // it.
104 using node_impl_t =
105 Kokkos::Impl::GraphNodeImpl<ExecutionSpace, Kernel, Predecessor>;
106 std::shared_ptr<node_impl_t> m_node_impl;
107
108 // </editor-fold> end Private Data Members }}}2
109 //----------------------------------------------------------------------------
110
111 //----------------------------------------------------------------------------
112 // <editor-fold desc="Implementation detail accessors"> {{{2
113
114 // Internally, use shallow constness
115 node_impl_t& get_node_impl() const { return *m_node_impl.get(); }
116 std::shared_ptr<node_impl_t> const& get_node_ptr() const& {
117 return m_node_impl;
118 }
119 std::shared_ptr<node_impl_t> get_node_ptr() && {
120 return std::move(m_node_impl);
121 }
122 std::weak_ptr<graph_impl_t> get_graph_weak_ptr() const {
123 return m_graph_impl;
124 }
125
126 // </editor-fold> end Implementation detail accessors }}}2
127 //----------------------------------------------------------------------------
128
129 // TODO kernel name propagation and exposure
130
131 template <class NextKernelDeduced>
132 auto _then_kernel(NextKernelDeduced&& arg_kernel) const {
133 // readability note:
134 // std::remove_cvref_t<NextKernelDeduced> is a specialization of
135 // Kokkos::Impl::GraphNodeKernelImpl:
136 static_assert(Kokkos::Impl::is_specialization_of<
137 Kokkos::Impl::remove_cvref_t<NextKernelDeduced>,
138 Kokkos::Impl::GraphNodeKernelImpl>::value,
139 "Kokkos internal error");
140
141 auto graph_ptr = m_graph_impl.lock();
142 KOKKOS_EXPECTS(bool(graph_ptr))
143
144 using next_kernel_t = Kokkos::Impl::remove_cvref_t<NextKernelDeduced>;
145
146 using return_t = GraphNodeRef<ExecutionSpace, next_kernel_t, GraphNodeRef>;
147
148 auto rv = Kokkos::Impl::GraphAccess::make_graph_node_ref(
149 m_graph_impl,
150 Kokkos::Impl::GraphAccess::make_node_shared_ptr<
151 typename return_t::node_impl_t>(
152 m_node_impl->execution_space_instance(),
153 Kokkos::Impl::_graph_node_kernel_ctor_tag{},
154 (NextKernelDeduced &&) arg_kernel,
155 // *this is the predecessor
156 Kokkos::Impl::_graph_node_predecessor_ctor_tag{}, *this));
157
158 // Add the node itself to the backend's graph data structure, now that
159 // everything is set up.
160 graph_ptr->add_node(rv.m_node_impl);
161 // Add the predecessaor we stored in the constructor above in the backend's
162 // data structure, now that everything is set up.
163 graph_ptr->add_predecessor(rv.m_node_impl, *this);
164 KOKKOS_ENSURES(bool(rv.m_node_impl))
165 return rv;
166 }
167
168 //----------------------------------------------------------------------------
169 // <editor-fold desc="Private constructors"> {{{2
170
171 GraphNodeRef(std::weak_ptr<graph_impl_t> arg_graph_impl,
172 std::shared_ptr<node_impl_t> arg_node_impl)
173 : m_graph_impl(std::move(arg_graph_impl)),
174 m_node_impl(std::move(arg_node_impl)) {}
175
176 // </editor-fold> end Private constructors }}}2
177 //----------------------------------------------------------------------------
178
179 public:
180 //----------------------------------------------------------------------------
181 // <editor-fold desc="Constructors, destructors, and assignment"> {{{2
182
183 //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
184 // <editor-fold desc="rule of 6 ctors"> {{{3
185
186 // Copyable and movable (basically just shared_ptr semantics
187 GraphNodeRef() noexcept = default;
188 GraphNodeRef(GraphNodeRef const&) = default;
189 GraphNodeRef(GraphNodeRef&&) noexcept = default;
190 GraphNodeRef& operator=(GraphNodeRef const&) = default;
191 GraphNodeRef& operator=(GraphNodeRef&&) noexcept = default;
192 ~GraphNodeRef() = default;
193
194 // </editor-fold> end rule of 6 ctors }}}3
195 //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
196
197 //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
198 // <editor-fold desc="Type-erasing converting ctor and assignment"> {{{3
199
200 template <
201 class OtherKernel, class OtherPredecessor,
202 std::enable_if_t<
203 // Not a copy/move constructor
204 !std::is_same<GraphNodeRef, GraphNodeRef<execution_space, OtherKernel,
205 OtherPredecessor>>::value &&
206 // must be an allowed type erasure of the kernel
207 Kokkos::Impl::is_compatible_type_erasure<OtherKernel,
208 graph_kernel>::value &&
209 // must be an allowed type erasure of the predecessor
210 Kokkos::Impl::is_compatible_type_erasure<
211 OtherPredecessor, graph_predecessor>::value,
212 int> = 0>
213 /* implicit */
214 GraphNodeRef(
215 GraphNodeRef<execution_space, OtherKernel, OtherPredecessor> const& other)
216 : m_graph_impl(other.m_graph_impl), m_node_impl(other.m_node_impl) {}
217
218 // Note: because this is an implicit conversion (as is supposed to be the
219 // case with most type-erasing wrappers like this), we don't also need
220 // a converting assignment operator.
221
222 // </editor-fold> end Type-erasing converting ctor and assignment }}}3
223 //- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
224
225 // </editor-fold> end Constructors, destructors, and assignment }}}2
226 //----------------------------------------------------------------------------
227
228 //----------------------------------------------------------------------------
229 // <editor-fold desc="then_parallel_for"> {{{2
230
231 template <
232 class Policy, class Functor,
233 std::enable_if_t<
234 // equivalent to:
235 // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
236 is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
237 // --------------------
238 int> = 0>
239 auto then_parallel_for(std::string arg_name, Policy&& arg_policy,
240 Functor&& functor) const {
241 //----------------------------------------
242 KOKKOS_EXPECTS(!m_graph_impl.expired())
243 KOKKOS_EXPECTS(bool(m_node_impl))
244 // TODO @graph restore this expectation once we add comparability to space
245 // instances
246 // KOKKOS_EXPECTS(
247 // arg_policy.space() == m_graph_impl->get_execution_space());
248
249 // needs to static assert constraint: DataParallelFunctor<Functor>
250
251 using policy_t = Kokkos::Impl::remove_cvref_t<Policy>;
252 // constraint check: same execution space type (or defaulted, maybe?)
253 static_assert(
254 std::is_same<typename policy_t::execution_space,
255 execution_space>::value,
256 // TODO @graph make defaulted execution space work
257 //|| policy_t::execution_space_is_defaulted,
258 "Execution Space mismatch between execution policy and graph");
259
260 auto policy = Experimental::require((Policy &&) arg_policy,
261 Kokkos::Impl::KernelInGraphProperty{});
262
263 using next_policy_t = decltype(policy);
264 using next_kernel_t =
265 Kokkos::Impl::GraphNodeKernelImpl<ExecutionSpace, next_policy_t,
266 std::decay_t<Functor>,
267 Kokkos::ParallelForTag>;
268 return this->_then_kernel(next_kernel_t{std::move(arg_name), policy.space(),
269 (Functor &&) functor,
270 (Policy &&) policy});
271 }
272
273 template <
274 class Policy, class Functor,
275 std::enable_if_t<
276 // equivalent to:
277 // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
278 is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
279 // --------------------
280 int> = 0>
281 auto then_parallel_for(Policy&& policy, Functor&& functor) const {
282 // needs to static assert constraint: DataParallelFunctor<Functor>
283 return this->then_parallel_for("", (Policy &&) policy,
284 (Functor &&) functor);
285 }
286
287 template <class Functor>
288 auto then_parallel_for(std::string name, std::size_t n,
289 Functor&& functor) const {
290 // needs to static assert constraint: DataParallelFunctor<Functor>
291 return this->then_parallel_for(std::move(name),
293 (Functor &&) functor);
294 }
295
296 template <class Functor>
297 auto then_parallel_for(std::size_t n, Functor&& functor) const {
298 // needs to static assert constraint: DataParallelFunctor<Functor>
299 return this->then_parallel_for("", n, (Functor &&) functor);
300 }
301
302 // </editor-fold> end then_parallel_for }}}2
303 //----------------------------------------------------------------------------
304
305 //----------------------------------------------------------------------------
306 // <editor-fold desc="then_parallel_reduce"> {{{2
307
308 template <
309 class Policy, class Functor, class ReturnType,
310 std::enable_if_t<
311 // equivalent to:
312 // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
313 is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
314 // --------------------
315 int> = 0>
316 auto then_parallel_reduce(std::string arg_name, Policy&& arg_policy,
317 Functor&& functor,
318 ReturnType&& return_value) const {
319 auto graph_impl_ptr = m_graph_impl.lock();
320 KOKKOS_EXPECTS(bool(graph_impl_ptr))
321 KOKKOS_EXPECTS(bool(m_node_impl))
322 // TODO @graph restore this expectation once we add comparability to space
323 // instances
324 // KOKKOS_EXPECTS(
325 // arg_policy.space() == m_graph_impl->get_execution_space());
326
327 // needs static assertion of constraint:
328 // DataParallelReductionFunctor<Functor, ReturnType>
329
330 using policy_t = std::remove_cv_t<std::remove_reference_t<Policy>>;
331 static_assert(
332 std::is_same<typename policy_t::execution_space,
333 execution_space>::value,
334 // TODO @graph make defaulted execution space work
335 // || policy_t::execution_space_is_defaulted,
336 "Execution Space mismatch between execution policy and graph");
337
338 // This is also just an expectation, but it's one that we expect the user
339 // to interact with (even in release mode), so we should throw an exception
340 // with an explanation rather than just doing a contract assertion.
341 // We can't static_assert this because of the way that Reducers store
342 // whether or not they point to a View as a runtime boolean rather than part
343 // of the type.
344 if (Kokkos::Impl::parallel_reduce_needs_fence(
345 graph_impl_ptr->get_execution_space(), return_value)) {
346 Kokkos::Impl::throw_runtime_exception(
347 "Parallel reductions in graphs can't operate on Reducers that "
348 "reference a scalar because they can't complete synchronously. Use a "
349 "Kokkos::View instead and keep in mind the result will only be "
350 "available once the graph is submitted (or in tasks that depend on "
351 "this one).");
352 }
353
354 //----------------------------------------
355 // This is a disaster, but I guess it's not a my disaster to fix right now
356 using return_type_remove_cvref =
357 std::remove_cv_t<std::remove_reference_t<ReturnType>>;
358 static_assert(Kokkos::is_view<return_type_remove_cvref>::value ||
360 "Output argument to parallel reduce in a graph must be a "
361 "View or a Reducer");
362 using return_type =
363 // Yes, you do really have to do this...
364 std::conditional_t<Kokkos::is_reducer<return_type_remove_cvref>::value,
365 return_type_remove_cvref,
366 const return_type_remove_cvref>;
367 using functor_type = Kokkos::Impl::remove_cvref_t<Functor>;
368 // see Kokkos_Parallel_Reduce.hpp for how these details are used there;
369 // we're just doing the same thing here
370 using return_value_adapter =
371 Kokkos::Impl::ParallelReduceReturnValue<void, return_type,
372 functor_type>;
373 // End of Kokkos reducer disaster
374 //----------------------------------------
375
376 auto policy = Experimental::require((Policy &&) arg_policy,
377 Kokkos::Impl::KernelInGraphProperty{});
378
379 using next_policy_t = decltype(policy);
380 using next_kernel_t = Kokkos::Impl::GraphNodeKernelImpl<
381 ExecutionSpace, next_policy_t, functor_type, Kokkos::ParallelReduceTag,
382 typename return_value_adapter::reducer_type>;
383
384 return this->_then_kernel(next_kernel_t{
385 std::move(arg_name), graph_impl_ptr->get_execution_space(),
386 (Functor &&) functor, (Policy &&) policy,
387 return_value_adapter::return_value(return_value, functor)});
388 }
389
390 template <
391 class Policy, class Functor, class ReturnType,
392 std::enable_if_t<
393 // equivalent to:
394 // requires Kokkos::ExecutionPolicy<remove_cvref_t<Policy>>
395 is_execution_policy<Kokkos::Impl::remove_cvref_t<Policy>>::value,
396 // --------------------
397 int> = 0>
398 auto then_parallel_reduce(Policy&& arg_policy, Functor&& functor,
399 ReturnType&& return_value) const {
400 return this->then_parallel_reduce("", (Policy &&) arg_policy,
401 (Functor &&) functor,
402 (ReturnType &&) return_value);
403 }
404
405 template <class Functor, class ReturnType>
406 auto then_parallel_reduce(std::string label,
407 typename execution_space::size_type idx_end,
408 Functor&& functor,
409 ReturnType&& return_value) const {
410 return this->then_parallel_reduce(
411 std::move(label), Kokkos::RangePolicy<execution_space>{0, idx_end},
412 (Functor &&) functor, (ReturnType &&) return_value);
413 }
414
415 template <class Functor, class ReturnType>
416 auto then_parallel_reduce(typename execution_space::size_type idx_end,
417 Functor&& functor,
418 ReturnType&& return_value) const {
419 return this->then_parallel_reduce("", idx_end, (Functor &&) functor,
420 (ReturnType &&) return_value);
421 }
422
423 // </editor-fold> end then_parallel_reduce }}}2
424 //----------------------------------------------------------------------------
425
426 // TODO @graph parallel scan, deep copy, etc.
427};
428
429} // end namespace Experimental
430} // end namespace Kokkos
431
432#endif // KOKKOS_KOKKOS_GRAPHNODE_HPP
A thread safe view to a bitset.
ReturnType