// Companion to cute_make_tiled_mma_walkthrough_zh.md.
// Host coordinate checks only: no MMA instruction or GPU kernel is executed.
// nvcc -std=c++17 -arch=sm_80 -Iinclude cute_make_tiled_mma_trace.cu \
// -o /tmp/cute_make_tiled_mma_trace
#include <cute/tensor.hpp>
#include <cute/atom/mma_traits_sm80.hpp>
#include <cstdio>
#include <cstdlib>
#include <vector>
using namespace cute;
using Op = SM80_16x8x16_F16F16F16F16_TN;
void require(bool ok, char const* message) {
if (!ok) {
std::fprintf(stderr, "FAILED: %s\n", message);
std::exit(1);
}
template<class T>
void show(char const* label, T const& x) {
std::printf("%s = ", label); print(x); print("\n");
template<class MMA>
void show_c_steps(MMA const& mma) {
using AtomTV = typename MMA::AtomLayoutC_TV;
auto C = make_layout(make_shape(_32{},_16{}), make_stride(_16{},_1{}));
auto t = logical_divide(C, make_tile(mma.template permutation_mnk<0>(),
mma.template permutation_mnk<1>()));
auto c = zipped_divide(t, make_tile(make_layout(_16{}), make_layout(_8{})));
auto tv = c.compose(AtomTV{}, _);
auto th = zipped_divide(tv, make_tile(_, make_tile(make_layout(_2{}), make_layout(_2{}))));
print("\n[C partition construction, row-major 32x16]\n");
show("input C", C);
show("logical_divide", t);
show("Atom zipped_divide", c);
show("compose Atom TV", tv);
show("thread zipped_divide", th);
show("partition_C thread 0", mma.get_slice(0).partition_C(make_tensor((half_t*)nullptr,C)).layout());
// Independent expectations from the SM80 lane/value rules, including reuse.
template<int M, int N, int K, int TM, int TN, int TK,
bool PermuteM = false, bool ReverseWarps = false, class MMA>
void verify(char const* label, MMA const& mma) {
auto A = make_identity_tensor(make_shape(Int<M>{}, Int<K>{}));
auto B = make_identity_tensor(make_shape(Int<N>{}, Int<K>{}));
auto C = make_identity_tensor(make_shape(Int<M>{}, Int<N>{}));
constexpr int threads = 32*TM*TN*TK;
std::vector<int> coverA(M*K), coverB(N*K), coverC(M*N);
int checks = 0;
auto map_m = [](int m) {
if constexpr (PermuteM) return (m/32)*32 + 4*(m%8) + (m%32)/8;
else return m;
};
std::printf("\n[%s]\n", label);
show("ThrLayoutVMNK", mma.get_thr_layout_vmnk());
show("tile MNK", tile_shape(mma));
show("layoutA_TV", mma.get_layoutA_TV());
show("layoutB_TV", mma.get_layoutB_TV());
show("layoutC_TV", mma.get_layoutC_TV());
auto tvA = mma.get_layoutA_TV();
auto tvB = mma.get_layoutB_TV();
auto tvC = mma.get_layoutC_TV();
int tileM = int(tile_size<0>(mma));
int tileN = int(tile_size<1>(mma));
for (int t = 0; t < threads; ++t) {
int lane = t%32, warp = t/32;
int wm = ReverseWarps ? warp/TN : warp%TM;
int wn = ReverseWarps ? warp%TN : (warp/TM)%TN;
int wk = warp/(TM*TN);
auto thr = mma.get_slice(t);
auto a = thr.partition_A(A);
auto b = thr.partition_B(B);
auto c = thr.partition_C(C);
if (t == 0) {
show("partition_A shape", shape(a));
show("partition_B shape", shape(b));
show("partition_C shape", shape(c));
auto rA = thr.partition_fragment_A(make_tensor((half_t*)nullptr, make_layout(shape(A))));
auto rB = thr.partition_fragment_B(make_tensor((half_t*)nullptr, make_layout(shape(B))));
auto rC = thr.partition_fragment_C(C);
show("fragment_A layout", rA.layout());
show("fragment_B layout", rB.layout());
show("fragment_C layout", rC.layout());
require(int(size(a)) == 8*(M/(16*TM))*(K/(16*TK)), "A size");
require(int(size(b)) == 4*(N/(8*TN))*(K/(16*TK)), "B size");
require(int(size(c)) == 4*(M/(16*TM))*(N/(8*TN)), "C size");
for (int rm = 0; rm < M/(16*TM); ++rm)
for (int rk = 0; rk < K/(16*TK); ++rk)
for (int v = 0; v < 8; ++v) {
int m = map_m(lane/4 + 8*((v/2)%2) + 16*wm + 16*TM*rm);
int k = 2*(lane%4) + v%2 + 8*(v/4) + 16*wk + 16*TK*rk;
auto coord = a(v,rm,rk);
require(int(get<0>(coord)) == m && int(get<1>(coord)) == k, "A coordinate");
++coverA[m+M*k]; ++checks;
if (M == tileM && K == int(tile_size<2>(mma)))
require(int(tvA(t,make_coord(v,make_coord(rm,rk)))) == m+M*k, "A TV");
for (int rn = 0; rn < N/(8*TN); ++rn)
for (int v = 0; v < 4; ++v) {
int n = lane/4 + 8*wn + 8*TN*rn;
int k = 2*(lane%4) + v%2 + 8*(v/2) + 16*wk + 16*TK*rk;
auto coord = b(v,rn,rk);
require(int(get<0>(coord)) == n && int(get<1>(coord)) == k, "B coordinate");
++coverB[n+N*k]; ++checks;
if (N == tileN && K == int(tile_size<2>(mma)))
require(int(tvB(t,make_coord(v,make_coord(rn,rk)))) == n+N*k, "B TV");
int m = map_m(lane/4 + 8*(v/2) + 16*wm + 16*TM*rm);
int n = 2*(lane%4) + v%2 + 8*wn + 8*TN*rn;
auto coord = c(v,rm,rn);
require(int(get<0>(coord)) == m && int(get<1>(coord)) == n, "C coordinate");
++coverC[m+M*n]; ++checks;
if (M == tileM && N == tileN)
require(int(tvC(t,make_coord(v,make_coord(rm,rn)))) == m+M*n, "C TV");
for (int count : coverA) require(count == TN, "A coverage multiplicity");
for (int count : coverB) require(count == TM, "B coverage multiplicity");
for (int count : coverC) require(count == TK, "C coverage multiplicity");
std::printf("Verified %d coordinates; A/B/C multiplicity = %d/%d/%d.\n", checks, TN, TM, TK);
int main() {
auto base = make_tiled_mma(Op{});
auto four = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>>{});
auto reverse = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>,Stride<_2,_1,_0>>{});
auto larger = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>>{}, Tile<_64,_32,_16>{});
using PermM = Layout<Shape<_8,_4>,Stride<_4,_1>>;
auto perm = make_tiled_mma(Op{}, Layout<Shape<_2,_2,_1>>{}, Tile<PermM,Underscore,Underscore>{});
auto split = make_tiled_mma(Op{}, Layout<Shape<_1,_1,_2>>{});
verify<16,8,16,1,1,1>("one warp", base);
verify<32,16,16,2,2,1>("four warps", four);
verify<64,32,32,2,2,1>("larger input", four);
verify<32,16,16,2,2,1,false,true>("reordered warps", reverse);
verify<64,32,16,2,2,1>("explicit tile extent", larger);
verify<32,16,16,2,2,1,true>("permuted M", perm);
verify<16,8,32,1,1,2>("K thread tiling", split);
show_c_steps(four);
// Separate real storage offsets from the identity-coordinate checks above.
auto rowC = make_tensor((half_t*)nullptr,
make_layout(make_shape(_64{},_32{}), make_stride(_32{},_1{})));
show("large row-major partition_C", four.get_slice(0).partition_C(rowC).layout());
print("All make_tiled_mma layout checks passed. No GPU kernel was launched.\n");