transform_view.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/core/compiler_attributes.hpp>
#include <uni20/core/types.hpp>
#include <uni20/mdspan/concepts.hpp>
#include <uni20/mdspan/mdspan.hpp>
#include <uni20/mdspan/zip_layout.hpp>
 
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace uni20
{
namespace detail
{
 
template <class Function, class... Arguments>
concept ConstDirectlyCallable = requires(Function const& function) { function(std::declval<Arguments>()...); };
 
template <class Function, class... Arguments>
using direct_call_result_t = decltype(std::declval<Function const&>()(std::declval<Arguments>()...));
 
template <class Function, MdspanLike Span> class unary_transform_accessor {
  public:
    using function_type = Function;
    using span_type = Span;
    using wrapped_accessor_type = typename span_type::accessor_type;
    using data_handle_type = typename span_type::data_handle_type;
    using offset_type = span_offset_t<wrapped_accessor_type>;
    using reference = direct_call_result_t<function_type, typename span_type::reference>;
    using element_type = logical_value_t<remove_proxy_reference_t<reference>> const;
    using offset_policy = unary_transform_accessor;
 
    template <class FwdFunction>
    UNI20_HOST_DEVICE constexpr unary_transform_accessor(FwdFunction&& function, span_type const& span)
        : function_(std::forward<FwdFunction>(function)), accessor_(span.accessor())
    {}
 
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto offset(data_handle_type const& handle,
                                                          offset_type const& offset) const -> data_handle_type
    {
      return accessor_.offset(handle, offset);
    }
 
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto access(data_handle_type const& handle,
                                                          offset_type const& offset) const -> reference
    {
      return function_(accessor_.access(handle, offset));
    }
 
  private:
    [[no_unique_address]] function_type function_;
    [[no_unique_address]] wrapped_accessor_type accessor_;
};
 
template <class Function, MdspanLike... Spans> class transform_accessor {
  public:
    static_assert(sizeof...(Spans) >= 2);
 
    using function_type = Function;
    using accessor_tuple = std::tuple<typename Spans::accessor_type...>;
    using data_handle_type = std::tuple<typename Spans::data_handle_type...>;
    using offset_type = std::tuple<span_offset_t<typename Spans::accessor_type>...>;
    using reference = direct_call_result_t<function_type, typename Spans::reference...>;
    using element_type = logical_value_t<remove_proxy_reference_t<reference>> const;
    using offset_policy = transform_accessor;
 
    template <class FwdFunction>
    UNI20_HOST_DEVICE constexpr transform_accessor(FwdFunction&& function, Spans const&... spans)
        : function_(std::forward<FwdFunction>(function)), accessors_(spans.accessor()...)
    {}
 
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto offset(data_handle_type const& handles,
                                                          offset_type const& offsets) const -> data_handle_type
    {
      return offset_impl(handles, offsets, std::index_sequence_for<Spans...>{});
    }
 
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto access(data_handle_type const& handles,
                                                          offset_type const& offsets) const -> reference
    {
      return access_impl(handles, offsets, std::index_sequence_for<Spans...>{});
    }
 
  private:
    template <std::size_t... Index>
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto offset_impl(data_handle_type const& handles,
                                                               offset_type const& offsets,
                                                               std::index_sequence<Index...>) const -> data_handle_type
    {
      return {std::get<Index>(accessors_).offset(std::get<Index>(handles), std::get<Index>(offsets))...};
    }
 
    template <std::size_t... Index>
    [[nodiscard]] UNI20_HOST_DEVICE constexpr auto access_impl(data_handle_type const& handles,
                                                               offset_type const& offsets,
                                                               std::index_sequence<Index...>) const -> reference
    {
      return function_(std::get<Index>(accessors_).access(std::get<Index>(handles), std::get<Index>(offsets))...);
    }
 
    [[no_unique_address]] function_type function_;
    [[no_unique_address]] accessor_tuple accessors_;
};
 
} // namespace detail
 
template <class Function, MdspanLike Span, class Domain>
inline constexpr bool enable_accessor_in_domain<detail::unary_transform_accessor<Function, Span>, Domain> =
    enable_accessor_in_domain<typename Span::accessor_type, Domain>;
 
template <class Function, MdspanLike... Spans>
inline constexpr bool enable_accessor_in_domain<detail::transform_accessor<Function, Spans...>, host_access_domain> =
    (enable_accessor_in_domain<typename Spans::accessor_type, host_access_domain> && ...);
 
template <class Function, MdspanLike... Spans>
inline constexpr bool enable_accessor_in_domain<detail::transform_accessor<Function, Spans...>, cuda_access_domain> =
    (enable_accessor_in_domain<typename Spans::accessor_type, cuda_access_domain> && ...);
 
template <class Function, MdspanLike Span>
  requires detail::ConstDirectlyCallable<std::decay_t<Function>, typename Span::reference>
[[nodiscard]] constexpr auto transform_view(Function&& function, Span const& span)
{
  using function_type = std::decay_t<Function>;
  using accessor_type = detail::unary_transform_accessor<function_type, Span>;
  using element_type = typename accessor_type::element_type;
  using extents_type = typename Span::extents_type;
  using layout_type = typename Span::layout_type;
 
  return stdex::mdspan<element_type, extents_type, layout_type, accessor_type>{
      span.data_handle(), span.mapping(), accessor_type(std::forward<Function>(function), span)};
}
 
template <class Function, MdspanLike First, MdspanLike Second, MdspanLike... Rest>
  requires detail::ConstDirectlyCallable<std::decay_t<Function>, typename First::reference, typename Second::reference,
                                         typename Rest::reference...>
[[nodiscard]] constexpr auto transform_view(Function&& function, First const& first, Second const& second,
                                            Rest const&... rest)
{
  static_assert(First::rank() == Second::rank() && ((First::rank() == Rest::rank()) && ...));
 
  using function_type = std::decay_t<Function>;
  using extents_type = common_extents_t<First, Second, Rest...>;
  using layout_type = zip_layout_t<First, Second, Rest...>;
  using mapping_type = typename layout_type::template mapping<extents_type>;
  using accessor_type = detail::transform_accessor<function_type, First, Second, Rest...>;
  using element_type = typename accessor_type::element_type;
 
  auto extents = make_common_extents(first, second, rest...);
  return stdex::mdspan<element_type, extents_type, layout_type, accessor_type>{
      std::tuple{first.data_handle(), second.data_handle(), rest.data_handle()...},
      mapping_type{extents, first.mapping(), second.mapping(), rest.mapping()...},
      accessor_type(std::forward<Function>(function), first, second, rest...)};
}
 
} // namespace uni20