leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/layout.hpp
+131 -73
View File
@@ -418,8 +418,8 @@ make_layout_like(Layout<Shape,Stride> const& layout)
// Make a compact layout with the same shape as @a layout
// and strides following the order induced by @a layout.stride(),
// except mode-0 is always stride-1 and generated column-major.
// The 0th mode is commonly used for MMA_Atoms or Copy_Atoms
// so this generates the 0th mode with LayoutLeft regardless of the reference layout.
// The 0th mode is commonly used for MMA_Atoms or Copy_Atoms so this
// generates the 0th mode with LayoutLeft (preserving stride-0s) regardless of the reference layout
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
@@ -427,7 +427,8 @@ make_fragment_like(Layout<Shape,Stride> const& layout)
{
constexpr int R = Layout<Shape,Stride>::rank;
if constexpr (R > 1 && is_static<Shape>::value) {
return tiled_product(make_layout(shape<0>(layout)),
return tiled_product(make_layout(get<0>(layout.shape()),
compact_col_major(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride())));
} else {
return make_layout(layout.shape());
@@ -757,7 +758,8 @@ bw_coalesce(OldShape const& old_shape, OldStride const& old_stride,
} else if constexpr (is_constant<1, NewShape>::value) {
// Replace our shape-1 with anything (Can only happen on input new_shape/new_stride)
return bw_coalesce<I-1>(old_shape, old_stride, get<I>(old_shape), get<I>(old_stride));
} else if constexpr (is_constant<true, decltype(get<I>(old_shape) * get<I>(old_stride) == get<0>(new_stride))>::value) {
} else if constexpr (is_static<decltype(get<0>(new_shape))>::value &&
is_constant<true, decltype(get<I>(old_shape) * get<I>(old_stride) == get<0>(new_stride))>::value) {
// Merge modes because the shapes and strides match
return bw_coalesce<I-1>(old_shape, old_stride,
replace_front(new_shape, get<I>(old_shape) * get<0>(new_shape)),
@@ -772,6 +774,45 @@ bw_coalesce(OldShape const& old_shape, OldStride const& old_stride,
CUTE_GCC_UNREACHABLE;
}
// cute::coalesce promises to not change the Layout as a function from integers to codomain.
// It accomplishes this inside of the Layout's domain, but not always outside of the domain.
// Example: (_4,_1):(_1,_0) coalesces to _4:_1.
// detail::coalesce_x preserves the Layout function inside its domain and outside.
//
// @post depth(@a result) <= 1
// @post for all i, 0 <= i, @a layout(i) == @a result(i)
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
coalesce_x(Layout<Shape,Stride> const& layout)
{
auto flat_shape = flatten(layout.shape());
auto flat_stride = flatten(layout.stride());
constexpr int R = decltype(rank(flat_shape))::value;
if constexpr (is_constant<1, decltype(get<R-1>(flat_shape))>::value) {
return detail::bw_coalesce<R-2>(flat_shape, flat_stride, Int<2>{}, get<R-1>(flat_stride));
} else {
return detail::bw_coalesce<R-2>(flat_shape, flat_stride, get<R-1>(flat_shape), get<R-1>(flat_stride));
}
}
// Apply coalesce_x at the terminals of trg_profile
template <class Shape, class Stride, class IntTuple>
CUTE_HOST_DEVICE constexpr
auto
coalesce_x(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
{
if constexpr (is_tuple<IntTuple>::value) {
static_assert(tuple_size<IntTuple>::value <= Layout<Shape,Stride>::rank);
return cute::transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return coalesce_x(l,t); });
} else {
return coalesce_x(layout);
}
CUTE_GCC_UNREACHABLE;
}
} // end namespace detail
// "Simplify" the layout by combining modes that are possible to combine
@@ -807,6 +848,25 @@ coalesce(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
CUTE_GCC_UNREACHABLE;
}
// Combine static and dynamic modes of a shape.
// @post size(@a result) == size(@a shape)
// @post depth(@a result) <= 1
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
coalesce(Shape const& shape)
{
static_assert(is_integral<Shape>::value || is_tuple<Shape>::value);
return cute::fold_first(flatten(shape), [](auto const& init, auto const& a) {
if constexpr (is_static<decltype(back(init))>::value == is_static<decltype(a)>::value) {
return replace_back(init, back(init) * a); // Both static or both dynamic, coalesce and replace
} else {
return append(init, a); // Can't coalesce, so append
}
});
}
// Replace the modes in layout that have a 0-stride with a 1-size
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
@@ -918,70 +978,64 @@ template <class LShape, class LStride,
class RShape, class RStride>
CUTE_HOST_DEVICE constexpr
auto
composition_impl(Layout<LShape,LStride> const& lhs,
composition_impl(LShape const& lhs_shape, LStride const& lhs_stride,
RShape const& rhs_shape, RStride const& rhs_stride)
{
if constexpr (is_tuple<RShape>::value) {
// Apply the right-distributivity of Layout composition
return transform_layout(rhs_shape, rhs_stride, [&](auto const& s, auto const& d) { return composition_impl(lhs, s, d); });
return transform_layout(rhs_shape, rhs_stride, [&](auto const& s, auto const& d) {
return composition_impl(lhs_shape, lhs_stride, s, d);
});
} else
if constexpr (is_scaled_basis<RStride>::value) {
// Special case for a ScaledBasis stride
return composition_impl(get<RStride::mode()>(lhs), rhs_shape, rhs_stride.value());
return composition_impl(basis_get(rhs_stride, lhs_shape), basis_get(rhs_stride, lhs_stride),
rhs_shape, basis_value(rhs_stride));
} else
if constexpr (is_integral<RStride>::value) {
// Integral Rstride (and RShape)
if constexpr (is_constant<0, RStride>::value) {
// Special case shortcut for any static stride-0
return Layout<RShape, RStride>{rhs_shape, rhs_stride};
} else
if constexpr (is_integral<decltype(lhs_shape)>::value) {
// Special case shortcut for any integral LShape
return Layout{rhs_shape, rhs_stride * lhs_stride};
} else
if constexpr (is_constant<1, RStride>::value) {
// Special case shortcut for any static stride-1
constexpr int R = rank_v<LShape>;
auto result_shape_0 = take<0,R-1>(lhs_shape);
// NOTE: Should only flatten once for efficiency
auto flat_shape = flatten(lhs.shape());
[[maybe_unused]] auto flat_stride = flatten(lhs.stride());
[[maybe_unused]] constexpr int R = rank(flat_shape);
// Mod out the rhs_shape from the lhs_shape
auto const [result_shape_1, rest_shape] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_shape),
[] (auto const& init, auto const& si) {
return cute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
});
if constexpr (is_constant<0, RStride>::value) {
// Special case shortcut for any static stride-0
return Layout<RShape, RStride>{rhs_shape, rhs_stride};
} else
if constexpr (is_integral<decltype(flat_shape)>::value) {
// Special case shortcut for any integral LShape
auto result_stride = rhs_stride * flat_stride;
return Layout<RShape, decltype(result_stride)>{rhs_shape, result_stride};
} else
if constexpr (is_constant<1, RStride>::value) {
// Special case shortcut for any static stride-1
auto result_shape_0 = take<0,R-1>(flat_shape);
// Jump into coalesce and append (rest_shape, get<R-1>(lhs_stride))
return detail::bw_coalesce<R-2>(result_shape_1, lhs_stride, rest_shape, get<R-1>(lhs_stride));
} else {
// General case: integral RShape and RStride, tuple LShape and LStride
constexpr int R = rank_v<LShape>;
auto result_shape_0 = take<0,R-1>(lhs_shape);
auto result_stride_0 = take<0,R-1>(lhs_stride);
// Mod out the rhs_shape from the lhs.shape()
auto const [result_shape_1, rest_shape] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_shape),
[] (auto const& init, auto const& si) {
return cute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
});
// Divide out the rhs_stride from the lhs_shape
auto const [result_shape_1, rest_stride] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_stride),
[] (auto const& init, auto const& di) {
return cute::make_tuple(append(get<0>(init), shape_div(di, get<1>(init))), shape_div(get<1>(init), di));
});
// Jump into coalesce and append (rest_shape, get<R-1>(lhs.stride())
return detail::bw_coalesce<R-2>(result_shape_1, flat_stride, rest_shape, get<R-1>(flat_stride));
} else
{
// General case
auto result_shape_0 = take<0,R-1>(flat_shape);
auto result_stride_0 = take<0,R-1>(flat_stride);
// Apply any lhs_shape changes to the stride
auto result_stride_1 = elem_scale(result_stride_0, shape_div(result_shape_0, result_shape_1));
// Divide out the rhs_stride from the lhs.shape()
auto const [result_shape_1, rest_stride] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_stride),
[] (auto const& init, auto const& di) {
return cute::make_tuple(append(get<0>(init), shape_div(di, get<1>(init))), shape_div(get<1>(init), di));
});
// Mod out the rhs_shape from the lhs_shape
auto const [result_shape_2, rest_shape] = fold(result_shape_1, cute::make_tuple(cute::make_tuple(), rhs_shape),
[] (auto const& init, auto const& si) {
return cute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
});
// Apply any lhs.shape() changes to the stride
auto result_stride_1 = elem_scale(result_stride_0, shape_div(result_shape_0, result_shape_1));
// Mod out the rhs_shape from the lhs.shape()
auto const [result_shape_2, rest_shape] = fold(result_shape_1, cute::make_tuple(cute::make_tuple(), rhs_shape),
[] (auto const& init, auto const& si) {
return cute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
});
// Jump into coalesce and append (rest_shape, rest_stride * get<R-1>(lhs.stride())
return detail::bw_coalesce<R-2>(result_shape_2, result_stride_1, rest_shape, rest_stride * get<R-1>(flat_stride));
}
// Jump into coalesce and append (rest_shape, rest_stride * get<R-1>(lhs_stride))
return detail::bw_coalesce<R-2>(result_shape_2, result_stride_1, rest_shape, rest_stride * get<R-1>(lhs_stride));
}
CUTE_GCC_UNREACHABLE;
@@ -996,7 +1050,9 @@ auto
composition(Layout<LShape,LStride> const& lhs,
Layout<RShape,RStride> const& rhs)
{
return detail::composition_impl(lhs, rhs.shape(), rhs.stride());
auto coprofile = repeat_like(decltype(coshape(rhs)){}, Int<0>{});
auto flat_lhs = detail::coalesce_x(lhs, coprofile);
return detail::composition_impl(flat_lhs.shape(), flat_lhs.stride(), rhs.shape(), rhs.stride());
}
template <class LShape, class LStride, class Tiler>
@@ -1012,7 +1068,8 @@ composition(Layout<LShape,LStride> const& lhs,
} else if constexpr (is_underscore<Tiler>::value) {
return lhs;
} else if constexpr (is_integral<Tiler>::value) {
return detail::composition_impl(lhs, rhs, Int<1>{});
auto flat_lhs = detail::coalesce_x(lhs);
return detail::composition_impl(flat_lhs.shape(), flat_lhs.stride(), rhs, Int<1>{});
}
CUTE_GCC_UNREACHABLE;
@@ -1032,14 +1089,14 @@ composition(Layout<LShape,LStride> const& lhs,
namespace detail {
// @pre @a layout has been filtered (flattened and no stride-0 or size-1 modes).
template <class Shape, class Stride, class CoSizeHi>
template <class Shape, class Stride, class CoTarget>
CUTE_HOST_DEVICE constexpr
auto
complement(Shape const& shape, Stride const& stride, CoSizeHi const& cosize_hi)
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(cosize_hi);
return make_layout(coalesce(cotarget));
} else {
// General case
constexpr int R = rank_v<Shape>;
@@ -1055,28 +1112,30 @@ complement(Shape const& shape, Stride const& stride, CoSizeHi const& cosize_hi)
{
auto [shape, stride, result_shape, result_stride] = init;
auto min_stride = cute::min(stride);
auto min_idx = find(stride, min_stride);
auto min_idx = cute::find(stride, min_stride);
auto new_shape = min_stride / get<i>(result_stride);
auto new_stride = get<min_idx>(shape) * min_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 = curr_shape * min_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);
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); // new shape = min_stride / last_stride
auto result_shape = append(result_shape_, new_shape);
// Compute the rest_shape and rest_stride
auto rest_stride = get<0>(shape_) * get<0>(stride_);
auto rest_shape = ceil_div(cosize_hi, 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_col_major(rest_shape, new_stride);
// Jump into coalesce and append (rest_shape, rest_stride)
return detail::bw_coalesce<R-1>(result_shape, result_stride, rest_shape, rest_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;
@@ -1084,14 +1143,13 @@ complement(Shape const& shape, Stride const& stride, CoSizeHi const& cosize_hi)
} // end namespace detail
template <class Shape, class Stride, class CoSizeHi>
template <class Shape, class Stride, class CoTarget>
CUTE_HOST_DEVICE constexpr
auto
complement(Layout<Shape,Stride> const& layout, CoSizeHi const& cosize_hi)
complement(Layout<Shape,Stride> const& layout, CoTarget const& cotarget)
{
static_assert(cute::is_integral<CoSizeHi>::value, "Expected integral codomain size in complement.");
auto filter_layout = filter(layout);
return detail::complement(filter_layout.shape(), filter_layout.stride(), cosize_hi);
return detail::complement(filter_layout.shape(), filter_layout.stride(), shape(cotarget));
}
template <class Shape, class Stride>
@@ -1365,7 +1423,7 @@ auto
logical_divide(Layout<LShape,LStride> const& layout,
Layout<TShape,TStride> const& tiler)
{
return composition(layout, make_layout(tiler, complement(tiler, size(layout))));
return composition(layout, make_layout(tiler, complement(tiler, shape(layout))));
}
template <class LShape, class LStride, class Tiler>