triton-syntax

Triton编程语言

概述

Triton是一种专为GPU编程设计的Python DSL，目标是让GPU编程像NumPy一样简单。

核心特性

1. Python-like语法

@triton.jit
def add_kernel(
    x_ptr, y_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr
):
    # 程序ID
    pid = tl.program_id(axis=0)
    
    # 块偏移
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    
    # Mask处理边界
    mask = offsets < n_elements
    
    # 加载数据
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    
    # 计算
    output = x + y
    
    # 存储结果
    tl.store(output_ptr + offsets, output, mask=mask)

2. 自动优化

Triton编译器自动进行：

• 内存合并优化
• 共享内存管理
• 寄存器分配
• 指令调度

3. Block编程模型

# 每个program处理一个block的数据
BLOCK_SIZE = 1024
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)

add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=BLOCK_SIZE)

数据类型

基本类型

tl.float32    # 32位浮点
tl.float16    # 半精度浮点
tl.bfloat16   # Brain浮点
tl.int32      # 32位整数
tl.int64      # 64位整数

Tensor类型

# 1D tensor
x = tl.arange(0, BLOCK_SIZE)  # shape: [BLOCK_SIZE]

# 2D tensor
rows = tl.arange(0, BLOCK_M)[:, None]  # shape: [BLOCK_M, 1]
cols = tl.arange(0, BLOCK_N)[None, :]  # shape: [1, BLOCK_N]

内存操作

加载（Load）

# 基本加载
data = tl.load(ptr + offsets)

# 带mask加载（处理边界）
data = tl.load(ptr + offsets, mask=mask, other=0.0)

# 带cache hint
data = tl.load(ptr + offsets, cache_modifier=".ca")  # cache all

存储（Store）

# 基本存储
tl.store(ptr + offsets, data)

# 带mask存储
tl.store(ptr + offsets, data, mask=mask)

计算操作

算术运算

# 逐元素运算
c = a + b
c = a * b
c = a / b
c = tl.exp(a)
c = tl.log(a)
c = tl.sqrt(a)

规约操作

# Sum
total = tl.sum(x, axis=0)

# Max
maximum = tl.max(x, axis=0)

# Min
minimum = tl.min(x, axis=0)

矩阵运算

# Dot product
c = tl.dot(a, b)  # 高度优化的矩阵乘法

# 支持混合精度
c = tl.dot(a.to(tl.float16), b.to(tl.float16), acc=tl.float32)