Tensor Memory Accelerator: Load a 2D Tile

Experimental — correctness is graded; the performance target awaits GPU calibration. The starter spends many thread instructions calculating addresses and copying a 32 by 32 tile into shared memory. On Hopper or newer, let one elected thread issue a descriptor-driven TMA transfer and let the block wait on a transaction barrier.

Languages: CUDA C++, Rust

Lesson

Hopper's Tensor Memory Accelerator transfers multidimensional tiles between global and shared memory from a tensor descriptor. A shared transaction barrier tracks both thread arrivals and the bytes completed by the asynchronous transfer. This problem requires compute capability 9.0 or newer and is disabled on the L4 worker.

The harness fixes effective grid = (N/32, N/32, 1), with 32 by 32 tiles. launch configuration grid and dynamic shared-memory values are ignored by this specialized runner; block dimensions still apply. This is a launch-visibility constraint, not a demonstrated execution failure.

What you will learn

Understand which tensor layout information lives in CUtensorMap.
Initialize a block-scoped transaction barrier in shared memory.
Have one elected thread announce the expected bytes and issue a 2D bulk tensor copy.
Wait for the TMA transaction before the block consumes shared memory.
Separate TMA data movement from Tensor Core matrix arithmetic.

Starter code

    __shared__ alignas(128) float tile[kTmaTile][kTmaTile];
    int tileRow = static_cast<int>(blockIdx.y) * kTmaTile;
    int tileCol = static_cast<int>(blockIdx.x) * kTmaTile;

    for (int i = static_cast<int>(threadIdx.x); i < kTmaTile * kTmaTile; i += blockDim.x) {
        int row = i / kTmaTile;
        int col = i % kTmaTile;
        tile[row][col] = input[(tileRow + row) * N + tileCol + col];
    }
    __syncthreads();
    for (int i = static_cast<int>(threadIdx.x); i < kTmaTile * kTmaTile; i += blockDim.x) {
        int row = i / kTmaTile;
        int col = i % kTmaTile;
        output[(tileRow + row) * N + tileCol + col] = tile[row][col] * 2.0f + 1.0f;
    }