leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 30 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/layout_composed.hpp
+50 -39
View File
@@ -37,7 +37,7 @@
/* This implements a ComposedLayout of the form
* LayoutA o Offset o LayoutB
* and is useful in cases where composition() does not or cannot apply to LayoutA and LayoutB.
* For example, then the "divisibility condition" in shape_div is violated in composition(LayoutA, LayoutB).
* For example, when the "divisibility condition" in shape_div is violated in composition(LayoutA, LayoutB).
*
* This ComposedLayout provides similar functionality to Layout including tiling, partitioning,
* coordinate-to-index mapping and layout manipulations, but is not considered a "normal" layout.
@@ -357,12 +357,11 @@ composition(LayoutA const& layoutA,
return ComposedLayout<LayoutA, Offset, LayoutB>{layoutA, offset, layoutB};
}
template <class A, class O, class B,
class LayoutOrTile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
composition(ComposedLayout<A,O,B> const& a,
LayoutOrTile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), composition(a.layout_b(), b));
}
@@ -433,92 +432,101 @@ zip(ComposedLayout<A,O,B> const& a)
// Partitions
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
logical_divide(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), logical_divide(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
tile_unzip(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), tile_unzip(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
tiled_divide(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), tiled_divide(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
zipped_divide(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), zipped_divide(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
flat_divide(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), flat_divide(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
logical_product(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), logical_product(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
zipped_product(ComposedLayout<A,O,B> const& a,
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), zipped_product(a.layout_b(), b));
}
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
tiled_product(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), tiled_product(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
flat_product(ComposedLayout<A,O,B> const& a,
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), flat_product(a.layout_b(), b));
}
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
blocked_product(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), blocked_product(a.layout_b(), b));
}
template <class A, class O, class B,
class Tile>
template <class A, class O, class B, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
raked_product(ComposedLayout<A,O,B> const& a,
Tile const& b)
Tiler const& b)
{
return composition(a.layout_a(), a.offset(), raked_product(a.layout_b(), b));
}
@@ -585,16 +593,19 @@ CUTE_HOST_DEVICE constexpr
auto
recast_layout(ComposedLayout<A,O,B> const& layout)
{
if constexpr (sizeof(NewType) == sizeof(OldType)) {
using scale = decltype(trait_ratio(sizeof_bits<NewType>{}, sizeof_bits<OldType>{}));
if constexpr (scale::num == 1 && scale::den == 1) {
return layout;
} else if constexpr (sizeof(NewType) > sizeof(OldType)) {
static_assert(sizeof(NewType) % sizeof(OldType) == 0, "NewType must be a multiple of OldType");
return upcast<sizeof(NewType)/sizeof(OldType)>(layout);
} else if constexpr (sizeof(NewType) < sizeof(OldType)) {
static_assert(sizeof(OldType) % sizeof(NewType) == 0, "NewType must be a divisor of OldType");
return downcast<sizeof(OldType)/sizeof(NewType)>(layout);
}
else if constexpr (scale::num == 1) {
return downcast<scale::den>(layout);
}
else if constexpr (scale::den == 1) {
return upcast<scale::num>(layout);
}
else {
static_assert(dependent_false<scale>, "Recast not supported.");
}
CUTE_GCC_UNREACHABLE;
}