Tritonの紹介 - ニューラルネットワークのためのGPUプログラミング

BLOCK = 512

# これはNumbaのGPUカーネルです
# この関数の様々なインスタンスが並行して実行される場合があります
@jit
def add(X, Y, Z, N):
    # Numba / CUDAでは、各カーネルインスタンス自体がSIMT実行モデルを使用します
    # ここで、threadIdxの様々な値に対して命令が並行して実行されます
 　  tid = threadIdx.x
 　  bid = blockIdx.x
 　  # スカラーインデックス
 　  idx = bid * BLOCK + tid
 　  if id < N:
 　      # Numbaにはポインタがありません
 　      # Z、X、Yは密なテンソルです
 　      Z[idx] = X[idx] + Y[idx]


...
grid = (ceil_div(N, BLOCK),)
block = (BLOCK,)
add[grid, block](x, y, z, x.shape[0])

BLOCK = 512

# これはTritonのGPUカーネルです。
# この関数の様々なインスタンスが並行して実行される場合があります
@jit
def add(X, Y, Z, N):
    # Tritonでは、各カーネルインスタンスが単一のスレッドでブロック操作を実行します
    # threadIdxに類似した構造はありません。
    pid = program_id(0)
    # block of indices
    idx = pid * BLOCK + arange(BLOCK)
    mask = idx < N
    # Tritonは、インデックス演算子ではなくポインタ演算を使用します
    x = load(X + idx, mask=mask)
    y = load(Y + idx, mask=mask)
    store(Z + idx, x + y, mask=mask)


...
grid = (ceil_div(N, BLOCK),)
# スレッドブロックなし
add[grid](x, y, z, x.shape[0])

import triton
import triton.language as tl

@triton.jit
def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):
    # 行インデックス
    m = tl.program_id(0)
    # 列インデックス
    # この特定のカーネルは、BLOCK_SIZE列より少ない行列に対してのみ機能します
    BLOCK_SIZE = 1024
    n = tl.arange(0, BLOCK_SIZE)
    # ロードしたいすべての要素のメモリアドレスは次のように計算できます
    X = X + m * stride_xm + n * stride_xn
    # 入力データをロードします。 範囲外の要素を0で埋めます
    x = tl.load(X, mask=n < N, other=-float('inf'))
    # 数値的に安定したソフトマックスを計算する
    z = x - tl.max(x, axis=0)
    num = tl.exp(z)
    denom = tl.sum(num, axis=0)
    y = num / denom
    # Yに書き戻す
    Y = Y + m * stride_ym + n * stride_yn
    tl.store(Y, y, mask=n < N)

import torch
# 入力/出力テンソルを割り当てます
X = torch.normal(0, 1, size=(583, 931), device='cuda')
Y = torch.empty_like(X)
# SPMD起動グリッド
grid = (X.shape[0], )
# GPUカーネルをエンキューします
softmax[grid](Y, Y.stride(0), Y.stride(1), 
    X, X.stride(0), X.stride(1),
    X.shape[0]    , X.shape[1])

@torch.jit.script
def softmax(x):
    x_max = x.max(dim=1)[0]
    z = x - x_max[:, None]
    numerator = torch.exp(x)
    denominator = numerator.sum(dim=1)
    return numerator / denominator[:, None]

@triton.jit
def matmul(A, B, C, M, N, K, stride_am, stride_ak, 
    stride_bk, stride_bn, stride_cm, stride_cn,
    **META):
    # メタパラメータを抽出する
    BLOCK_M, GROUP_M = META['BLOCK_M'], META['GROUP_M']
    BLOCK_N = META['BLOCK_N']
    BLOCK_K = META['BLOCK_K']
    # プログラムは、L2ヒット率を向上させるためにグループ化されます
    _pid_m = tl.program_id(0)
    _pid_n = tl.program_id(1)
    pid_m = _pid_m // GROUP_M
    pid_n = (_pid_n * GROUP_M) + (_pid_m % GROUP_M)
    # rm (またはrn)は、Cの行（または列）のインデックスの範囲を示します
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    # rkは、A（またはB）の列（または行）のインデックスの範囲を示します
    rk = tl.arange(0, BLOCK_K)
    # AとBの最初のブロックの要素のメモリアドレスは、
    # numpyスタイルのブロードキャストを使用して計算できます
    A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
    B = B + (rk [:, None] * stride_bk  + rn[None, :] * stride_bn)
    # アキュムレータを初期化して繰り返し更新します
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(K, 0, -BLOCK_K):
        a = tl.load(A)
        b = tl.load(B)
        # ブロックレベルの行列乗算
        acc += tl.dot(a, b)
        # 次の反復中にAとBの次のブロックがロードされるようにポインタをインクリメントします
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk
    # fuse leaky ReLU if desired
    # acc = tl.where(acc >= 0, acc, alpha * acc)
    # 結果を書き戻す
    C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
    mask = (rm[:, None] < M) & (rn[None, :] < N)
    tl.store(C, acc, mask=mask)