uni20/linalg_2async_2contract_8hpp_source.html

#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/contract.hpp>


#include <array>

#include <cstddef>

#include <memory>

#include <tuple>

#include <utility>


namespace uni20::linalg

{

namespace detail

{

template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor,

          uni20::TensorView RhsTensor, class AlphaAwaiter, class BetaAwaiter, std::size_t LhsRank, std::size_t RhsRank,

          std::size_t ContractedRank>

async::AsyncTask co_contract(BackendSelector const selector, async::WriteBuffer<OutputTensor> output,

                             async::ReadBuffer<LhsTensor> lhs, async::ReadBuffer<RhsTensor> rhs, AlphaAwaiter alpha,

                             BetaAwaiter beta, ContractionAxes<LhsRank, RhsRank, ContractedRank> axes)

{

  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);

  using scalar_type = uni20::tensor_element_t<OutputTensor>;

  scalar_type const alpha_value = static_cast<scalar_type>(std::get<3>(awaited));

  scalar_type const beta_value = static_cast<scalar_type>(std::get<4>(awaited));


  auto const required_extents = contraction_output_extents(lhs_value, rhs_value, axes);

  uni20::require_output(output_value, required_extents);

  auto output_descriptor = uni20::mdspec_of(output_value);

  auto lhs_descriptor = uni20::mdspec_of(lhs_value);

  auto rhs_descriptor = uni20::mdspec_of(rhs_value);

  auto operation = contract_op<LhsRank, RhsRank, ContractedRank>{.axes = std::move(axes)};

  co_await co_dispatch_kernel(selector, std::move(operation), output_descriptor, alpha_value, lhs_descriptor,

                              rhs_descriptor, beta_value);

  co_return;

}


template <class BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor,

          uni20::TensorView RhsTensor, class AlphaAwaiter, class BetaAwaiter, std::size_t LhsRank, std::size_t RhsRank,

          std::size_t ContractedRank>

void schedule_async_contract(BackendSelector selector, async::Async<OutputTensor>& output,

                             async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs, AlphaAwaiter alpha,

                             BetaAwaiter beta, ContractionAxes<LhsRank, RhsRank, ContractedRank> axes)

{

  auto const* const output_queue = std::addressof(output.queue());

  ERROR_IF(output_queue == std::addressof(lhs.queue()) || output_queue == std::addressof(rhs.queue()),

           "async contraction output must not share an epoch queue with an input");

  auto task = co_contract(std::move(selector), output.write(), lhs.read(), rhs.read(), std::move(alpha),

                          std::move(beta), std::move(axes));

  task.debug_name("contract");

  async::schedule(std::move(task));

}

} // namespace detail


template <KernelBackendSelector BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor,

          uni20::TensorView RhsTensor, class Alpha, class Beta, std::size_t LhsRank, std::size_t RhsRank,

          std::size_t ContractedRank>

  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&

           uni20::MutableRankedTensorView<OutputTensor,

                                          ContractionAxes<LhsRank, RhsRank, ContractedRank>::output_rank> &&

           uni20::RankedTensorView<LhsTensor, LhsRank> && uni20::RankedTensorView<RhsTensor, RhsRank> &&

           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&

           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>


void contract(BackendSelector selector, async::Async<OutputTensor>& output, Alpha&& alpha,

              async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs,

              ContractionAxes<LhsRank, RhsRank, ContractedRank> axes, Beta&& beta)

{

  detail::schedule_async_contract(std::move(selector), output, lhs, rhs, async::read(std::forward<Alpha>(alpha)),

                                  async::read(std::forward<Beta>(beta)), std::move(axes));

}


template <uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor, uni20::TensorView RhsTensor, class Alpha,

          class Beta, std::size_t LhsRank, std::size_t RhsRank, std::size_t ContractedRank>

  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&

           uni20::MutableRankedTensorView<OutputTensor,

                                          ContractionAxes<LhsRank, RhsRank, ContractedRank>::output_rank> &&

           uni20::RankedTensorView<LhsTensor, LhsRank> && uni20::RankedTensorView<RhsTensor, RhsRank> &&

           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&

           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>


void contract(async::Async<OutputTensor>& output, Alpha&& alpha, async::Async<LhsTensor> const& lhs,

              async::Async<RhsTensor> const& rhs, ContractionAxes<LhsRank, RhsRank, ContractedRank> axes, Beta&& beta)

