Update Kokkos library in LAMMPS to v3.2

lib/kokkos/core/src/Kokkos_UniqueToken.hpp

@@ -46,6 +46,8 @@
 #define KOKKOS_UNIQUE_TOKEN_HPP
 
 #include <Kokkos_Macros.hpp>
+#include <Kokkos_MemoryTraits.hpp>
+#include <Kokkos_Core_fwd.hpp>
 
 namespace Kokkos {
 namespace Experimental {
@@ -82,6 +84,95 @@ class UniqueToken {
   void release(size_type) const;
 };
 
+/// \brief Instance scope UniqueToken allows for a max size other than
+/// execution_space::concurrency()
+///
+/// This object should behave like a ref-counted object, so that when the last
+/// instance is destroyed, resources are free if needed
+template <typename ExecutionSpace>
+class UniqueToken<ExecutionSpace, UniqueTokenScope::Instance>
+    : public UniqueToken<ExecutionSpace, UniqueTokenScope::Global> {
+ public:
+  using execution_space = ExecutionSpace;
+  using size_type       = typename execution_space::size_type;
+
+  /// \brief Create object with specified size
+  ///
+  /// It is required that max_size is >= the maximum number of concurrent
+  /// threads that will attempt to acquire the UniqueToken. This constructor is
+  /// most commonly useful when you:
+  ///   1) Have a loop bound that may be smaller than
+  ///   execution_space::concurrency().
+  ///   2) Want a per-team unique token in the range [0,
+  ///   execution_space::concurrency() / team_size)
+  UniqueToken(size_type max_size, execution_space const& = execution_space());
+};
+
+// NOTE There was an agreement amongst developers that "AcquireUniqueToken" is a
+// bad name but at this time no one has suggested a better alternative.
+
+/// \brief RAII helper for per-thread unique token values.
+///
+/// The token value will be acquired at construction and automatically
+/// released at destruction.
+template <typename ExecutionSpace,
+          UniqueTokenScope TokenScope = UniqueTokenScope::Instance>
+class AcquireUniqueToken {
+ public:
+  using exec_space = ExecutionSpace;
+  using size_type  = typename exec_space::size_type;
+  using token_type = UniqueToken<exec_space, TokenScope>;
+
+ private:
+  token_type my_token;
+  size_type my_acquired_val;
+
+ public:
+  KOKKOS_FUNCTION AcquireUniqueToken(token_type t)
+      : my_token(t), my_acquired_val(my_token.acquire()) {}
+
+  KOKKOS_FUNCTION ~AcquireUniqueToken() { my_token.release(my_acquired_val); }
+
+  KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
+};
+
+/// \brief RAII helper for per-team unique token values.
+///
+/// The token value will be acquired at construction and automatically
+/// released at destruction. All threads in a team will share the same
+/// token value.
+template <typename TeamPolicy>
+class AcquireTeamUniqueToken {
+ public:
+  using exec_space       = typename TeamPolicy::execution_space;
+  using token_type       = UniqueToken<exec_space>;
+  using size_type        = typename token_type::size_type;
+  using team_member_type = typename TeamPolicy::member_type;
+  using scratch_view =
+      Kokkos::View<size_type, typename exec_space::scratch_memory_space,
+                   Kokkos::MemoryUnmanaged>;
+
+ private:
+  token_type my_token;
+  size_type my_acquired_val;
+  scratch_view my_team_acquired_val;
+  team_member_type my_team;
+
+ public:
+  // NOTE The implementations of the constructor and destructor use
+  // `Kokkos::single()` which is an inline function defined in each backend.
+  // This creates circular dependency issues.  Moving them to a separate header
+  // is less than ideal and should be revisited later.  Having a `UniqueToken`
+  // forward declaration was considered but the non-type template parameter
+  // makes things complicated because it would require moving the definition of
+  // `UniqueTokenScope` enumeration type and its enumerators away which would
+  // hurt readability.
+  KOKKOS_FUNCTION AcquireTeamUniqueToken(token_type t, team_member_type team);
+  KOKKOS_FUNCTION ~AcquireTeamUniqueToken();
+  KOKKOS_FUNCTION size_type value() const { return my_acquired_val; }
+  static std::size_t shmem_size() { return scratch_view::shmem_size(); }
+};
+
 }  // namespace Experimental
 }  // namespace Kokkos