shared.h

Go to the documentation of this file.

 
#pragma once
 
#include <atomic>
#include <cstddef>  // std::ptrdiff_t
#include <type_traits>
 
namespace detail
{
 
// define logic for reference counting (on the host)
// - this default implementation uses built-in atomic operations provided with
//   the C++ standard library
// - this logic has been pull out of ControlBlock in case we want to take a
//   crack at implementing an alternative version in terms of OpenMP constructs
//   (this requires the use of OpenMP 3.1 or newer)
// - for context, the use of std::atomic with OpenMP **PROBABLY** works as you
//   would expect, but that is not strictly required by the OpenMP standard. In
//   practice compiler writers generally try to "do the right thing" and make
//   things work correctly. Furthermore, the kind of logic where this plays a
//   role is typically executed at the very start and end of Cholla's execution
//   (i.e. outside of all OpenMP blocks)
//
// optimization opportunity: the memory ordering is a little stricter than
//   necessary (out of an abundance of cautious), which introduces additional
//   overhead. In practice, this shouldn't matter very much in cholla since we
//   increment/decrement reference counts relatively infrequently; this happens
//   when initializing relevant datastructures and when they go out of scope
//   (which is almost entirely constrained to startup and shutdown)
 
using RefCountType = std::atomic<long>;
inline long ref_count_increment(RefCountType& count) noexcept { return count.fetch_add(1L, std::memory_order_seq_cst); }
inline long ref_count_decrement(RefCountType& count) noexcept
{
  return count.fetch_add(-1L, std::memory_order_seq_cst);
}
 
class ControlBlock
{
  RefCountType ref_count_;
 
 protected:
  virtual ~ControlBlock() noexcept {}
 
 public:
  ControlBlock() noexcept : ref_count_{1L} {}
 
  void increment_count() noexcept
  {
    long preincrement_val = ref_count_increment(ref_count_);
    CHOLLA_ASSERT(preincrement_val >= 1L, "invariant is violated");  // <- sanity check!
  }
 
  void decrement_count() noexcept
  {
    long predecrement_val = ref_count_decrement(ref_count_);
    CHOLLA_ASSERT(predecrement_val > 0L, "invariant is violated");  // <- sanity check!
    if (predecrement_val == 1L) delete this;
  }
};
 
template <typename HandleOrPtrType, typename Deleter>
class ControlBlockImpl : public ControlBlock
{
  HandleOrPtrType managed_;
  Deleter deleter_;
 
 public:
  ControlBlockImpl(HandleOrPtrType& m, Deleter& d) : ControlBlock(), managed_{m}, deleter_{d} {}
 
  ~ControlBlockImpl() { deleter_(managed_); }
};
 
}  // namespace detail
 
// down below, we move on to actually defining SharedHandle and SharedDevPtr.
 
// define the CALL_INCREMENT_COUNT and CALL_DECREMENT_COUNT macros
// -> these macros forward onto detail::ControlBlock::increment_count and
//    detail::ControlBlock::decrement_count when invoked on the host and do nothing
//    when invoked on the host
// -> this is the desired behavior we want (it's explained in more detail at the top of
//    this page)
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
  #define CALL_INCREMENT_COUNT(cb_ptr) /* DOES NOTHING ON DEVICE */
  #define CALL_DECREMENT_COUNT(cb_ptr) /* DOES NOTHING ON DEVICE */
#else
  #define CALL_INCREMENT_COUNT(cb_ptr) (cb_ptr)->increment_count()
  #define CALL_DECREMENT_COUNT(cb_ptr) (cb_ptr)->decrement_count()
#endif
 
