// Host-only companion to cute_ldmatrix_mma_coordinate_walkthrough_zh.md.
// No ldmatrix or mma.sync instruction is executed.
// From the CUTLASS root:
// g++ -x c++ -std=c++17 -Iinclude -I/usr/local/cuda/include \
// -I/usr/local/cuda/include/cccl cute_ldmatrix_mma_coordinate_trace.cu \
// -o /tmp/cute_ldmatrix_mma_coordinate_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");
int main() {
alignas(16) half_t storage[16*16];
auto sA = make_tensor(make_smem_ptr(storage),
Layout<Shape<_16,_16>, Stride<_16,_1>>{});
for (int m = 0; m < 16; ++m) {
for (int k = 0; k < 16; ++k) {
sA(m,k) = half_t(16*m+k);
// Keep the transported value distinct from element ids and address offsets.
sA(0,9) = half_t(37);
auto mma = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{});
auto load = Copy_Atom<SM75_U32x4_LDSM_N, half_t>{};
auto tiled_copy = make_tiled_copy_A(load, mma);
using Copy = decltype(tiled_copy);
using Src = typename Copy::AtomLayoutSrc;
using Dst = typename Copy::AtomLayoutDst;
using Ref = typename Copy::AtomLayoutRef;
auto F = mma.get_layoutA_TV();
auto src_to_ref = right_inverse(Ref{}).compose(Src{});
auto dst_to_src = right_inverse(Src{}).compose(Dst{});
auto coordinates = make_identity_tensor(shape(sA));
show("MMA ALayout", typename decltype(mma)::AtomLayoutA_TV{});
show("TiledLayout_TV", F);
show("AtomLayoutSrc (half units)", Src{});
show("AtomLayoutDst (half units)", Dst{});
show("AtomLayoutRef (half units)", Ref{});
show("source TV -> reference TV", src_to_ref);
int target_q = int(F(0,5));
int copy_destination_q = int(tiled_copy.get_layoutD_TV()(0,5));
int route_e = int(Dst{}(0,5));
int source_tv = int(dst_to_src(0,5));
require(target_q == 144 && copy_destination_q == target_q && route_e == 129 &&
source_tv%32 == 16 && source_tv/32 == 1,
"Destination slot 5 requires A(0,9) and traces back to source lane 16 half 1");
std::printf("\nMMA target (0,5): q=%d; Copy final destination q=%d\n",
target_q, copy_destination_q);
std::printf("Copy route: destination (0,5) -> e=%d -> source (%d,%d)\n",
route_e, source_tv%32, source_tv/32);
std::puts("\nsource TV | e | reference TV | tv | q | A(m,k) | half offset | example byte address");
auto src16_coordinates = tiled_copy.get_slice(16).partition_S(coordinates);
for (int p = 0; p < 2; ++p) {
int e = int(Src{}(16,p));
int tv = int(src_to_ref(16,p));
int ref_lane = tv % 32;
int ref_value = tv / 32;
int q = int(F(ref_lane,ref_value));
auto coord = src16_coordinates(p,0,0);
int m = int(get<0>(coord));
int k = int(get<1>(coord));
int offset = int(sA.layout()(coord));
require(e == 128+p && ref_lane == 0 && ref_value == 4+p,
"The two entries of source lane 16 recover reference values 4 and 5");
require(q == 128+16*p && m == 0 && k == 8+p && offset == 8+p,
"The example distinguishes matrix encoding from memory offset");
std::printf("(16,%d) | %d | (%d,%d) | %d | %d | A(%d,%d) | %d | 0x%x\n",
p, e, ref_lane, ref_value, tv, q, m, k, offset, 0x1000+2*offset);
// Check the non-transposed m8n8.x4 hardware distribution independently
// of the inverse-layout calculation. A source lane supplies one row.
int visited[32][8] = {};
int route_count = 0;
std::puts("\nSource row addresses inferred by partition_S:");
for (int src_lane = 0; src_lane < 32; ++src_lane) {
auto src_coords = tiled_copy.get_slice(src_lane).partition_S(coordinates);
auto src_values = tiled_copy.get_slice(src_lane).partition_S(sA);
auto first = src_coords(0,0,0);
int first_offset = int(sA.layout()(first));
require((2*first_offset)%16 == 0, "Every source row is 16-byte aligned");
int block = src_lane / 8;
int row = src_lane % 8;
int expected_m = row + 8*(block%2);
int expected_k_base = 8*(block/2);
require(first == make_coord(expected_m,expected_k_base),
"Each group of eight address lanes describes the expected A quadrant");
if (row == 0) {
std::printf("lanes %2d..%2d: A rows %d..%d, K %d..%d\n",
8*block, 8*block+7, expected_m, expected_m+7,
expected_k_base, expected_k_base+7);
for (int col = 0; col < 8; ++col) {
int dst_lane = 4*row + col/2;
int dst_value = 2*block + col%2;
auto source_coord = src_coords(col,0,0);
auto mma_coord = mma.get_slice(dst_lane).partition_A(coordinates)(dst_value,0,0);
require(source_coord == make_coord(expected_m,expected_k_base+col),
"The source row is contiguous in K");
require(int(sA.layout()(source_coord)) == first_offset+col,
"The actual source address agrees with the row offset");
require(int(Src{}(src_lane,col)) == int(Dst{}(dst_lane,dst_value)),
"The Copy traits describe the hardware's source-to-destination route");
int tv = int(src_to_ref(src_lane,col));
int q = int(F(tv%32,tv/32));
require(q == int(get<0>(source_coord)) + 16*int(get<1>(source_coord)),
"Source TV reaches its matrix coordinate through reference TV");
require(source_coord == mma_coord,
"The hardware route reaches the slot requiring this matrix element");
require(float(src_values(col,0,0)) == float(sA(mma_coord)),
"The address holds the value required by the destination MMA slot");
++visited[dst_lane][dst_value];
++route_count;
int slot_count = 0;
int operand_count = 0;
for (int lane = 0; lane < 32; ++lane) {
auto thr_mma = mma.get_slice(lane);
auto tCrA = thr_mma.partition_fragment_A(sA);
auto thr_copy = tiled_copy.get_slice(lane);
auto tXrA = thr_copy.retile_D(tCrA);
auto expected_coords = thr_mma.partition_A(coordinates);
require(tXrA.data() == tCrA.data(), "retile_D preserves the data iterator");
for (int v = 0; v < 8; ++v) {
require(visited[lane][v] == 1, "Every destination slot is reached exactly once");
require(&tXrA(v,0,0) == &tCrA(v,0,0),
"Copy value v aliases MMA value v for this x4 configuration");
int src_tv = int(dst_to_src(lane,v));
int source_lane = src_tv%32;
int source_value = src_tv/32;
auto source = tiled_copy.get_slice(source_lane).partition_S(sA);
// Model the fixed hardware route on the host, then read through the MMA view.
tXrA(v,0,0) = source(source_value,0,0);
auto coord = expected_coords(v,0,0);
require(float(tCrA(v,0,0)) == float(sA(coord)),
"The Copy route fills the correct MMA value slot");
++slot_count;
auto mma_operands = recast<uint32_t>(tCrA(_,0,0));
auto copy_operands = recast<uint32_t>(tXrA(_,0,0));
require(int(size(mma_operands)) == 4 && int(size(copy_operands)) == 4,
"Copy and MMA both expose four 32-bit A 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(©_operands(j) == &mma_operands(j),
"Copy output j and MMA input j refer to the same packed storage");
require(mma_operands(j) == expected, "The low and high half fields retain their bits");
++operand_count;
if (lane == 0) {
show("\nlane 0 tCrA layout", tCrA.layout());
show("lane 0 tXrA layout", tXrA.layout());
std::puts("lane 0 operands: initial labels A(m,k)=16*m+k, with A(0,9) changed to 37:");
std::printf("dst%d / a%d: low=%.0f, high=%.0f\n", j, j,
float(tCrA(2*j,0,0)), float(tCrA(2*j+1,0,0)));
require(float(tCrA(5,0,0)) == 37.0f, "The value 37 at A(0,9) reaches lane 0 slot 5");
std::printf("\nPASS: %d hardware routes, %d fragment slots, %d packed operands verified (host only).\n",
route_count, slot_count, operand_count);