handle.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./allocators.hpp"
#include "./memcpy.hpp"
#include "../concepts.hpp"
#include "../macros.hpp"
 
#include <array>
#include <memory>
#include <type_traits>
#include <utility>
 
namespace nda::mem {
 
  static constexpr bool init_dcmplx = true;
 
  template <typename T, int Al>
  struct alignas(Al) aligner {
    T x;
 
    [[nodiscard]] T &get() noexcept { return x; }
 
    [[nodiscard]] T const &get() const noexcept { return x; }
  };
 
  struct do_not_initialize_t {};
 
  inline static constexpr do_not_initialize_t do_not_initialize{};
 
  struct init_zero_t {};
 
  inline static constexpr init_zero_t init_zero{};
 
  // Forward declaration.
  template <typename T, AddressSpace AdrSp = Host>
  struct handle_shared;
 
  template <typename T, Allocator A = mallocator<>>
  struct handle_heap {
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_heap requires the value_type to have a non-throwing destructor");
 
    private:
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != nullptr.
    size_t _size = 0;
 
    // Allocator to use.
#ifndef NDA_DEBUG_LEAK_CHECK
    static inline A allocator; // NOLINT (allocator is not specific to a single instance)
#else
    static inline leak_check<A> allocator; // NOLINT (allocator is not specific to a single instance)
#endif
 
    // For shared ownership (points to a blk_T_t).
    mutable std::shared_ptr<void> sptr;
 
    // Type of the memory block, i.e. a pointer to the data and its size.
    using blk_T_t = std::pair<T *, size_t>;
 
    // Release the handled memory (data pointer and size are not set to null here).
    static void destruct(blk_T_t b) noexcept {
      auto [data, size] = b;
 
      // do nothing if the data is null
      if (data == nullptr) return;
 
      // if needed, call the destructors of the objects stored
      if constexpr (A::address_space == Host and !(std::is_trivial_v<T> or nda::is_complex_v<T>)) {
        for (size_t i = 0; i < size; ++i) data[i].~T();
      }
 
      // deallocate the memory block
      allocator.deallocate({(char *)data, size * sizeof(T)});
    }
 
    // Deleter for the shared pointer.
    static void deleter(void *p) noexcept { destruct(*((blk_T_t *)p)); }
 
    public:
    using value_type = T;
 
    using allocator_type = A;
 
    static constexpr auto address_space = allocator_type::address_space;
 
    std::shared_ptr<void> get_sptr() const {
      if (not sptr) sptr.reset(new blk_T_t{_data, _size}, deleter);
      return sptr;
    }
 
    ~handle_heap() noexcept {
      if (not sptr and not(is_null())) destruct({_data, _size});
    }
 
    handle_heap() = default;
 
    handle_heap(handle_heap &&h) noexcept : _data(h._data), _size(h._size), sptr(std::move(h.sptr)) {
      h._data = nullptr;
      h._size = 0;
    }
 
    handle_heap &operator=(handle_heap &&h) noexcept {
      // release current resources if they are not shared and not null
      if (not sptr and not(is_null())) destruct({_data, _size});
 
      // move the resources from the source handle
      _data = h._data;
      _size = h._size;
      sptr  = std::move(h.sptr);
 
      // reset the source handle to a null state
      h._data = nullptr;
      h._size = 0;
      return *this;
    }
 
    explicit handle_heap(handle_heap const &h) : handle_heap(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, address_space>(_data, h.data(), h.size() * sizeof(T));
      } else {
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }
 
    handle_heap &operator=(handle_heap const &h) {
      *this = handle_heap{h};
      return *this;
    }
 
    template <OwningHandle<value_type> H>
    explicit handle_heap(H const &h) : handle_heap(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, H::address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        static_assert(address_space == H::address_space,
                      "Constructing an nda::mem::handle_heap from a handle of a different address space requires a trivially copyable value_type");
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }
 
    template <Allocator AS>
    handle_heap &operator=(handle_heap<T, AS> const &h) {
      *this = handle_heap{h};
      return *this;
    }
 
