CUDA GPU プログラミング上級編: Warp最適化、Tensor Core、Tritonカーネル作成まで

#define TILE_SIZE 16

__global__ void tiledMatMul(
    float* A, float* B, float* C,
    int M, int N, int K
) {
    __shared__ float tileA[TILE_SIZE][TILE_SIZE];
    __shared__ float tileB[TILE_SIZE][TILE_SIZE];

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
    float sum = 0.0f;

    int numTiles = (K + TILE_SIZE - 1) / TILE_SIZE;

    for (int t = 0; t < numTiles; t++) {
        int aCol = t * TILE_SIZE + threadIdx.x;
        int bRow = t * TILE_SIZE + threadIdx.y;

        // タイルをshared memoryに協調ロード
        tileA[threadIdx.y][threadIdx.x] =
            (row < M && aCol < K) ? A[row * K + aCol] : 0.0f;
        tileB[threadIdx.y][threadIdx.x] =
            (bRow < K && col < N) ? B[bRow * N + col] : 0.0f;

        __syncthreads();

        for (int k = 0; k < TILE_SIZE; k++) {
            sum += tileA[threadIdx.y][k] * tileB[k][threadIdx.x];
        }

        __syncthreads();
    }

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

__constant__ float convWeights[256];

__global__ void convKernel(float* input, float* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float result = 0.0f;
        for (int i = 0; i < 256; i++) {
            result += input[idx + i] * convWeights[i]; // ブロードキャストアクセス
        }
        output[idx] = result;
    }
}

// 悪い例: ワープ内でdivergenceが発生
__global__ void badKernel(float* data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx % 2 == 0) {
        data[idx] = data[idx] * 2.0f;
    } else {
        data[idx] = data[idx] + 1.0f;
    }
}

// 良い例: ワープ単位で分岐を揃える
__global__ void goodKernel(float* data, int n) {
    int warpId = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int idx    = blockIdx.x * blockDim.x + threadIdx.x;
    if (warpId % 2 == 0) {   // ワープ単位の分岐 — divergenceなし
        data[idx] = data[idx] * 2.0f;
    } else {
        data[idx] = data[idx] + 1.0f;
    }
}

// 悪い例: 列優先ストライドアクセス (stride = N)
__global__ void columnAccess(float* matrix, int M, int N) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    for (int row = 0; row < M; row++) {
        sum += matrix[row * N + col]; // stride-Nアクセス
    }
}

// 良い例: 行優先コアレスドアクセス
__global__ void rowAccess(float* matrix, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    for (int col = 0; col < N; col++) {
        sum += matrix[row * N + col]; // 連続アクセス
    }
}

// 悪い例: 転置でbank conflictが発生
// smem[0][0], smem[1][0], smem[2][0]... はすべてbank 0

// 良い例: 1列のパディングでbank conflictを解消
__global__ void matTransposeNoBankConflict(float* in, float* out, int N) {
    __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; // +1パディング

    int x = blockIdx.x * TILE_SIZE + threadIdx.x;
    int y = blockIdx.y * TILE_SIZE + threadIdx.y;

    if (x < N && y < N)
        tile[threadIdx.y][threadIdx.x] = in[y * N + x];

    __syncthreads();

    x = blockIdx.y * TILE_SIZE + threadIdx.x;
    y = blockIdx.x * TILE_SIZE + threadIdx.y;

    if (x < N && y < N)
        out[y * N + x] = tile[threadIdx.x][threadIdx.y];
}

// __shfl_down_syncによるwarpレベルのsum reduction
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void blockReduceSum(float* input, float* output, int n) {
    __shared__ float shared[32];

    int tid   = threadIdx.x;
    int idx   = blockIdx.x * blockDim.x + tid;
    float val = (idx < n) ? input[idx] : 0.0f;

    // Step 1: 各ワープ内でreduction
    val = warpReduceSum(val);

    // Step 2: 各ワープの結果をshared memoryに書き込み
    if (tid % 32 == 0) shared[tid / 32] = val;
    __syncthreads();

    // Step 3: 最初のワープで最終reduction
    if (tid < 32) {
        val = (tid < blockDim.x / 32) ? shared[tid] : 0.0f;
        val = warpReduceSum(val);
        if (tid == 0) atomicAdd(output, val);
    }
}

