access.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/tensor/concepts.hpp>
 
#include <concepts>
#include <cstddef>
#include <optional>
#include <type_traits>
#include <utility>
 
namespace uni20
{
template <class State>
concept TensorAccessState = std::move_constructible<State> && std::is_nothrow_move_constructible_v<State> &&
                            requires(State& state) {
  { state.release() } noexcept;
};
 
namespace detail
{
 
struct immediate_mdspan_access_state
{
    void release() noexcept {}
};
 
template <ImmediateTensorView Tensor> class borrowed_read_tensor_lease {
  public:
    using tensor_type = std::remove_cvref_t<Tensor>;
    using mdspan_type = std::remove_cvref_t<decltype(std::declval<tensor_type const&>().mdspan())>;
    using storage_policy = tensor_storage_policy_t<tensor_type>;
    using extents_type = typename mdspan_type::extents_type;
    using index_type = typename mdspan_type::index_type;
 
    explicit borrowed_read_tensor_lease(tensor_type const& tensor) noexcept : tensor_(&tensor) {}
 
    borrowed_read_tensor_lease(borrowed_read_tensor_lease const&) = delete;
    borrowed_read_tensor_lease& operator=(borrowed_read_tensor_lease const&) = delete;
 
    borrowed_read_tensor_lease(borrowed_read_tensor_lease&& other) noexcept
        : tensor_(std::exchange(other.tensor_, nullptr))
    {}
 
    borrowed_read_tensor_lease& operator=(borrowed_read_tensor_lease&& other) noexcept
    {
      if (this != &other) tensor_ = std::exchange(other.tensor_, nullptr);
      return *this;
    }
 
    [[nodiscard]] decltype(auto) mdspan() const& { return this->source().mdspan(); }
 
    void mdspan() && = delete;
 
    [[nodiscard]] decltype(auto) backend_selector() const { return this->source().backend_selector(); }
 
    [[nodiscard]] decltype(auto) storage() const
      requires requires(tensor_type const& tensor) { tensor.storage(); }
    {
      return this->source().storage();
    }
 
    [[nodiscard]] decltype(auto) extents() const { return this->source().extents(); }
 
    [[nodiscard]] index_type extent(std::size_t axis) const { return this->source().extent(axis); }
 
    [[nodiscard]] static constexpr std::size_t rank() noexcept { return mdspan_type::rank(); }
 
    void release() noexcept { tensor_ = nullptr; }
 
  private:
    [[nodiscard]] tensor_type const& source() const
    {
      CHECK(tensor_ != nullptr, "cannot use a released borrowed tensor lease");
      return *tensor_;
    }
 
    tensor_type const* tensor_ = nullptr;
};
 
template <MutableImmediateTensorView Tensor> class borrowed_write_tensor_lease {
  public:
    using tensor_type = std::remove_cvref_t<Tensor>;
    using mdspan_type = std::remove_cvref_t<decltype(std::declval<tensor_type&>().mdspan())>;
    using const_mdspan_type = std::remove_cvref_t<decltype(std::declval<tensor_type const&>().mdspan())>;
    using storage_policy = tensor_storage_policy_t<tensor_type>;
    using extents_type = typename mdspan_type::extents_type;
    using index_type = typename mdspan_type::index_type;
 
    explicit borrowed_write_tensor_lease(tensor_type& tensor) noexcept : tensor_(&tensor) {}
 
    borrowed_write_tensor_lease(borrowed_write_tensor_lease const&) = delete;
    borrowed_write_tensor_lease& operator=(borrowed_write_tensor_lease const&) = delete;
 
    borrowed_write_tensor_lease(borrowed_write_tensor_lease&& other) noexcept
        : tensor_(std::exchange(other.tensor_, nullptr))
    {}
 
    borrowed_write_tensor_lease& operator=(borrowed_write_tensor_lease&& other) noexcept
    {
      if (this != &other) tensor_ = std::exchange(other.tensor_, nullptr);
      return *this;
    }
 
    [[nodiscard]] decltype(auto) mdspan() & { return this->source().mdspan(); }
 
    [[nodiscard]] decltype(auto) mdspan() const& { return std::as_const(this->source()).mdspan(); }
 
    void mdspan() && = delete;
 
    [[nodiscard]] decltype(auto) backend_selector() const { return this->source().backend_selector(); }
 
    [[nodiscard]] decltype(auto) extents() const { return this->source().extents(); }
 
    [[nodiscard]] index_type extent(std::size_t axis) const { return this->source().extent(axis); }
 
    [[nodiscard]] static constexpr std::size_t rank() noexcept { return mdspan_type::rank(); }
 
    void release() noexcept { tensor_ = nullptr; }
 
  private:
    [[nodiscard]] tensor_type& source() const
    {
      CHECK(tensor_ != nullptr, "cannot use a released borrowed tensor lease");
      return *tensor_;
    }
 
    tensor_type* tensor_ = nullptr;
};
 
} // namespace detail
 
template <MdspanLike Mdspan, TensorAccessState AccessState>
  requires std::is_const_v<typename Mdspan::element_type> && std::is_nothrow_move_constructible_v<Mdspan>
class read_mdspan_lease {
  public:
    using mdspan_type = Mdspan;
    using access_state_type = AccessState;
 
