tensor/copy.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/mdspan/concepts.hpp>
#include <uni20/tensor/copy_into.hpp>
#include <uni20/tensor/output.hpp>
#include <uni20/tensor/tensor.hpp>
 
#include <concepts>
#include <cstddef>
#include <type_traits>
#include <utility>
 
namespace uni20
{
template <class RequestedLayout = void, class BackendSelector, MdspanLike InputMdspan>
[[nodiscard]] auto make_tensor(BackendSelector&& selector, InputMdspan&& input)
{
  using input_type = std::remove_cvref_t<InputMdspan>;
  using layout_type = detail::materialized_layout_t<RequestedLayout, input_type>;
  using result_type =
      Tensor<std::remove_cv_t<typename input_type::element_type>, input_type::rank(), HostStorage, layout_type>;
 
  result_type result(uninitialized, convert_tensor_extents<typename result_type::extents_type>(input.extents()));
  auto output_span = result.mdspan();
  copy(std::forward<BackendSelector>(selector), output_span, std::forward<InputMdspan>(input));
  return result;
}
 
template <class RequestedLayout = void, TensorView InputTensor> [[nodiscard]] auto make_tensor(InputTensor const& input)
{
  using input_mdspan = tensor_mdspec_t<InputTensor>;
  using layout_type = detail::materialized_layout_t<RequestedLayout, input_mdspan>;
  using result_type = Tensor<tensor_element_t<InputTensor>, input_mdspan::rank(), HostStorage, layout_type>;
 
  return result_type(input);
}
 
#if UNI20_BACKEND_CUDA
template <TensorView InputTensor>
  requires(std::same_as<tensor_storage_policy_t<InputTensor>, CudaStorage> &&
           detail::CanonicalReshapeLayout<typename tensor_mdspec_t<InputTensor>::layout_type>)
[[nodiscard]] auto to_host(InputTensor const& input)
{
  using input_mdspan = tensor_mdspec_t<InputTensor>;
  using layout_type = typename input_mdspan::layout_type;
  using result_type = Tensor<tensor_element_t<InputTensor>, input_mdspan::rank(), HostStorage, layout_type>;
 
  result_type result(uninitialized, convert_tensor_extents<typename result_type::extents_type>(input.extents()));
  copy(result, input);
  return result;
}
 
template <TensorView InputTensor>
  requires detail::CanonicalReshapeLayout<typename tensor_mdspec_t<InputTensor>::layout_type>
[[nodiscard]] auto to_device(InputTensor const& input, cuda::DeviceResources& resources)
{
  using input_mdspan = tensor_mdspec_t<InputTensor>;
  using layout_type = typename input_mdspan::layout_type;
  using result_type = Tensor<tensor_element_t<InputTensor>, input_mdspan::rank(), CudaStorage, layout_type>;
 
  result_type result(uninitialized, resources,
                     convert_tensor_extents<typename result_type::extents_type>(input.extents()));
  copy(result, input);
  return result;
}
 
template <TensorView InputTensor>
  requires detail::CanonicalReshapeLayout<typename tensor_mdspec_t<InputTensor>::layout_type>
[[nodiscard]] auto to_device(InputTensor const& input, int device)
{
  return to_device(input, cuda::device_resources(device));
}
 
template <TensorView InputTensor>
  requires detail::CanonicalReshapeLayout<typename tensor_mdspec_t<InputTensor>::layout_type>
[[nodiscard]] auto to_device(InputTensor const& input, cuda::Device device)
{
  return to_device(input, device.ordinal());
}
#endif
 
template <class RequestedLayout = void, TensorView InputTensor, std::integral... Extents>
  requires(!OwningTensor<InputTensor> &&
           (std::is_void_v<RequestedLayout> || detail::CanonicalReshapeLayout<RequestedLayout>) &&
           (!std::is_void_v<RequestedLayout> || !StridedTensorView<InputTensor> ||
            detail::CanonicalReshapeLayout<typename tensor_mdspec_t<InputTensor>::layout_type>))
[[nodiscard]] auto reshape(InputTensor const& input, Extents... requested_extents)
{
  using input_layout = typename tensor_mdspec_t<InputTensor>::layout_type;
  if constexpr (ImmediateTensorView<InputTensor> && StridedTensorView<InputTensor> &&
                detail::CanonicalReshapeLayout<input_layout>)
  {
    auto view = reshape_view(input, requested_extents...);
    return make_tensor<RequestedLayout>(view);
  }
  else
  {
    auto materialized = make_tensor<RequestedLayout>(input);
    return reshape(std::move(materialized), requested_extents...);
  }
}
 
} // namespace uni20