Generated Triton code, GEMM examples and ROCm warp-size pitfalls

Examining these three practical mechanisms—inspecting torch.compile's generated Triton code, comparing tiled matrix multiplication (GEMM) code, and navigating warp-size pitfalls when porting CUDA to AMD ROCm—illustrates how high-level AI frameworks translate directly into hardware instructions.

1. Inspecting Triton Code Generated by torch.compile

When you wrap a standard PyTorch operation like RMSNorm (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight) with @torch.compile, TorchInductor eliminates four separate HBM round-trips by generating a single fused Triton kernel behind the scenes.

Running your script with TORCH_LOGS="output_code" python script.py prints the actual compiler-generated Triton code, which looks like this:

import triton
import triton.language as tl

@triton.jit
def triton_per_fused_mul_pow_rsqrt_0(in_ptr0, in_ptr1, out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel

    # 1. Single coalesced read of input tensor (x) and weight vector from HBM
    x = tl.load(in_ptr0 + (rindex + (4096 * xindex)), rmask & xmask, other=0.0).to(tl.float32)
    w = tl.load(in_ptr1 + (rindex), rmask, eviction_policy='evict_last').to(tl.float32)

    # 2. Compute square, row-wise mean, and reciprocal square root entirely in registers
    x_sq = x * x
    row_sum = tl.sum(x_sq, 1)[:, None]
    rstd = tl.rsqrt((row_sum / 4096.0) + 1e-06)

    # 3. Multiply by inverse standard deviation and learned weight, then write once
    out = (x * rstd) * w
    tl.store(out_ptr0 + (rindex + (4096 * xindex)), out, rmask & xmask)

• What the Compiler Discovered Automatically: Notice eviction_policy='evict_last' on the weight vector w. TorchInductor automatically detected that w is reused across every row (xindex), instructing the GPU's L2 cache to keep w pinned on-chip while streaming x through once.

2. Side-by-Side Matrix Multiplication (GEMM) Code: CUDA C++ vs. Triton

Matrix multiplication (C = A × B) highlights the stark difference between manual thread synchronization in CUDA C++ and declarative tile math in Triton.

A. CUDA C++ (Explicit Shared Memory Tiling & Sync Barriers)

In CUDA C++, every single thread computes one scalar output element (acc), while threads cooperate manually to load tiles into __shared__ memory:

#define TILE 16

__global__ void gemm_cuda(const float* A, const float* B, float* C, int M, int N, int K) {
    __shared__ float sA[TILE][TILE];
    __shared__ float sB[TILE][TILE];

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float acc = 0.0f;

    for (int t = 0; t < (K + TILE - 1) / TILE; ++t) {
        // 1. Each thread manually loads 1 element into Shared Memory
        sA[threadIdx.y][threadIdx.x] = (row < M && t * TILE + threadIdx.x < K)
                                       ? A[row * K + t * TILE + threadIdx.x] : 0.0f;
        sB[threadIdx.y][threadIdx.x] = (t * TILE + threadIdx.y < K && col < N)
                                       ? B[(t * TILE + threadIdx.y) * N + col] : 0.0f;
        __syncthreads(); // Mandatory barrier before reading shared tiles

        // 2. Inner dot product across the 16x16 tile
        for (int k = 0; k < TILE; ++k) {
            acc += sA[threadIdx.y][k] * sB[k][threadIdx.x];
        }
        __syncthreads(); // Mandatory barrier before loading next tile
    }

    if (row < M && col < N) C[row * N + col] = acc;
}

B. OpenAI Triton (Declarative 2D Tile Accumulation & Tensor Cores)

In Triton, a single program block computes an entire BLOCK_M × BLOCK_N output tile using tl.dot(), which automatically lowers to hardware Tensor Cores:

@triton.jit
def gemm_triton(A_ptr, B_ptr, C_ptr, M, N, K,
                stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    # Accumulate in FP32 registers while multiplying FP16 tiles on Tensor Cores
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (k + offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(k + offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)  # Hardware Tensor Core instruction
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    c_ptrs = C_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

3. Warp-Size Pitfalls When Porting CUDA C++ to AMD ROCm (HIP)

When porting low-level CUDA C++ kernels to AMD Instinct GPUs (such as the MI300X or MI325X) via AMD's HIP interface, syntax translation (cudaMalloc → hipMalloc) is trivial. However, silent logical bugs frequently occur because NVIDIA warps contain 32 threads, whereas AMD CDNA wavefronts contain 64 threads.

• Why Triton Avoids All Four Pitfalls: In OpenAI Triton, operations like tl.sum() and tl.max() operate across entire logical tensor dimensions rather than hardcoded 32-thread warps. During compilation, Triton's NVPTX backend automatically emits 32-thread shuffle trees for NVIDIA GPUs, while its AMDGPU backend emits 64-thread wavefront reductions for AMD Instinct GPUs—eliminating warp-size bugs by design.

We have now examined the full lifecycle of modern GPU kernels—from PyTorch bytecode capture and Triton code generation down to CUDA vs. AMD wavefront mechanics.