linalg/async/transform.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/async/async.hpp>
#include <uni20/async/awaiters.hpp>
#include <uni20/async/debug_scheduler.hpp>
#include <uni20/common/trace.hpp>
#include <uni20/linalg/async/detail/output.hpp>
#include <uni20/tensor/output.hpp>
#include <uni20/tensor/transform.hpp>
 
#include <concepts>
#include <cstddef>
#include <tuple>
#include <utility>
 
namespace uni20
{
namespace detail
{
 
template <class OutputTensor, class... InputTensors>
void validate_async_transform_aliasing(async::Async<OutputTensor> const& output,
                                       async::Async<InputTensors> const&... inputs)
{
  if constexpr (sizeof...(InputTensors) > 0)
  {
    auto const* const output_queue = std::addressof(output.queue());
    auto validate_input = [&](auto const& input) {
      ERROR_IF(output_queue == std::addressof(input.queue()),
               "async elementwise transform output must not share an epoch queue with an input");
    };
    (validate_input(inputs), ...);
  }
}
 
template <AsyncTensorOutput OutputTensor, TensorView FirstInputTensor>
[[nodiscard]] OutputTensor& prepare_async_assign_transform_output(async::shared_storage<OutputTensor>& storage,
                                                                  FirstInputTensor const& first_input)
{
  if (storage.constructed()) return *storage;
 
  using extents_type = tensor_extents_t<OutputTensor>;
  if constexpr (std::constructible_from<OutputTensor, uninitialized_t, extents_type const&>)
  {
    auto const extents = convert_tensor_extents<extents_type>(first_input.extents());
    return storage.emplace(uninitialized, extents);
  }
  else if constexpr (std::constructible_from<OutputTensor, extents_type const&>)
  {
    auto const extents = convert_tensor_extents<extents_type>(first_input.extents());
    return storage.emplace(extents);
  }
  else
  {
    throw async::buffer_write_uninitialized{};
  }
}
 
template <class BackendSelector, class Function, AsyncTensorOutput OutputTensor, class Awaited, std::size_t... Input>
void invoke_async_assign_transform(BackendSelector const& selector, linalg::transform_op<Function>& operation,
                                   OutputTensor& output, Awaited& awaited, std::index_sequence<Input...>)
{
  uni20::assign_transform(selector, output, std::move(operation.function), std::get<Input + 1>(awaited)...);
}
 
template <class BackendSelector, class Function, AsyncTensorOutput OutputTensor, class Awaited, std::size_t... Input>
void invoke_async_transform_inplace(BackendSelector const& selector, linalg::transform_inplace_op<Function>& operation,
                                    OutputTensor& output, Awaited& awaited, std::index_sequence<Input...>)
{
  uni20::transform_inplace(selector, output, std::move(operation.function), std::get<Input + 1>(awaited)...);
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
async::AsyncTask co_assign_transform(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,
                                     linalg::transform_op<Function> operation,
                                     async::ReadBuffer<InputTensors>... inputs)
{
  if constexpr (async::is_async_alias_v<OutputTensor>)
  {
    AsyncAliasWriteDescriptorAwaiter output_descriptor(output);
    auto awaited = co_await async::all(output_descriptor, inputs...);
    auto mutable_output = std::get<0>(awaited);
    require_output(mutable_output, std::get<1>(awaited).extents());
    invoke_async_assign_transform(selector, operation, mutable_output, awaited,
                                  std::index_sequence_for<InputTensors...>{});
  }
  else
  {
    auto output_storage = output.storage();
    auto awaited = co_await async::all(output_storage, inputs...);
    auto& storage = std::get<0>(awaited);
    auto& output_value = prepare_async_assign_transform_output(storage, std::get<1>(awaited));
    invoke_async_assign_transform(selector, operation, output_value, awaited,
                                  std::index_sequence_for<InputTensors...>{});
  }
  co_return;
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
async::AsyncTask co_transform_inplace(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,
                                      linalg::transform_inplace_op<Function> operation,
                                      async::ReadBuffer<InputTensors>... inputs)
{
  if constexpr (async::is_async_alias_v<OutputTensor>)
  {
    AsyncAliasWriteDescriptorAwaiter output_descriptor(output);
    auto awaited = co_await async::all(output_descriptor, inputs...);
    auto mutable_output = std::get<0>(awaited);
    invoke_async_transform_inplace(selector, operation, mutable_output, awaited,
                                   std::index_sequence_for<InputTensors...>{});
  }
  else
  {
    auto output_storage = output.storage();
    auto awaited = co_await async::all(output_storage, inputs...);
    auto& storage = std::get<0>(awaited);
    if (!storage.constructed()) throw async::buffer_write_uninitialized{};
    invoke_async_transform_inplace(selector, operation, *storage, awaited, std::index_sequence_for<InputTensors...>{});
  }
  co_return;
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function>
async::AsyncTask co_assign_transform_existing(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,
                                              linalg::transform_op<Function> operation)
{
  if constexpr (async::is_async_alias_v<OutputTensor>)
  {
    AsyncAliasWriteDescriptorAwaiter output_descriptor(output);
    auto awaited = co_await async::all(output_descriptor);
    auto mutable_output = std::get<0>(awaited);
    uni20::assign_transform(selector, mutable_output, std::move(operation.function));
  }
  else
  {
    auto output_storage = output.storage();
    auto awaited = co_await async::all(output_storage);
    auto& storage = std::get<0>(awaited);
    if (!storage.constructed()) throw async::buffer_write_uninitialized{};
    uni20::assign_transform(selector, *storage, std::move(operation.function));
  }
  co_return;
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
void schedule_async_assign_transform(BackendSelector selector, async::Async<OutputTensor>& output,
                                     linalg::transform_op<Function> operation,
                                     async::Async<InputTensors> const&... inputs)
{
  validate_async_transform_aliasing(output, inputs...);
  auto task = co_assign_transform(std::move(selector), output.write(), std::move(operation), inputs.read()...);
  task.debug_name("assign_transform");
  async::schedule(std::move(task));
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
void schedule_async_transform_inplace(BackendSelector selector, async::Async<OutputTensor>& output,
                                      linalg::transform_inplace_op<Function> operation,
                                      async::Async<InputTensors> const&... inputs)
{
  validate_async_transform_aliasing(output, inputs...);
  auto task = co_transform_inplace(std::move(selector), output.write(), std::move(operation), inputs.read()...);
  task.debug_name("transform_inplace");
  async::schedule(std::move(task));
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function>
void schedule_async_assign_transform_existing(BackendSelector selector, async::Async<OutputTensor>& output,
                                              linalg::transform_op<Function> operation)
{
  auto task = co_assign_transform_existing(std::move(selector), output.write(), std::move(operation));
  task.debug_name("assign_transform");
  async::schedule(std::move(task));
}
 
} // namespace detail
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function>
void assign_transform(BackendSelector selector, async::Async<OutputTensor>& output, Function&& function)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  detail::schedule_async_assign_transform_existing(std::move(selector), output, std::move(operation));
}
 
template <AsyncTensorOutput OutputTensor, class Function>
void assign_transform(async::Async<OutputTensor>& output, Function&& function)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend_for<OutputTensor>(operation);
  detail::schedule_async_assign_transform_existing(std::move(selector), output, std::move(operation));
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Value>
void fill(BackendSelector selector, async::Async<OutputTensor>& output, Value&& value)
{
  assign_transform(std::move(selector), output, linalg::constant{std::forward<Value>(value)});
}
 
template <AsyncTensorOutput OutputTensor, class Value> void fill(async::Async<OutputTensor>& output, Value&& value)
{
  assign_transform(output, linalg::constant{std::forward<Value>(value)});
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView FirstInputTensor,
          TensorView... RestInputTensors>
  requires detail::OverwriteTransformTensors<OutputTensor, FirstInputTensor, RestInputTensors...>
void assign_transform(BackendSelector selector, async::Async<OutputTensor>& output, Function&& function,
                      async::Async<FirstInputTensor> const& first_input,
                      async::Async<RestInputTensors> const&... rest_inputs)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  detail::schedule_async_assign_transform(std::move(selector), output, std::move(operation), first_input,
                                          rest_inputs...);
}
 
template <AsyncTensorOutput OutputTensor, class Function, TensorView FirstInputTensor, TensorView... RestInputTensors>
  requires detail::OverwriteTransformTensors<OutputTensor, FirstInputTensor, RestInputTensors...>
void assign_transform(async::Async<OutputTensor>& output, Function&& function,
                      async::Async<FirstInputTensor> const& first_input,
                      async::Async<RestInputTensors> const&... rest_inputs)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend_for<OutputTensor, FirstInputTensor, RestInputTensors...>(operation);
  detail::schedule_async_assign_transform(std::move(selector), output, std::move(operation), first_input,
                                          rest_inputs...);
}
 
template <class BackendSelector, AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
  requires detail::UpdateTransformTensors<OutputTensor, InputTensors...>
void transform_inplace(BackendSelector selector, async::Async<OutputTensor>& output, Function&& function,
                       async::Async<InputTensors> const&... inputs)
{
  auto operation = linalg::transform_inplace_op{std::forward<Function>(function)};
  detail::schedule_async_transform_inplace(std::move(selector), output, std::move(operation), inputs...);
}
 
template <AsyncTensorOutput OutputTensor, class Function, TensorView... InputTensors>
  requires detail::UpdateTransformTensors<OutputTensor, InputTensors...>
void transform_inplace(async::Async<OutputTensor>& output, Function&& function,
                       async::Async<InputTensors> const&... inputs)
{
  auto operation = linalg::transform_inplace_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend_for<OutputTensor, InputTensors...>(operation);
  detail::schedule_async_transform_inplace(std::move(selector), output, std::move(operation), inputs...);
}
 
} // namespace uni20