// Companion to cute_partition_s_copy_unpack_walkthrough_zh.md.
// Host: g++ -x c++ -std=c++17 -Iinclude -I/usr/local/cuda/include \
// -I/usr/local/cuda/include/cccl cute_partition_s_copy_unpack_trace.cu \
// -o /tmp/cute_partition_s_copy_unpack_trace
// PTX only: nvcc -std=c++17 -Iinclude -arch=sm_80 --ptx \
// cute_partition_s_copy_unpack_trace.cu -o /tmp/cute_partition_s_copy_unpack_trace.ptx
#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm75.hpp>
#include <cute/atom/mma_traits_sm80.hpp>
#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cstring>
using namespace cute;
using MmaOp = SM80_16x8x16_F16F16F16F16_TN;
#if !defined(__CUDACC__)
void require(bool condition, char const* message) {
if (!condition) {
std::fprintf(stderr, "FAILED: %s\n", message);
std::exit(1);
}
// A host-only operation with the same argument types as ldmatrix.x2.
// The actual copy_unpack below expands its tensors into these references.
// It checks bindings without executing or emulating a GPU instruction.
struct CaptureLdsmX2 {
using SRegisters = uint128_t[1];
using DRegisters = uint32_t[2];
inline static void const* expected_source = nullptr;
inline static void const* expected_outputs[2] = {};
inline static int calls = 0;
static void copy(uint128_t const& src, uint32_t& d0, uint32_t& d1) {
require(static_cast<void const*>(&src) == expected_source,
"copy_unpack must pass the source row-start reference");
require(static_cast<void const*>(&d0) == expected_outputs[0],
"copy_unpack dst0 must bind the first half pair of this group");
require(static_cast<void const*>(&d1) == expected_outputs[1],
"copy_unpack dst1 must bind the second half pair of this group");
++calls;
};
namespace cute {
template <>
struct Copy_Traits<::CaptureLdsmX2> : Copy_Traits<SM75_U32x2_LDSM_N> {};
int main() {
auto mma = make_tiled_mma(MmaOp{}, Layout<Shape<_2,_2,_1>>{});
auto tiled_copy = make_tiled_copy_A(
Copy_Atom<SM75_U32x2_LDSM_N,half_t>{}, mma);
using Copy = decltype(tiled_copy);
using S = typename Copy::AtomLayoutSrc;
using D = typename Copy::AtomLayoutDst;
alignas(16) half_t storage[512];
auto sA = make_tensor(make_smem_ptr(storage),
Layout<Shape<_32,_16>,Stride<_16,_1>>{});
for (int m=0; m<32; ++m)
for (int k=0; k<16; ++k) sA(m,k) = half_t(16*m+k);
sA(8,0) = half_t(23);
sA(0,9) = half_t(37);
auto coordinates = make_identity_tensor(shape(sA));
int rows = 0;
int slots = 0;
int packed_outputs = 0;
for (int t=0; t<128; ++t) {
int lane = t%32;
int warp_base = t-lane;
int wm = (t/32)%2;
int source_m = 16*wm + lane%16;
auto thr_copy = tiled_copy.get_slice(t);
auto tXsA = thr_copy.partition_S(sA);
auto source_coordinates = thr_copy.partition_S(coordinates);
auto fragment = mma.get_slice(t).partition_fragment_A(sA);
auto tXrA = thr_copy.retile_D(fragment);
auto mma_coordinates = mma.get_slice(t).partition_A(coordinates);
auto destination_coordinates = thr_copy.partition_D(coordinates);
require(tXrA.data() == fragment.data(), "retile_D must retain fragment storage");
for (int g=0; g<2; ++g) {
auto src_atom = tXsA(make_coord(_,g),0,0);
auto dst_atom = tXrA(make_coord(_,g),0,0);
require(int(size(src_atom)) == 8 && int(size(dst_atom)) == 4,
"One x2 call must expose eight source and four destination halves");
for (int p=0; p<8; ++p) {
require(source_coordinates(make_coord(p,g),0,0) == make_coord(source_m,8*g+p),
"partition_S must select the expected source row coordinates");
require(&src_atom(p) == &sA(source_m,8*g+p),
"Each source atom must expose eight physically contiguous halves");
require(reinterpret_cast<uintptr_t>(&src_atom(0))%16 == 0,
"Each row-start address must be 16-byte aligned");
++rows;
CaptureLdsmX2::expected_source = &sA(source_m,8*g);
CaptureLdsmX2::expected_outputs[0] = &fragment(4*g,0,0);
CaptureLdsmX2::expected_outputs[1] = &fragment(4*g+2,0,0);
copy_unpack(Copy_Traits<CaptureLdsmX2>{}, src_atom, dst_atom);
// Independent non-transposed ldmatrix route, expressed for a destination lane:
// row=lane/4, col=2*(lane%4)+a%2, matrix=a/2, source lane=8*matrix+row.
for (int a=0; a<4; ++a) {
int row = lane/4;
int col = 2*(lane%4) + a%2;
int source_lane = 8*(a/2) + row;
int source_thread = warp_base + source_lane;
int v = a+4*g;
auto routed_source = tiled_copy.get_slice(source_thread).partition_S(sA);
auto routed_coordinates = tiled_copy.get_slice(source_thread).partition_S(coordinates);
require(S{}(source_lane,col) == D{}(lane,a),
"Copy traits must match the independently specified hardware route");
require(routed_coordinates(make_coord(col,g),0,0) == mma_coordinates(v,0,0),
"The hardware source must match the MMA destination coordinate");
require(destination_coordinates(make_coord(a,g),0,0) == mma_coordinates(v,0,0),
"Copy partition_D must match the MMA coordinate");
require(&dst_atom(a) == &fragment(v,0,0),
"retile_D must bind the corresponding original fragment slot");
// Host model of the route; no GPU instruction is executed.
dst_atom(a) = routed_source(make_coord(col,g),0,0);
require(fragment(v,0,0).raw() == sA(mma_coordinates(v,0,0)).raw(),
"The expected half bits must reach the original fragment slot");
++slots;
auto packed = recast<uint32_t>(dst_atom);
for (int j=0; j<2; ++j) {
uint32_t actual = 0;
std::memcpy(&actual, &packed(j), sizeof(actual));
uint32_t expected = uint32_t(dst_atom(2*j).raw()) |
(uint32_t(dst_atom(2*j+1).raw()) << 16);
require(actual == expected, "Recast must preserve the packed low/high half bits");
++packed_outputs;
if (t==0) {
std::printf("tXsA = "); print(tXsA.layout()); std::printf("\n");
std::printf("tXrA = "); print(tXrA.layout()); std::printf("\n");
std::printf("lane0 slot2 = %.0f; slot5 = %.0f\n",
float(fragment(2,0,0)),float(fragment(5,0,0)));
if (t==0 || t==1 || t==8 || t==16 || t==32 || t==64)
std::printf("thread%d row starts: A(%d,0), A(%d,8)\n",t,source_m,source_m);
std::printf("PASS: %d source rows, %d copy_unpack calls, %d destination half slots, "
"%d packed outputs; host only.\n",
rows,CaptureLdsmX2::calls,slots,packed_outputs);
#endif
#if defined(__CUDACC__)
// Compile-only kernel. If launched, use exactly one block of 128 threads.
// The caller supplies 512 input halves and storage for 1024 output halves.
extern "C" __global__ void source_to_fragment_x2(
half_t const* input, half_t* output) {
__shared__ __align__(16) half_t shared_a[512];
int t = int(threadIdx.x);
for (int i=t; i<512; i+=128) shared_a[i] = input[i];
__syncthreads();
auto sA = make_tensor(make_smem_ptr(shared_a),
Copy_Atom<SM75_U32x2_LDSM_N,half_t>{},mma);
copy(tiled_copy,tXsA,tXrA);
for (int v=0; v<8; ++v) output[t*8+v] = fragment(v,0,0);