__global__ void exclusiveScan(int* input, int* output, int n) {
    __shared__ int temp[1024 * 2];
    int tid    = threadIdx.x;
    int offset = 1;

    temp[2 * tid]     = (2 * tid < n)     ? input[2 * tid]     : 0;
    temp[2 * tid + 1] = (2 * tid + 1 < n) ? input[2 * tid + 1] : 0;

    // Up-sweep (reduce) フェーズ
    for (int d = n >> 1; d > 0; d >>= 1) {
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            temp[bi] += temp[ai];
        }
        offset <<= 1;
    }

    if (tid == 0) temp[n - 1] = 0;

    // Down-sweep フェーズ
    for (int d = 1; d < n; d <<= 1) {
        offset >>= 1;
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            int t  = temp[ai];
            temp[ai]  = temp[bi];
            temp[bi] += t;
        }
    }

    __syncthreads();
    if (2 * tid < n)     output[2 * tid]     = temp[2 * tid];
    if (2 * tid + 1 < n) output[2 * tid + 1] = temp[2 * tid + 1];
}

#include <mma.h>
using namespace nvcuda;

#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 16

__global__ void tensorCoreMatMul(
    half* A, half* B, float* C,
    int M, int N, int K
) {
    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::col_major> bFrag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K,
                   float> cFrag;

    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int warpN = (blockIdx.y * blockDim.y + threadIdx.y);

    wmma::fill_fragment(cFrag, 0.0f);

    for (int i = 0; i < K; i += WMMA_K) {
        int aRow = warpM * WMMA_M;
        int aCol = i;
        int bRow = i;
        int bCol = warpN * WMMA_N;

        if (aRow < M && aCol < K && bRow < K && bCol < N) {
            wmma::load_matrix_sync(aFrag, A + aRow * K + aCol, K);
            wmma::load_matrix_sync(bFrag, B + bRow * N + bCol, N);
            wmma::mma_sync(cFrag, aFrag, bFrag, cFrag);
        }
    }

    int cRow = warpM * WMMA_M;
    int cCol = warpN * WMMA_N;
    if (cRow < M && cCol < N) {
        wmma::store_matrix_sync(C + cRow * N + cCol, cFrag, N,
                                wmma::mem_row_major);
    }
}

__device__ float fastSqrt(float x) {
    float result;
    asm("sqrt.approx.f32 %0, %1;" : "=f"(result) : "f"(x));
    return result;
}

__device__ int ldgInt(const int* ptr) {
    int val;
    // キャッシュヒントを使ったグローバルメモリロード
    asm volatile("ld.global.cg.s32 %0, [%1];"
                 : "=r"(val) : "l"(ptr));
    return val;
}

標準Attention HBM I/O:
  Q, K, V 読み込み      → O(Nd)
  S = QK^T 書き込み    → O(N^2)   <-- ボトルネック
  S 読み込み、P 書き込み → O(N^2)
  P 読み込み、O 書き込み → O(Nd)
  合計: O(N^2 * d)

Flash Attention HBM I/O:
  Q, K, V 読み込み → O(Nd)
  O 書き込み       → O(Nd)
  合計: O(Nd)  — N^2項なし!

// Flash Attention フォワードパス (簡略版)
__global__ void flashAttentionKernel(
    const float* Q, const float* K, const float* V,
    float* O, float* l, float* m,
    int seqLen, int headDim, float scale
) {
    int tileSize = blockDim.x;
    int qStart   = blockIdx.x * tileSize;

    extern __shared__ float smem[];
    float* sQ = smem;
    float* sK = smem + tileSize * headDim;
    float* sV = smem + 2 * tileSize * headDim;

    // QタイルをSRAMにロード
    for (int d = 0; d < headDim; d++) {
        sQ[threadIdx.x * headDim + d] =
            Q[(qStart + threadIdx.x) * headDim + d];
    }

    float mi  = -INFINITY;  // 行の最大値 (数値安定性)
    float li  = 0.0f;       // softmax分母の累積
    float oi[64] = {};      // 出力の累積

    for (int kStart = 0; kStart < seqLen; kStart += tileSize) {
        // K, VタイルをSRAMにロード
        for (int d = 0; d < headDim; d++) {
            sK[threadIdx.x * headDim + d] =
                K[(kStart + threadIdx.x) * headDim + d];
            sV[threadIdx.x * headDim + d] =
                V[(kStart + threadIdx.x) * headDim + d];
        }
        __syncthreads();

        for (int j = 0; j < tileSize; j++) {
            float sij = 0.0f;
            for (int d = 0; d < headDim; d++) {
                sij += sQ[threadIdx.x * headDim + d]
                     * sK[j * headDim + d];
            }
            sij *= scale;

            // オンラインsoftmax更新
            float miNew = fmaxf(mi, sij);
            float liNew = li * expf(mi - miNew) + expf(sij - miNew);

            for (int d = 0; d < headDim; d++) {
                oi[d] = oi[d] * (li / liNew) * expf(mi - miNew)
                      + (expf(sij - miNew) / liNew) * sV[j * headDim + d];
            }
            mi = miNew;
            li = liNew;
        }
        __syncthreads();
    }

    for (int d = 0; d < headDim; d++) {
        O[(qStart + threadIdx.x) * headDim + d] = oi[d];
    }
    l[qStart + threadIdx.x] = li;
    m[qStart + threadIdx.x] = mi;
}

