tensor/transform.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/linalg/backend_selector.hpp>
#include <uni20/linalg/backends/cpu/transform.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/elementwise_functions.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/mdspan/concepts.hpp>
#include <uni20/tensor/access.hpp>
#include <uni20/tensor/concepts.hpp>
#include <uni20/tensor/output.hpp>
 
#if UNI20_BACKEND_CUDA
#include <uni20/linalg/backends/cuda/transform.hpp>
#endif
 
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace uni20
{
namespace detail
{
 
template <class Output, class... Inputs>
concept OverwriteTransformSpans =
    MutableMdspanLike<Output> && (MdspanLike<Inputs> && ...) &&
    ((std::remove_cvref_t<Output>::rank() == std::remove_cvref_t<Inputs>::rank()) && ...);
 
template <class Output, class... Inputs>
concept UpdateTransformSpans = MutableMdspanLike<Output> && (MdspanLike<Inputs> && ...) &&
                               ((std::remove_cvref_t<Output>::rank() == std::remove_cvref_t<Inputs>::rank()) && ...);
 
template <class Output, class... Inputs>
concept OverwriteTransformTensors =
    MutableTensorView<Output> && (sizeof...(Inputs) >= 1) && (TensorView<Inputs> && ...) &&
    ((tensor_mdspec_t<Output>::rank() == tensor_mdspec_t<Inputs>::rank()) && ...);
 
template <class Output, class... Inputs>
concept UpdateTransformTensors = MutableTensorView<Output> && (TensorView<Inputs> && ...) &&
                                 ((tensor_mdspec_t<Output>::rank() == tensor_mdspec_t<Inputs>::rank()) && ...);
 
template <class Reference, class... Others>
void require_transform_extents(Reference const& reference, Others const&... others)
{
  if constexpr (sizeof...(Others) > 0)
  {
    constexpr std::size_t Rank = [] {
      if constexpr (TensorView<Reference>)
        return tensor_mdspec_t<Reference>::rank();
      else
        return std::remove_cvref_t<Reference>::rank();
    }();
    if constexpr (Rank > 0)
    {
      auto require_equal = [&](auto const& other) {
        for (std::size_t axis = 0; axis < Rank; ++axis)
          ERROR_IF(reference.extent(axis) != other.extent(axis),
                   "elementwise transform operands have different extents");
      };
      (require_equal(others), ...);
    }
  }
}
 
template <class BackendSelector, class Operation, class OutputMdspan, class... InputMdspans>
void dispatch_transform(BackendSelector&& selector, Operation operation, OutputMdspan&& output,
                        InputMdspans&&... inputs)
{
  linalg::dispatch_kernel(std::forward<BackendSelector>(selector), std::move(operation),
                          std::forward<OutputMdspan>(output), std::forward<InputMdspans>(inputs)...);
}
 
} // namespace detail
 
template <class BackendSelector, class OutputMdspan, class Function, class... InputMdspans>
  requires detail::OverwriteTransformSpans<OutputMdspan, InputMdspans...>
void assign_transform(BackendSelector&& selector, OutputMdspan&& output, Function&& function, InputMdspans&&... inputs)
{
  detail::require_transform_extents(output, inputs...);
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto output_descriptor = std::forward<OutputMdspan>(output);
  auto input_descriptors = std::tuple{make_const_mdspan(inputs)...};
  std::apply(
      [&](auto&... input_descriptor) {
        detail::dispatch_transform(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor,
                                   input_descriptor...);
      },
      input_descriptors);
}
 
template <class BackendSelector, class OutputMdspan, class Function, class... InputMdspans>
  requires detail::UpdateTransformSpans<OutputMdspan, InputMdspans...>
void transform_inplace(BackendSelector&& selector, OutputMdspan&& output, Function&& function, InputMdspans&&... inputs)
{
  detail::require_transform_extents(output, inputs...);
  auto operation = linalg::transform_inplace_op{std::forward<Function>(function)};
  auto output_descriptor = std::forward<OutputMdspan>(output);
  auto input_descriptors = std::tuple{make_const_mdspan(inputs)...};
  std::apply(
      [&](auto&... input_descriptor) {
        detail::dispatch_transform(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor,
                                   input_descriptor...);
      },
      input_descriptors);
}
 
template <class BackendSelector, MutableTensorView OutputTensor, class Function>
void assign_transform(BackendSelector&& selector, OutputTensor&& output, Function&& function)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto output_descriptor = mdspec_of(output);
  detail::dispatch_transform(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor);
}
 
