// Companion to cute_mma_ldmatrix_ab_retile_walkthrough_zh.md.
// Host build validates layouts and models the hardware route; it executes no GPU instruction.
// g++ -x c++ -std=c++17 -Iinclude -I/usr/local/cuda/include \
// -I/usr/local/cuda/include/cccl cute_mma_ldmatrix_ab_retile_trace.cu \
// -o /tmp/cute_mma_ldmatrix_ab_retile_trace
// Compile the three one-warp GPU kernels to PTX without launching them:
// nvcc -std=c++17 -Iinclude -arch=sm_80 --ptx \
// cute_mma_ldmatrix_ab_retile_trace.cu -o /tmp/cute_mma_ldmatrix_ab_retile_trace.ptx
#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm75.hpp>
#include <cute/atom/mma_traits_sm80.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cstdio>
#include <cstdlib>
using namespace cute;
using MmaOp = SM80_16x8x16_F16F16F16F16_TN;
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<bool IsA, class ThrMma, class Tensor>
auto partition_operand(ThrMma const& thr, Tensor const& tensor) {
if constexpr (IsA) return thr.partition_A(tensor);
else return thr.partition_B(tensor);
auto make_operand_fragment(ThrMma const& thr, Tensor const& tensor) {
if constexpr (IsA) return thr.partition_fragment_A(tensor);
else return thr.partition_fragment_B(tensor);
template<class Copy, class Fragment>
void show_retile_steps(Fragment const& fragment) {
constexpr int Rank = Fragment::rank;
using NumThreads = typename Copy::TiledNumThr;
auto V = size<0>(fragment);
auto inverse = right_inverse(typename Copy::TiledLayout_TV{})
.with_shape(shape(typename Copy::Tiler_MN{}));
auto fmn = upcast<NumThreads{} * V>(inverse);
auto product = logical_product(make_layout(V), right_inverse(fmn));
auto fv = zipped_divide(product, make_layout(typename Copy::AtomNumVal{}));
auto tiler = prepend(product_each(shape(fmn)), V);
auto tt = zipped_divide(fragment.layout(), tiler);
auto vt = tt.compose(fv, _);
show("frg_layout_mn", fmn);
show("gather extents", product_each(shape(fmn)));
show("logical_product", product);
show("frg_layout_v", fv);
show("retile tiler", tiler);
show("t_tensor", tt);
show("v_tensor", vt);
show("retile result", vt(_, append<Rank>(Int<0>{}, _)));
template<bool IsA, class Load, class Tensor, class Mma>
void verify(char const* label, Tensor const& source, Mma const& mma,
bool print_steps = false) {
auto tiled_copy = [&] {
if constexpr (IsA) return make_tiled_copy_A(Copy_Atom<Load, half_t>{}, mma);
else return make_tiled_copy_B(Copy_Atom<Load, half_t>{}, mma);
}();
using Copy = decltype(tiled_copy);
using Src = typename Copy::AtomLayoutSrc;
using Dst = typename Copy::AtomLayoutDst;
constexpr int values = IsA ? 8 : 4;
constexpr int atom_values = int(typename Copy::AtomNumVal{});
constexpr int matrices = atom_values / 2;
constexpr int groups = values / atom_values;
constexpr int address_lanes = 8 * matrices;
auto coords = make_identity_tensor(shape(source));
int coverage[32][values] = {};
int source_checks = 0;
int slot_checks = 0;
int packed_checks = 0;
std::printf("\n[%s] atom half values=%d; instruction groups=%d\n",
label, atom_values, groups);
show("AtomLayoutSrc", Src{});
show("AtomLayoutDst", Dst{});
// Independent non-transposed ldmatrix route:
// source row address lane = 8*local_matrix + row
// destination lane = 4*row + col/2
// destination atom half = 2*local_matrix + col%2
// Upper lanes have repeated source views for x1/x2; only the first
// address_lanes lanes provide the distinct rows consumed by the instruction.
for (int g = 0; g < groups; ++g) {
for (int ts = 0; ts < 32; ++ts) {
auto source_coords = tiled_copy.get_slice(ts).partition_S(coords);
int local_matrix = (ts % address_lanes) / 8;
int row = ts % 8;
int global_matrix = local_matrix + matrices * g;
auto first = source_coords(make_coord(0,g),0,0);
int first_offset = int(source.layout()(first));
require((2 * first_offset) % 16 == 0, "Every source row must be 16-byte aligned");
for (int col = 0; col < 8; ++col) {
int td = 4*row + col/2;
int av = 2*local_matrix + col%2;
int v = av + atom_values*g;
auto src_coord = source_coords(make_coord(col,g),0,0);
int mn = IsA ? row + 8*(global_matrix%2) : row;
int k = col + 8*(IsA ? global_matrix/2 : global_matrix);
require(src_coord == make_coord(mn,k), "The inferred source quadrant must match");
require(int(source.layout()(src_coord)) == first_offset+col,
"A source row must contain eight contiguous half values");
auto expected = partition_operand<IsA>(mma.get_slice(td),coords)(v,0,0);
require(src_coord == expected, "Hardware route must reach the required MMA coordinate");
require(int(Src{}(ts,col)) == int(Dst{}(td,av)),
"Copy traits must agree with the independent hardware route");
if (ts < address_lanes) ++coverage[td][v];
++source_checks;
for (int ts = 0; ts < address_lanes; ts += 8) {
auto sc = tiled_copy.get_slice(ts).partition_S(coords);
std::printf("group %d address lane %2d starts at ",g,ts);
print(sc(make_coord(0,g),0,0)); print("\n");
for (int lane = 0; lane < 32; ++lane) {
auto thr_mma = mma.get_slice(lane);
auto fragment = make_operand_fragment<IsA>(thr_mma,source);
auto thr_copy = tiled_copy.get_slice(lane);
auto retiled = thr_copy.retile_D(fragment);
auto expected = partition_operand<IsA>(thr_mma,coords);
auto destination_coords = thr_copy.partition_D(coords);
require(retiled.data() == fragment.data(), "retile_D must preserve the storage iterator");
for (int av = 0; av < atom_values; ++av) {
int row = lane/4;
int col = 2*(lane%4) + av%2;
int ts = 8*(av/2) + row;
auto src = tiled_copy.get_slice(ts).partition_S(source);
require(coverage[lane][v] == 1, "Each destination half must be reached exactly once");
require(destination_coords(make_coord(av,g),0,0) == expected(v,0,0),
"Composed Copy destination must equal the MMA coordinate");
require(&retiled(make_coord(av,g),0,0) == &fragment(v,0,0),
"Copy groups must bind the expected fragment half slots");
// Host model of the fixed hardware route; no GPU instruction is executed.
retiled(make_coord(av,g),0,0) = src(make_coord(col,g),0,0);
require(fragment(v,0,0).raw() == source(expected(v,0,0)).raw(),
"The actual source half bits must reach the required fragment slot");
++slot_checks;
auto copy_regs = recast<uint32_t>(retiled(make_coord(_,g),0,0));
auto mma_regs = recast<uint32_t>(fragment(_,0,0));
require(int(size(copy_regs)) == matrices, "Each instruction must expose its declared outputs");
for (int j = 0; j < matrices; ++j) {
int output = j + matrices*g;
require(©_regs(j) == &mma_regs(output),
"A Copy output and the corresponding MMA input must share packed storage");
uint32_t bits = uint32_t(fragment(2*output,0,0).raw()) |
(uint32_t(fragment(2*output+1,0,0).raw()) << 16);
require(copy_regs(j) == bits, "Packed low/high half bits must be preserved");
++packed_checks;
if (lane == 0) {
show("partition source-backed layout", partition_operand<IsA>(thr_mma,source).layout());
show("MMA fragment", fragment.layout());
show("retile_D view", retiled.layout());
if (print_steps) show_retile_steps<Copy>(fragment);
for (int v = 0; v < values; ++v) {
std::printf("lane0 slot%d = %.0f from ", v, float(fragment(v,0,0)));
print(expected(v,0,0)); print("\n");
int target_v = IsA ? 5 : 3;
int g = target_v/atom_values;
int av = target_v%atom_values;
int ts = 8*(av/2);
int e = int(Dst{}(0,av));
auto f = typename Copy::TiledLayout_TV{};
int q = int(f(0,target_v));
std::printf("target: group%d source(%d,1) -> e%d -> q%d -> lane0 slot%d = %.0f\n",
g,ts,e,q,target_v,float(fragment(target_v,0,0)));
std::printf("PASS %s: %d source-view entries, %d fragment slots, %d packed bindings.\n",
label,source_checks,slot_checks,packed_checks);
int main() {
alignas(16) half_t storage_a[256];
alignas(16) half_t storage_b[128];
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>>{});
for (int m=0; m<16; ++m)
for (int k=0; k<16; ++k) sA(m,k)=half_t(16*m+k);
for (int n=0; n<8; ++n)
for (int k=0; k<16; ++k) sB(n,k)=half_t(16*n+k);
sA(0,9)=half_t(37);
sB(0,9)=half_t(53);
auto mma = make_tiled_mma(MmaOp{});
show("MMA ALayout", typename decltype(mma)::AtomLayoutA_TV{});
show("MMA BLayout", typename decltype(mma)::AtomLayoutB_TV{});
verify<true,SM75_U32x4_LDSM_N>("A/x4",sA,mma);
verify<true,SM75_U32x2_LDSM_N>("A/x2",sA,mma,true);
verify<false,SM75_U32x2_LDSM_N>("B/x2",sB,mma);
verify<false,SM75_U32x1_LDSM_N>("B/x1",sB,mma,true);
std::puts("PASS: 4 configurations, 768 fragment slots and 384 packed bindings; host only.");
#if defined(__CUDACC__)
// Compile-only kernels. If launched, each requires exactly one warp per block.
// Inputs use the same physical layouts as the article; output is row-major 16x8.
template<class LoadA, class LoadB>
__device__ __forceinline__ void ab_body(
half_t const* a, half_t const* b, half_t* output,
half_t* shared_a, half_t* shared_b) {
int lane = int(threadIdx.x);
for (int i=lane; i<256; i+=32) shared_a[i]=a[i];
for (int i=lane; i<128; i+=32) shared_b[i]=b[i];
__syncthreads();
auto sA = make_tensor(make_smem_ptr(shared_a),
auto sB = make_tensor(make_smem_ptr(shared_b),
auto gC = make_tensor(make_gmem_ptr(output),
Layout<Shape<_16,_8>,Stride<_8,_1>>{});
auto copy_a = make_tiled_copy_A(Copy_Atom<LoadA,half_t>{},mma);
auto copy_b = make_tiled_copy_B(Copy_Atom<LoadB,half_t>{},mma);
auto thr_a = copy_a.get_slice(lane);
auto thr_b = copy_b.get_slice(lane);
auto tCrA = thr_mma.partition_fragment_A(sA);
auto tCrB = thr_mma.partition_fragment_B(sB);
auto tCrC = thr_mma.partition_fragment_C(gC);
auto tXsA = thr_a.partition_S(sA);
auto tXsB = thr_b.partition_S(sB);
auto tXrA = thr_a.retile_D(tCrA);
auto tXrB = thr_b.retile_D(tCrB);
copy(copy_a,tXsA,tXrA);
copy(copy_b,tXsB,tXrB);
clear(tCrC);
gemm(mma,tCrA,tCrB,tCrC);
copy(tCrC,thr_mma.partition_C(gC));
extern "C" __global__ void ab_x4_x2(half_t const* a, half_t const* b, half_t* output) {
__shared__ __align__(16) half_t shared_a[256];
__shared__ __align__(16) half_t shared_b[128];
ab_body<SM75_U32x4_LDSM_N,SM75_U32x2_LDSM_N>(a,b,output,shared_a,shared_b);
extern "C" __global__ void ab_x2_x2(half_t const* a, half_t const* b, half_t* output) {
ab_body<SM75_U32x2_LDSM_N,SM75_U32x2_LDSM_N>(a,b,output,shared_a,shared_b);
extern "C" __global__ void ab_x2_x1(half_t const* a, half_t const* b, half_t* output) {
ab_body<SM75_U32x2_LDSM_N,SM75_U32x1_LDSM_N>(a,b,output,shared_a,shared_b);
#endif