    read_mdspan_lease(access_state_type state, mdspan_type mdspan) noexcept
        : state_(std::move(state)), mdspan_(std::move(mdspan))
    {}
 
    read_mdspan_lease(read_mdspan_lease const&) = delete;
    read_mdspan_lease& operator=(read_mdspan_lease const&) = delete;
 
    read_mdspan_lease(read_mdspan_lease&& other) noexcept
        : state_(std::move(other.state_)), mdspan_(std::move(other.mdspan_))
    {
      other.mdspan_.reset();
    }
 
    read_mdspan_lease&
    operator=(read_mdspan_lease&& other) noexcept
      requires(std::is_nothrow_move_assignable_v<access_state_type> && std::is_nothrow_move_assignable_v<mdspan_type>)
    {
      if (this != &other)
      {
        this->release();
        state_ = std::move(other.state_);
        mdspan_ = std::move(other.mdspan_);
        other.mdspan_.reset();
      }
      return *this;
    }
 
    ~read_mdspan_lease() { this->release(); }
 
    [[nodiscard]] mdspan_type& mdspan() &
    {
      CHECK(mdspan_.has_value(), "cannot use an inactive mdspan lease");
      return *mdspan_;
    }
 
    [[nodiscard]] mdspan_type const& mdspan() const&
    {
      CHECK(mdspan_.has_value(), "cannot use an inactive mdspan lease");
      return *mdspan_;
    }
 
    mdspan_type mdspan() && = delete;
 
    void release() noexcept
    {
      if (!mdspan_) return;
      mdspan_.reset();
      state_.release();
    }
 
  private:
    [[no_unique_address]] access_state_type state_;
    std::optional<mdspan_type> mdspan_;
};
 
template <MutableMdspanLike Mdspan, TensorAccessState AccessState>
  requires std::is_nothrow_move_constructible_v<Mdspan>
class write_mdspan_lease {
  public:
    using mdspan_type = Mdspan;
    using access_state_type = AccessState;
 
    write_mdspan_lease(access_state_type state, mdspan_type mdspan) noexcept
        : state_(std::move(state)), mdspan_(std::move(mdspan))
    {}
 
    write_mdspan_lease(write_mdspan_lease const&) = delete;
    write_mdspan_lease& operator=(write_mdspan_lease const&) = delete;
 
    write_mdspan_lease(write_mdspan_lease&& other) noexcept
        : state_(std::move(other.state_)), mdspan_(std::move(other.mdspan_))
    {
      other.mdspan_.reset();
    }
 
    write_mdspan_lease&
    operator=(write_mdspan_lease&& other) noexcept
      requires(std::is_nothrow_move_assignable_v<access_state_type> && std::is_nothrow_move_assignable_v<mdspan_type>)
    {
      if (this != &other)
      {
        this->release();
        state_ = std::move(other.state_);
        mdspan_ = std::move(other.mdspan_);
        other.mdspan_.reset();
      }
      return *this;
    }
 
    ~write_mdspan_lease() { this->release(); }
 
    [[nodiscard]] mdspan_type& mdspan() &
    {
      CHECK(mdspan_.has_value(), "cannot use an inactive mdspan lease");
      return *mdspan_;
    }
 
    [[nodiscard]] mdspan_type const& mdspan() const&
    {
      CHECK(mdspan_.has_value(), "cannot use an inactive mdspan lease");
      return *mdspan_;
    }
 
    mdspan_type mdspan() && = delete;
 
    void release() noexcept
    {
      if (!mdspan_) return;
      mdspan_.reset();
      state_.release();
    }
 
  private:
    [[no_unique_address]] access_state_type state_;
    std::optional<mdspan_type> mdspan_;
};
 
template <MdspanLike Mdspan, TensorAccessState AccessState, class BackendSelector, class StoragePolicy>
  requires std::is_const_v<typename Mdspan::element_type> && std::is_nothrow_move_constructible_v<Mdspan> &&
           std::is_nothrow_move_constructible_v<BackendSelector>
class read_tensor_lease {
  public:
    using mdspan_type = Mdspan;
    using access_state_type = AccessState;
    using backend_selector_type = BackendSelector;
    using storage_policy = StoragePolicy;
    using extents_type = typename mdspan_type::extents_type;
    using index_type = typename mdspan_type::index_type;
 
