// Host-only CuTe Layout exercises. No CUDA kernel is launched.
// Starter:
// nvcc -std=c++17 -I include cute_layout_exercises.cu -o /tmp/cute_layout_exercises
// Reference answers:
// nvcc -std=c++17 -I include -DCUTE_LAYOUT_EXERCISES_REFERENCE \
// cute_layout_exercises.cu -o /tmp/cute_layout_exercises_reference
// Run all exercises, or select one: /tmp/cute_layout_exercises 4
// Edit the six TODO branches below; the default branches intentionally fail checks.
#include <cute/tensor.hpp>
#include <array>
#include <cstdio>
#include <cstdlib>
using namespace cute;
// -------------------- Fill in these six functions --------------------
auto exercise1_layout()
{
#ifdef CUTE_LAYOUT_EXERCISES_REFERENCE
return make_layout(Shape<_3,_4>{}, Stride<_8,_1>{});
#else
// TODO(1): Keep a 3x4 logical matrix; start each row eight elements apart.
return make_layout(Shape<_3,_4>{}, Stride<_4,_1>{});
#endif
}
template <class Src, class Dst>
void exercise2_copy(Src const& src, Dst const& dst)
// The tensors have the same shape, so the same integer coordinate denotes
// the same (m,n), even when their physical offsets differ.
for (int i = 0; i < int(size(src)); ++i) {
dst(i) = src(i);
// TODO(2): Preserve matrix coordinates when copying row-major to column-major.
dst.data()[i] = src.data()[i];
auto exercise3_layout()
return make_layout(Shape<_2,Shape<_2,_3>>{},
Stride<_6,Stride<_3,_1>>{});
// TODO(3): Make h(m,(a,b)) = 6*m + 3*a + b, with a<2 and b<3.
Stride<_6,Stride<_1,_2>>{});
template <class Tensor>
auto exercise4_transpose(Tensor const& a)
auto transposed_shape = make_shape(size<1>(a), size<0>(a));
auto transposed_stride = make_stride(stride<1>(a), stride<0>(a));
// TODO(4): A transpose view swaps both shape modes and stride modes.
auto transposed_stride = make_stride(_3{}, _1{});
return make_tensor(a.data(), make_layout(transposed_shape, transposed_stride));
auto exercise5_tv()
return make_layout(Shape<_8,_2>{}, Stride<_1,_8>{});
// TODO(5): Eight logical threads, two values each; q = t + 8*v.
return make_layout(Shape<_8,_2>{}, Stride<_2,_1>{});
auto exercise6_mma_coord(int lane, int v)
int a = v % 2;
int b = (v / 2) % 2;
int c = v / 4;
return make_coord(lane / 4 + 8*b, 2*(lane % 4) + a + 8*c);
// TODO(6), optional: Decode all eight values of each SM80 MMA A lane.
return make_coord(lane / 4, 2*(lane % 4));
// -------------------- Acceptance checks: leave these unchanged --------------------
namespace {
struct CheckFailure {};
void expect_equal(char const* label, int actual, int expected)
if (actual != expected) {
std::printf(" FAIL: %s; got %d, expected %d\n", label, actual, expected);
throw CheckFailure{};
void expect_true(char const* label, bool condition)
expect_equal(label, condition ? 1 : 0, 1);
template <class Layout>
void print_offsets(Layout const& layout)
for (int m = 0; m < int(size<0>(layout)); ++m) {
for (int n = 0; n < int(size<1>(layout)); ++n) {
std::printf("%4d", int(layout(m,n)));
std::printf("\n");
void check1()
auto padded = exercise1_layout();
print_offsets(padded);
expect_equal("logical size", int(size(padded)), 12);
expect_equal("cosize (no trailing padding required)", int(cosize(padded)), 20);
for (int m = 0; m < 3; ++m) {
for (int n = 0; n < 4; ++n) {
expect_equal("padded(m,n)", int(padded(m,n)), 8*m+n);
void check2()
std::array<int,12> source{};
std::array<int,12> destination{};
auto src = make_tensor(source.data(), make_layout(Shape<_3,_4>{}, Stride<_4,_1>{}));
auto dst = make_tensor(destination.data(), make_layout(Shape<_3,_4>{}, Stride<_1,_3>{}));
src(m,n) = 10*m+n;
exercise2_copy(src, dst);
std::printf("destination storage:");
for (int value : destination) {
std::printf(" %d", value);
expect_equal("src.layout()(5)", int(src.layout()(5)), 9);
expect_equal("dst.layout()(5)", int(dst.layout()(5)), 5);
expect_equal("dst(m,n) preserves matrix value", dst(m,n), 10*m+n);
void check3()
auto h = exercise3_layout();
print_offsets(h);
expect_equal("rank", int(rank(h)), 2);
expect_equal("depth", int(depth(h)), 2);
expect_equal("size", int(size(h)), 12);
auto flat = flatten(h);
expect_true("group restores hierarchy", group<1,3>(flat) == h);
for (int m = 0; m < 2; ++m) {
for (int j = 0; j < 6; ++j) {
int a = j % 2;
int b = j / 2;
int expected = 6*m+3*a+b;
expect_equal("h(m,j)", int(h(m,j)), expected);
expect_equal("h(m,(a,b))", int(h(make_coord(m,make_coord(a,b)))), expected);
expect_equal("h(i)", int(h(m+2*j)), expected);
expect_equal("flatten(h)(i)", int(flat(m+2*j)), expected);
void check4()
std::array<int,20> storage;
storage.fill(-99);
auto a = make_tensor(storage.data(), make_layout(Shape<_3,_4>{}, Stride<_8,_1>{}));
a(m,n) = 10*m+n;
auto b = exercise4_transpose(a);
expect_equal("transpose rows", int(size<0>(b)), 4);
expect_equal("transpose columns", int(size<1>(b)), 3);
expect_true("transpose shares the data pointer", b.data() == a.data());
// Check every computed offset before dereferencing the view.
expect_equal("transpose offset", int(b.layout()(n,m)), 8*m+n);
expect_equal("b(n,m) == a(m,n)", b(n,m), a(m,n));
auto row = b(2,_);
expect_true("b(2,_) starts at storage+2", row.data() == storage.data()+2);
expect_equal("b(2,_) stride", int(stride<0>(row)), 8);
row(1) = 777;
expect_equal("write through transpose row updates a(1,2)", a(1,2), 777);
for (int offset = 0; offset < 20; ++offset) {
int expected = offset % 8 < 4 ? 10*(offset/8)+offset%8 : -99;
if (offset == 10) {
expected = 777;
expect_equal("only the selected element changed; padding preserved", storage[offset], expected);
std::printf("b(2,1) = %d; a(1,2) = %d; storage[10] = %d\n",
b(2,1), a(1,2), storage[10]);
void check5()
auto data_layout = make_layout(Shape<_4,_4>{}, Stride<_6,_1>{});
auto tv = exercise5_tv();
auto address = composition(data_layout, tv);
std::array<int,22> storage;
std::array<int,16> visits{};
auto output = make_tensor(storage.data(), data_layout);
for (int t = 0; t < 8; ++t) {
for (int v = 0; v < 2; ++v) {
int q = int(tv(t,v));
expect_equal("tv(t,v) is a logical position", q, t+8*v);
int m = q % 4;
int n = q / 4;
int offset = int(address(t,v));
expect_equal("composition yields the physical offset", offset, 6*m+n);
expect_equal("logical position is written exactly once", ++visits[q], 1);
storage[offset] = 100*t+v;
for (int q = 0; q < 16; ++q) {
expect_equal("complete tile coverage", visits[q], 1);
expect_equal("matrix value identifies its owner", output(q), 100*(q%8)+q/8);
for (int offset = 0; offset < 22; ++offset) {
if (offset % 6 >= 4) {
expect_equal("padding not written", storage[offset], -99);
std::printf("owner values (100*t+v):\n");
for (int m = 0; m < 4; ++m) {
std::printf("%5d", output(m,n));
void check6()
auto mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{});
auto tv = mma.get_layoutA_TV();
auto data_layout = make_layout(Shape<_16,_16>{}, Stride<_16,_1>{});
std::array<int,256> visits{};
for (int lane = 0; lane < 32; ++lane) {
for (int v = 0; v < 8; ++v) {
auto coord = exercise6_mma_coord(lane,v);
int m = int(get<0>(coord));
int k = int(get<1>(coord));
expect_true("MMA coordinates in [0,16)", m >= 0 && m < 16 && k >= 0 && k < 16);
int q = m+16*k;
expect_equal("coordinate agrees with actual MMA A TV", q, int(tv(lane,v)));
expect_equal("each A element has one owner", ++visits[q], 1);
if (lane == 4) {
std::printf("lane=4 v=%d -> A(%d,%d) q=%d physical_offset=%d\n",
v,m,k,q,int(data_layout(m,k)));
for (int count : visits) {
expect_equal("all 256 A elements covered", count, 1);
struct Exercise {
char const* name;
void (*run)();
};
} // namespace
int main(int argc, char** argv)
int selected = 0;
if (argc > 2) {
std::fprintf(stderr, "Usage: %s [1-6]\n", argv[0]);
return 2;
if (argc == 2) {
char* end = nullptr;
long parsed = std::strtol(argv[1], &end, 10);
if (end == argv[1] || *end != '\0' || parsed < 1 || parsed > 6) {
std::fprintf(stderr, "Choose an exercise from 1 to 6.\n");
selected = int(parsed);
Exercise exercises[] = {
{"padded matrix", check1},
{"copy between storage layouts", check2},
{"hierarchical coordinates", check3},
{"transpose view and slicing", check4},
{"thread/value ownership", check5},
{"SM80 MMA A coordinates (optional)", check6},
std::printf("Reference answers; host execution only.\n");
std::printf("Starter; fill TODO(1)-TODO(6). Initial failures are expected.\n");
int attempted = 0;
int passed = 0;
for (int i = 0; i < 6; ++i) {
if (selected != 0 && selected != i+1) {
continue;
++attempted;
std::printf("\nExercise %d: %s\n", i+1, exercises[i].name);
try {
exercises[i].run();
++passed;
std::printf(" PASS\n");
} catch (CheckFailure const&) {
// Continue so each exercise has its own feedback.
std::printf("\nPassed %d/%d exercises.\n", passed, attempted);
return passed == attempted ? 0 : 1;