// Host-only layout inspection. No ldmatrix or mma.sync is executed.
// Compile from the CUTLASS root:
// g++ -x c++ -std=c++17 -Iinclude -I/usr/local/cuda/include \
// -I/usr/local/cuda/include/cccl cute_mma_ldmatrix_registers_trace.cu \
// -o /tmp/cute_mma_ldmatrix_registers_trace
#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm75.hpp>
#include <cute/atom/mma_traits_sm80.hpp>
#include <cstdio>
#include <cstdlib>
using namespace cute;
void require(bool condition, char const* message) {
if (!condition) {
std::fprintf(stderr, "FAILED: %s\n", message);
std::exit(1);
}
template<class T>
void show(char const* name, T const& value) {
std::printf("%s = ", name);
print(value);
print("\n");
template<class Layout>
void show_fragment_steps(Layout const& layout) {
constexpr int R = Layout::rank;
auto block = make_layout(get<0>(layout.shape()),
compact_major<LayoutLeft>(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape()))));
auto rest = make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride()));
show("block", block);
show("rest", rest);
show("fragment layout", tiled_product(block, rest));
void inspect_fragment_zero_strides() {
std::puts("=== Why the original Atom stride is inspected ===");
auto atom_shape = Shape<_2,_2,_2>{};
auto original_stride = Stride<_1,_128,_8>{};
auto original_compact_shape = filter_zeros(original_stride, atom_shape);
auto original_block = make_layout(atom_shape,
compact_major<LayoutLeft>(original_compact_shape));
show("nonzero strides: compact shape", original_compact_shape);
show("nonzero strides: block", original_block);
auto broadcast_stride = Stride<_1,_0,_8>{};
auto broadcast_compact_shape = filter_zeros(broadcast_stride, atom_shape);
auto broadcast_block = make_layout(atom_shape,
compact_major<LayoutLeft>(broadcast_compact_shape));
show("static zero: compact shape", broadcast_compact_shape);
show("static zero: block", broadcast_block);
auto unfiltered_block = make_layout(atom_shape, compact_major<LayoutLeft>(atom_shape));
auto dynamic_zero_stride = make_stride(_1{}, int(0), _8{});
auto dynamic_compact_shape = filter_zeros(dynamic_zero_stride, atom_shape);
auto dynamic_block = make_layout(atom_shape,
compact_major<LayoutLeft>(dynamic_compact_shape));
show("dynamic zero: compact shape", dynamic_compact_shape);
show("dynamic zero: block", dynamic_block);
require(broadcast_block(0,0,0) == broadcast_block(0,1,0),
"Static broadcast coordinates retain their alias");
require(int(size(broadcast_block)) == 8 && int(cosize(broadcast_block)) == 4,
"Eight logical positions use four half slots");
require(unfiltered_block(0,0,0) != unfiltered_block(0,1,0),
"Using only the original shape loses the broadcast alias");
require(dynamic_block(0,0,0) != dynamic_block(0,1,0),
"A runtime zero is not a static broadcast marker");
auto input = make_layout(make_shape(atom_shape, _1{}, _1{}),
make_stride(broadcast_stride, _0{}, _0{}));
auto fragment = make_fragment_like(input);
for (int v = 0; v < 8; ++v) {
require(fragment(v,0,0) == broadcast_block(v),
"The broadcast block agrees with make_fragment_like");
template<class Copy>
void inspect_copy_atom_cut(char const* label, Copy const&) {
std::printf("\n--- %s: split the TV input domain ---\n", label);
using TV = typename Copy::TiledLayout_TV;
using Ref = typename Copy::AtomLayoutRef;
using Dst = typename Copy::AtomLayoutDst;
constexpr int atom_threads = int(typename Copy::AtomNumThr{});
constexpr int atom_values = int(typename Copy::AtomNumVal{});
constexpr int tiled_threads = int(typename Copy::TiledNumThr{});
constexpr int tiled_values = int(typename Copy::TiledNumVal{});
auto cut = zipped_divide(TV{}, make_shape(
typename Copy::AtomNumThr{}, typename Copy::AtomNumVal{}));
auto dst_to_ref = right_inverse(Ref{}).compose(Dst{});
auto target = cut.compose(dst_to_ref, _);
std::printf("Tiled TV size = (%d,%d), Copy Atom size = (%d,%d)\n",
tiled_threads, tiled_values, atom_threads, atom_values);
show("Copy Atom reference layout (half units)", Ref{});
show("Copy Atom destination layout (half units)", Dst{});
show("MMA TV divided into Copy-sized groups", cut);
show("destination TV -> reference TV", dst_to_ref);
for (int group_t = 0; group_t < tiled_threads/atom_threads; ++group_t) {
for (int group_v = 0; group_v < tiled_values/atom_values; ++group_v) {
for (int t = 0; t < atom_threads; ++t) {
for (int v = 0; v < atom_values; ++v) {
auto rest = make_coord(group_t,group_v);
int tv = int(dst_to_ref(t,v));
int ref_t = tv % atom_threads;
int ref_v = tv / atom_threads;
require(tv == t + atom_threads*v, "ldmatrix reference equals destination");
require(cut(make_coord(t,v),rest) ==
TV{}(t + atom_threads*group_t, v + atom_values*group_v),
"Copy-sized cuts preserve the original MMA TV mapping");
require(target(make_coord(t,v),rest) ==
TV{}(ref_t + atom_threads*group_t, ref_v + atom_values*group_v),
"The local Copy conversion composes into the MMA TV domain");
if constexpr (atom_values == 8) {
std::printf("(t,v)=(0,5): Copy element id=%d, reference TV index=%d, MMA matrix index=%d\n",
int(Dst{}(0,5)), int(dst_to_ref(0,5)), int(TV{}(0,5)));
void inspect_one_warp() {
std::puts("=== One warp: A(16,16), B(8,16), no swizzle ===");
alignas(16) cutlass::half_t storage_a[16*16];
alignas(16) cutlass::half_t storage_b[8*16];
auto sA = make_tensor(make_smem_ptr(storage_a), Layout<Shape<_16,_16>,Stride<_16,_1>>{});
auto sB = make_tensor(make_smem_ptr(storage_b), Layout<Shape<_8,_16>,Stride<_16,_1>>{});
auto mma = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{});
auto load_a = Copy_Atom<SM75_U32x4_LDSM_N, cutlass::half_t>{};
auto load_b = Copy_Atom<SM75_U32x2_LDSM_N, cutlass::half_t>{};
auto copy_a = make_tiled_copy_A(load_a, mma);
auto copy_a_x1 = make_tiled_copy_A(
Copy_Atom<SM75_U32x1_LDSM_N, cutlass::half_t>{}, mma);
auto copy_a_x2 = make_tiled_copy_A(
Copy_Atom<SM75_U32x2_LDSM_N, cutlass::half_t>{}, mma);
auto copy_b = make_tiled_copy_B(load_b, mma);
auto coords_a = make_identity_tensor(shape(sA));
auto coords_b = make_identity_tensor(shape(sB));
using CopyA = decltype(copy_a);
using SrcA = typename CopyA::AtomLayoutSrc;
using DstA = typename CopyA::AtomLayoutDst;
using RefA = typename CopyA::AtomLayoutRef;
auto src_to_ref_a = right_inverse(RefA{}).compose(SrcA{});
auto src_to_dst_a = right_inverse(DstA{}).compose(SrcA{});
auto dst_to_src_a = right_inverse(SrcA{}).compose(DstA{});
for (int m = 0; m < 16; ++m) {
for (int k = 0; k < 16; ++k) {
storage_a[16*m+k] = cutlass::half_t(16*m+k);
inspect_copy_atom_cut("A + ldmatrix.x4", copy_a);
inspect_copy_atom_cut("A + ldmatrix.x2", copy_a_x2);
inspect_copy_atom_cut("A + ldmatrix.x1", copy_a_x1);
require(int(size<1>(mma.get_layoutB_TV())) == 4,
"A single MMA B TV has four values per lane");
require(int(size<1>(mma.get_layoutB_TV())) % int(decltype(load_a)::NumValDst) != 0,
"A single MMA B TV cannot be cut into x4 groups of eight values");
for (int lane = 0; lane < 32; ++lane) {
auto thr_mma = mma.get_slice(lane);
auto tCrA = thr_mma.partition_fragment_A(sA);
auto tCrB = thr_mma.partition_fragment_B(sB);
auto thr_a = copy_a.get_slice(lane);
auto thr_b = copy_b.get_slice(lane);
auto tXrA = thr_a.retile_D(tCrA);
auto tXrA_x1 = copy_a_x1.get_slice(lane).retile_D(tCrA);
auto tXrB = thr_b.retile_D(tCrB);
auto mma_a = thr_mma.partition_A(coords_a);
auto mma_b = thr_mma.partition_B(coords_b);
auto copy_dst_a = thr_a.partition_D(coords_a);
auto copy_dst_b = thr_b.partition_D(coords_b);
auto source_a = thr_a.partition_S(coords_a);
auto source_values_a = thr_a.partition_S(sA);
auto first_source_a = source_a(0,0,0);
require((2*int(sA.layout()(first_source_a))) % 16 == 0,
"Each x4 source row begins at a 16-byte-aligned address");
for (int p = 0; p < 8; ++p) {
auto source_coord = source_a(p,0,0);
require(get<0>(source_coord) == get<0>(first_source_a) &&
get<1>(source_coord) == get<1>(first_source_a) + p,
"Each x4 source address describes eight contiguous half values in K");
int element = int(SrcA{}(lane,p));
int ref_tv = int(src_to_ref_a(lane,p));
int dst_tv = int(src_to_dst_a(lane,p));
int dst_lane = dst_tv % 32;
int dst_value = dst_tv / 32;
auto dst_coord = mma.get_slice(dst_lane).partition_A(coords_a)(dst_value,0,0);
require(int(DstA{}(dst_lane,dst_value)) == element,
"Copy source and destination identify the same internal element");
require(int(CopyA::TiledLayout_TV{}(ref_tv%32,ref_tv/32)) ==
int(get<0>(source_coord)) + 16*int(get<1>(source_coord)),
"Source TV reaches the matrix coordinate through reference TV");
require(source_coord == dst_coord,
"Equal Copy element ids reach equal source and MMA matrix coordinates");
require(float(source_values_a(p,0,0)) ==
16*int(get<0>(dst_coord)) + int(get<1>(dst_coord)),
"The source address holds the value needed by the destination MMA slot");
require(tXrA.data() == tCrA.data(), "A view retains the fragment data iterator");
auto mma_coord = make_coord(1,0,1);
auto copy_coord = make_coord(5,0);
require(&tCrA(mma_coord,0,0) == &tXrA(copy_coord,0,0),
"Different nested coordinates select the same half slot");
tXrA(copy_coord,0,0) = cutlass::half_t(7);
require(float(tCrA(mma_coord,0,0)) == 7.0f,
"Writing the Copy view changes the value read through the MMA view");
int m = lane/4 + 8*((v/2)%2);
int k = 2*(lane%4) + v%2 + 8*(v/4);
require(mma_a(v,0,0) == make_coord(m,k), "MMA A coordinates");
require(copy_dst_a(v,0,0) == mma_a(v,0,0), "Copy and MMA A coordinates agree");
require(&tCrA(v,0,0) - &tCrA(0,0,0) == v, "A values are compact");
require(&tXrA(v,0,0) == &tCrA(v,0,0), "A retile aliases the original fragment");
require(&tXrA_x1(make_coord(v%2,v/2),0,0) == &tCrA(v,0,0),
"x1 groups the same A storage into four pairs");
require(4*(m%8) + (k%8)/2 == lane &&
k%2 + 2*(m/8) + 4*(k/8) == v,
"A matrix coordinates recover the hardware lane and half position");
int src_tv = int(dst_to_src_a(lane,v));
auto src_values = copy_a.get_slice(src_tv%32).partition_S(sA);
// Model the Copy_Traits route on the host; no GPU instruction is executed.
tXrA(v,0,0) = src_values(src_tv/32,0,0);
require(float(tCrA(v,0,0)) == 16*m+k,
"Following the Copy route fills the intended MMA fragment slot");
for (int v = 0; v < 4; ++v) {
int n = lane/4;
int k = 2*(lane%4) + v%2 + 8*(v/2);
require(mma_b(v,0,0) == make_coord(n,k), "MMA B coordinates");
require(copy_dst_b(v,0,0) == mma_b(v,0,0), "Copy and MMA B coordinates agree");
require(&tXrB(v,0,0) == &tCrB(v,0,0), "B retile aliases the original fragment");
require(4*n + (k%8)/2 == lane && k%2 + 2*(k/8) == v,
"B matrix coordinates recover the hardware lane and half position");
tCrB(v,0,0) = cutlass::half_t(16*n + k);
auto registers_a = recast<uint32_t>(tCrA(_,0,0));
auto registers_b = recast<uint32_t>(tCrB(_,0,0));
require(int(size(registers_a)) == 4 && int(size(registers_b)) == 2,
"mma_unpack exposes four A and two B 32-bit operands");
for (int j = 0; j < 4; ++j) {
uint32_t expected = uint32_t(tCrA(2*j,0,0).raw()) |
(uint32_t(tCrA(2*j+1,0,0).raw()) << 16);
require(registers_a(j) == expected, "A operand contains the correct low and high half fields");
for (int j = 0; j < 2; ++j) {
uint32_t expected = uint32_t(tCrB(2*j,0,0).raw()) |
(uint32_t(tCrB(2*j+1,0,0).raw()) << 16);
require(registers_b(j) == expected, "B operand contains the correct low and high half fields");
if (lane == 0) {
show("Copy Atom source layout (half units)", SrcA{});
int example_e = int(SrcA{}(16,1));
int example_tv = int(src_to_ref_a(16,1));
int example_q = int(CopyA::TiledLayout_TV{}(example_tv%32,example_tv/32));
require(example_e == 129 && example_tv == 160 && example_q == 144,
"The A(0,9) source-to-reference-to-MMA example");
require(copy_a.get_slice(16).partition_S(coords_a)(1,0,0) == make_coord(0,9),
"Source lane 16 half 1 reads A(0,9)");
std::printf("source (16,1) -> Copy e=%d -> reference tv=%d -> MMA q=%d -> destination (0,5)\n",
example_e, example_tv, example_q);
std::puts("A(0,9): source row base=8 halves, row offset=1, memory offset=9; destination a2 high half.");
show("MMA A TV", mma.get_layoutA_TV());
show("Copy reference TV", decltype(copy_a)::TiledLayout_TV{});
show("Copy destination TV", copy_a.get_layoutD_TV());
auto part_a = thr_mma.partition_A(sA);
show("partition_A layout", part_a.layout());
show_fragment_steps(part_a.layout());
show("tCrA layout", tCrA.layout());
show("tCrB layout", tCrB.layout());
show("tXrA layout", tXrA.layout());
show("tXrA x1 layout", tXrA_x1.layout());
show("tXrB layout", tXrB.layout());
std::puts("A nested aliases: MMA (1,0,1) = x4 Copy (5,0) = x1 Copy (1,2) = slot 5");
std::puts("Copy-view write/MMA-view readback verified before filling matrix labels.");
std::puts("Matrix labels below use A(m,k)=16*m+k and B(n,k)=16*n+k.");
std::printf("lane 0 a%d: low half=%.0f, high half=%.0f\n", j,
float(tCrA(2*j,0,0)), float(tCrA(2*j+1,0,0)));
std::printf("lane 0 b%d: low half=%.0f, high half=%.0f\n", j,
float(tCrB(2*j,0,0)), float(tCrB(2*j+1,0,0)));
print("lane 0 A source coordinates = ");
for (int v = 0; v < 8; ++v) { print(source_a(v,0,0)); print(" "); }
print("lane 0 A(v=5) = "); print(mma_a(5,0,0)); print("\n");
void inspect_official_cta() {
std::puts("\n=== Official CTA: 128x128x64, 2x2 warps ===");
auto atom = composition(Swizzle<3,3,3>{},
Layout<Shape<_8,Shape<_8,_8>>,Stride<_8,Stride<_1,_64>>>{});
auto shared_layout = tile_to_shape(atom, Shape<_128,_64,_3>{});
alignas(16) cutlass::half_t shared_stub[128*64*3];
auto s = make_tensor(make_smem_ptr(shared_stub), shared_layout);
auto mma = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{},
Layout<Shape<_2,_2>>{}, Tile<_32,_32,_16>{});
auto thr_mma = mma.get_slice(0);
auto tCrA = thr_mma.partition_fragment_A(s(_,_,0));
auto tCrB = thr_mma.partition_fragment_B(s(_,_,0));
auto copy_b = make_tiled_copy_B(Copy_Atom<SM75_U32x4_LDSM_N,cutlass::half_t>{}, mma);
auto tXrB = copy_b.get_slice(0).retile_D(tCrB);
auto part_a = thr_mma.partition_A(s(_,_,0));
for (int k = 0; k < 4; ++k) {
require(&tXrB(v,j,k) == &tCrB(v%4,2*j+v/4,k), "B Copy group joins two MMA fragments");
require(int(size(tCrA)) == 128 && int(size(tCrB)) == 128, "Official A/B fragment sizes");
int main() {
inspect_fragment_zero_strides();
inspect_one_warp();
inspect_official_cta();
std::puts("PASS: layouts, coordinates, and aliases verified; no GPU instruction executed.");