{

  auto operation = contract_op<LhsRank, RhsRank, ContractedRank>{.axes = axes};

  auto selector = select_backend_for<OutputTensor, LhsTensor, RhsTensor>(operation);

  contract(std::move(selector), output, std::forward<Alpha>(alpha), lhs, rhs, std::move(axes),

           std::forward<Beta>(beta));

}


template <KernelBackendSelector BackendSelector, uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor,

          uni20::TensorView RhsTensor, class Alpha, class Beta, std::size_t ContractedRank>

  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&

           (ContractedRank <= uni20::tensor_mdspec_t<LhsTensor>::rank()) &&

           (ContractedRank <= uni20::tensor_mdspec_t<RhsTensor>::rank()) &&

           (uni20::tensor_mdspec_t<OutputTensor>::rank() == uni20::tensor_mdspec_t<LhsTensor>::rank() +

                                                                uni20::tensor_mdspec_t<RhsTensor>::rank() -

                                                                2 * ContractedRank) &&

           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&

           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>


void contract(BackendSelector selector, async::Async<OutputTensor>& output, Alpha&& alpha,

              async::Async<LhsTensor> const& lhs, async::Async<RhsTensor> const& rhs,

              std::array<std::pair<std::size_t, std::size_t>, ContractedRank> requested_axes, Beta&& beta)

{

  constexpr std::size_t lhs_rank = uni20::tensor_mdspec_t<LhsTensor>::rank();

  constexpr std::size_t rhs_rank = uni20::tensor_mdspec_t<RhsTensor>::rank();

  auto axes = make_contraction_axes<lhs_rank, rhs_rank>(std::move(requested_axes));

  contract(std::move(selector), output, std::forward<Alpha>(alpha), lhs, rhs, std::move(axes),

           std::forward<Beta>(beta));

}


template <uni20::AsyncTensorOutput OutputTensor, uni20::TensorView LhsTensor, uni20::TensorView RhsTensor, class Alpha,

          class Beta, std::size_t ContractedRank>

  requires detail::CompatibleContractionTensors<OutputTensor, LhsTensor, RhsTensor> &&

           (ContractedRank <= uni20::tensor_mdspec_t<LhsTensor>::rank()) &&

           (ContractedRank <= uni20::tensor_mdspec_t<RhsTensor>::rank()) &&

           (uni20::tensor_mdspec_t<OutputTensor>::rank() == uni20::tensor_mdspec_t<LhsTensor>::rank() +

                                                                uni20::tensor_mdspec_t<RhsTensor>::rank() -

                                                                2 * ContractedRank) &&

           AsyncOperationScalar<Alpha, uni20::tensor_element_t<OutputTensor>> &&

           AsyncOperationScalar<Beta, uni20::tensor_element_t<OutputTensor>>


void contract(async::Async<OutputTensor>& output, Alpha&& alpha, async::Async<LhsTensor> const& lhs,

              async::Async<RhsTensor> const& rhs,

              std::array<std::pair<std::size_t, std::size_t>, ContractedRank> requested_axes, Beta&& beta)

{

  constexpr std::size_t lhs_rank = uni20::tensor_mdspec_t<LhsTensor>::rank();

  constexpr std::size_t rhs_rank = uni20::tensor_mdspec_t<RhsTensor>::rank();

  auto axes = make_contraction_axes<lhs_rank, rhs_rank>(std::move(requested_axes));

  contract(output, std::forward<Alpha>(alpha), lhs, rhs, std::move(axes), std::forward<Beta>(beta));

}


} // namespace uni20::linalg

async.hpp
The Async<T> container: coroutine‐safe asynchronous read/write.

dispatch.hpp
Coroutine-aware kernel dispatch with an ordinary blocking fallback.

awaiters.hpp
Utilities for awaiting on Async<T>: co_await overload, all(), try_await().

uni20::async::AsyncTask
Canonical host-oriented coroutine task type.
Definition async_task.hpp:196

uni20::async::Async
Async<T> is a container for coroutine-safe asynchronous read/write.
Definition async/async.hpp:93

uni20::async::Async::read
ReadBuffer< T > read() const
Begin an asynchronous read of the value.
Definition async/async.hpp:339

uni20::async::Async::queue
EpochQueue const & queue() const
Inspect the shared epoch queue.
Definition async/async.hpp:423

uni20::async::Async::write
WriteBuffer< T > write()
Begin exclusive asynchronous access to the payload.
Definition async/async.hpp:346

