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_VIEW_TRACKER_HPP
      18             : #define KOKKOS_VIEW_TRACKER_HPP
      19             : 
      20             : namespace Kokkos {
      21             : 
      22             : template <class DataType, class... Properties>
      23             : class View;
      24             : 
      25             : namespace Impl {
      26             : 
      27             : /*
      28             :  * \class ViewTracker
      29             :  * \brief template class to wrap the shared allocation tracker
      30             :  *
      31             :  * \section This class is templated on the View and provides
      32             :  * constructors that match the view.  The constructors and assignments
      33             :  * from view will externalize the logic needed to enable/disable
      34             :  * ref counting to provide a single gate to enable further developments
      35             :  * which may hinge on the same logic.
      36             :  *
      37             :  */
      38             : template <class ParentView>
      39             : struct ViewTracker {
      40             :   using track_type  = Kokkos::Impl::SharedAllocationTracker;
      41             :   using view_traits = typename ParentView::traits;
      42             : 
      43             :   track_type m_tracker;
      44             : 
      45             :   KOKKOS_INLINE_FUNCTION
      46             :   ViewTracker() : m_tracker() {}
      47             : 
      48             :   KOKKOS_INLINE_FUNCTION
      49       33714 :   ViewTracker(const ViewTracker& vt) noexcept
      50       33714 :       : m_tracker(vt.m_tracker, view_traits::is_managed) {}
      51             : 
      52             :   KOKKOS_INLINE_FUNCTION
      53             :   explicit ViewTracker(const ParentView& vt) noexcept : m_tracker() {
      54             :     assign(vt);
      55             :   }
      56             : 
      57             :   template <class RT, class... RP>
      58             :   KOKKOS_INLINE_FUNCTION explicit ViewTracker(
      59             :       const View<RT, RP...>& vt) noexcept
      60             :       : m_tracker() {
      61             :     assign(vt);
      62             :   }
      63             : 
      64             :   template <class RT, class... RP>
      65             :   KOKKOS_INLINE_FUNCTION void assign(const View<RT, RP...>& vt) noexcept {
      66             :     if (this == reinterpret_cast<const ViewTracker*>(&vt.m_track)) return;
      67             :     KOKKOS_IF_ON_HOST((
      68             :         if (view_traits::is_managed && Kokkos::Impl::SharedAllocationRecord<
      69             :                                            void, void>::tracking_enabled()) {
      70             :           m_tracker.assign_direct(vt.m_track.m_tracker);
      71             :         } else { m_tracker.assign_force_disable(vt.m_track.m_tracker); }))
      72             : 
      73             :     KOKKOS_IF_ON_DEVICE((m_tracker.assign_force_disable(vt.m_track.m_tracker);))
      74             :   }
      75             : 
      76             :   KOKKOS_INLINE_FUNCTION ViewTracker& operator=(
      77             :       const ViewTracker& rhs) noexcept {
      78             :     if (this == &rhs) return *this;
      79             :     KOKKOS_IF_ON_HOST((
      80             :         if (view_traits::is_managed && Kokkos::Impl::SharedAllocationRecord<
      81             :                                            void, void>::tracking_enabled()) {
      82             :           m_tracker.assign_direct(rhs.m_tracker);
      83             :         } else { m_tracker.assign_force_disable(rhs.m_tracker); }))
      84             : 
      85             :     KOKKOS_IF_ON_DEVICE((m_tracker.assign_force_disable(rhs.m_tracker);))
      86             :     return *this;
      87             :   }
      88             : 
      89             :   KOKKOS_INLINE_FUNCTION
      90             :   explicit ViewTracker(const track_type& tt) noexcept
      91             :       : m_tracker(tt, view_traits::is_managed) {}
      92             : };
      93             : 
      94             : }  // namespace Impl
      95             : 
      96             : }  // namespace Kokkos
      97             : 
      98             : #endif  // KOKKOS_VIEW_TRACKER_HPP