tensor/output.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/async/shared_storage.hpp>
#include <uni20/common/initialization.hpp>
#include <uni20/common/trace.hpp>
#include <uni20/core/runtime_precision.hpp>
#include <uni20/tensor/concepts.hpp>
 
#include <concepts>
#include <cstddef>
#include <type_traits>
#include <utility>
 
namespace uni20
{
template <class Extents>
concept TensorExtentsLike = requires(Extents const& extents) {
  { std::remove_cvref_t<Extents>::rank() } -> std::convertible_to<std::size_t>;
  extents.extent(std::size_t{});
};
 
template <TensorExtentsLike LhsExtents, TensorExtentsLike RhsExtents>
[[nodiscard]] constexpr bool tensor_extents_equal(LhsExtents const& lhs, RhsExtents const& rhs) noexcept
{
  if constexpr (std::remove_cvref_t<LhsExtents>::rank() != std::remove_cvref_t<RhsExtents>::rank())
  {
    return false;
  }
  else
  {
    for (std::size_t axis = 0; axis < std::remove_cvref_t<LhsExtents>::rank(); ++axis)
    {
      if (lhs.extent(axis) != rhs.extent(axis)) return false;
    }
    return true;
  }
}
 
namespace detail
{
template <class TargetExtents, TensorExtentsLike SourceExtents, std::size_t... Axis>
[[nodiscard]] TargetExtents convert_tensor_extents_impl(SourceExtents const& source, std::index_sequence<Axis...>)
{
  for (std::size_t axis = 0; axis < TargetExtents::rank(); ++axis)
  {
    auto const fixed_extent = TargetExtents::static_extent(axis);
    ERROR_IF(fixed_extent != stdex::dynamic_extent && !std::cmp_equal(source.extent(axis), fixed_extent),
             "tensor extent does not match the destination's static extent");
  }
  using index_type = typename TargetExtents::index_type;
  return TargetExtents(static_cast<index_type>(source.extent(Axis))...);
}
} // namespace detail
 
template <class TargetExtents, TensorExtentsLike SourceExtents>
[[nodiscard]] TargetExtents convert_tensor_extents(SourceExtents const& source)
{
  static_assert(TargetExtents::rank() == std::remove_cvref_t<SourceExtents>::rank(),
                "tensor output rank does not match the required shape");
  return detail::convert_tensor_extents_impl<TargetExtents>(source, std::make_index_sequence<TargetExtents::rank()>{});
}
 
template <class T>
concept ResizableTensorOutput =
    MutableTensorView<T> &&
    requires(std::remove_reference_t<T>& output, tensor_extents_t<T> const& extents) { output.reset_shape(extents); };
 
template <TensorView Output, TensorExtentsLike RequiredExtents>
void require_output(Output const& output, RequiredExtents const& required)
{
  static_assert(tensor_extents_t<Output>::rank() == std::remove_cvref_t<RequiredExtents>::rank(),
                "tensor output rank does not match the required shape");
  ERROR_IF(!tensor_extents_equal(output.extents(), required),
           "tensor output shape does not match the required extents");
}
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents>
Output& prepare_output(Output& output, RequiredExtents const& required)
{
  static_assert(tensor_extents_t<Output>::rank() == std::remove_cvref_t<RequiredExtents>::rank(),
                "tensor output rank does not match the required shape");
  if (tensor_extents_equal(output.extents(), required)) return output;
 
  if constexpr (ResizableTensorOutput<Output>)
  {
    auto const converted = convert_tensor_extents<tensor_extents_t<Output>>(required);
    if constexpr (requires { output.reset_shape(uninitialized, converted); })
      output.reset_shape(uninitialized, converted);
    else
      output.reset_shape(converted);
    CHECK(tensor_extents_equal(output.extents(), required));
  }
  else
  {
    require_output(output, required);
  }
  return output;
}
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents, class Placement>
  requires requires(Output& output, tensor_extents_t<Output> const& extents, Placement const& placement) {
    { output.storage_is_compatible(placement) } -> std::convertible_to<bool>;
    output.reset_shape(extents);
    output.replace(extents, placement);
  }
Output& prepare_output(Output& output, RequiredExtents const& required, Placement const& placement)
{
  static_assert(tensor_extents_t<Output>::rank() == std::remove_cvref_t<RequiredExtents>::rank(),
                "tensor output rank does not match the required shape");
  if (tensor_extents_equal(output.extents(), required) && output.storage_is_compatible(placement)) return output;
 
  auto const converted = convert_tensor_extents<tensor_extents_t<Output>>(required);
  if (output.storage_is_compatible(placement))
  {
    if constexpr (requires { output.reset_shape(uninitialized, converted); })
      output.reset_shape(uninitialized, converted);
    else
      output.reset_shape(converted);
  }
  else
  {
    if constexpr (requires { output.replace(uninitialized, converted, placement); })
      output.replace(uninitialized, converted, placement);
    else
      output.replace(converted, placement);
  }
 
  CHECK(tensor_extents_equal(output.extents(), required));
  CHECK(output.storage_is_compatible(placement));
  return output;
}
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents>
Output& prepare_output(async::shared_storage<Output>& storage, RequiredExtents const& required)
{
  if (storage.constructed()) return prepare_output(*storage, required);
 
  if constexpr (std::constructible_from<Output, tensor_extents_t<Output> const&>)
  {
    auto const converted = convert_tensor_extents<tensor_extents_t<Output>>(required);
    if constexpr (std::constructible_from<Output, uninitialized_t, tensor_extents_t<Output> const&>)
      return storage.emplace(uninitialized, converted);
    else
      return storage.emplace(converted);
  }
  else
  {
    ERROR("unconstructed tensor output cannot be initialized from the required extents");
  }
}
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents, class Placement>
  requires std::constructible_from<Output, Placement const&, tensor_extents_t<Output> const&> &&
           requires(Output& output, tensor_extents_t<Output> const& extents, Placement const& placement) {
             prepare_output(output, extents, placement);
           }
Output& prepare_output(async::shared_storage<Output>& storage, RequiredExtents const& required,
                       Placement const& placement)
{
  if (storage.constructed()) return prepare_output(*storage, required, placement);
 
  auto const converted = convert_tensor_extents<tensor_extents_t<Output>>(required);
  if constexpr (std::constructible_from<Output, uninitialized_t, Placement const&, tensor_extents_t<Output> const&>)
    return storage.emplace(uninitialized, placement, converted);
  else
    return storage.emplace(placement, converted);
}
 
namespace detail
{
// Only outputs with independent writable defaults record operation precision.
// Structural views preserve their parent's metadata; no elements are converted.
template <MutableTensorView Output> void record_output_precision(Output& output, Precision p)
{
  require_precision<tensor_element_t<Output>>(p);
  if constexpr (requires { output.default_precision(p); }) output.default_precision(p);
}
} // namespace detail
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents>
Output& prepare_output(Output& output, RequiredExtents const& required, Precision p)
{
  detail::record_output_precision(output, p);
  return prepare_output(output, required);
}
 
template <MutableTensorView Output, TensorExtentsLike RequiredExtents>
  requires(!has_runtime_precision_v<tensor_element_t<Output>> ||
           std::constructible_from<Output, uninitialized_t, tensor_extents_t<Output> const&, Precision>)
Output& prepare_output(async::shared_storage<Output>& storage, RequiredExtents const& required, Precision p)
{
  require_precision<tensor_element_t<Output>>(p);
  if (storage.constructed()) return prepare_output(*storage, required, p);
  if constexpr (has_runtime_precision_v<tensor_element_t<Output>>)
  {
    auto const converted = convert_tensor_extents<tensor_extents_t<Output>>(required);
    return storage.emplace(uninitialized, converted, p);
  }
  else
    return prepare_output(storage, required);
}
 
} // namespace uni20