async/matrix_product.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/concepts.hpp>
#include <uni20/linalg/async/detail/output.hpp>
#include <uni20/linalg/async/dispatch.hpp>
#include <uni20/linalg/ops/matrix_product.hpp>
#include <uni20/tensor/output.hpp>
 
#if UNI20_BACKEND_CUBLAS
#include <uni20/linalg/backends/cublas/gemm_task.hpp>
#endif
 
#include <concepts>
#include <memory>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
namespace detail
{
template <class OutputTensor, class LhsTensor, class RhsTensor>
void validate_async_matrix_product_aliasing(async::Async<OutputTensor> const& output,
                                            async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs)
{
  auto const* const output_queue = std::addressof(output.queue());
  ERROR_IF(output_queue == std::addressof(lhs.queue()) || output_queue == std::addressof(rhs.queue()),
           "async matrix product output must not share an epoch queue with an input");
}
 
template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class AlphaAwaiter>
async::AsyncTask co_assign_product(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,
                                   async::ReadBuffer<LhsTensor> lhs, async::ReadBuffer<RhsTensor> rhs,
                                   AlphaAwaiter alpha, std::optional<Precision> requested)
{
  auto output_awaiter = uni20::detail::mutable_async_tensor_awaiter(output);
  auto awaited = co_await async::all(output_awaiter, lhs, rhs, alpha);
  auto const& lhs_value = std::get<1>(awaited);
  auto const& rhs_value = std::get<2>(awaited);
  auto const& alpha_value = std::get<3>(awaited);
 
  using scalar_type = uni20::tensor_element_t<OutputTensor>;
  scalar_type const alpha_scalar = static_cast<scalar_type>(alpha_value);
  auto lhs_descriptor = uni20::mdspec_of(lhs_value);
  auto rhs_descriptor = uni20::mdspec_of(rhs_value);
  auto p = requested ? *requested : common_default_precision(lhs_value, rhs_value);
  require_precision<scalar_type>(p);
#if UNI20_HAS_MPREAL
  if constexpr (has_runtime_precision_v<scalar_type>)
  {
    if constexpr (async::is_async_alias_v<OutputTensor>)
    {
      auto output_value = uni20::detail::mutable_async_tensor_value<OutputTensor>(std::get<0>(awaited));
      uni20::linalg::assign_product(selector, output_value, lhs_value, rhs_value, alpha_scalar, p);
    }
    else
    {
      auto& result = prepare_output(std::get<0>(awaited), matrix_product_shape(lhs_value, rhs_value), p);
      uni20::linalg::assign_product(selector, result, lhs_value, rhs_value, alpha_scalar, p);
    }
  }
  else
#endif
      if constexpr (async::is_async_alias_v<OutputTensor>)
  {
    auto output_value = uni20::detail::mutable_async_tensor_value<OutputTensor>(std::get<0>(awaited));
    co_await co_dispatch_kernel(selector, assign_product_op{}, output_value, alpha_scalar, lhs_descriptor,
                                rhs_descriptor);
  }
  else
  {
    auto& storage = std::get<0>(awaited);
    co_await co_dispatch_kernel(selector, assign_product_op{}, storage, alpha_scalar, lhs_descriptor, rhs_descriptor);
  }
  co_return;
}
 
template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class AlphaAwaiter, class BetaAwaiter>
async::AsyncTask co_gemm(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,
                         async::ReadBuffer<LhsTensor> lhs, async::ReadBuffer<RhsTensor> rhs, AlphaAwaiter alpha,
                         BetaAwaiter beta, std::optional<Precision> requested)
{
  auto output_awaiter = uni20::detail::mutable_async_tensor_awaiter(output);
  auto awaited = co_await async::all(output_awaiter, lhs, rhs, alpha, beta);
  decltype(auto) output_value = uni20::detail::mutable_async_tensor_value<OutputTensor>(std::get<0>(awaited));
  auto const& lhs_value = std::get<1>(awaited);
  auto const& rhs_value = std::get<2>(awaited);
  auto const& alpha_value = std::get<3>(awaited);
  auto const& beta_value = std::get<4>(awaited);
 
  auto const shape = matrix_product_shape(lhs_value, rhs_value);
  uni20::require_output(output_value, shape);
  using scalar_type = uni20::tensor_element_t<OutputTensor>;
  scalar_type const alpha_scalar = static_cast<scalar_type>(alpha_value);
  scalar_type const beta_scalar = static_cast<scalar_type>(beta_value);
  if (requested) require_precision<scalar_type>(*requested);
  auto output_span = uni20::mdspec_of(output_value);
  auto lhs_span = uni20::mdspec_of(lhs_value);
  auto rhs_span = uni20::mdspec_of(rhs_value);
#if UNI20_HAS_MPREAL
  if constexpr (has_runtime_precision_v<scalar_type>)
    uni20::linalg::gemm(selector, output_value, alpha_scalar, lhs_value, rhs_value, beta_scalar, requested);
  else
#endif
    co_await co_dispatch_kernel(selector, gemm_op{}, output_span, alpha_scalar, lhs_span, rhs_span, beta_scalar);
  co_return;
}
 
template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class AlphaAwaiter>
void schedule_async_assign_product(BackendSelector selector, async::Async<OutputTensor>& output,
                                   async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs,
                                   AlphaAwaiter alpha, std::optional<Precision> p = {})
{
  validate_async_matrix_product_aliasing(output, lhs, rhs);
  auto task = co_assign_product(std::move(selector), output.write(), lhs.read(), rhs.read(), std::move(alpha), p);
  task.debug_name("assign_product");
  async::schedule(std::move(task));
}
 
template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class AlphaAwaiter, class BetaAwaiter>
void schedule_async_gemm(BackendSelector selector, async::Async<OutputTensor>& output,
                         async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs, AlphaAwaiter alpha,
                         BetaAwaiter beta, std::optional<Precision> p = {})
{
  validate_async_matrix_product_aliasing(output, lhs, rhs);
  auto task =
      co_gemm(std::move(selector), output.write(), lhs.read(), rhs.read(), std::move(alpha), std::move(beta), p);
  task.debug_name("gemm");
  async::schedule(std::move(task));
}
} // namespace detail
 
