group_indices.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./layout/idx_map.hpp"
#include "./layout/permutation.hpp"
#include "./stdutil/array.hpp"
#include "./traits.hpp"
 
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <numeric>
#include <stdexcept>
 
namespace nda {
 
  namespace detail {
 
    // Check if given integer arrays correspond to a partition of the set `J = {0, 1, ..., R-1}`.
    template <size_t R, size_t... Rs>
    constexpr bool is_partition_of_indices(std::array<int, Rs> const &...grps) {
      // check that the total number of array elements is equal to R
      constexpr auto Rstot = (Rs + ...);
      static_assert(Rstot == R, "Error in nda::detail::is_partition_of_indices: Number of indices has to be equal to the rank");
 
      // check that every element of J is in one and only one group
      auto count = stdutil::make_initialized_array<R>(0); // FIXME : not necessary in C++20
      auto l     = [&](auto &&grp) mutable -> void {
        for (int u = 0; u < grp.size(); ++u) {
          if (grp[u] < 0) throw std::runtime_error("Error in nda::detail::is_partition_of_indices: Negative indices are not allowed");
          if (grp[u] >= R) throw std::runtime_error("Error in nda::detail::is_partition_of_indices: Indices larger than the rank are not allowed");
          count[grp[u]]++;
        }
      };
      (l(grps), ...);
 
      return (count == stdutil::make_initialized_array<R>(1));
    }
 
    // Given an original stride order and a partition of the indices, determine the new stride order of the grouped
    // nda::idx_map.
    template <size_t R, size_t... Rs>
    constexpr std::array<int, sizeof...(Rs)> stride_order_of_grouped_idx_map(std::array<int, R> const &stride_order,
                                                                             std::array<int, Rs> const &...grps) {
      // inverse permutation of the stride order
      auto mem_pos = permutations::inverse(stride_order);
 
      // check if the indices/dimensions in a group are consecutive in memory and find the slowest varying index within
      // the group
      auto min_mem_pos = [&mem_pos](auto &&grp) {
        int min = R, max = 0;
        for (int idx : grp) {
          int v = mem_pos[idx];
          if (v > max) max = v;
          if (v < min) min = v;
        }
        bool idx_are_consecutive_in_memory = (max - min + 1 == grp.size());
        if (!idx_are_consecutive_in_memory)
          throw std::runtime_error("Error in nda::detail::stride_order_of_grouped_idx_map: Indices are not consecutive in memory");
        return min;
      };
 
      // number of groups in the partition
      constexpr int Ngrps = sizeof...(Rs);
 
      // array containing the slowest varying index within each group
      std::array<int, Ngrps> min_mem_positions{min_mem_pos(grps)...};
 
      // rank each group by its slowest varying index, i.e. the slowest group gets rank 0 and the fastest rank Ngrps-1
      std::array<int, Ngrps> mem_pos_out = stdutil::make_initialized_array<Ngrps>(0); // FIXME : not necessary in C++20
      for (int u = 0; u < Ngrps; ++u) {
        for (int i = 0; i < Ngrps; ++i) {
          if (min_mem_positions[i] < min_mem_positions[u]) ++mem_pos_out[u];
        }
      }
 
      // the new stride order is the inverse of mem_pos_out
      return permutations::inverse(mem_pos_out);
    }
 
  } // namespace detail
 
  template <int... Is>
  struct idx_group_t {
    static constexpr std::array<int, sizeof...(Is)> as_std_array{Is...};
  };
 
  template <int... Is>
  inline idx_group_t<Is...> idx_group = {}; // NOLINT (global variable template is what we want here)
 
  template <int Rank, uint64_t StaticExtents, uint64_t StrideOrder, layout_prop_e LayoutProp, typename... IdxGrps>
  auto group_indices_layout(idx_map<Rank, StaticExtents, StrideOrder, LayoutProp> const &idxm, IdxGrps...) {
    // compile-time checks
    static_assert(LayoutProp == layout_prop_e::contiguous, "Error in nda::group_indices_layout: Only contiguous layouts are supported");
    static_assert(detail::is_partition_of_indices<Rank>(IdxGrps::as_std_array...),
                  "Error in nda::group_indices_layout: Given index groups are not a valid partition");
 
    // decoded original stride order
    static constexpr auto stride_order = idx_map<Rank, StaticExtents, StrideOrder, LayoutProp>::stride_order;
 
    // number of elements in a group
    auto total_len_of_a_grp = [&idxm](auto &&grp) {
      return std::accumulate(grp.begin(), grp.end(), 1L, [&idxm](long l, long u) { return l * idxm.lengths()[u]; });
    };
 
    // minimum stride of a group
    auto min_stride_of_a_grp = [&idxm](auto &&grp) {
      return std::accumulate(grp.begin(), grp.end(), idxm.size(), [&idxm](long s, long u) { return std::min(s, idxm.strides()[u]); });
    };
 
    // new rank, shape, strides and stride order
    static constexpr int new_rank = sizeof...(IdxGrps);
    std::array<long, new_rank> new_extents{total_len_of_a_grp(IdxGrps::as_std_array)...};
    std::array<long, new_rank> new_strides{min_stride_of_a_grp(IdxGrps::as_std_array)...};
    constexpr auto new_stride_order = encode(detail::stride_order_of_grouped_idx_map(stride_order, IdxGrps::as_std_array...));
 
    return idx_map<new_rank, 0, new_stride_order, layout_prop_e::contiguous>{new_extents, new_strides};
  }
 
} // namespace nda