    read_tensor_lease(access_state_type state, mdspan_type mdspan, backend_selector_type selector) noexcept
        : state_(std::move(state)), mdspan_(std::move(mdspan)), selector_(std::move(selector))
    {}
 
    read_tensor_lease(read_tensor_lease const&) = delete;
    read_tensor_lease& operator=(read_tensor_lease const&) = delete;
 
    read_tensor_lease(read_tensor_lease&& other) noexcept
        : state_(std::move(other.state_)), mdspan_(std::move(other.mdspan_)), selector_(std::move(other.selector_))
    {
      other.mdspan_.reset();
    }
 
    read_tensor_lease&
    operator=(read_tensor_lease&& other) noexcept
      requires(std::is_nothrow_move_assignable_v<access_state_type> && std::is_nothrow_move_assignable_v<mdspan_type> &&
               std::is_nothrow_move_assignable_v<backend_selector_type>)
    {
      if (this != &other)
      {
        this->release();
        state_ = std::move(other.state_);
        mdspan_ = std::move(other.mdspan_);
        selector_ = std::move(other.selector_);
        other.mdspan_.reset();
      }
      return *this;
    }
 
    ~read_tensor_lease() { this->release(); }
 
    [[nodiscard]] mdspan_type& mdspan() &
    {
      CHECK(mdspan_.has_value(), "cannot resolve mdspan through an inactive tensor lease");
      return *mdspan_;
    }
 
    [[nodiscard]] mdspan_type const& mdspan() const&
    {
      CHECK(mdspan_.has_value(), "cannot resolve mdspan through an inactive tensor lease");
      return *mdspan_;
    }
 
    mdspan_type mdspan() && = delete;
 
    [[nodiscard]] backend_selector_type const& backend_selector() const noexcept { return selector_; }
 
    [[nodiscard]] decltype(auto) storage() const
      requires requires(access_state_type const& state) { state.storage(); }
    {
      return state_.storage();
    }
 
    [[nodiscard]] extents_type const& extents() const { return this->mdspan().extents(); }
 
    [[nodiscard]] index_type extent(std::size_t axis) const { return this->mdspan().extent(axis); }
 
    [[nodiscard]] static constexpr std::size_t rank() noexcept { return mdspan_type::rank(); }
 
    void release() noexcept
    {
      if (!mdspan_) return;
      mdspan_.reset();
      state_.release();
    }
 
  private:
    [[no_unique_address]] access_state_type state_;
    std::optional<mdspan_type> mdspan_;
    [[no_unique_address]] backend_selector_type selector_;
};
 
template <MutableMdspanLike MutableMdspan, MdspanLike ConstMdspan, TensorAccessState AccessState, class BackendSelector,
          class StoragePolicy>
  requires std::is_const_v<typename ConstMdspan::element_type> &&
           std::is_nothrow_move_constructible_v<MutableMdspan> &&
           std::is_nothrow_move_constructible_v<ConstMdspan> &&
           std::is_nothrow_move_constructible_v<BackendSelector>
class write_tensor_lease {
  public:
    using mdspan_type = MutableMdspan;
    using const_mdspan_type = ConstMdspan;
    using access_state_type = AccessState;
    using backend_selector_type = BackendSelector;
    using storage_policy = StoragePolicy;
    using extents_type = typename mdspan_type::extents_type;
    using index_type = typename mdspan_type::index_type;
 
    write_tensor_lease(access_state_type state, mdspan_type mdspan, const_mdspan_type const_mdspan,
                       backend_selector_type selector) noexcept
        : state_(std::move(state)), mdspan_(std::move(mdspan)), const_mdspan_(std::move(const_mdspan)),
          selector_(std::move(selector))
    {}
 
    write_tensor_lease(write_tensor_lease const&) = delete;
    write_tensor_lease& operator=(write_tensor_lease const&) = delete;
 
    write_tensor_lease(write_tensor_lease&& other) noexcept
        : state_(std::move(other.state_)), mdspan_(std::move(other.mdspan_)),
          const_mdspan_(std::move(other.const_mdspan_)), selector_(std::move(other.selector_))
    {
      other.const_mdspan_.reset();
      other.mdspan_.reset();
    }
 