template <class BackendSelector, uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha, class Beta>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&
           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>
void gemm(BackendSelector selector, async::Async<OutputTensor>& output, Alpha&& alpha,
          async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs, Beta&& beta,
          std::optional<Precision> p = {})
{
  detail::schedule_async_gemm(std::move(selector), output, lhs, rhs, async::read(std::forward<Alpha>(alpha)),
                              async::read(std::forward<Beta>(beta)), p);
}
 
template <uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha, class Beta>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&
           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>
void gemm(async::Async<OutputTensor>& output, Alpha&& alpha, async::Async<LhsTensor> const& lhs,
          async::Async<RhsTensor> const& rhs, Beta&& beta, std::optional<Precision> p = {})
{
  auto selector = select_backend_for<OutputTensor, LhsTensor, RhsTensor>(gemm_op{});
  gemm(std::move(selector), output, std::forward<Alpha>(alpha), lhs, rhs, std::forward<Beta>(beta), p);
}
 
template <class BackendSelector, uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha = uni20::tensor_element_t<OutputTensor>>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>>
void assign_product(BackendSelector selector, async::Async<OutputTensor>& output, async::Async<LhsTensor> const& lhs,
                    async::Async<RhsTensor> const& rhs, Alpha&& alpha, std::optional<Precision> p = {})
{
  detail::schedule_async_assign_product(std::move(selector), output, lhs, rhs, async::read(std::forward<Alpha>(alpha)),
                                        p);
}
 
template <uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha = uni20::tensor_element_t<OutputTensor>>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>>
void assign_product(async::Async<OutputTensor>& output, async::Async<LhsTensor> const& lhs,
                    async::Async<RhsTensor> const& rhs, Alpha&& alpha, std::optional<Precision> p = {})
{
  auto selector = select_backend_for<OutputTensor, LhsTensor, RhsTensor>(assign_product_op{});
  detail::schedule_async_assign_product(std::move(selector), output, lhs, rhs, async::read(std::forward<Alpha>(alpha)),
                                        p);
}
 
template <class BackendSelector, uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha = uni20::tensor_element_t<OutputTensor>>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>>
void add_product(BackendSelector selector, async::Async<OutputTensor>& output, async::Async<LhsTensor> const& lhs,
                 async::Async<RhsTensor> const& rhs, Alpha&& alpha, std::optional<Precision> p = {})
{
  gemm(std::move(selector), output, std::forward<Alpha>(alpha), lhs, rhs, uni20::tensor_element_t<OutputTensor>{1}, p);
}
 
template <uni20::MutableRankedTensorView<2> OutputTensor, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor, class Alpha = uni20::tensor_element_t<OutputTensor>>
  requires uni20::AsyncTensorOutput<OutputTensor> &&
           detail::CompatibleMatrixProductTensors<OutputTensor, LhsTensor, RhsTensor> &&
           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>>
void add_product(async::Async<OutputTensor>& output, async::Async<LhsTensor> const& lhs,
                 async::Async<RhsTensor> const& rhs, Alpha&& alpha, std::optional<Precision> p = {})
{
  gemm(output, std::forward<Alpha>(alpha), lhs, rhs, uni20::tensor_element_t<OutputTensor>{1}, p);
}
 
template <KernelBackendSelector Selector, uni20::AsyncTensorOutput Output, RankedTensorView<2> A, RankedTensorView<2> B>
  requires detail::CompatibleMatrixProductTensors<Output, A, B>
void assign_product(Selector selector, async::Async<Output>& output, async::Async<A> const& a, async::Async<B> const& b,
                    std::optional<Precision> p = {})
{
  assign_product(std::move(selector), output, a, b, tensor_element_t<Output>{1}, p);
}
 
template <uni20::AsyncTensorOutput Output, RankedTensorView<2> A, RankedTensorView<2> B>
  requires detail::CompatibleMatrixProductTensors<Output, A, B>
void assign_product(async::Async<Output>& output, async::Async<A> const& a, async::Async<B> const& b,
                    std::optional<Precision> p = {})
{
  assign_product(select_backend_for<Output, A, B>(assign_product_op{}), output, a, b, tensor_element_t<Output>{1}, p);
}
template <KernelBackendSelector Selector, uni20::AsyncTensorOutput Output, RankedTensorView<2> A, RankedTensorView<2> B>
  requires detail::CompatibleMatrixProductTensors<Output, A, B>
void add_product(Selector selector, async::Async<Output>& output, async::Async<A> const& a, async::Async<B> const& b,
                 std::optional<Precision> p = {})
{
  add_product(std::move(selector), output, a, b, tensor_element_t<Output>{1}, p);
}
 
template <uni20::AsyncTensorOutput Output, RankedTensorView<2> A, RankedTensorView<2> B>
  requires detail::CompatibleMatrixProductTensors<Output, A, B>
void add_product(async::Async<Output>& output, async::Async<A> const& a, async::Async<B> const& b,
                 std::optional<Precision> p = {})
{
  add_product(select_backend_for<Output, A, B>(gemm_op{}), output, a, b, tensor_element_t<Output>{1}, p);
}
 
} // namespace uni20::linalg