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

Updates for 3.4 release. (#1305)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2024-01-16 13:42:51 -05:00
committed by GitHub
parent acba5beee5
commit 2f589ffa76
166 changed files with 5996 additions and 4702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/stride.hpp
+22 -24
View File
@@ -37,7 +37,8 @@
namespace cute
{
/** crd2idx maps a coordinate within <Shape,Stride> to an index
/** crd2idx(c,s,d) maps a coordinate within <Shape,Stride> to an index
*
* This is computed as follows:
* [coord, shape, and stride are all integers => step forward by stride]
* op(c, s, d) => c * d
@@ -46,7 +47,6 @@ namespace cute
* [coord, shape, and stride are all tuples => consider each mode independently]
* op((c,C), (s,S), (d,D)) => op(c, s, d) + op((C), (S), (D))
*/
template <class Coord, class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
@@ -115,10 +115,6 @@ crd2idx(Coord const& coord,
CUTE_GCC_UNREACHABLE;
}
//
// If we know Stride is default [CompactColMajor], then we can take shortcuts
//
namespace detail {
template <class CTuple, class STuple, int I0, int... Is>
@@ -138,26 +134,31 @@ crd2idx_horner(CTuple const& coord,
} // end namespace detail
/** crd2idx(c,s) maps a coordinate within Shape to an index
* via a colexicographical enumeration of coordinates in Shape.
* i = c0 + s0 * (c1 + s1 * (c2 + s2 * ...))
*/
template <class Coord, class Shape>
CUTE_HOST_DEVICE constexpr
auto
crd2idx(Coord const& coord,
Shape const& shape)
{
static_assert(decltype(congruent(coord,shape))::value, "Mismatched Ranks");
if constexpr (is_tuple<Shape>::value) {
// Flatten and apply Horner's method
auto flat_coord = flatten(coord);
auto flat_shape = flatten(shape);
return detail::crd2idx_horner(flat_coord, flat_shape, tuple_seq<decltype(flat_shape)>{});
} else {
if constexpr (is_integral<Coord>::value) { // Coord is already an index
return coord;
} else if constexpr (is_integral<Shape>::value) {
static_assert(dependent_false<Shape>, "Invalid parameters");
} else { // Make congruent, flatten, and apply Horner's method
static_assert(tuple_size<Coord>::value == tuple_size<Shape>::value, "Mismatched Ranks");
auto flat_coord = flatten(coord);
auto flat_shape = flatten(product_like(shape, coord));
return detail::crd2idx_horner(flat_coord, flat_shape, tuple_seq<decltype(flat_shape)>{});
}
CUTE_GCC_UNREACHABLE;
}
/** idx2crd splits an index to a coordinate within <Shape,Stride>.
/** idx2crd(i,s,d) splits an index into a coordinate within <Shape,Stride>.
*
* This is computed as follows:
* [index, shape, and stride are all integers => determine 1D coord]
@@ -170,7 +171,6 @@ crd2idx(Coord const& coord,
* NOTE: This only works for compact shape+stride layouts. A more general version would
* apply to all surjective layouts
*/
template <class Index, class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
@@ -207,15 +207,13 @@ idx2crd(Index const& idx,
CUTE_GCC_UNREACHABLE;
}
//
// If we know Stride is default [CompactColMajor], then we can take shortcuts
//
//(idx / 1) % s0
//(idx / s0) % s1
//(idx / (s0 * s1)) % s2
//...
/** idx2crd(i,s) splits an index into a coordinate within Shape
* via a colexicographical enumeration of coordinates in Shape.
* c0 = (idx / 1) % s0
* c1 = (idx / s0) % s1
* c2 = (idx / (s0 * s1)) % s2
* ...
*/
template <class Index, class Shape>
CUTE_HOST_DEVICE constexpr
auto