Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_UniqueToken.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_UNIQUE_TOKEN_HPP
23#define KOKKOS_UNIQUE_TOKEN_HPP
24
25#include <Kokkos_Macros.hpp>
26#include <Kokkos_MemoryTraits.hpp>
27#include <Kokkos_Core_fwd.hpp>
28
29namespace Kokkos {
30namespace Experimental {
31
32enum class UniqueTokenScope : int { Instance, Global };
33
39template <typename ExecutionSpace = Kokkos::DefaultExecutionSpace,
40 UniqueTokenScope = UniqueTokenScope::Instance>
42 public:
43 using execution_space = ExecutionSpace;
44 using size_type = typename execution_space::size_type;
45
49 UniqueToken(execution_space const& = execution_space());
50
52 KOKKOS_INLINE_FUNCTION
53 size_type size() const;
54
56 KOKKOS_INLINE_FUNCTION
57 size_type acquire() const;
58
60 KOKKOS_INLINE_FUNCTION
61 void release(size_type) const;
62};
63
69template <typename ExecutionSpace>
70class UniqueToken<ExecutionSpace, UniqueTokenScope::Instance>
71 : public UniqueToken<ExecutionSpace, UniqueTokenScope::Global> {
72 public:
73 using execution_space = ExecutionSpace;
74 using size_type = typename execution_space::size_type;
75
85 UniqueToken(size_type max_size, execution_space const& = execution_space());
86};
87
88// NOTE There was an agreement amongst developers that "AcquireUniqueToken" is a
89// bad name but at this time no one has suggested a better alternative.
90
95template <typename ExecutionSpace,
96 UniqueTokenScope TokenScope = UniqueTokenScope::Instance>
98 public:
99 using exec_space = ExecutionSpace;
100 using size_type = typename exec_space::size_type;
102
103 private:
104 token_type my_token;
105 size_type my_acquired_val;
106
107 public:
108 KOKKOS_FUNCTION AcquireUniqueToken(token_type t)
109 : my_token(t), my_acquired_val(my_token.acquire()) {}
110
111 KOKKOS_FUNCTION ~AcquireUniqueToken() { my_token.release(my_acquired_val); }
112
113 KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
114};
115
121template <typename TeamPolicy>
123 public:
124 using exec_space = typename TeamPolicy::execution_space;
126 using size_type = typename token_type::size_type;
127 using team_member_type = typename TeamPolicy::member_type;
128 using scratch_view =
129 Kokkos::View<size_type, typename exec_space::scratch_memory_space,
130 Kokkos::MemoryUnmanaged>;
131
132 private:
133 token_type my_token;
134 size_type my_acquired_val;
135 scratch_view my_team_acquired_val;
136 team_member_type my_team;
137
138 public:
139 // NOTE The implementations of the constructor and destructor use
140 // `Kokkos::single()` which is an inline function defined in each backend.
141 // This creates circular dependency issues. Moving them to a separate header
142 // is less than ideal and should be revisited later. Having a `UniqueToken`
143 // forward declaration was considered but the non-type template parameter
144 // makes things complicated because it would require moving the definition of
145 // `UniqueTokenScope` enumeration type and its enumerators away which would
146 // hurt readability.
147 KOKKOS_FUNCTION AcquireTeamUniqueToken(token_type t, team_member_type team);
148 KOKKOS_FUNCTION ~AcquireTeamUniqueToken();
149 KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
150 static std::size_t shmem_size() { return scratch_view::shmem_size(); }
151};
152
153} // namespace Experimental
154} // namespace Kokkos
155
156#endif // KOKKOS_UNIQUE_TOKEN_HPP
RAII helper for per-team unique token values.
RAII helper for per-thread unique token values.
UniqueToken(size_type max_size, execution_space const &=execution_space())
Create object with specified size.
class to generate unique ids base on the required amount of concurrency
KOKKOS_INLINE_FUNCTION size_type acquire() const
acquire value such that 0 <= value < size()
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 void release(size_type) const
release a value acquired by generate
Parallel execution of a functor calls the functor once with each member of the execution policy.