template <MutableTensorView OutputTensor, class Function>
void assign_transform(OutputTensor&& output, Function&& function)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend(operation, output);
  auto output_descriptor = mdspec_of(output);
  detail::dispatch_transform(std::move(selector), std::move(operation), output_descriptor);
}
 
template <class BackendSelector, class Output, class Value>
  requires(detail::OverwriteTransformSpans<Output> || MutableTensorView<Output>)
void fill(BackendSelector&& selector, Output&& output, Value&& value)
{
  assign_transform(std::forward<BackendSelector>(selector), std::forward<Output>(output),
                   linalg::constant{std::forward<Value>(value)});
}
 
template <MutableTensorView Output, class Value> void fill(Output&& output, Value&& value)
{
  assign_transform(std::forward<Output>(output), linalg::constant{std::forward<Value>(value)});
}
 
template <class BackendSelector, class OutputTensor, class Function, class FirstInputTensor, class... RestInputTensors>
  requires detail::OverwriteTransformTensors<OutputTensor, FirstInputTensor, RestInputTensors...>
void assign_transform(BackendSelector&& selector, OutputTensor&& output, Function&& function,
                      FirstInputTensor const& first_input, RestInputTensors const&... rest_inputs)
{
  detail::require_transform_extents(first_input, rest_inputs...);
  prepare_output(output, first_input.extents());
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto output_descriptor = mdspec_of(output);
  auto input_descriptors = std::tuple{mdspec_of(first_input), mdspec_of(rest_inputs)...};
  std::apply(
      [&](auto&... inputs) {
        detail::dispatch_transform(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor,
                                   inputs...);
      },
      input_descriptors);
}
 
template <class OutputTensor, class Function, class FirstInputTensor, class... RestInputTensors>
  requires detail::OverwriteTransformTensors<OutputTensor, FirstInputTensor, RestInputTensors...>
void assign_transform(OutputTensor&& output, Function&& function, FirstInputTensor const& first_input,
                      RestInputTensors const&... rest_inputs)
{
  auto operation = linalg::transform_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend(operation, output, first_input, rest_inputs...);
  detail::require_transform_extents(first_input, rest_inputs...);
  prepare_output(output, first_input.extents());
  auto output_descriptor = mdspec_of(output);
  auto input_descriptors = std::tuple{mdspec_of(first_input), mdspec_of(rest_inputs)...};
  std::apply(
      [&](auto&... inputs) {
        detail::dispatch_transform(selector, std::move(operation), output_descriptor, inputs...);
      },
      input_descriptors);
}
 
template <class BackendSelector, class OutputTensor, class Function, class... InputTensors>
  requires detail::UpdateTransformTensors<OutputTensor, InputTensors...>
void transform_inplace(BackendSelector&& selector, OutputTensor&& output, Function&& function,
                       InputTensors const&... inputs)
{
  detail::require_transform_extents(output, inputs...);
  auto operation = linalg::transform_inplace_op{std::forward<Function>(function)};
  auto output_descriptor = mdspec_of(output);
  auto input_descriptors = std::tuple{mdspec_of(inputs)...};
  std::apply(
      [&](auto&... input_descriptors) {
        detail::dispatch_transform(std::forward<BackendSelector>(selector), std::move(operation), output_descriptor,
                                   input_descriptors...);
      },
      input_descriptors);
}
 
template <class OutputTensor, class Function, class... InputTensors>
  requires detail::UpdateTransformTensors<OutputTensor, InputTensors...>
void transform_inplace(OutputTensor&& output, Function&& function, InputTensors const&... inputs)
{
  auto operation = linalg::transform_inplace_op{std::forward<Function>(function)};
  auto selector = linalg::select_backend(operation, output, inputs...);
  detail::require_transform_extents(output, inputs...);
  auto output_descriptor = mdspec_of(output);
  auto input_descriptors = std::tuple{mdspec_of(inputs)...};
  std::apply(
      [&](auto&... input_descriptors) {
        detail::dispatch_transform(selector, std::move(operation), output_descriptor, input_descriptors...);
      },
      input_descriptors);
}
 
} // namespace uni20