    handle_heap(long size, do_not_initialize_t) {
      if (size == 0) return;
      auto b = allocator.allocate(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
    }
 
    handle_heap(long size, init_zero_t) {
      if (size == 0) return;
      auto b = allocator.allocate_zero(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
    }
 
    handle_heap(long size) {
      if (size == 0) return;
      blk_t b;
      if constexpr (is_complex_v<T> && init_dcmplx)
        b = allocator.allocate_zero(size * sizeof(T));
      else
        b = allocator.allocate(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
 
      // call placement new for non trivial and non complex types
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < size; ++i) new (_data + i) T();
      }
    }
 
    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }
 
    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }
 
    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      // check the invariants in debug mode
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }
 
    [[nodiscard]] T *data() const noexcept { return _data; }
 
    [[nodiscard]] long size() const noexcept { return _size; }
  };
 
  template <typename T, size_t Size>
  struct handle_stack {
    static_assert(std::is_copy_constructible_v<T>, "nda::mem::handle_stack requires the value_type to be copy constructible");
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_stack requires the value_type to have a non-throwing destructor");
 
    private:
    // Memory buffer on the stack to store the data.
    std::array<char, sizeof(T) * Size> buffer;
 
    public:
    using value_type = T;
 
    static constexpr auto address_space = Host;
 
    ~handle_stack() noexcept {
      if constexpr (!std::is_trivial_v<T>) {
        // explicitly call the destructor for each non-trivial object
        for (size_t i = 0; i < Size; ++i) data()[i].~T();
      }
    }
 
    handle_stack() = default;
 
    handle_stack(handle_stack &&h) noexcept { operator=(h); }
 
    handle_stack &operator=(handle_stack &&h) noexcept {
      operator=(h);
      return *this;
    }
 
    handle_stack(handle_stack const &h) noexcept { operator=(h); }
 
    handle_stack &operator=(handle_stack const &h) {
      for (size_t i = 0; i < Size; ++i) new (data() + i) T(h[i]);
      return *this;
    }
 
    handle_stack(long /*size*/, do_not_initialize_t) {}
 
    handle_stack(long /*size*/, init_zero_t) {
      static_assert(std::is_scalar_v<T> or is_complex_v<T>, "nda::mem::handle_stack can only be initialized to zero for scalar and complex types");
      for (size_t i = 0; i < Size; ++i) data()[i] = 0;
    }
 
    handle_stack(long /*size*/) : handle_stack{} {
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < Size; ++i) new (data() + i) T();
      }
    }
 
    [[nodiscard]] T &operator[](long i) noexcept { return data()[i]; }
 
    [[nodiscard]] T const &operator[](long i) const noexcept { return data()[i]; }
 
    static constexpr bool is_null() noexcept { return false; }
 
    [[nodiscard]] T *data() const noexcept { return (T *)buffer.data(); }
 
    static constexpr long size() noexcept { return Size; }
  };
 
  template <typename T, size_t Size>
  struct handle_sso { // struct alignas(alignof(T)) handle_sso {
    static_assert(Size > 0, "Size == 0 makes no sense in nda::mem::handle_sso");
    static_assert(std::is_copy_constructible_v<T>, "nda::mem::handle_sso requires the value_type to be copy constructible");
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_sso requires the value_type to have a non-throwing destructor");
 
    private:
    // Memory buffer on the stack to store the data.
    std::array<char, sizeof(T) * Size> buffer;
 
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != nullptr.
    size_t _size = 0;
 
    // Release the handled memory.
    void clean() noexcept {
      if (is_null()) return;
      if constexpr (!std::is_trivial_v<T>) {
        // explicitly call the destructor for each non-trivial object
        for (size_t i = 0; i < _size; ++i) data()[i].~T();
      }
      if (on_heap()) mallocator<>::deallocate({(char *)_data, _size * sizeof(T)});
      _data = nullptr;
      _size = 0;
    }
 
    public:
    using value_type = T;
 
    static constexpr auto address_space = Host;
 
    handle_sso(){}; // NOLINT (user-defined constructor to avoid value initialization of the buffer)
 
    ~handle_sso() noexcept { clean(); }
 
    handle_sso(handle_sso &&h) noexcept : _size(h._size) {
      if (on_heap()) {
        _data = h._data;
      } else {
        if (_size != 0) {
          _data = (T *)buffer.data();
          for (size_t i = 0; i < _size; ++i) new (data() + i) T(h[i]);
        }
      }
      h._data = nullptr;
      h._size = 0;
    }
 
    handle_sso &operator=(handle_sso &&h) noexcept {
      clean();
      _size = h._size;
      if (on_heap()) {
        _data = h._data;
      } else {
        if (_size != 0) {
          _data = (T *)buffer.data();
          for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
        }
      }
      h._data = nullptr;
      h._size = 0;
      return *this;
    }
 
    handle_sso(handle_sso const &h) : handle_sso(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }
 
    handle_sso &operator=(handle_sso const &h) {
      if (this == &h) return *this; // self assignment
      clean();
      _size = h._size;
      if (_size == 0) return *this;
      if (on_heap()) {
        auto b = mallocator<>::allocate(_size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      } else {
        _data = (T *)buffer.data();
      }
      for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      return *this;
    }
 
    template <OwningHandle<value_type> H>
    explicit handle_sso(H const &h) : handle_sso(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, H::address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        static_assert(address_space == H::address_space,
                      "Constructing an nda::mem::handle_sso from a handle of a different address space requires a trivially copyable value_type");
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }
 
    handle_sso(long size, do_not_initialize_t) {
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
      } else {
        auto b = mallocator<>::allocate(size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
    }
 
    handle_sso(long size, init_zero_t) {
      static_assert(std::is_scalar_v<T> or is_complex_v<T>, "nda::mem::handle_sso can only be initialized to zero for scalar and complex types");
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
        for (size_t i = 0; i < _size; ++i) data()[i] = 0;
      } else {
        auto b = mallocator<>::allocate_zero(size * sizeof(T)); //, alignof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
    }
 
    handle_sso(long size) {
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
      } else {
        blk_t b;
        if constexpr (is_complex_v<T> && init_dcmplx)
          b = mallocator<>::allocate_zero(size * sizeof(T));
        else
          b = mallocator<>::allocate(size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
 
      // call placement new for non trivial and non complex types
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < size; ++i) new (_data + i) T();
      }
    }
 
    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }
 
    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }
 
    [[nodiscard]] bool on_heap() const { return _size > Size; }
 
    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }
 
    [[nodiscard]] T *data() const noexcept { return _data; }
 
    [[nodiscard]] long size() const noexcept { return _size; }
  };
 
  template <typename T, AddressSpace AdrSp>
  struct handle_shared {
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_shared requires the value_type to have a non-throwing destructor");
 
    private:
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != 0.
    size_t _size = 0;
 
    // Type of the memory block, i.e. a pointer to the data and its size.
    using blk_t = std::pair<T *, size_t>;
 
    // For shared ownership (points to a blk_T_t).
    std::shared_ptr<void> sptr;
 
    public:
    using value_type = T;
 
    static constexpr auto address_space = AdrSp;
 
    handle_shared() = default;
 
    handle_shared(T *data, size_t size, void *foreign_handle, void (*foreign_decref)(void *)) noexcept
       : _data(data), _size(size), sptr{foreign_handle, foreign_decref} {}
 
    template <Allocator A>
    handle_shared(handle_heap<T, A> const &h) noexcept
      requires(A::address_space == address_space)
       : _data(h.data()), _size(h.size()) {
      if (not h.is_null()) sptr = h.get_sptr();
    }
 
    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }
 
    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }
 
    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      // Check the Invariants in Debug Mode
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }
 
    [[nodiscard]] long refcount() const noexcept { return sptr.use_count(); }
 
    [[nodiscard]] T *data() const noexcept { return _data; }
 
    [[nodiscard]] long size() const noexcept { return _size; }
  };
 
  template <typename T, AddressSpace AdrSp = Host>
  struct handle_borrowed {
    private:
    // Value type of the data with const removed.
    using T0 = std::remove_const_t<T>;
 
    // Parent handle (required for regular -> shared promotion in Python Converter).
    handle_heap<T0> const *_parent = nullptr;
 
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    public:
    using value_type = T;
 
    static constexpr auto address_space = AdrSp;
 
    handle_borrowed() = default;
 
    handle_borrowed &operator=(handle_borrowed &&) = default;
 
    handle_borrowed(handle_borrowed const &) = default;
 
    handle_borrowed &operator=(handle_borrowed const &) = default;
 
    handle_borrowed(T *ptr) noexcept : _data(ptr) {}
 
    template <Handle H>
      requires(address_space == H::address_space and (std::is_const_v<value_type> or !std::is_const_v<typename H::value_type>)
               and std::is_same_v<const value_type, const typename H::value_type>)
    handle_borrowed(H const &h, long offset = 0) noexcept : _data(h.data() + offset) {
      if constexpr (std::is_same_v<H, handle_heap<T0>>) _parent = &h;
    }
 
    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }
 
    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }
 
    [[nodiscard]] bool is_null() const noexcept { return _data == nullptr; }
 
    [[nodiscard]] handle_heap<T0> const *parent() const { return _parent; }
 
    [[nodiscard]] T *data() const noexcept { return _data; }
  };
 
} // namespace nda::mem