Line data    Source code

       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_HOST_SHARED_PTR_HPP
      18             : #define KOKKOS_IMPL_HOST_SHARED_PTR_HPP
      19             : 
      20             : #include <Kokkos_Macros.hpp>
      21             : #include <Kokkos_Atomic.hpp>
      22             : #include <impl/Kokkos_Error.hpp>
      23             : 
      24             : #include <functional>
      25             : 
      26             : namespace Kokkos {
      27             : namespace Impl {
      28             : 
      29             : template <typename T>
      30             : class HostSharedPtr {
      31             :  public:
      32             :   using element_type = T;
      33             : 
      34             :   KOKKOS_DEFAULTED_FUNCTION constexpr HostSharedPtr() = default;
      35             :   KOKKOS_FUNCTION constexpr HostSharedPtr(std::nullptr_t) {}
      36             : 
      37             :   explicit HostSharedPtr(T* element_ptr)
      38             :       : HostSharedPtr(element_ptr, [](T* const t) { delete t; }) {}
      39             : 
      40             :   template <class Deleter>
      41             :   HostSharedPtr(T* element_ptr, const Deleter& deleter)
      42             :       : m_element_ptr(element_ptr) {
      43             :     static_assert(std::is_invocable_v<Deleter, T*> &&
      44             :                   std::is_copy_constructible_v<Deleter>);
      45             :     if (element_ptr) {
      46             :       try {
      47             :         m_control = new Control{deleter, 1};
      48             :       } catch (...) {
      49             :         deleter(element_ptr);
      50             :         throw;
      51             :       }
      52             :     }
      53             :   }
      54             : 
      55             :   KOKKOS_FUNCTION HostSharedPtr(HostSharedPtr&& other) noexcept
      56             :       : m_element_ptr(other.m_element_ptr), m_control(other.m_control) {
      57             :     other.m_element_ptr = nullptr;
      58             :     other.m_control     = nullptr;
      59             :   }
      60             : 
      61        4944 :   KOKKOS_FUNCTION HostSharedPtr(const HostSharedPtr& other) noexcept
      62        4944 :       : m_element_ptr(other.m_element_ptr), m_control(other.m_control) {
      63        4944 :     KOKKOS_IF_ON_HOST(
      64             :         (if (m_control) Kokkos::atomic_add(&(m_control->m_counter), 1);))
      65             :     KOKKOS_IF_ON_DEVICE(m_control = nullptr;)
      66        4944 :   }
      67             : 
      68             :   KOKKOS_FUNCTION HostSharedPtr& operator=(HostSharedPtr&& other) noexcept {
      69             :     if (&other != this) {
      70             :       cleanup();
      71             :       m_element_ptr       = other.m_element_ptr;
      72             :       other.m_element_ptr = nullptr;
      73             :       m_control           = other.m_control;
      74             :       other.m_control     = nullptr;
      75             :     }
      76             :     return *this;
      77             :   }
      78             : 
      79             :   KOKKOS_FUNCTION HostSharedPtr& operator=(
      80             :       const HostSharedPtr& other) noexcept {
      81             :     if (&other != this) {
      82             :       cleanup();
      83             :       m_element_ptr = other.m_element_ptr;
      84             :       m_control     = other.m_control;
      85             :       KOKKOS_IF_ON_HOST(
      86             :           (if (m_control) Kokkos::atomic_add(&(m_control->m_counter), 1);))
      87             :       KOKKOS_IF_ON_DEVICE(m_control = nullptr;)
      88             :     }
      89             :     return *this;
      90             :   }
      91             : 
      92        6180 :   KOKKOS_FUNCTION ~HostSharedPtr() { cleanup(); }
      93             : 
      94             :   // returns the stored pointer
      95        1236 :   KOKKOS_FUNCTION T* get() const noexcept { return m_element_ptr; }
      96             :   // dereferences the stored pointer
      97             :   KOKKOS_FUNCTION T& operator*() const noexcept {
      98             :     KOKKOS_EXPECTS(bool(*this));
      99             :     return *get();
     100             :   }
     101             :   // dereferences the stored pointer
     102             :   KOKKOS_FUNCTION T* operator->() const noexcept {
     103             :     KOKKOS_EXPECTS(bool(*this));
     104             :     return get();
     105             :   }
     106             : 
     107             :   // checks if the stored pointer is not null
     108             :   KOKKOS_FUNCTION explicit operator bool() const noexcept {
     109             :     return get() != nullptr;
     110             :   }
     111             : 
     112             :   // returns the number of HostSharedPtr instances managing the current object
     113             :   // or 0 if there is no managed object.
     114             :   int use_count() const noexcept {
     115             :     return m_control ? m_control->m_counter : 0;
     116             :   }
     117             : 
     118             :  private:
     119        7416 :   KOKKOS_FUNCTION void cleanup() noexcept {
     120        9888 :     KOKKOS_IF_ON_HOST((
     121             :         // If m_counter is set, then this instance is responsible for managing
     122             :         // the object pointed to by m_counter and m_element_ptr.
     123             :         if (m_control) {
     124             :           int const count =
     125             :               Kokkos::atomic_fetch_sub(&(m_control->m_counter), 1);
     126             :           // atomic_fetch_sub might have memory order relaxed, so we need to
     127             :           // force synchronization to avoid multiple threads doing the cleanup.
     128             :           Kokkos::memory_fence();
     129             :           if (count == 1) {
     130             :             (m_control->m_deleter)(m_element_ptr);
     131             :             m_element_ptr = nullptr;
     132             :             delete m_control;
     133             :             m_control = nullptr;
     134             :           }
     135             :         }))
     136        7416 :   }
     137             : 
     138        2472 :   struct Control {
     139             :     std::function<void(T*)> m_deleter;
     140             :     int m_counter;
     141             :   };
     142             : 
     143             :   T* m_element_ptr   = nullptr;
     144             :   Control* m_control = nullptr;
     145             : };
     146             : }  // namespace Impl
     147             : }  // namespace Kokkos
     148             : 
     149             : #endif