Function complement

include/cute/layout.hpp:1178–1227

// @pre @a layout has been filtered (flattened and no stride-0 or size-1 modes).
template <class Shape, class Stride, class CoTarget>
CUTE_HOST_DEVICE constexpr
auto
complement(Shape const& shape, Stride const& stride, CoTarget const& cotarget)
{
  if constexpr (is_constant<0, Stride>::value) {
    // Special case for irreducible rank-1 stride-0 layout
    return make_layout(coalesce(cotarget));
  } else {
    // General case
    constexpr int R = rank_v<Shape>;
    static_assert(R == 1 || is_static<Stride>::value,
                  "Dynamic-stride complement only for rank-1 layouts");

    // Should just be a sort and a fold...
    // Then we could even handle dynamic strides (but they would destroy all static strides)
    auto [shape_, stride_, result_shape_, result_stride] =
      fold(make_seq<R-1>{},
           cute::make_tuple(shape, stride, cute::make_tuple(), cute::make_tuple(Int<1>{})),
           [](auto const& init, auto i)
           {
              auto [shape, stride, result_shape, result_stride] = init;
              auto min_stride = cute::min(stride);
              auto min_idx    = cute::find(stride, min_stride);
              auto new_shape  = min_stride / get<i>(result_stride);
              auto new_stride = min_stride * get<min_idx>(shape);
              static_assert(not is_constant<0, decltype(new_shape)>::value, "Non-injective Layout detected in complement.");

              return cute::make_tuple(remove<min_idx>(shape),              // Remove the min_idx from shape
                                      remove<min_idx>(stride),             // Remove the min_idx from stride
                                      append(result_shape , new_shape ),   // new shape  = min_stride / last_stride
                                      append(result_stride, new_stride));  // new stride = min_stride * curr_shape
            });

    // Append the last shape mode
    auto new_shape    = get<0>(stride_) / get<R-1>(result_stride);         // new shape  = min_stride / last_stride
    static_assert(not is_constant<0, decltype(new_shape)>::value, "Non-injective Layout detected in complement.");
    auto result_shape = append(result_shape_, new_shape);

    // Compute the rest_shape and rest_stride
    auto new_stride  = get<0>(stride_) * get<0>(shape_);                   // new stride = min_stride * curr_shape
    auto rest_shape  = coalesce(ceil_div(cotarget, new_stride));
    auto rest_stride = compact_major<LayoutLeft>(rest_shape, new_stride);

    // Coalesce and append (rest_shape, rest_stride)
    return coalesce(make_layout(make_shape (result_shape , rest_shape ),
                                make_stride(result_stride, rest_stride)));
  }

  CUTE_GCC_UNREACHABLE;
}

} // end namespace detail