#define DEFINE_COMMON_METHODS(KLASS)                                                                             \
  /* The lint that we disable just below tells us that we should enclose every occurrence of KLASS within   */   \
  /* parentheses. Basically, we'd write it as (KLASS). I'm not actually sure that the result of the macro   */   \
  /* expansion would be valid code and it would certainly make things harder to read. This check is most    */   \
  /* useful when you would interject a macro call inside a C++ expression like a function. Since this macro */   \
  /* is used in a fundamentally different way (i.e. to define functions) its okay to disable the check      */   \
  /* NOLINTBEGIN(bugprone-macro-parentheses) */                                                                  \
                                                                                                                 \
  /* implement the copy constructor */                                                                           \
  template <typename T>                                                                                          \
  __host__ __device__ KLASS<T>::KLASS(const KLASS<T>& other) noexcept : wrapped_{other.wrapped_}, cb_{other.cb_} \
  {                                                                                                              \
    if (cb_ != nullptr) {                                                                                        \
      /* The lint that we disable just below tells us that we should enclose every is warning about using */     \
      /* memory after it is freed. However, the fact that we always initialize the cb_ to nullptr and     */     \
      /* that we only decrement the reference count if we overwrite cb_ immediately afterwards means that */     \
      /* it's impossible for this scenario to arise                                                       */     \
      /* NOLINTBEGIN(clang-analyzer-cplusplus.NewDelete,-warnings-as-errors) */                                  \
      CALL_INCREMENT_COUNT(cb_);                                                                                 \
      /* NOLINTEND(clang-analyzer-cplusplus.NewDelete,-warnings-as-errors) */                                    \
    }                                                                                                            \
  }                                                                                                              \
                                                                                                                 \
  /* implement the move constructor */                                                                           \
  template <typename T>                                                                                          \
  __host__ __device__ KLASS<T>::KLASS(KLASS<T>&& other) noexcept : wrapped_{other.wrapped_}, cb_{other.cb_}      \
  {                                                                                                              \
    other.set_empty_wrapped_();                                                                                  \
    other.cb_ = nullptr;                                                                                         \
  }                                                                                                              \
                                                                                                                 \
  /* implement the copy assignment operation */                                                                  \
  template <typename T>                                                                                          \
  __host__ __device__ KLASS<T>& KLASS<T>::operator=(const KLASS<T>& other) noexcept                              \
  {                                                                                                              \
    /* care needs to be taken for self-assignment when `this` is the ONLY owner of a resource (it would be  */   \
    /* bad if a naive implementation decremented the reference count to zero before trying to increment it) */   \
    /* -> at the time of writing, the rest of the implementation actually makes this branch unnecessary.    */   \
    /*    For now, we explicitly handle this case, since its more idiomatic (plus it addresses the          */   \
    /*    associated clang-tidy warning)                                                                    */   \
    if (this == &other) return *this;                                                                            \
                                                                                                                 \
    /* the different_control_blocks check is a minor optimization. It lets us avoid an unnecessary pair of  */   \
    /* atomic increments and decrements (the fact they are atomic makes the operation more expensive). This */   \
    /* is a no-brainer since we already branch based on whether cb_ or other.cb_ are null pointers          */   \
    bool different_control_blocks = cb_ != other.cb_;                                                            \
    if (different_control_blocks and (other.cb_ != nullptr)) CALL_INCREMENT_COUNT(other.cb_);                    \
    if (different_control_blocks and (cb_ != nullptr)) CALL_DECREMENT_COUNT(cb_);                                \
    cb_      = other.cb_;                                                                                        \
    wrapped_ = other.wrapped_;                                                                                   \
    return *this;                                                                                                \
  }                                                                                                              \
                                                                                                                 \
  /* implement the move assignment operation */                                                                  \
  template <typename T>                                                                                          \
  __host__ __device__ KLASS<T>& KLASS<T>::operator=(KLASS<T>&& other) noexcept                                   \
  {                                                                                                              \
    reset();                                             \
    swap(other);                                 \
    return *this;                                                                                                \
  }                                                                                                              \
                                                                                                                 \
  template <typename T>                                                                                          \
  __host__ __device__ void KLASS<T>::swap(KLASS<T>& other) noexcept                                              \
  {                                                                                                              \
    decltype(wrapped_) tmp_wrapped = wrapped_;                                                                   \
    wrapped_                       = other.wrapped_;                                                             \
    other.wrapped_                 = tmp_wrapped;                                                                \
                                                                                                                 \
    detail::ControlBlock* tmp_cb = cb_;                                                                          \
    cb_                          = other.cb_;                                                                    \
    other.cb_                    = tmp_cb;                                                                       \
  }                                                                                                              \
                                                                                                                 \
  template <typename T>                                                                                          \
  __host__ __device__ void KLASS<T>::reset() noexcept                                                            \
  {                                                                                                              \
    /* reminder: if the reference count hits 0, the control block automatically: */                              \
    /*    - destroys wrapped_ (a copy of wrapped is tracked within the control   */                              \
    /*      block for this purpose)                                              */                              \
    /*    - calls delete on itself                                               */                              \
    if (cb_ != nullptr) {                                                                                        \
      /* NOLINTBEGIN(clang-analyzer-unix.DynamicMemoryModeling, clang-analyzer-cplusplus.NewDelete) */           \
      CALL_DECREMENT_COUNT(cb_);                                                                                 \
      /* NOLINTEND(clang-analyzer-unix.DynamicMemoryModeling, clang-analyzer-cplusplus.NewDelete) */             \
      cb_ = nullptr;                                                                                             \
      set_empty_wrapped_();                                                                                      \
    }                                                                                                            \
  }                                                                                                              \
  /* NOLINTEND(bugprone-macro-parentheses) */
 
