GPU Architecture · Code example
Computes the same matrix multiply as scalar FP32 on the CUDA cores and as 16x16 tiles on the Tensor Cores through the WMMA API, checks both against a CPU reference, and times them so the Tensor Core speedup is a measured number.
// 04-wmma-matmul.cu — the same C = A x B computed two ways:
// (1) naiveMatmul — one thread per output element, scalar FP32 FMAs on the
// CUDA cores (the baseline every matmul starts as), and
// (2) wmmaMatmul — 16x16 output tiles produced by the Tensor Cores through
// the WMMA API, one warp per tile.
// Both are checked against a CPU reference; then we time them so the Tensor
// Core speedup is a number, not a claim.
//
// The Tensor Core does exactly one thing: D = A*B + C on small matrix TILES,
// FP16 inputs accumulated in FP32. A whole warp issues one mma_sync and the
// Tensor Cores in its sub-partition chew through the tile cooperatively.
//
// Build: nvcc -arch=sm_120 04-wmma-matmul.cu -o wmma (Tensor Cores need sm_70+)
// Run: ./wmma
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>
using namespace nvcuda;
#define CHECK(call) do { \
cudaError_t _e = (call); \
if (_e != cudaSuccess) { \
printf("CUDA error at line %d: %s\n", __LINE__, \
cudaGetErrorString(_e)); \
return 1; \
} \
} while (0)
// Square matrices, a multiple of 64 so every 16x16 tile is full (no masking —
// boundary handling would only distract from the WMMA mechanism).
static const int N = 1024;
// The one fixed tile shape the Tensor Core works in: a 16x16x16 MMA. The
// fragment template parameters below MUST match a shape the hardware supports.
static const int WMMA_M = 16, WMMA_N = 16, WMMA_K = 16;
// ---- (1) Baseline: scalar FP32 on the CUDA cores --------------------------
// One thread computes one C[row][col] as a dot product — K scalar FMAs, each on
// a single CUDA-core lane. This is what a matmul is before you specialize it.
__global__ void naiveMatmul(const float* A, const float* B, float* C, int n) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
int row = blockIdx.y * blockDim.y + threadIdx.y;
if (row < n && col < n) {
float acc = 0.0f;
for (int k = 0; k < n; ++k)
acc += A[row * n + k] * B[k * n + col];
C[row * n + col] = acc;
}
}
// ---- (2) Tensor Core: 16x16 tiles via WMMA --------------------------------
// One WARP owns one 16x16 output tile. It sweeps across K in steps of 16,
// loading a 16x16 tile of A and of B each step and accumulating their product
// into acc. Every thread in the warp calls these together — the *_sync suffix
// means the 32 lanes cooperate to feed one Tensor Core operation.
__global__ void wmmaMatmul(const half* A, const half* B, float* C, int n) {
// Global warp coordinates → which 16x16 output tile this warp produces.
int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
int rowTile = warpM * WMMA_M; // top row of this warp's output tile
int colTile = warpN * WMMA_N; // left column
if (rowTile >= n || colTile >= n) return;
// Opaque, warp-owned register fragments. You never index them by hand — the
// hardware decides how the tile's elements are spread across the 32 lanes.
wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> aFrag;
wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> bFrag;
wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc;
wmma::fill_fragment(acc, 0.0f); // C tile starts at zero
for (int k = 0; k < n; k += WMMA_K) {
// Leading dimension = row stride of the full matrix (elements), not the tile.
wmma::load_matrix_sync(aFrag, A + rowTile * n + k, n); // A tile: rows rowTile.., col k
wmma::load_matrix_sync(bFrag, B + k * n + colTile, n); // B tile: row k.., cols colTile..
wmma::mma_sync(acc, aFrag, bFrag, acc); // acc += aFrag * bFrag (one MMA)
}
wmma::store_matrix_sync(C + rowTile * n + colTile, acc, n, wmma::mem_row_major);
}
static float timeKernel(void (*launch)(), int iters) {
cudaEvent_t beg, end;
cudaEventCreate(&beg); cudaEventCreate(&end);
launch(); // one warm-up (not timed)
cudaDeviceSynchronize();
cudaEventRecord(beg);
for (int i = 0; i < iters; ++i) launch();
cudaEventRecord(end);
cudaEventSynchronize(end);
float ms = 0.0f; cudaEventElapsedTime(&ms, beg, end);
cudaEventDestroy(beg); cudaEventDestroy(end);
return ms / iters;
}
// File-scope pointers so the timed launch thunks can see them.
static float *dAf, *dBf, *dCf;
static half *dAh, *dBh;
static float *dCw;
static void launchNaive() {
dim3 block(16, 16);
dim3 grid((N + 15) / 16, (N + 15) / 16);
naiveMatmul<<<grid, block>>>(dAf, dBf, dCf, N);
}
static void launchWmma() {
// 128x4 threads = 4x4 = 16 warps per block → each block makes a 64x64 output.
dim3 block(128, 4);
dim3 grid((N + 63) / 64, (N + 63) / 64);
wmmaMatmul<<<grid, block>>>(dAh, dBh, dCw, N);
}
int main() {
size_t nf = (size_t)N * N;
size_t bytesF = nf * sizeof(float);
size_t bytesH = nf * sizeof(half);
float *hA = (float*)malloc(bytesF);
float *hB = (float*)malloc(bytesF);
float *hRef = (float*)malloc(bytesF); // CPU reference
float *hC = (float*)malloc(bytesF); // device results copied back
half *hAh = (half*)malloc(bytesH);
half *hBh = (half*)malloc(bytesH);
// Small values so FP16's ~3-decimal-digit precision stays sane. Store both a
// float copy (for the naive kernel + reference) and the half copy WMMA reads.
for (size_t i = 0; i < nf; ++i) {
float a = (float)((i * 7 + 1) % 13) / 13.0f - 0.5f;
float b = (float)((i * 5 + 3) % 11) / 11.0f - 0.5f;
hA[i] = a; hB[i] = b;
hAh[i] = __float2half(a); hBh[i] = __float2half(b);
}
// CPU reference from the SAME values the GPU sees (half cast back to float),
// so we measure only arithmetic differences, not input rounding.
for (int r = 0; r < N; ++r)
for (int c = 0; c < N; ++c) {
float acc = 0.0f;
for (int k = 0; k < N; ++k)
acc += __half2float(hAh[r * N + k]) * __half2float(hBh[k * N + c]);
hRef[r * N + c] = acc;
}
CHECK(cudaMalloc(&dAf, bytesF)); CHECK(cudaMalloc(&dBf, bytesF)); CHECK(cudaMalloc(&dCf, bytesF));
CHECK(cudaMalloc(&dAh, bytesH)); CHECK(cudaMalloc(&dBh, bytesH)); CHECK(cudaMalloc(&dCw, bytesF));
CHECK(cudaMemcpy(dAf, hA, bytesF, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(dBf, hB, bytesF, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(dAh, hAh, bytesH, cudaMemcpyHostToDevice));
CHECK(cudaMemcpy(dBh, hBh, bytesH, cudaMemcpyHostToDevice));
float naiveMs = timeKernel(launchNaive, 20);
CHECK(cudaGetLastError());
float wmmaMs = timeKernel(launchWmma, 20);
CHECK(cudaGetLastError());
// Correctness: max relative error of each kernel vs the CPU reference.
auto maxRelErr = [&](float* dOut) {
CHECK(cudaMemcpy(hC, dOut, bytesF, cudaMemcpyDeviceToHost));
double worst = 0.0;
for (size_t i = 0; i < nf; ++i) {
double denom = fabs(hRef[i]) + 1e-3;
double rel = fabs((double)hC[i] - hRef[i]) / denom;
if (rel > worst) worst = rel;
}
return worst;
};
double naiveErr = maxRelErr(dCf);
double wmmaErr = maxRelErr(dCw);
double gflop = 2.0 * N * N * (double)N / 1e9; // 2 ops per MAC
printf("Matmul %d x %d x %d (%.2f GFLOP per multiply)\n\n", N, N, N, gflop);
printf("naive FP32 (CUDA cores): %7.3f ms %7.1f GFLOP/s max rel err %.1e\n",
naiveMs, gflop / (naiveMs / 1e3), naiveErr);
printf("WMMA FP16 (Tensor Cores): %7.3f ms %7.1f GFLOP/s max rel err %.1e\n",
wmmaMs, gflop / (wmmaMs / 1e3), wmmaErr);
printf("\nTensor Cores: %.1fx faster on the same multiply.\n", naiveMs / wmmaMs);
printf("Both agree with the CPU reference (err is FP16 rounding, not a bug).\n");
free(hA); free(hB); free(hRef); free(hC); free(hAh); free(hBh);
cudaFree(dAf); cudaFree(dBf); cudaFree(dCf);
cudaFree(dAh); cudaFree(dBh); cudaFree(dCw);
return 0;
}