void pipelinedProcessing(float* h_input, float* h_output,
                         int totalSize, int chunkSize) {
    const int numStreams = 4;
    cudaStream_t streams[numStreams];
    float *d_input[numStreams], *d_output[numStreams];

    for (int i = 0; i < numStreams; i++) {
        cudaStreamCreate(&streams[i]);
        cudaMalloc(&d_input[i],  chunkSize * sizeof(float));
        cudaMalloc(&d_output[i], chunkSize * sizeof(float));
    }

    // パイプライン: HtoD → カーネル → DtoH をストリーム間でオーバーラップ
    for (int chunk = 0; chunk < totalSize / chunkSize; chunk++) {
        int s = chunk % numStreams;
        float* hIn  = h_input  + chunk * chunkSize;
        float* hOut = h_output + chunk * chunkSize;

        cudaMemcpyAsync(d_input[s], hIn,
                        chunkSize * sizeof(float),
                        cudaMemcpyHostToDevice, streams[s]);

        dim3 block(256);
        dim3 grid((chunkSize + 255) / 256);
        processKernel<<<grid, block, 0, streams[s]>>>(
            d_input[s], d_output[s], chunkSize);

        cudaMemcpyAsync(hOut, d_output[s],
                        chunkSize * sizeof(float),
                        cudaMemcpyDeviceToHost, streams[s]);
    }

    for (int i = 0; i < numStreams; i++) {
        cudaStreamSynchronize(streams[i]);
        cudaStreamDestroy(streams[i]);
        cudaFree(d_input[i]);
        cudaFree(d_output[i]);
    }
}

cudaEvent_t startEvent, stopEvent;
cudaEventCreate(&startEvent);
cudaEventCreate(&stopEvent);

cudaEventRecord(startEvent, stream1);
kernel1<<<grid, block, 0, stream1>>>(d_data);

// stream2はstream1のkernel1完了を待機
cudaStreamWaitEvent(stream2, startEvent, 0);
kernel2<<<grid, block, 0, stream2>>>(d_data);

cudaEventRecord(stopEvent, stream2);
cudaEventSynchronize(stopEvent);

float ms = 0;
cudaEventElapsedTime(&ms, startEvent, stopEvent);
printf("経過時間: %.3f ms\n", ms);

// GPU 0からGPU 1へNVLink/PCIe P2P経由で直接コピー
int canAccess;
cudaDeviceCanAccessPeer(&canAccess, 0, 1);
if (canAccess) {
    cudaSetDevice(0);
    cudaDeviceEnablePeerAccess(1, 0);

    float* d_src; // GPU 0上
    float* d_dst; // GPU 1上

    cudaSetDevice(0);
    cudaMalloc(&d_src, size);

    cudaSetDevice(1);
    cudaMalloc(&d_dst, size);

    // ホストを経由せず直接コピー
    cudaMemcpyPeer(d_dst, 1, d_src, 0, size);
}

#include <nccl.h>

void ncclAllReduceExample(int numGPUs, float** d_buffers, size_t count) {
    ncclComm_t comms[numGPUs];
    int devList[numGPUs];
    for (int i = 0; i < numGPUs; i++) devList[i] = i;

    ncclCommInitAll(comms, numGPUs, devList);

    cudaStream_t streams[numGPUs];
    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        cudaStreamCreate(&streams[i]);
    }

    // 全GPUのバッファを加算し、結果を全GPUに配布
    ncclGroupStart();
    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        ncclAllReduce(d_buffers[i], d_buffers[i], count,
                      ncclFloat, ncclSum, comms[i], streams[i]);
    }
    ncclGroupEnd();

    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        cudaStreamSynchronize(streams[i]);
        ncclCommDestroy(comms[i]);
    }
}

import triton
import triton.language as tl
import torch

@triton.jit
def matmul_kernel(
    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)

    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) & (offs_k[None, :] < K - k),
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=(offs_k[:, None] < K - k) & (offs_n[None, :] < N),
                    other=0.0)
        acc += tl.dot(a, b)
        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, :]
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc, mask=c_mask)


