linalg/async/detail/output.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/async/async.hpp>
#include <uni20/async/debug_scheduler.hpp>
#include <uni20/common/initialization.hpp>
#include <uni20/linalg/async/concepts.hpp>
 
#include <concepts>
#include <memory>
#include <type_traits>
#include <utility>
 
namespace uni20
{
namespace detail
{
 
template <class Alias> class AsyncAliasWriteDescriptorAwaiter {
  public:
    explicit AsyncAliasWriteDescriptorAwaiter(async::WriteBuffer<Alias>& output) : output_(std::addressof(output)) {}
 
    [[nodiscard]] bool await_ready() const noexcept { return output_->await_ready(); }
 
    void await_suspend(async::BasicTask task) noexcept { output_->await_suspend(std::move(task)); }
 
    [[nodiscard]] Alias const& await_resume() { return output_->await_resume().get(); }
 
#if UNI20_DEBUG_DAG
    [[nodiscard]] auto node() const { return output_->node(); }
    static constexpr auto debug_task_role() noexcept { return TaskRegistry::EpochTaskRole::Writer; }
#endif
 
  private:
    async::WriteBuffer<Alias>* output_;
};
 
template <AsyncTensorOutput Tensor> [[nodiscard]] auto mutable_async_tensor_awaiter(async::WriteBuffer<Tensor>& output)
{
  if constexpr (async::is_async_alias_v<Tensor>)
    return AsyncAliasWriteDescriptorAwaiter<Tensor>{output};
  else
    return output.storage();
}
 
template <AsyncTensorOutput Tensor, class Awaited>
[[nodiscard]] decltype(auto) mutable_async_tensor_value(Awaited&& awaited)
{
  if constexpr (async::is_async_alias_v<Tensor>)
  {
    return Tensor(std::forward<Awaited>(awaited));
  }
  else
  {
    if (!awaited.constructed()) throw async::buffer_write_uninitialized{};
    return *awaited;
  }
}
 
template <AsyncTensorOutput Tensor>
[[nodiscard]] Tensor& prepare_async_scalar_output(async::shared_storage<Tensor>& storage)
{
  if (storage.constructed()) return *storage;
  if constexpr (std::constructible_from<Tensor, uninitialized_t>)
    return storage.emplace(uninitialized);
  else if constexpr (std::default_initializable<Tensor>)
    return storage.emplace();
  else
    throw async::buffer_write_uninitialized{};
}
 
} // namespace detail
} // namespace uni20