block_tensor_mapped_view.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <uni20/mdspan/mdspec.hpp>
#include <uni20/symmetry/block_tensor.hpp>
 
#include <algorithm>
#include <array>
#include <concepts>
#include <cstddef>
#include <memory>
#include <optional>
#include <span>
#include <stdexcept>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace uni20::detail
{
 
template <std::size_t Offset, class Tuple, std::size_t... I>
auto tuple_slice_impl(Tuple const& tuple, std::index_sequence<I...>)
{
  return std::tuple{std::get<Offset + I>(tuple)...};
}
 
template <std::size_t Offset, std::size_t Count, class Tuple> auto tuple_slice(Tuple const& tuple)
{
  return tuple_slice_impl<Offset>(tuple, std::make_index_sequence<Count>{});
}
 
template <class Tuple> auto domain_from_tuple(Tuple&& tuple)
{
  return std::apply([](auto&&... spaces) { return Domain{std::forward<decltype(spaces)>(spaces)...}; },
                    std::forward<Tuple>(tuple));
}
 
template <class Tuple> auto codomain_from_tuple(Tuple&& tuple)
{
  return std::apply([](auto&&... spaces) { return Codomain{std::forward<decltype(spaces)>(spaces)...}; },
                    std::forward<Tuple>(tuple));
}
 
template <class Tensor, std::size_t Axis,
          bool DomainAxis = Axis<std::remove_cvref_t<Tensor>::domain_type::size()> struct TensorAxisSpace;
 
template <class Tensor, std::size_t Axis> struct TensorAxisSpace<Tensor, Axis, true>
{
    using tensor_type = std::remove_cvref_t<Tensor>;
    using type = typename tensor_type::domain_type::template space_type<Axis>;
};
 
template <class Tensor, std::size_t Axis> struct TensorAxisSpace<Tensor, Axis, false>
{
    using tensor_type = std::remove_cvref_t<Tensor>;
    static constexpr std::size_t codomain_axis = Axis - tensor_type::domain_type::size();
    using type = typename tensor_type::codomain_type::template space_type<codomain_axis>;
};
 
template <class Tensor, std::size_t Axis> using tensor_axis_space_t = typename TensorAxisSpace<Tensor, Axis>::type;
 
template <std::size_t Size>
consteval auto is_valid_axis_permutation(std::array<std::size_t, Size> const& permutation) -> bool
{
  std::array<bool, Size> seen{};
  for (std::size_t const axis : permutation)
  {
    if (axis >= Size || seen[axis]) return false;
    seen[axis] = true;
  }
  return true;
}
 
template <std::size_t Size> consteval auto make_identity_axis_permutation()
{
  std::array<std::size_t, Size> permutation{};
  for (std::size_t axis = 0; axis < Size; ++axis)
    permutation[axis] = axis;
  return permutation;
}
 
template <class Tensor, std::size_t... Axis> consteval auto key_axis_flags_impl(std::index_sequence<Axis...>)
{
  return std::array<bool, sizeof...(Axis)>{
      BlockTensorSpaceTraits<tensor_axis_space_t<Tensor, Axis>>::has_block_coordinate...};
}
 
template <class Tensor> consteval auto key_axis_flags()
{
  return key_axis_flags_impl<Tensor>(std::make_index_sequence<std::remove_cvref_t<Tensor>::order()>{});
}
 
template <class Tensor, std::size_t... Axis> consteval auto dense_axis_flags_impl(std::index_sequence<Axis...>)
{
  return std::array<bool, sizeof...(Axis)>{
      BlockTensorSpaceTraits<tensor_axis_space_t<Tensor, Axis>>::has_dense_axis...};
}
 
template <class Tensor> consteval auto dense_axis_flags()
{
  return dense_axis_flags_impl<Tensor>(std::make_index_sequence<std::remove_cvref_t<Tensor>::order()>{});
}
 
template <class Tensor, auto FactorPermutation, bool DenseAxes> consteval auto make_filtered_axis_permutation()
{
  using tensor_type = std::remove_cvref_t<Tensor>;
  constexpr std::size_t order = tensor_type::order();
  constexpr std::size_t filtered_count =
      DenseAxes ? tensor_type::dense_block_order() : tensor_type::key_coordinate_count();
  static_assert(std::tuple_size_v<decltype(FactorPermutation)> == order);
  static_assert(is_valid_axis_permutation(FactorPermutation));
 
  constexpr auto flags = [] {
    if constexpr (DenseAxes)
      return dense_axis_flags<tensor_type>();
    else
      return key_axis_flags<tensor_type>();
  }();
 
  std::array<std::size_t, order> source_filtered_axis{};
  std::size_t source_count = 0;
  for (std::size_t axis = 0; axis < order; ++axis)
  {
    if (flags[axis]) source_filtered_axis[axis] = source_count++;
  }
 
  std::array<std::size_t, filtered_count> permutation{};
  std::size_t output = 0;
  for (std::size_t factor = 0; factor < order; ++factor)
  {
    std::size_t const source_factor = FactorPermutation[factor];
    if (flags[source_factor]) permutation[output++] = source_filtered_axis[source_factor];
  }
  if (output != filtered_count) throw "invalid filtered BlockTensor axis permutation";
  return permutation;
}
 
template <std::size_t CoordinateCount>
auto permute_key(BlockKey<CoordinateCount> const& source, std::array<std::size_t, CoordinateCount> const& permutation)
    -> BlockKey<CoordinateCount>
{
  std::array<std::size_t, CoordinateCount> coordinates{};
  for (std::size_t axis = 0; axis < CoordinateCount; ++axis)
  {
    coordinates[axis] = source.coordinate(permutation[axis]);
  }
  return BlockKey<CoordinateCount>{coordinates};
}
 
template <class Block, bool Immediate = MdspanLike<std::remove_cvref_t<Block>>> struct StridedBlock;
 
template <class Block> struct StridedBlock<Block, true>
{
    using source_type = std::remove_cvref_t<Block>;
    using extents_type = stdex::dextents<typename source_type::index_type, source_type::rank()>;
    using type = stdex::mdspan<typename source_type::element_type, extents_type, stdex::layout_stride,
                               typename source_type::accessor_type>;
};
 
template <class Block> struct StridedBlock<Block, false>
{
    using source_type = std::remove_cvref_t<Block>;
    static_assert(MdspecLike<source_type>);
    using extents_type = stdex::dextents<typename source_type::index_type, source_type::rank()>;
    using type = mdspec<typename source_type::element_type, extents_type, stdex::layout_stride,
                        typename source_type::accessor_type, typename source_type::data_descriptor_type>;
};
 
template <class Block> using strided_block_t = typename StridedBlock<Block>::type;
 
template <MdspecLike Block, std::size_t DenseBlockOrder>
  requires(DenseBlockOrder == 0 || StridedMdspecLike<Block>)
auto permute_block_mdspec(Block block, std::array<std::size_t, DenseBlockOrder> const& permutation)
    -> strided_block_t<Block>
{
  using result_type = strided_block_t<Block>;
  using extents_type = typename result_type::extents_type;
  using mapping_type = typename result_type::mapping_type;
  using index_type = typename extents_type::index_type;
 
  std::array<std::size_t, DenseBlockOrder> extents{};
  std::array<index_type, DenseBlockOrder> strides{};
  if constexpr (DenseBlockOrder > 0)
  {
    for (std::size_t axis = 0; axis < DenseBlockOrder; ++axis)
    {
      extents[axis] = static_cast<std::size_t>(block.extent(permutation[axis]));
      strides[axis] = static_cast<index_type>(block.stride(permutation[axis]));
    }
  }
  auto const result_extents = make_extents<extents_type>(extents);
  if constexpr (MdspanLike<Block>)
  {
    return result_type(block.data_handle(), mapping_type(result_extents, strides), block.accessor());
  }
  else
  {
    return result_type(std::move(block.data_descriptor()), mapping_type(result_extents, strides), block.accessor());
  }
}
 
template <TensorView Block, std::size_t DenseBlockOrder>
  requires(DenseBlockOrder == tensor_mdspec_t<Block>::rank() && (DenseBlockOrder == 0 || StridedTensorView<Block>))
auto permute_block(Block block, std::array<std::size_t, DenseBlockOrder> const& permutation)
{
  auto mutable_mdspec = permute_block_mdspec(mdspec_of(block), permutation);
  auto const_mdspec = permute_block_mdspec(mdspec_of(std::as_const(block)), permutation);
  using mutable_mdspec_type = decltype(mutable_mdspec);
  using const_mdspec_type = decltype(const_mdspec);
  using storage_policy = tensor_storage_policy_t<Block>;
  return MdspecTensorView<mutable_mdspec_type, const_mdspec_type, storage_policy>{std::move(mutable_mdspec),
                                                                                  std::move(const_mdspec)};
}
 
template <TensorView Block> auto make_read_only_block_view(Block const& block)
{
  auto const_mdspec = mdspec_of(block);
  using mdspec_type = std::remove_cvref_t<decltype(const_mdspec)>;
  using storage_policy = tensor_storage_policy_t<Block>;
  return MdspecTensorView<mdspec_type, mdspec_type, storage_policy>{const_mdspec, const_mdspec};
}
 
// An identity map preserves the original mapping/accessor, including compressed
// diagonal blocks which cannot be represented by a layout_stride descriptor.
template <auto Permutation, TensorView Block> auto map_block(Block block)
{
  if constexpr (Permutation == make_identity_axis_permutation<std::tuple_size_v<decltype(Permutation)>>())
    return block;
  else
    return permute_block(std::move(block), Permutation);
}
 
template <class SourceTensor, class = void> class BlockTensorAllocationContextBinding {
  public:
    explicit BlockTensorAllocationContextBinding(SourceTensor&) noexcept {}
};
 
template <class SourceTensor>
class BlockTensorAllocationContextBinding<
    SourceTensor, std::void_t<decltype(std::declval<SourceTensor const&>().allocation_context())>> {
  public:
    using context_type = std::remove_reference_t<decltype(std::declval<SourceTensor const&>().allocation_context())>;
 
    explicit BlockTensorAllocationContextBinding(SourceTensor& source) noexcept
        : context_(std::addressof(source.allocation_context()))
    {}
 
    [[nodiscard]] auto get() const noexcept -> context_type& { return *context_; }
 
  private:
    context_type* context_;
};
 
template <class SourceTensor, class DomainType, class CodomainType, auto KeyPermutation, auto DensePermutation>
class BlockTensorMappedView {
  public:
    using source_type = SourceTensor;
    using source_tensor_type = std::remove_cvref_t<SourceTensor>;
    using borrowed_block_tensor_view_tag = void;
    using element_type = typename source_tensor_type::element_type;
    using value_type = typename source_tensor_type::value_type;
    using storage_policy = typename source_tensor_type::storage_policy;
    using storage_type = typename source_tensor_type::storage_type;
    using backend_selector_type = typename source_tensor_type::backend_selector_type;
    using key_type = typename source_tensor_type::key_type;
    using domain_type = DomainType;
    using codomain_type = CodomainType;
    using source_mutable_access_block_type =
        decltype(std::declval<SourceTensor&>().block_by_ordinal(std::size_t{}));
    using mutable_block_type = decltype(map_block<DensePermutation>(std::declval<source_mutable_access_block_type>()));
    using const_block_type = decltype(make_read_only_block_view(std::declval<mutable_block_type const&>()));
    using allocation_context_binding_type = BlockTensorAllocationContextBinding<SourceTensor>;
 
    static constexpr std::size_t static_order = source_tensor_type::order();
    static constexpr std::size_t static_key_coordinate_count = source_tensor_type::key_coordinate_count();
    static constexpr std::size_t static_dense_block_order = source_tensor_type::dense_block_order();
 
  private:
    static constexpr bool source_blocks_writable =
        !std::is_const_v<typename std::remove_cvref_t<source_mutable_access_block_type>::element_type>;
 
    static_assert(IsDomain<domain_type>::value);
    static_assert(IsCodomain<codomain_type>::value);
    static_assert(domain_type::size() + codomain_type::size() == static_order);
    static_assert(BoundaryBlockShape<domain_type>::key_coordinate_count +
                      BoundaryBlockShape<codomain_type>::key_coordinate_count ==
                  static_key_coordinate_count);
    static_assert(BoundaryBlockShape<domain_type>::dense_block_order +
                      BoundaryBlockShape<codomain_type>::dense_block_order ==
                  static_dense_block_order);
    static_assert(std::tuple_size_v<decltype(KeyPermutation)> == static_key_coordinate_count);
    static_assert(std::tuple_size_v<decltype(DensePermutation)> == static_dense_block_order);
    static_assert(is_valid_axis_permutation(KeyPermutation));
    static_assert(is_valid_axis_permutation(DensePermutation));
 
  public:
    explicit BlockTensorMappedView(SourceTensor& source, domain_type domain, codomain_type codomain)
        : symmetry_(source.symmetry()), domain_(std::move(domain)), codomain_(std::move(codomain)),
          allocation_context_(source)
    {
      this->build_key_index(source);
    }
 
    BlockTensorMappedView(BlockTensorMappedView const&) = default;
 
    BlockTensorMappedView(BlockTensorMappedView&&) = default;
 
    auto operator=(BlockTensorMappedView const&) -> BlockTensorMappedView& = delete;
 
    auto operator=(BlockTensorMappedView&&) -> BlockTensorMappedView& = delete;
 
    auto symmetry() const -> Symmetry { return symmetry_; }
 
    auto domain() const -> domain_type const& { return domain_; }
 
    auto codomain() const -> codomain_type const& { return codomain_; }
 
    static constexpr auto order() noexcept -> std::size_t { return static_order; }
 
    static constexpr auto key_coordinate_count() noexcept -> std::size_t { return static_key_coordinate_count; }
 
    static constexpr auto dense_block_order() noexcept -> std::size_t { return static_dense_block_order; }
 
    auto stored_block_count() const noexcept -> std::size_t { return keys_.size(); }
 
    auto legal_block_count() const -> std::size_t
    {
      std::size_t count = 0;
      this->for_each_possible_key([&](key_type const& key) {
        if (this->is_legal(key)) ++count;
      });
      return count;
    }
 
    auto has_all_legal_blocks() const -> bool { return this->stored_block_count() == this->legal_block_count(); }
 
    auto is_legal(key_type const& key) const -> bool
    {
      return is_legal_block_key(this->symmetry(), domain_, codomain_, key);
    }
 
    auto contains(key_type const& key) const -> bool { return this->ordinal(key).has_value(); }
 
    auto find_block(key_type const& key) -> std::optional<mutable_block_type>
      requires source_blocks_writable
    {
      auto const found = this->ordinal(key);
      if (!found) return std::nullopt;
      return blocks_[*found];
    }
 
    auto find_block(key_type const& key) const -> std::optional<const_block_type>
    {
      auto const found = this->ordinal(key);
      if (!found) return std::nullopt;
      return make_read_only_block_view(std::as_const(blocks_[*found]));
    }
 
    auto block(key_type const& key) -> mutable_block_type
      requires source_blocks_writable
    {
      auto const found = this->find_block(key);
      if (!found) throw std::out_of_range("mapped BlockTensor block key is not stored");
      return *found;
    }
 
    auto block(key_type const& key) const -> const_block_type
    {
      auto const found = this->find_block(key);
      if (!found) throw std::out_of_range("mapped BlockTensor block key is not stored");
      return *found;
    }
 
    auto block_by_ordinal(std::size_t ordinal) -> mutable_block_type
      requires source_blocks_writable
    {
      if (ordinal >= keys_.size()) throw std::out_of_range("mapped BlockTensor block ordinal is out of range");
      return blocks_[ordinal];
    }
 
    auto block_by_ordinal(std::size_t ordinal) const -> const_block_type
    {
      if (ordinal >= keys_.size()) throw std::out_of_range("mapped BlockTensor block ordinal is out of range");
      return make_read_only_block_view(std::as_const(blocks_[ordinal]));
    }
 
    auto stored_keys() const noexcept -> std::span<key_type const> { return keys_; }
 
    static constexpr auto backend_selector() noexcept -> backend_selector_type
    {
      return source_tensor_type::backend_selector();
    }
 
    decltype(auto) allocation_context() const noexcept
      requires requires(allocation_context_binding_type const& binding) { binding.get(); }
    {
      return allocation_context_.get();
    }
 
  private:
    struct KeyBinding
    {
        key_type key;
        std::size_t source_ordinal;
    };
 
    void build_key_index(SourceTensor& source)
    {
      std::vector<KeyBinding> bindings;
      bindings.reserve(source.stored_block_count());
      for (std::size_t ordinal = 0; ordinal < source.stored_block_count(); ++ordinal)
      {
        bindings.push_back({permute_key(source.stored_keys()[ordinal], KeyPermutation), ordinal});
      }
      std::ranges::sort(bindings, {}, &KeyBinding::key);
 
      keys_.reserve(bindings.size());
      blocks_.reserve(bindings.size());
      for (auto const& binding : bindings)
      {
        keys_.push_back(binding.key);
        auto source_block = source.block_by_ordinal(binding.source_ordinal);
        blocks_.push_back(map_block<DensePermutation>(std::move(source_block)));
      }
    }
 
    template <class Function> void for_each_possible_key(Function&& function) const
    {
      auto const sizes = boundary_factor_sizes<static_key_coordinate_count>(domain_, codomain_);
      if (std::ranges::find(sizes, std::size_t{0}) != sizes.end()) return;
 
      if constexpr (static_key_coordinate_count == 0)
      {
        function(key_type{});
      }
      else
      {
        std::array<std::size_t, static_key_coordinate_count> coordinates{};
        while (true)
        {
          function(key_type{coordinates});
          std::size_t axis = static_key_coordinate_count;
          while (axis > 0)
          {
            --axis;
            if (++coordinates[axis] < sizes[axis]) break;
            coordinates[axis] = 0;
          }
          if (axis == 0 && coordinates[0] == 0) return;
        }
      }
    }
 
    auto ordinal(key_type const& key) const -> std::optional<std::size_t>
    {
      auto const found = std::ranges::lower_bound(keys_, key);
      if (found == keys_.end() || *found != key) return std::nullopt;
      return static_cast<std::size_t>(found - keys_.begin());
    }
 
    Symmetry symmetry_;
    domain_type domain_;
    codomain_type codomain_;
    [[no_unique_address]] allocation_context_binding_type allocation_context_;
    std::vector<key_type> keys_;
    std::vector<mutable_block_type> blocks_;
};
 
} // namespace uni20::detail