// Host-only layout checks for cutlass_cute_flash_attention_walkthrough_zh.md.
// nvcc -std=c++17 -Iinclude cute_flash_attention_trace.cu -o /tmp/cute_flash_attention_trace
#include <cute/tensor.hpp>
#include <cstdio>
#include <cstdlib>
using namespace cute;
void require(bool ok, char const* what) {
if (!ok) { std::fprintf(stderr, "FAILED: %s\n", what); std::exit(1); }
}
template<class T> void show(char const* name, T const& x) {
std::printf("%s: ", name); print(x); print("\n");
template<class L> auto rowcol(L l) {
auto d = logical_divide(l, Shape<_2>{});
return make_layout(make_layout(get<1>(get<0>(d)), get<1>(d)),
make_layout(get<0>(get<0>(d)), get<2>(d)));
template<class L> auto aregs(L l) {
using X = Underscore;
auto d = logical_divide(l, Shape<X, Shape<X, _2>>{});
return make_layout(make_layout(get<0>(get<1>(d)), get<0>(get<0>(d)),
get<0>(get<1>(get<1>(d)))),
get<1>(get<0>(d)), get<1>(get<1>(get<1>(d))));
template<class Copy, class Fragment> void steps(Copy, Fragment const& f) {
auto v = size<0>(f);
using NumThr = typename Copy::TiledNumThr;
auto mn = upcast<NumThr{} * v>(
right_inverse(typename Copy::TiledLayout_TV{})
.with_shape(shape(typename Copy::Tiler_MN{})));
auto fv = zipped_divide(logical_product(make_layout(v), right_inverse(mn)),
make_layout(typename Copy::AtomNumVal{}));
show("frg_layout_mn", mn);
show("gather extents", product_each(shape(mn)));
show("frg_layout_v", fv);
int main() {
using Mma = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<_4,_1,_1>>, Tile<_64,_16,_16>>;
using Smem = decltype(tile_to_shape(composition(Swizzle<3,3,3>{},
Layout<Shape<_8,_64>, Stride<_64,_1>>{}), Shape<_64,_64>{}));
using SmemVt = decltype(tile_to_shape(composition(Swizzle<3,3,3>{},
Layout<Shape<_64,_64>, Stride<_1,_64>>{}), Shape<_64,_64>{}));
Mma mma;
alignas(16) half_t storage[64*64]{};
auto smem = make_tensor(make_smem_ptr(storage), Smem{});
auto coords = make_identity_tensor(Shape<_64,_64>{});
auto ca = make_tiled_copy_A(Copy_Atom<SM75_U32x4_LDSM_N,half_t>{}, mma);
auto cb = make_tiled_copy_B(Copy_Atom<SM75_U32x4_LDSM_N,half_t>{}, mma);
int qa = 0, kb = 0, pa = 0, cc = 0;
for (int tid = 0; tid < 128; ++tid) {
auto thr = mma.get_slice(tid);
auto a = thr.partition_fragment_A(smem);
auto b = thr.partition_fragment_B(smem);
auto c = partition_fragment_C(mma, Shape<_64,_64>{});
auto xa = ca.get_slice(tid).retile_D(a);
auto xb = cb.get_slice(tid).retile_D(b);
require(xa.data() == a.data() && xb.data() == b.data(), "retile aliases storage");
auto ac = thr.partition_A(coords);
auto bc = thr.partition_B(coords);
auto dc = thr.partition_C(coords);
auto axc = ca.get_slice(tid).retile_D(ac);
auto bxc = cb.get_slice(tid).retile_D(bc);
auto ad = ca.get_slice(tid).partition_D(coords);
auto bd = cb.get_slice(tid).partition_D(coords);
for (int i = 0; i < size(ad); ++i) {
require(ad(i) == axc(i), "Q Copy destination matches MMA"); ++qa;
for (int i = 0; i < size(bd); ++i) {
require(bd(i) == bxc(i), "K Copy destination matches MMA"); ++kb;
auto plain = make_tensor(make_smem_ptr(storage),
Layout<Shape<_64,_64>, Stride<_64,_1>>{});
auto bp = thr.partition_fragment_B(plain);
auto xbp = cb.get_slice(tid).retile_D(bp);
for (int k = 0; k < 4; ++k)
for (int n = 0; n < 4; ++n)
for (int v = 0; v < 8; ++v) {
require(int(xb.layout()(make_coord(v,0),n,k)) ==
v + 8*n + 32*k, "swizzled K x4 slot formula");
require(int(xbp.layout()(make_coord(v,0),n,k)) ==
v%4 + 16*(2*n+v/4) + 4*k, "plain K x4 slot formula");
int slots[32];
for (int i = 0; i < 32; ++i) slots[i] = i;
auto rc = make_tensor(make_rmem_ptr(&slots[0]), rowcol(c.layout()));
auto pc = make_tensor(make_rmem_ptr(&slots[0]), aregs(rc.layout()));
for (int i = 0; i < size(ac); ++i) {
require(dc(pc(i)) == ac(i), "C to A stays in the same lane"); ++pa;
for (int j = 0; j < 16; ++j)
for (int r = 0; r < 2; ++r) {
require(dc(rc(r,j)) == make_coord(16*(tid/32)+(tid%32)/4+8*r,
2*(tid%4)+j%2+8*(j/2)), "score coordinates");
++cc;
if (tid == 0) {
show("tCrQ", a.layout()); show("tCrK", b.layout());
show("tCrS", c.layout()); show("tXrQ", xa.layout());
show("tXrK", xb.layout()); show("tCrS_rowcol", rc.layout());
show("tCrP", pc.layout());
show("plain tCrK", bp.layout()); show("plain tXrK", xbp.layout());
print("Q retile internals\n"); steps(ca,a);
print("K retile internals\n"); steps(cb,b);
std::printf("K x4 value %d -> actual slot %d, plain slot %d -> ", v,
int(xb.layout()(make_coord(v,0),0,0)),
int(xbp.layout()(make_coord(v,0),0,0)));
print(bxc(make_coord(v,0),0,0)); print("\n");
for (int key = 0; key < 64; ++key)
for (int dim = 0; dim < 64; ++dim)
require(Smem{}(key,dim) == SmemVt{}(dim,key), "V transpose aliases same half");
std::printf("PASS: Q %d, K %d, C-to-A %d, score coordinates %d, V addresses 4096\n",
qa,kb,pa,cc);
std::puts("Host layout checks only; no GPU instruction executed.");