strides.hpp

Go to the documentation of this file.

#pragma once
 
#include "concepts.hpp"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <functional>
#include <initializer_list>
#include <limits>
#include <uni20/common/static_vector.hpp>
#include <uni20/common/trace.hpp>
#include <uni20/core/types.hpp>
#include <uni20/mdspan/mdspan.hpp>
 
namespace uni20
{
 
template <std::size_t N> struct extent_strides
{
    index_type extent;
    std::array<std::ptrdiff_t, N> strides;
 
    constexpr extent_strides() = default;
 
    template <typename Ext>
      requires std::is_integral_v<Ext>
    constexpr extent_strides(Ext e, std::initializer_list<std::ptrdiff_t> s) : extent(static_cast<index_type>(e))
    {
      assert(s.size() == N);
      std::copy(s.begin(), s.end(), strides.begin());
    }
 
    template <typename Ext, typename Str>
      requires std::is_integral_v<Ext> && std::is_integral_v<Str>
    constexpr extent_strides(Ext e, std::array<Str, N> s) : extent(static_cast<index_type>(e))
    {
      for (std::size_t i = 0; i < N; ++i)
      {
        strides[i] = static_cast<std::ptrdiff_t>(s[i]);
      }
    }
 
    template <typename Ext, typename... Strides>
      requires(std::is_integral_v<Ext> && sizeof...(Strides) == N)
    constexpr extent_strides(Ext e, Strides... s)
        : extent(static_cast<index_type>(e)), strides{static_cast<std::ptrdiff_t>(s)...}
    {}
 
    [[nodiscard]] constexpr bool can_merge_with_inner(extent_strides inner) const noexcept
    {
      if (extent < 0 || inner.extent < 0) return false;
      if (extent == 1 || inner.extent == 1) return true;
 
      auto const inner_extent = static_cast<std::ptrdiff_t>(inner.extent);
      for (std::size_t i = 0; i < N; ++i)
      {
        auto const inner_stride = inner.strides[i];
        if (inner_extent != 0 &&
            ((inner_stride > 0 && inner_stride > std::numeric_limits<std::ptrdiff_t>::max() / inner_extent) ||
             (inner_stride < 0 && inner_stride < std::numeric_limits<std::ptrdiff_t>::min() / inner_extent)))
          return false;
        if (strides[i] != inner_stride * inner_extent) return false;
      }
      return true;
    }
 
    constexpr void merge_with_inner(extent_strides inner) noexcept
    {
      if (extent == 1)
      {
        extent = inner.extent;
        if (inner.extent == 1)
          strides.fill(1);
        else
          strides = inner.strides;
        return;
      }
      if (inner.extent == 1) return;
 
      extent *= inner.extent;
      strides = inner.strides;
    }
};
 
namespace detail
{
 
template <std::size_t N, std::size_t R>
void remove_singleton_dimensions(static_vector<extent_strides<N>, R>& dimensions)
{
  if (dimensions.empty()) return;
 
  std::size_t write = 0;
  for (std::size_t read = 0; read < dimensions.size(); ++read)
  {
    if (dimensions[read].extent != 1) dimensions[write++] = dimensions[read];
  }
 
  if (write == 0)
  {
    dimensions[0].extent = 1;
    dimensions[0].strides.fill(1);
    dimensions.resize(1);
    return;
  }
  dimensions.resize(write);
}
 
template <std::size_t N, std::size_t R> void sort_and_merge_left(static_vector<extent_strides<N>, R>& out)
{
  remove_singleton_dimensions(out);
  if (out.size() <= 1) return;
 
  std::sort(out.begin(), out.end(),
            [](auto const& lhs, auto const& rhs) { return std::abs(lhs.strides[0]) < std::abs(rhs.strides[0]); });
 
  extent_strides<N> current = out[0];
  std::size_t write = 0;
  for (std::size_t i = 1; i < out.size(); ++i)
  {
    auto candidate = out[i];
    if (candidate.can_merge_with_inner(current))
    {
      candidate.merge_with_inner(current);
      current = candidate;
    }
    else
    {
      out[write++] = current;
      current = candidate;
    }
  }
 
  out[write++] = current;
  out.resize(write);
}
 
template <std::size_t N, std::size_t R> void sort_and_merge_right(static_vector<extent_strides<N>, R>& out)
{
  remove_singleton_dimensions(out);
  if (out.size() <= 1) return;
 
  std::sort(out.begin(), out.end(),
            [](auto const& lhs, auto const& rhs) { return std::abs(lhs.strides[0]) > std::abs(rhs.strides[0]); });
 
  extent_strides<N> current = out[0];
  std::size_t write = 0;
  for (std::size_t i = 1; i < out.size(); ++i)
  {
    auto inner = out[i];
    if (current.can_merge_with_inner(inner))
    {
      current.merge_with_inner(inner);
    }
    else
    {
      out[write++] = current;
      current = inner;
    }
  }
 
  out[write++] = current;
  out.resize(write);
}
 
} // namespace detail
 
template <typename Extents, std::size_t R>
  requires(Extents::rank() == R)
inline static_vector<extent_strides<2>, R> merge_strides_left(Extents const& ext,
                                                              std::array<std::ptrdiff_t, R> const& Stride1,
                                                              std::array<std::ptrdiff_t, R> const& Stride2)
{
  static_vector<extent_strides<2>, R> out;
  for (std::size_t i = 0; i < R; ++i)
  {
    out.emplace_back({ext[i], {Stride1[i], Stride2[i]}});
  }
 
  detail::sort_and_merge_left(out);
 
  return out;
}
 
template <typename Extents, std::size_t R>
  requires(Extents::rank() == R)
inline static_vector<extent_strides<2>, R> merge_strides_right(Extents const& ext,
                                                               std::array<std::ptrdiff_t, R> const& Stride1,
                                                               std::array<std::ptrdiff_t, R> const& Stride2)
{
  static_vector<extent_strides<2>, R> out;
  for (std::size_t i = 0; i < R; ++i)
  {
    out.emplace_back({ext[i], {Stride1[i], Stride2[i]}});
  }
 
  detail::sort_and_merge_right(out);
 
  return out;
}
 
template <std::size_t N, std::size_t R> void merge_strides_left(static_vector<extent_strides<N>, R>& out)
{
  detail::sort_and_merge_left(out);
}
 
template <std::size_t N, std::size_t R> void merge_strides_right(static_vector<extent_strides<N>, R>& out)
{
  detail::sort_and_merge_right(out);
}
 
template <StridedMdspanLike AType, StridedMdspanLike BType, StridedMdspanLike CType, std::size_t N>
auto extract_strides(AType const& A, BType const& B,
                     std::array<std::pair<std::size_t, std::size_t>, N> const& contractDims, CType const& C)
{
  constexpr std::size_t MR = AType::rank() - N; // rank of the M group (A/C legs that are not contracted over)
  constexpr std::size_t NR = BType::rank() - N; // rank of the N group (B/C legs that are not contracted over)
  constexpr std::size_t KR = N;                 // rank of the K group (A/Blegs that are contracted over)
 
  static_vector<extent_strides<2>, MR> Mgroup;
  static_vector<extent_strides<2>, NR> Ngroup;
  static_vector<extent_strides<2>, KR> Kgroup;
 
  // Assemble the K array of the contracted legs and mark which legs of A and B are contracted over
  std::array<bool, AType::rank()> AContracted{};
  std::array<bool, BType::rank()> BContracted{};
  for (std::size_t i = 0; i < N; ++i)
  {
    // check that the extents match
    auto [ai, bi] = contractDims[i];
    ERROR_IF(A.extent(ai) != B.extent(bi), "Extent along tensor contraction dimension does not match", ai, bi);
    AContracted[ai] = true;
    BContracted[bi] = true;
    Kgroup.emplace_back(A.extent(ai), A.stride(ai), B.stride(bi));
  }
  // Now fill out the uncontracted dimensions and verify that they match with C
  std::size_t ci = 0;
  for (std::size_t ai = 0; ai < AType::rank(); ++ai)
  {
    if (!AContracted[ai])
    {
      ERROR_IF(A.extent(ai) != C.extent(ci), "Extent along uncontracted dimension does not match", ai, ci);
      Mgroup.emplace_back(A.extent(ai), A.stride(ai), C.stride(ci));
      ++ci;
    }
  }
  for (std::size_t bi = 0; bi < BType::rank(); ++bi)
  {
    if (!BContracted[bi])
    {
      ERROR_IF(B.extent(bi) != C.extent(ci), "Extent along uncontracted dimension does not match", bi, ci);
      Ngroup.emplace_back(B.extent(bi), B.stride(bi), C.stride(ci));
      ++ci;
    }
  }
 
  merge_strides_right(Mgroup);
  merge_strides_right(Ngroup);
  merge_strides_right(Kgroup);
 
  return std::tuple{Mgroup, Ngroup, Kgroup};
}
 
} // namespace uni20