    write_tensor_lease&
    operator=(write_tensor_lease&& other) noexcept
      requires(std::is_nothrow_move_assignable_v<access_state_type> && std::is_nothrow_move_assignable_v<mdspan_type> &&
               std::is_nothrow_move_assignable_v<const_mdspan_type> &&
               std::is_nothrow_move_assignable_v<backend_selector_type>)
    {
      if (this != &other)
      {
        this->release();
        state_ = std::move(other.state_);
        mdspan_ = std::move(other.mdspan_);
        const_mdspan_ = std::move(other.const_mdspan_);
        selector_ = std::move(other.selector_);
        other.const_mdspan_.reset();
        other.mdspan_.reset();
      }
      return *this;
    }
 
    ~write_tensor_lease() { this->release(); }
 
    [[nodiscard]] mdspan_type& mdspan() &
    {
      CHECK(mdspan_.has_value(), "cannot resolve mdspan through an inactive tensor lease");
      return *mdspan_;
    }
 
    [[nodiscard]] const_mdspan_type const& mdspan() const&
    {
      CHECK(const_mdspan_.has_value(), "cannot resolve mdspan through an inactive tensor lease");
      return *const_mdspan_;
    }
 
    mdspan_type mdspan() && = delete;
 
    [[nodiscard]] backend_selector_type const& backend_selector() const noexcept { return selector_; }
 
    [[nodiscard]] extents_type const& extents() const { return this->mdspan().extents(); }
 
    [[nodiscard]] index_type extent(std::size_t axis) const { return this->mdspan().extent(axis); }
 
    [[nodiscard]] static constexpr std::size_t rank() noexcept { return mdspan_type::rank(); }
 
    void release() noexcept
    {
      if (!mdspan_) return;
      const_mdspan_.reset();
      mdspan_.reset();
      state_.release();
    }
 
