copy_into.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/linalg/backends/cpu/copy.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/mdspan/concepts.hpp>
#include <uni20/tensor/output.hpp>
 
#if UNI20_BACKEND_CUDA
#include <uni20/linalg/backends/cuda/copy.hpp>
#include <uni20/storage/cuda_storage.hpp>
#include <uni20/storage/host_storage.hpp>
#endif
 
#include <cstddef>
#include <type_traits>
#include <utility>
 
namespace uni20
{
namespace detail
{
template <class Output, class Input>
concept CopySpans = MutableMdspanLike<Output> && MdspanLike<Input> &&
                    (std::remove_cvref_t<Output>::rank() == std::remove_cvref_t<Input>::rank());
 
template <class Output, class Input>
concept CopyTensors = MutableTensorView<Output> && TensorView<Input> &&
                      (tensor_mdspec_t<Output>::rank() == tensor_mdspec_t<Input>::rank());
 
#if UNI20_BACKEND_CUDA
template <class Output, class Input>
inline constexpr bool is_pageable_cuda_transfer = (std::same_as<tensor_storage_policy_t<Output>, CudaStorage> &&
                                                   std::same_as<tensor_storage_policy_t<Input>, HostStorage>) ||
                                                  (std::same_as<tensor_storage_policy_t<Output>, HostStorage> &&
                                                   std::same_as<tensor_storage_policy_t<Input>, CudaStorage>);
#endif
 
template <MdspecLike Output, MdspecLike Input>
[[nodiscard]] constexpr bool copy_extents_match(Output const& output, Input const& input) noexcept
{
  if constexpr (std::remove_cvref_t<Output>::rank() != std::remove_cvref_t<Input>::rank())
  {
    return false;
  }
  else
  {
    constexpr std::size_t rank = std::remove_cvref_t<Output>::rank();
    if constexpr (rank > 0)
    {
      for (std::size_t axis = 0; axis < rank; ++axis)
      {
        if (output.extent(axis) != input.extent(axis)) return false;
      }
    }
    return true;
  }
}
 
} // namespace detail
 
template <class BackendSelector, class OutputMdspan, class InputMdspan>
  requires detail::CopySpans<OutputMdspan, InputMdspan>
void copy(BackendSelector&& selector, OutputMdspan&& output, InputMdspan&& input)
{
  ERROR_IF(!detail::copy_extents_match(output, input), "copy output shape does not match input shape");
  auto output_span = std::forward<OutputMdspan>(output);
  auto input_span = make_const_mdspan(input);
  linalg::dispatch_kernel(std::forward<BackendSelector>(selector), linalg::copy_op{}, output_span, input_span);
}
 
template <class BackendSelector, class OutputTensor, class InputTensor>
  requires detail::CopyTensors<OutputTensor, InputTensor>
void copy(BackendSelector&& selector, OutputTensor&& output, InputTensor const& input)
{
  prepare_output(output, input.extents());
  auto output_span = mdspec_of(output);
  auto input_span = mdspec_of(input);
  linalg::dispatch_kernel(std::forward<BackendSelector>(selector), linalg::copy_op{}, output_span, input_span);
}
 
template <class OutputTensor, class InputTensor>
  requires detail::CopyTensors<OutputTensor, InputTensor>
void copy(OutputTensor&& output, InputTensor const& input)
{
#if UNI20_BACKEND_CUDA
  if constexpr (detail::is_pageable_cuda_transfer<OutputTensor, InputTensor>)
  {
    copy(linalg::CudaReferenceBackend{}, output, input);
  }
  else
#endif
  {
    auto selector = linalg::select_backend(linalg::copy_op{}, output, input);
    copy(selector, output, input);
  }
}
 
template <MutableTensorView Output, TensorView Input>
  requires(tensor_mdspec_t<Output>::rank() == tensor_mdspec_t<Input>::rank())
void assign_through(Output& output, Input const& input)
{
  copy(output, input);
}
 
} // namespace uni20