uni20::async::ReadBuffer
RAII handle for reading the value from an Async container at a given epoch.
Definition buffers.hpp:53

uni20::async::WriteBuffer
Exclusive mutable-access buffer for one Async<T> epoch.
Definition buffers.hpp:859

uni20::AsyncTensorOutput
Mutable Tensor output supported by async operation wrappers.
Definition linalg/async/concepts.hpp:25

uni20::MutableRankedTensorView
MutableTensorView with a specified static rank.
Definition tensor/concepts.hpp:166

uni20::RankedTensorView
TensorView whose readable mdspec has a specified static rank.
Definition tensor/concepts.hpp:149

uni20::TensorView
Tensor-level object exposing readable multidimensional metadata and a backend selector.
Definition tensor/concepts.hpp:82

debug_scheduler.hpp
Coroutine scheduler for AsyncTask.

concepts.hpp
Capabilities shared by asynchronous Tensor operation wrappers.

output.hpp
Shared output capabilities for Async Tensor operation wrappers.

contract.hpp
Fixed-output pairwise Tensor contraction front end.

uni20::async::schedule
void schedule(AsyncTask &&task)
Schedule a task on the currently configured global scheduler.
Definition debug_scheduler.hpp:237

uni20::async::all
auto all(Aw &&... aw) noexcept
Build an awaitable that waits for all of the provided awaitables.
Definition awaiters.hpp:200

uni20::async::read
ReadBuffer< T > read(Async< T > const &a)
Convenience helper that forwards to Async<T>::read().
Definition async/async.hpp:540

uni20::linalg
Definition linalg/async.hpp:28

uni20::linalg::contract
void contract(BackendSelector selector, async::Async< OutputTensor > &output, Alpha &&alpha, async::Async< LhsTensor > const &lhs, async::Async< RhsTensor > const &rhs, ContractionAxes< LhsRank, RhsRank, ContractedRank > axes, Beta &&beta)
Schedule a fixed-output pairwise contraction through an explicit selector.
Definition linalg/async/contract.hpp:84

uni20::linalg::contraction_output_extents
auto contraction_output_extents(Lhs const &lhs, Rhs const &rhs, ContractionAxes< LhsRank, RhsRank, ContractedRank > const &axes) -> stdex::dextents< uni20::index_type, ContractionAxes< LhsRank, RhsRank, ContractedRank >::output_rank >
Validate paired extents and return the canonical contraction output shape.
Definition contraction_axes.hpp:141

uni20::linalg::co_dispatch_kernel
async::AsyncTask co_dispatch_kernel(BackendSelector selector, Op op, Args &... args)
Dispatch a kernel from a coroutine without requiring coroutine wrappers for blocking backends.
Definition async/dispatch.hpp:106

uni20::require_output
void require_output(Output const &output, RequiredExtents const &required)
Require an existing tensor output to have the specified shape.
Definition tensor/output.hpp:84

uni20::mdspec_of
constexpr decltype(auto) mdspec_of(Tensor &tensor) noexcept(noexcept(tensor.mdspec()))
Return a tensor's immediate or descriptor-backed multidimensional metadata.
Definition tensor/concepts.hpp:50

uni20::tensor_element_t
std::remove_cv_t< typename tensor_mdspec_t< T >::element_type > tensor_element_t
Scalar value type exposed by a tensor-level object.
Definition tensor/concepts.hpp:126

uni20::beta
mpreal beta(mpreal const &a, mpreal const &x)
Real Beta function; positive rational/integer pairs retain exact values.
Definition mpreal_special.hpp:164

uni20::tensor_mdspec_t
std::remove_cvref_t< detail::normalized_const_mdspec_t< T > > tensor_mdspec_t
Normalized readable MdspecLike type exposed by a tensor-level object.
Definition tensor/concepts.hpp:123

uni20::linalg::ContractionAxes
Canonical contracted and surviving axes for two fixed-rank operands.
Definition contraction_axes.hpp:26

uni20::linalg::ContractionAxes::output_rank
static constexpr std::size_t output_rank
Definition contraction_axes.hpp:35

uni20::linalg::contract_op
Fixed-output pairwise tensor contraction operation.
Definition operation_tags.hpp:98

uni20::linalg::contract_op::axes
ContractionAxes< LhsRank, RhsRank, ContractedRank > axes
Definition operation_tags.hpp:100

trace.hpp

ERROR_IF
#define ERROR_IF(cond,...)
Definition trace.hpp:389