template <typename HandleT>
class SharedHandle
{
  // perform some sanity checks:
  static_assert(std::is_arithmetic_v<HandleT> or std::is_aggregate_v<HandleT> or std::is_pointer_v<HandleT>);
  static_assert(not std::is_const_v<HandleT>);
 
  // data members:
  HandleT wrapped_;           
  detail::ControlBlock* cb_;  
 
  // helper method used by DEFINE_COMMON_METHODS to set the value of wrapped_ to the
  // appropriate value when a SharedHandle is "empty"
  // -> this method does nothing since the value of wrapped_ in an empty SharedHandle
  //    instance is explicitly undefined
  // -> this is defined for parity with SharedDevPtr
  __host__ __device__ __forceinline__ void set_empty_wrapped_() const noexcept {}
 
 public:
  // used for testing purposes
  typedef HandleT wrapped_ref_type;
 
  __host__ __device__ SharedHandle() : wrapped_{}, cb_{nullptr} {};
 
  template <typename Deleter>
  __host__ SharedHandle(HandleT handle, Deleter d)
      // NOLINTNEXTLINE(clang-analyzer-unix.DynamicMemoryModeling, clang-analyzer-cplusplus.NewDelete)
      : wrapped_{handle}, cb_{new detail::ControlBlockImpl<HandleT, Deleter>(handle, d)}
  {
  }
 
  __host__ __device__ ~SharedHandle() noexcept { reset(); }
 
  // copy/move construction and assignment
  __host__ __device__ SharedHandle(const SharedHandle& other) noexcept;
  __host__ __device__ SharedHandle(SharedHandle&& other) noexcept;
  __host__ __device__ SharedHandle& operator=(const SharedHandle& other) noexcept;
  __host__ __device__ SharedHandle& operator=(SharedHandle&& other) noexcept;
 
  __host__ __device__ void swap(SharedHandle& o) noexcept;
 
  __host__ __device__ void reset() noexcept;
 
  __host__ __device__ __forceinline__ HandleT get() const noexcept { return wrapped_; }
 
  __host__ __device__ explicit operator bool() const noexcept { return cb_ != nullptr; }
};
 
// provide definitions for the remainder of methods
DEFINE_COMMON_METHODS(SharedHandle)
 
 
template <typename T>
class SharedDevPtr
{
  // data members:
  T* wrapped_;                
  detail::ControlBlock* cb_;  
 
  // helper method used by DEFINE_COMMON_METHODS to set the value of wrapped_ to the
  // appropriate value when a SharedDevPtr is "empty"
  __host__ __device__ __forceinline__ void set_empty_wrapped_() noexcept { wrapped_ = nullptr; }
 
 public:
  // used for testing purposes
  typedef T* wrapped_ref_type;
 
  __host__ __device__ SharedDevPtr() : wrapped_{nullptr}, cb_{nullptr} {}
 
  template <typename Deleter>
  __host__ SharedDevPtr(T* ptr, Deleter d)
      // NOLINTNEXTLINE(clang-analyzer-unix.DynamicMemoryModeling, clang-analyzer-cplusplus.NewDelete)
      : wrapped_{ptr}, cb_{new detail::ControlBlockImpl<T*, Deleter>(ptr, d)}
  {
  }
 
  __host__ __device__ ~SharedDevPtr() noexcept { reset(); }
 
  // copy/move construction and assignment
  __host__ __device__ SharedDevPtr(const SharedDevPtr& other) noexcept;
  __host__ __device__ SharedDevPtr(SharedDevPtr&& other) noexcept;
  __host__ __device__ SharedDevPtr& operator=(const SharedDevPtr& other) noexcept;
  __host__ __device__ SharedDevPtr& operator=(SharedDevPtr&& other) noexcept;
 
  __host__ __device__ void swap(SharedDevPtr& other) noexcept;
 
  __host__ __device__ void reset() noexcept;
 
  __host__ __device__ __forceinline__ T* get() const noexcept { return wrapped_; }
 
  __host__ __device__ explicit operator bool() const noexcept { return wrapped_ != nullptr; }
 
  __device__ __forceinline__ T& operator*() const noexcept { return *wrapped_; }
 
  __device__ __forceinline__ T& operator[](std::ptrdiff_t idx) const { return wrapped_[idx]; }
};
 
// provide definitions for the remainder of methods
DEFINE_COMMON_METHODS(SharedDevPtr)
 
// let's do some cleanup to avoid leaking of macros into other parts of the codebase
#undef CALL_INCREMENT_COUNT
#undef CALL_DECREMENT_COUNT
#undef DEFINE_COMMON_METHODS