  private:
    [[no_unique_address]] access_state_type state_;
    std::optional<mdspan_type> mdspan_;
    std::optional<const_mdspan_type> const_mdspan_;
    [[no_unique_address]] backend_selector_type selector_;
};
 
template <class Lease>
concept ReadMdspanLease =
    std::move_constructible<Lease> && (!std::copy_constructible<Lease>) && requires(Lease& lease) {
      requires MdspanLike<decltype(lease.mdspan())>;
      requires std::is_const_v<typename std::remove_cvref_t<decltype(lease.mdspan())>::element_type>;
      { lease.release() } noexcept;
    };
 
template <class Lease>
concept WriteMdspanLease =
    std::move_constructible<Lease> && (!std::copy_constructible<Lease>) && requires(Lease& lease) {
      requires MutableMdspanLike<decltype(lease.mdspan())>;
      { lease.release() } noexcept;
    };
 
template <class Lease>
concept HostReadMdspanLease = ReadMdspanLease<Lease> && HostAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept HostWriteMdspanLease =
    WriteMdspanLease<Lease> && HostAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept CudaReadMdspanLease = ReadMdspanLease<Lease> && CudaAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept CudaWriteMdspanLease =
    WriteMdspanLease<Lease> && CudaAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept ReadTensorLease = ImmediateTensorView<Lease> && std::move_constructible<Lease> &&
                          (!std::copy_constructible<Lease>) && requires(Lease& lease) {
                            { lease.release() } noexcept;
                          };
 
template <class Lease>
concept WriteTensorLease = ReadTensorLease<Lease> && MutableImmediateTensorView<Lease>;
 
template <class Lease>
concept HostReadTensorLease = ReadTensorLease<Lease> && HostAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept HostWriteTensorLease =
    WriteTensorLease<Lease> && HostAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept CudaReadTensorLease = ReadTensorLease<Lease> && CudaAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease>
concept CudaWriteTensorLease =
    WriteTensorLease<Lease> && CudaAccessibleMdspan<decltype(std::declval<Lease&>().mdspan())>;
 
template <class Lease> class ready_access {
  public:
    using lease_type = Lease;
 
    explicit ready_access(lease_type lease) : lease_(std::move(lease)) {}
 
    ready_access(ready_access const&) = delete;
    ready_access& operator=(ready_access const&) = delete;
    ready_access(ready_access&&) = default;
    ready_access& operator=(ready_access&&) = default;
 
    [[nodiscard]] bool await_ready() const noexcept { return true; }
 
    template <class Task> void await_suspend(Task&&) const noexcept {}
 
    [[nodiscard]] lease_type await_resume() { return std::move(lease_); }
 
  private:
    lease_type lease_;
};
 
template <HostAccessibleMdspan Mdspan>
  requires std::is_const_v<typename std::remove_cvref_t<Mdspan>::element_type>
[[nodiscard]] auto acquire_host_read_access_sync(Mdspan const& mdspan)
{
  using mdspan_type = std::remove_cvref_t<Mdspan>;
  return read_mdspan_lease<mdspan_type, detail::immediate_mdspan_access_state>{detail::immediate_mdspan_access_state{},
                                                                               mdspan};
}
 
template <MutableMdspanLike Mdspan>
  requires HostAccessibleMdspan<Mdspan>
[[nodiscard]] auto acquire_host_write_access_sync(Mdspan& mdspan)
{
  using mdspan_type = std::remove_cvref_t<Mdspan>;
  return write_mdspan_lease<mdspan_type, detail::immediate_mdspan_access_state>{detail::immediate_mdspan_access_state{},
                                                                                mdspan};
}
 
template <ImmediateTensorView Tensor>
  requires HostAccessibleMdspan<immediate_tensor_mdspan_t<Tensor>>
[[nodiscard]] auto acquire_host_read_access_sync(Tensor& tensor)
{
  using tensor_type = std::remove_cvref_t<Tensor>;
  return detail::borrowed_read_tensor_lease<tensor_type>{tensor};
}
 
template <MutableImmediateTensorView Tensor>
  requires HostAccessibleMdspan<mutable_immediate_tensor_mdspan_t<Tensor>>
[[nodiscard]] auto acquire_host_write_access_sync(Tensor& tensor)
{
  using tensor_type = std::remove_cvref_t<Tensor>;
  return detail::borrowed_write_tensor_lease<tensor_type>{tensor};
}
 
template <ImmediateTensorView Tensor>
  requires HostAccessibleMdspan<immediate_tensor_mdspan_t<Tensor>>
[[nodiscard]] auto acquire_host_read_access_async(Tensor& tensor)
{
  return ready_access{acquire_host_read_access_sync(tensor)};
}
 
template <MutableImmediateTensorView Tensor>
  requires HostAccessibleMdspan<mutable_immediate_tensor_mdspan_t<Tensor>>
[[nodiscard]] auto acquire_host_write_access_async(Tensor& tensor)
{
  return ready_access{acquire_host_write_access_sync(tensor)};
}
 
namespace detail
{
template <class Mdspec>
using host_read_mdspan_lease_t = decltype(acquire_host_read_access_sync(std::declval<Mdspec const&>()));
 
template <class Mdspec>
using host_write_mdspan_lease_t = decltype(acquire_host_write_access_sync(std::declval<Mdspec&>()));
} // namespace detail
 
template <class Mdspec>
using host_read_mdspan_t =
    std::remove_cvref_t<decltype(std::declval<detail::host_read_mdspan_lease_t<Mdspec>&>().mdspan())>;
 
template <class Mdspec>
using host_write_mdspan_t =
    std::remove_cvref_t<decltype(std::declval<detail::host_write_mdspan_lease_t<Mdspec>&>().mdspan())>;
 
template <class Mdspec>
concept HostReadableMdspec = MdspecLike<Mdspec> && requires(Mdspec const& mdspec) {
  { acquire_host_read_access_sync(mdspec) } -> HostReadMdspanLease;
};
 
template <class Mdspec>
concept HostWritableMdspec = MutableMdspecLike<Mdspec> && requires(Mdspec& mdspec) {
  { acquire_host_write_access_sync(mdspec) } -> HostWriteMdspanLease;
};
 
namespace detail
{
template <class Tensor>
using host_read_tensor_lease_t = decltype(acquire_host_read_access_sync(std::declval<Tensor const&>()));
 
template <class Tensor>
using host_write_tensor_lease_t = decltype(acquire_host_write_access_sync(std::declval<Tensor&>()));
} // namespace detail
 
template <class Tensor>
using host_read_tensor_mdspan_t =
    std::remove_cvref_t<decltype(std::declval<detail::host_read_tensor_lease_t<Tensor>&>().mdspan())>;
 
template <class Tensor>
using host_write_tensor_mdspan_t =
    std::remove_cvref_t<decltype(std::declval<detail::host_write_tensor_lease_t<Tensor>&>().mdspan())>;
 
template <class Tensor>
concept HostReadableTensor = TensorView<Tensor> && requires(Tensor const& tensor) {
  { acquire_host_read_access_sync(tensor) } -> HostReadTensorLease;
};
 
template <class Tensor>
concept HostWritableTensor = MutableTensorView<Tensor> && requires(Tensor& tensor) {
  { acquire_host_write_access_sync(tensor) } -> HostWriteTensorLease;
};
 
} // namespace uni20