idx_map.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-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, Dominik Kiese, Henri Menke, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./permutation.hpp"
#include "./range.hpp"
#include "./slice_static.hpp"
#include "../macros.hpp"
#include "../traits.hpp"
 
#include <algorithm>
#include <array>
#include <concepts>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <numeric>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace nda {
 
  template <int Rank>
  constexpr uint64_t Fortran_stride_order = nda::encode(nda::permutations::reverse_identity<Rank>());
 
  template <int Rank>
  constexpr uint64_t C_stride_order = nda::encode(nda::permutations::identity<Rank>());
 
  template <int Rank, uint64_t StaticExtents, uint64_t StrideOrder, layout_prop_e LayoutProp>
  class idx_map {
    static_assert(Rank < 16, "Error in nda::idx_map: Rank must be < 16");
    static_assert((StrideOrder != 0) or (Rank == 1), "Error in nda::idx_map: StrideOrder can only be zero for 1D arrays");
 
    // Extents of all dimensions (the shape of the map).
    std::array<long, Rank> len{};
 
    // Strides of all dimensions.
    std::array<long, Rank> str{};
 
    public:
    static constexpr uint64_t static_extents_encoded = StaticExtents;
 
    static constexpr std::array<int, Rank> static_extents = decode<Rank>(StaticExtents);
 
    static constexpr std::array<int, Rank> stride_order = (StrideOrder == 0 ? permutations::identity<Rank>() : decode<Rank>(StrideOrder));
 
    static constexpr uint64_t stride_order_encoded = encode(stride_order);
 
    static constexpr layout_prop_e layout_prop = LayoutProp;
 
    static constexpr layout_info_t layout_info = layout_info_t{stride_order_encoded, layout_prop};
 
    template <typename T>
    static constexpr int argument_is_allowed_for_call = std::is_constructible_v<long, T>;
 
    template <typename T>
    static constexpr int argument_is_allowed_for_call_or_slice =
       std::is_same_v<range, T> or std::is_same_v<range::all_t, T> or std::is_same_v<ellipsis, T> or std::is_constructible_v<long, T>;
 
    protected:
    static constexpr int n_dynamic_extents = []() {
      int r = 0;
      for (int u = 0; u < Rank; ++u) r += (static_extents[u] == 0 ? 1 : 0);
      return r;
    }();
 
    public:
    static constexpr int rank() noexcept { return Rank; }
 
    [[nodiscard]] long size() const noexcept { return std::accumulate(len.cbegin(), len.cend(), 1L, std::multiplies<>{}); }
 
    static constexpr long ce_size() noexcept {
      if constexpr (n_dynamic_extents != 0) {
        return 0;
      } else {
        return std::accumulate(static_extents.cbegin(), static_extents.cend(), 1L, std::multiplies<>{});
      }
    }
 
    [[nodiscard]] std::array<long, Rank> const &lengths() const noexcept { return len; }
 
    [[nodiscard]] std::array<long, Rank> const &strides() const noexcept { return str; }
 
    [[nodiscard]] long min_stride() const noexcept { return str[stride_order[Rank - 1]]; }
 
    [[nodiscard]] bool is_contiguous() const noexcept {
      auto s = size();
      if (s == 0) return true;
      return (std::abs(str[stride_order[0]] * len[stride_order[0]]) == s);
    }
 
    [[nodiscard]] bool has_positive_strides() const noexcept { return (*std::min_element(str.cbegin(), str.cend()) >= 0); }
 
    [[nodiscard]] bool is_strided_1d() const noexcept {
      auto s = size();
      if (s == 0) return true;
      int i = Rank - 1;
      while (len[stride_order[i]] == 1 and i > 0) --i;
      return (std::abs(str[stride_order[0]] * len[stride_order[0]]) == s * std::abs(str[stride_order[i]]));
    }
 
    static constexpr bool is_stride_order_C() { return (stride_order == permutations::identity<Rank>()); }
 
    static constexpr bool is_stride_order_Fortran() { return (stride_order == permutations::reverse_identity<Rank>()); }
 
    template <std::integral Int>
    [[nodiscard]] static bool is_stride_order_valid(Int *lenptr, Int *strptr) {
      auto dims_to_check = std::vector<int>{};
      dims_to_check.reserve(Rank);
      for (auto dim : stride_order)
        if (lenptr[dim] > 1) dims_to_check.push_back(dim);
      for (int n = 1; n < dims_to_check.size(); ++n)
        if (std::abs(strptr[dims_to_check[n - 1]]) < std::abs(strptr[dims_to_check[n]])) return false;
      return true;
    }
 
    [[nodiscard]] bool is_stride_order_valid() const { return is_stride_order_valid(len.data(), str.data()); }
 
    private:
    // Compute contiguous strides from the shape.
    void compute_strides_contiguous() {
      long s = 1;
      for (int v = rank() - 1; v >= 0; --v) {
        int u  = stride_order[v];
        str[u] = s;
        s *= len[u];
      }
      ENSURES(s == size());
    }
 
    // Check that the static extents and the shape are compatible.
    void assert_static_extents_and_len_are_compatible() const {
#ifdef NDA_ENFORCE_BOUNDCHECK
      if constexpr (n_dynamic_extents != Rank) {
#ifndef NDEBUG
        for (int u = 0; u < Rank; ++u)
          if (static_extents[u] != 0) EXPECTS(static_extents[u] == len[u]);
#endif
      }
#endif
    }
 
    // Should we check the stride order when constructing an idx_map from a shape and strides?
#ifdef NDA_DEBUG
    static constexpr bool check_stride_order = true;
#else
    static constexpr bool check_stride_order = false;
#endif
 
    // Combine static and dynamic extents.
    static std::array<long, Rank> merge_static_and_dynamic_extents(std::array<long, n_dynamic_extents> const &dynamic_extents) {
      std::array<long, Rank> extents;
      for (int u = 0, v = 0; u < Rank; ++u) extents[u] = (static_extents[u] == 0 ? dynamic_extents[v++] : static_extents[u]);
      return extents;
    }
 
    // FIXME ADD A CHECK layout_prop_e ... compare to stride and
 
    public:
    idx_map() {
      if constexpr (n_dynamic_extents == 0) {
        for (int u = 0; u < Rank; ++u) len[u] = static_extents[u];
        compute_strides_contiguous();
      }
    }
 
    template <layout_prop_e LP>
    idx_map(idx_map<Rank, StaticExtents, StrideOrder, LP> const &idxm) noexcept : len(idxm.lengths()), str(idxm.strides()) {
      // check strides and stride order of the constructed map
      EXPECTS(is_stride_order_valid());
 
      // check that the layout properties are compatible
      if constexpr (not layout_property_compatible(LP, layout_prop)) {
        if constexpr (has_contiguous(layout_prop)) {
          EXPECTS_WITH_MESSAGE(idxm.is_contiguous(), "Error in nda::idx_map: Constructing a contiguous from a non-contiguous layout");
        }
        if constexpr (has_strided_1d(layout_prop)) {
          EXPECTS_WITH_MESSAGE(idxm.is_strided_1d(), "Error in nda::idx_map: Constructing a strided_1d from a non-strided_1d layout");
        }
      }
    }
 
    template <uint64_t SE, layout_prop_e LP>
    idx_map(idx_map<Rank, SE, StrideOrder, LP> const &idxm) noexcept(false) : len(idxm.lengths()), str(idxm.strides()) {
      // check strides and stride order
      EXPECTS(is_stride_order_valid());
 
      // check that the layout properties are compatible
      if constexpr (not layout_property_compatible(LP, LayoutProp)) {
        if constexpr (has_contiguous(LayoutProp)) {
          EXPECTS_WITH_MESSAGE(idxm.is_contiguous(), "Error in nda::idx_map: Constructing a contiguous from a non-contiguous layout");
        }
        if constexpr (has_strided_1d(LayoutProp)) {
          EXPECTS_WITH_MESSAGE(idxm.is_strided_1d(), "Error in nda::idx_map: Constructing a strided_1d from a non-strided_1d layout");
        }
      }
 
      // check that the static extents and the shape are compatible
      assert_static_extents_and_len_are_compatible();
    }
 
    idx_map(std::array<long, Rank> const &shape, // NOLINT (only throws if check_stride_order is true)
            std::array<long, Rank> const &strides) noexcept(!check_stride_order)
       : len(shape), str(strides) {
      EXPECTS(std::all_of(shape.cbegin(), shape.cend(), [](auto const &i) { return i >= 0; }));
      if constexpr (check_stride_order) {
        if (not is_stride_order_valid()) throw std::runtime_error("Error in nda::idx_map: Incompatible strides, shape and stride order");
      }
    }
 
    template <std::integral Int = long>
    idx_map(std::array<Int, Rank> const &shape) noexcept : len(stdutil::make_std_array<long>(shape)) {
      EXPECTS(std::all_of(shape.cbegin(), shape.cend(), [](auto const &i) { return i >= 0; }));
      assert_static_extents_and_len_are_compatible();
      compute_strides_contiguous();
    }
 
    idx_map(std::array<long, n_dynamic_extents> const &shape) noexcept
      requires((n_dynamic_extents != Rank) and (n_dynamic_extents != 0))
       : idx_map(merge_static_and_dynamic_extents(shape)) {}
 
    template <uint64_t SE, uint64_t SO, layout_prop_e LP>
    idx_map(idx_map<Rank, SE, SO, LP> const &)
      requires(stride_order_encoded != SO)
    {
      static_assert((stride_order_encoded == SO), "Error in nda::idx_map: Incompatible stride orders");
    }
 
    template <int R>
    idx_map(std::array<long, R> const &)
      requires(R != Rank)
    {
      static_assert(R == Rank, "Error in nda::idx_map: Incompatible ranks");
    }
 
    idx_map(idx_map const &) = default;
 
    idx_map(idx_map &&) = default;
 
    idx_map &operator=(idx_map const &) = default;
 
    idx_map &operator=(idx_map &&) = default;
 
    private:
    // Get the contribution to the linear index in case of an nda::ellipsis argument.
    template <bool skip_stride, auto I>
    [[nodiscard]] FORCEINLINE long myget(ellipsis) const noexcept {
      // nda::ellipsis are skipped and do not contribute to the linear index
      return 0;
    }
 
    // Get the contribution to the linear index in case of a long argument.
    template <bool skip_stride, auto I>
    [[nodiscard]] FORCEINLINE long myget(long arg) const noexcept {
      if constexpr (skip_stride and (I == stride_order[Rank - 1])) {
        // optimize for the case when the fastest varying dimension is contiguous in memory
        return arg;
      } else {
        // otherwise multiply the argument by the stride of the current dimension
        return arg * std::get<I>(str);
      }
    }
 
    // Is the smallest stride equal to one, i.e. is the fastest varying dimension contiguous in memory?
    static constexpr bool smallest_stride_is_one = has_smallest_stride_is_one(layout_prop);
 
    // Implementation of the function call operator that takes a multi-dimensional index and maps it to a linear index.
    template <typename... Args, size_t... Is>
    [[nodiscard]] FORCEINLINE long call_impl(std::index_sequence<Is...>, Args... args) const noexcept {
      // set e_pos to the position of the ellipsis, otherwise to -1
      static constexpr int e_pos = ((std::is_same_v<Args, ellipsis> ? int(Is) + 1 : 0) + ...) - 1;
 
      if constexpr (e_pos == -1) {
        // no ellipsis present
        if constexpr (smallest_stride_is_one) {
          // optimize for the case that the fastest varying dimension is contiguous in memory
          return (myget<true, Is>(static_cast<long>(args)) + ...);
        } else {
          // arbitrary layouts
          return ((args * std::get<Is>(str)) + ...);
        }
      } else {
        // empty ellipsis is present and needs to be skipped
        return (myget<smallest_stride_is_one, (Is < e_pos ? Is : Is - 1)>(args) + ...);
      }
    }
 
    public:
    template <typename... Args>
    FORCEINLINE long operator()(Args const &...args) const
#ifdef NDA_ENFORCE_BOUNDCHECK
       noexcept(false) {
      assert_in_bounds(rank(), len.data(), args...);
#else
       noexcept(true) {
#endif
      return call_impl(std::make_index_sequence<sizeof...(Args)>{}, args...);
    }
 
    std::array<long, Rank> to_idx(long lin_idx) const {
      // compute residues starting from slowest index
      std::array<long, Rank> residues;
      residues[0] = lin_idx;
      for (auto i : range(1, Rank)) { residues[i] = residues[i - 1] % str[stride_order[i - 1]]; }
 
      // convert residues to indices, ordered from slowest to fastest
      std::array<long, Rank> idx;
      idx[Rank - 1] = residues[Rank - 1] / str[stride_order[Rank - 1]];
      for (auto i : range(Rank - 2, -1, -1)) { idx[i] = (residues[i] - residues[i + 1]) / str[stride_order[i]]; }
 
      // reorder indices according to stride order
      return permutations::apply_inverse(stride_order, idx);
    }
 
    template <typename... Args>
    auto slice(Args const &...args) const {
      return slice_static::slice_idx_map(*this, args...);
    }
 
    template <int R, uint64_t SE, uint64_t SO, layout_prop_e LP>
    bool operator==(idx_map<R, SE, SO, LP> const &rhs) const {
      return (Rank == R and len == rhs.lengths() and str == rhs.strides());
    }
 
    template <uint64_t Permutation>
    auto transpose() const {
      // Makes a new transposed idx_map with permutation P such that
      // denoting here A = this, A' = P A = returned_value
      // A'(i_k) = A(i_{P[k]})
      //
      // Note that this convention is the correct one to have a (left) action of the symmetric group on
      // a array and it may not be completely obvious.
      // Proof
      //  let's operate with P then Q, and denote A'' = Q A'. We want to show that A'' = (QP) A
      //   A'(i_k) = A(i_{P[k]})
      //   A''(j_k) = A'(j_{Q[k]})
      //   then i_k = j_{Q[k]} and  A''(j_k) =  A(i_{P[k]}) = A(j_{Q[P[k]]}) = A(j_{(QP)[k]}), q.e.d
      //
      // NB test will test this composition
      // Denoting this as A, an indexmap, calling it returns the linear index given by
      //
      // A(i_k) = sum_k i_k * S[k] (1)
      //
      // where S[k] denotes the strides.
      //
      // 1- S' : strides of A'
      //    A'(i_k) = sum_k i_{P[k]} * S[k] = sum_k i_k * S[P{^-1}[k]]
      //     so
      //         S'[k] = S[P{^-1}[k]]  (2)
      //    i.e. apply (inverse(P), S) or apply_inverse directly.
      //
      // 2- L' : lengths of A'
      //    if L[k] is the k-th length, then because of the definition of A', i.e. A'(i_k) = A(i_{P[k]})
      //    i_q in the lhs A is at position q' such that P[q'] = q  (A'(i0 i1 i2...) = A( i_P0 i_P1 i_P2....)
      //    hence L'[q] = L[q'] =  L[P^{-1}[q]]
      //    same for static length
      //
      // 3- stride_order: denoted in this paragraph as Q (and Q' for A').
      //    by definition Q is a permutation such that Q[0] is the slowest index, Q[Rank -1] the fastest
      //    hence S[Q[k]] is a strictly decreasing sequence (as checked by strides_compatible_to_stride_order)
      //    we want therefore Q' the permutation that will sort the S', i.e.
      //    S'[Q'[k]] = S[Q[k]]
      //    using (2), we have S[P{^-1}[Q'[k]]] = S[Q[k]]
      //    so the permutation Q' is such that  P{^-1}Q' = Q  or Q' = PQ (as permutation product/composition).
      //    NB : Q and P are permutations, so the operation must be a composition, not an apply (apply applies
      //    a P to any set, like L, S, not only a permutation) even though they are all std::array in the code ...
      //
      static constexpr std::array<int, Rank> permu              = decode<Rank>(Permutation);
      static constexpr std::array<int, Rank> new_stride_order   = permutations::compose(permu, stride_order);
      static constexpr std::array<int, Rank> new_static_extents = permutations::apply_inverse(permu, static_extents);
 
      return idx_map<Rank, encode(new_static_extents), encode(new_stride_order), LayoutProp>{permutations::apply_inverse(permu, lengths()),
                                                                                             permutations::apply_inverse(permu, strides())};
    }
  };
 
} // namespace nda