Kokkos_UniqueToken.hpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/
 
#ifndef KOKKOS_UNIQUE_TOKEN_HPP
#define KOKKOS_UNIQUE_TOKEN_HPP
 
#include <Kokkos_Macros.hpp>
#include <Kokkos_MemoryTraits.hpp>
#include <Kokkos_Core_fwd.hpp>
 
namespace Kokkos {
namespace Experimental {
 
enum class UniqueTokenScope : int { Instance, Global };
 
template <typename ExecutionSpace,
          UniqueTokenScope = UniqueTokenScope::Instance>
class UniqueToken {
 public:
  using execution_space = ExecutionSpace;
  using size_type       = typename execution_space::size_type;
 
  UniqueToken(execution_space const& = execution_space());
 
  KOKKOS_INLINE_FUNCTION
  size_type size() const;
 
  KOKKOS_INLINE_FUNCTION
  size_type acquire() const;
 
  KOKKOS_INLINE_FUNCTION
  void release(size_type) const;
};
 
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;
 
  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.
 
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; }
};
 
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
 
#endif  // KOKKOS_UNIQUE_TOKEN_HPP