Kokkos Core Kernels Package  Version of the Day
Kokkos_UniqueToken.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_UNIQUE_TOKEN_HPP
46 #define KOKKOS_UNIQUE_TOKEN_HPP
47 
48 #include <Kokkos_Macros.hpp>
49 #include <Kokkos_MemoryTraits.hpp>
50 #include <Kokkos_Core_fwd.hpp>
51 
52 namespace Kokkos {
53 namespace Experimental {
54 
55 enum class UniqueTokenScope : int { Instance, Global };
56 
62 template <typename ExecutionSpace,
63  UniqueTokenScope = UniqueTokenScope::Instance>
64 class UniqueToken {
65  public:
66  using execution_space = ExecutionSpace;
67  using size_type = typename execution_space::size_type;
68 
72  UniqueToken(execution_space const& = execution_space());
73 
75  KOKKOS_INLINE_FUNCTION
76  size_type size() const;
77 
79  KOKKOS_INLINE_FUNCTION
80  size_type acquire() const;
81 
83  KOKKOS_INLINE_FUNCTION
84  void release(size_type) const;
85 };
86 
92 template <typename ExecutionSpace>
93 class UniqueToken<ExecutionSpace, UniqueTokenScope::Instance>
94  : public UniqueToken<ExecutionSpace, UniqueTokenScope::Global> {
95  public:
96  using execution_space = ExecutionSpace;
97  using size_type = typename execution_space::size_type;
98 
108  UniqueToken(size_type max_size, execution_space const& = execution_space());
109 };
110 
111 // NOTE There was an agreement amongst developers that "AcquireUniqueToken" is a
112 // bad name but at this time no one has suggested a better alternative.
113 
118 template <typename ExecutionSpace,
119  UniqueTokenScope TokenScope = UniqueTokenScope::Instance>
121  public:
122  using exec_space = ExecutionSpace;
123  using size_type = typename exec_space::size_type;
125 
126  private:
127  token_type my_token;
128  size_type my_acquired_val;
129 
130  public:
131  KOKKOS_FUNCTION AcquireUniqueToken(token_type t)
132  : my_token(t), my_acquired_val(my_token.acquire()) {}
133 
134  KOKKOS_FUNCTION ~AcquireUniqueToken() { my_token.release(my_acquired_val); }
135 
136  KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
137 };
138 
144 template <typename TeamPolicy>
146  public:
147  using exec_space = typename TeamPolicy::execution_space;
149  using size_type = typename token_type::size_type;
150  using team_member_type = typename TeamPolicy::member_type;
151  using scratch_view =
152  Kokkos::View<size_type, typename exec_space::scratch_memory_space,
153  Kokkos::MemoryUnmanaged>;
154 
155  private:
156  token_type my_token;
157  size_type my_acquired_val;
158  scratch_view my_team_acquired_val;
159  team_member_type my_team;
160 
161  public:
162  // NOTE The implementations of the constructor and destructor use
163  // `Kokkos::single()` which is an inline function defined in each backend.
164  // This creates circular dependency issues. Moving them to a separate header
165  // is less than ideal and should be revisited later. Having a `UniqueToken`
166  // forward declaration was considered but the non-type template parameter
167  // makes things complicated because it would require moving the definition of
168  // `UniqueTokenScope` enumeration type and its enumerators away which would
169  // hurt readability.
170  KOKKOS_FUNCTION AcquireTeamUniqueToken(token_type t, team_member_type team);
171  KOKKOS_FUNCTION ~AcquireTeamUniqueToken();
172  KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
173  static std::size_t shmem_size() { return scratch_view::shmem_size(); }
174 };
175 
176 } // namespace Experimental
177 } // namespace Kokkos
178 
179 #endif // KOKKOS_UNIQUE_TOKEN_HPP
KOKKOS_INLINE_FUNCTION void release(size_type) const
release a value acquired by generate
RAII helper for per-thread unique token values.
class to generate unique ids base on the required amount of concurrency
KOKKOS_INLINE_FUNCTION size_type size() const
upper bound for acquired values, i.e. 0 <= value < size()
UniqueToken(execution_space const &=execution_space())
create object size for concurrency on the given instance
KOKKOS_INLINE_FUNCTION size_type acquire() const
acquire value such that 0 <= value < size()
RAII helper for per-team unique token values.
Definition: dummy.cpp:3
Parallel execution of a functor calls the functor once with each member of the execution policy...