code/day76-tma/tma_tile.cuThis is the source used by the lesson and its recorded evidence. Compile commands and expected output live in the directory README.
// SPDX-License-Identifier: MIT
//
// Day 76: load a tile with the Tensor Memory Accelerator and verify it.
//
// A 256x256 float matrix, cut into 64x64 tiles. Each block asks the TMA
// engine for its tile with a single cp.async.bulk.tensor issued by one
// thread, waits on an mbarrier whose completion is counted in bytes, adds
// 1.0f to every element in shared memory, and stores the result with plain
// per-thread writes. The host checks the round trip exactly: the copy and
// the add are both bit-deterministic, so any tolerance would hide a bug.
//
// The tensor map (CUtensorMap) is built on the host by
// cuTensorMapEncodeTiled, reached through the runtime's driver entry point
// query, so nothing here links against libcuda at build time. Every API
// used exists in CUDA 12.6: cudaGetDriverEntryPointByVersion (runtime),
// cuda::ptx::cp_async_bulk_tensor, cuda::ptx::fence_proxy_async and
// cuda::device::barrier_arrive_tx (all in the libcu++ 2.5.0 that 12.6
// ships).
//
// Needs compute capability 9.0 or newer: H100, B200, or any RTX 50. The
// program prints a clear message and exits nonzero on anything older.
//
// Build: nvcc -std=c++17 -O3 -arch=sm_90 -lineinfo -o tma_tile tma_tile.cu
// Run: ./tma_tile
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cudaTypedefs.h>
#include <cuda/barrier>
#include <cuda/ptx>
#include <cuda/std/utility>
// The one error macro. This file is standalone, so it carries its own
// verbatim copy. `err_` carries a trailing underscore so it cannot collide
// with a variable at the call site.
#define CUDA_CHECK(call) \
do { \
cudaError_t err_ = (call); \
if (err_ != cudaSuccess) { \
std::fprintf(stderr, "CUDA error %s:%d: %s: %s\n", __FILE__, \
__LINE__, #call, cudaGetErrorString(err_)); \
std::exit(EXIT_FAILURE); \
} \
} while (0)
// Matrix and tile shapes. The TMA encode call has hard rules the
// static_asserts below pin at compile time: each box dimension is at most
// 256, the innermost box row must span a multiple of 16 bytes, and the
// global row stride must be a multiple of 16 bytes. The grid covers the
// matrix exactly, so no tile hangs off the edge and the out-of-bounds
// fill path never runs.
constexpr size_t kGmemWidth = 256; // floats per row
constexpr size_t kGmemHeight = 256; // rows
constexpr int kTileWidth = 64; // floats per tile row
constexpr int kTileHeight = 64; // tile rows
constexpr int kThreadsPerBlock = 256;
static_assert(kTileWidth <= 256 && kTileHeight <= 256,
"cuTensorMapEncodeTiled caps each box dimension at 256");
static_assert(kTileWidth * sizeof(float) % 16 == 0,
"the innermost box dimension must span a multiple of 16 bytes");
static_assert(kGmemWidth * sizeof(float) % 16 == 0,
"the global row stride must be a multiple of 16 bytes");
static_assert(kGmemWidth % kTileWidth == 0 && kGmemHeight % kTileHeight == 0,
"the grid is sized to cover the matrix exactly");
static_assert(kThreadsPerBlock % 32 == 0,
"block size must be a whole number of warps");
using BlockBarrier = cuda::barrier<cuda::thread_scope_block>;
// Loads one 64x64 tile through the TMA engine, adds 1.0f to every element
// in shared memory, and stores the result with plain per-thread writes.
//
// One thread (thread 0) issues the whole 16384-byte tile copy; the other
// 255 threads issue no global load at all. In the store loop each thread
// owns 16 elements and consecutive threads write consecutive columns, so
// one warp's 32 addresses cover 128 contiguous bytes.
//
// Launch assumptions: blockDim.x == kThreadsPerBlock, and gridDim is
// (kGmemWidth / kTileWidth, kGmemHeight / kTileHeight) so every tile
// coordinate is in range.
__global__ void loadTileTma(const __grid_constant__ CUtensorMap tileMap,
float* __restrict__ out) {
// TMA writes into shared memory and requires the destination 128-byte
// aligned.
__shared__ alignas(128) float tile[kTileHeight][kTileWidth];
// A barrier in shared memory is dynamically initialized, which nvcc warns
// about for a __shared__ variable; the pragma accepts it, same as the
// programming guide's own TMA example.
#pragma nv_diag_suppress static_var_with_dynamic_init
__shared__ BlockBarrier bar;
if (threadIdx.x == 0) {
init(&bar, blockDim.x);
// The barrier was written through the generic proxy, and the TMA
// engine reads it through the async proxy. The fence publishes the
// initialized barrier to the engine; without it the copy can race
// with the init.
cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);
}
__syncthreads();
// snippet: issue-tma
// One thread asks for the whole tile. The coordinates are element
// offsets into the global tensor, x (the contiguous dimension) first.
// Everyone else just arrives; the engine, not any thread, completes
// the barrier by counting 16384 copied bytes against the expectation
// that barrier_arrive_tx registered.
BlockBarrier::arrival_token token;
if (threadIdx.x == 0) {
const int32_t coords[2] = {
static_cast<int32_t>(blockIdx.x) * kTileWidth,
static_cast<int32_t>(blockIdx.y) * kTileHeight};
cuda::ptx::cp_async_bulk_tensor(
cuda::ptx::space_cluster, cuda::ptx::space_global, &tile, &tileMap,
coords, cuda::device::barrier_native_handle(bar));
token = cuda::device::barrier_arrive_tx(bar, 1, sizeof(tile));
} else {
token = bar.arrive();
}
bar.wait(cuda::std::move(token));
// end snippet
// The wait makes the engine's shared-memory writes visible to every
// thread in the block; no further fence is needed on the load path.
// snippet: plain-store
const size_t colBase = blockIdx.x * static_cast<size_t>(kTileWidth);
const size_t rowBase = blockIdx.y * static_cast<size_t>(kTileHeight);
for (int k = static_cast<int>(threadIdx.x); k < kTileWidth * kTileHeight;
k += kThreadsPerBlock) {
const int r = k / kTileWidth;
const int c = k % kTileWidth;
out[(rowBase + r) * kGmemWidth + colBase + c] = tile[r][c] + 1.0f;
}
// end snippet
}
int main() {
const int device = 0;
CUDA_CHECK(cudaSetDevice(device));
cudaDeviceProp prop;
CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
std::printf("GPU: %s (compute capability %d.%d)\n", prop.name, prop.major,
prop.minor);
// Real gate, not an assert: TMA needs compute capability 9.0 or newer.
// 12.0 (RTX 50) passes; 8.9 (L4, RTX 40) and below do not.
if (prop.major < 9) {
std::fprintf(stderr,
"TMA needs compute capability 9.0 or newer; this GPU "
"is %d.%d. An H100, a B200 or any RTX 50 runs this "
"program.\n",
prop.major, prop.minor);
return EXIT_FAILURE;
}
// cuTensorMapEncodeTiled is a driver API function. The runtime hands
// out a pointer to the 12.0-era version of it, so the build needs no
// -lcuda and no driver headers beyond cudaTypedefs.h.
void* fn = nullptr;
cudaDriverEntryPointQueryResult qr = cudaDriverEntryPointSymbolNotFound;
CUDA_CHECK(cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &fn,
12000, cudaEnableDefault, &qr));
if (qr != cudaDriverEntryPointSuccess || fn == nullptr) {
std::fprintf(stderr,
"cuTensorMapEncodeTiled not found in this driver "
"(query result %d)\n",
static_cast<int>(qr));
return EXIT_FAILURE;
}
const PFN_cuTensorMapEncodeTiled_v12000 encodeTiled =
reinterpret_cast<PFN_cuTensorMapEncodeTiled_v12000>(fn);
// Distinct value per element, exactly representable in float, so the
// comparison below can demand bit equality.
const size_t elems = kGmemWidth * kGmemHeight;
const size_t bytes = elems * sizeof(float);
std::vector<float> h_in(elems);
std::vector<float> h_out(elems);
for (size_t i = 0; i < elems; ++i) {
h_in[i] = static_cast<float>(i);
}
float* d_in = nullptr;
float* d_out = nullptr;
CUDA_CHECK(cudaMalloc(&d_in, bytes));
CUDA_CHECK(cudaMalloc(&d_out, bytes));
CUDA_CHECK(cudaMemcpy(d_in, h_in.data(), bytes, cudaMemcpyHostToDevice));
// Zero the output so a tile the kernel never wrote reads as wrong
// (expected values start at 1.0f), instead of as leftover input.
CUDA_CHECK(cudaMemset(d_out, 0, bytes));
int status = EXIT_SUCCESS;
// snippet: encode-map
// The tensor map is the whole address story, built once on the host:
// element type, tensor rank, base pointer, tensor shape, row stride in
// bytes (one stride for rank 2: the innermost dimension is implicitly
// contiguous), tile shape, and no interleave, swizzle, L2 promotion or
// out-of-bounds fill. The kernel receives the map by value and never
// computes a global address.
CUtensorMap tileMap{};
const cuuint64_t gmemDim[2] = {kGmemWidth, kGmemHeight};
const cuuint64_t gmemStrideBytes[1] = {kGmemWidth * sizeof(float)};
const cuuint32_t boxDim[2] = {kTileWidth, kTileHeight};
const cuuint32_t elemStride[2] = {1, 1};
const CUresult mapResult = encodeTiled(
&tileMap, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 2, d_in, gmemDim,
gmemStrideBytes, boxDim, elemStride, CU_TENSOR_MAP_INTERLEAVE_NONE,
CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE,
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
// end snippet
if (mapResult != CUDA_SUCCESS) {
std::fprintf(stderr, "cuTensorMapEncodeTiled failed: CUresult %d\n",
static_cast<int>(mapResult));
status = EXIT_FAILURE;
}
if (status == EXIT_SUCCESS) {
const dim3 grid(kGmemWidth / kTileWidth, kGmemHeight / kTileHeight);
loadTileTma<<<grid, kThreadsPerBlock>>>(tileMap, d_out);
CUDA_CHECK(cudaGetLastError());
CUDA_CHECK(cudaDeviceSynchronize());
CUDA_CHECK(
cudaMemcpy(h_out.data(), d_out, bytes, cudaMemcpyDeviceToHost));
// Bit-exact check. in[i] + 1.0f is one float add of exactly
// representable values on both sides, so the first difference is a
// real bug (a tile landing in the wrong place, a race with the
// engine), never rounding.
size_t bad = elems;
for (size_t i = 0; i < elems; ++i) {
if (h_out[i] != h_in[i] + 1.0f) {
bad = i;
break;
}
}
if (bad != elems) {
std::fprintf(stderr,
"round trip wrong at %zu (row %zu, col %zu): got "
"%.1f, want %.1f\n",
bad, bad / kGmemWidth, bad % kGmemWidth,
static_cast<double>(h_out[bad]),
static_cast<double>(h_in[bad] + 1.0f));
status = EXIT_FAILURE;
} else {
std::printf(
"round trip: PASS. %zu elements, %d tiles of %d x %d "
"floats, one cp.async.bulk.tensor per tile.\n",
elems,
static_cast<int>((kGmemWidth / kTileWidth) *
(kGmemHeight / kTileHeight)),
kTileHeight, kTileWidth);
}
}
CUDA_CHECK(cudaFree(d_in));
CUDA_CHECK(cudaFree(d_out));
return status;
}