def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    K2, N = b.shape
    assert K == K2
    c = torch.empty((M, N), device=a.device, dtype=torch.float32)
    grid = lambda meta: (
        triton.cdiv(M, meta['BLOCK_M']),
        triton.cdiv(N, meta['BLOCK_N']),
    )
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_M=128, BLOCK_N=256, BLOCK_K=64,
    )
    return c

@triton.jit
def flash_attn_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    L_ptr, M_ptr,
    seqlen, headdim, scale,
    BLOCK_Q:  tl.constexpr,
    BLOCK_KV: tl.constexpr,
):
    start_q = tl.program_id(0) * BLOCK_Q
    offs_q  = start_q + tl.arange(0, BLOCK_Q)
    offs_d  = tl.arange(0, headdim)

    q = tl.load(Q_ptr + offs_q[:, None] * headdim + offs_d[None, :],
                mask=(offs_q[:, None] < seqlen), other=0.0)

    mi  = tl.full((BLOCK_Q,), value=-float('inf'), dtype=tl.float32)
    li  = tl.zeros((BLOCK_Q,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_Q, headdim), dtype=tl.float32)

    for start_kv in range(0, seqlen, BLOCK_KV):
        offs_kv = start_kv + tl.arange(0, BLOCK_KV)

        k = tl.load(K_ptr + offs_kv[:, None] * headdim + offs_d[None, :],
                    mask=(offs_kv[:, None] < seqlen), other=0.0)
        v = tl.load(V_ptr + offs_kv[:, None] * headdim + offs_d[None, :],
                    mask=(offs_kv[:, None] < seqlen), other=0.0)

        s      = tl.dot(q, tl.trans(k)) * scale
        mi_new = tl.maximum(mi, tl.max(s, axis=1))
        p      = tl.exp(s - mi_new[:, None])
        li_new = li * tl.exp(mi - mi_new) + tl.sum(p, axis=1)

        acc = acc * (li * tl.exp(mi - mi_new))[:, None] / li_new[:, None]
        acc += tl.dot(p, v) / li_new[:, None]

        mi = mi_new
        li = li_new

    tl.store(O_ptr + offs_q[:, None] * headdim + offs_d[None, :],
             acc, mask=(offs_q[:, None] < seqlen))
    tl.store(L_ptr + offs_q, li, mask=(offs_q < seqlen))
    tl.store(M_ptr + offs_q, mi, mask=(offs_q < seqlen))

# 全メトリクスのフルプロファイル
ncu --set full -o profile_report ./my_cuda_app

# 特定カーネルのみ、3回分プロファイル
ncu --kernel-name tiledMatMul --launch-count 3 ./my_cuda_app

# Rooflineモデル用のメモリ帯域幅メトリクス収集
ncu --metrics \
    l1tex__t_bytes_pipe_lsu_mem_global_op_ld.sum,\
    l1tex__t_bytes_pipe_lsu_mem_global_op_st.sum,\
    sm__cycles_elapsed.avg,\
    sm__cycles_elapsed.avg.per_second \
    ./my_cuda_app

# DRAMピーク帯域幅に対する使用率
ncu --metrics \
    gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed \
    ./my_cuda_app

# 完全なアプリケーショントレースを収集
nsys profile --trace=cuda,nvtx,osrt \
             --output=timeline_report \
             ./my_cuda_app

# 統計サマリを出力
nsys stats timeline_report.nsys-rep

# GPUトレース詳細をエクスポート
nsys analyze --report gputrace timeline_report.nsys-rep

#include <nvtx3/nvToolsExt.h>

void profiledFunction() {
    nvtxRangePush("DataPreprocessing");
    preprocessData();
    nvtxRangePop();

    nvtxRangePush("ForwardPass");
    forwardPass();
    nvtxRangePop();
}

算術強度 (Arithmetic Intensity) = FLOP / アクセスバイト数

A100 80GBの例:
  FP16 Tensor Core ピーク: 312 TFLOPS
  HBM帯域幅: 2 TB/s
  リッジポイント: 312 / 2 = 156 FLOP/Byte

  算術強度 < 156 → メモリ帯域幅バウンド
  算術強度 > 156 → 演算バウンド

# 算術強度の測定
ncu --metrics \
    sm__ops_path_tensor_src_fp16.sum,\
    dram__bytes.sum \
    ./matmul_app
# 算術強度 = (tensor_ops * 2) / dram_bytes