triton-cuda-elementwise

Element-wise 算子优化

适用于逐元素独立计算的算子

适用算子

算术运算: add, mul, div, sub, pow 激活函数: relu, sigmoid, tanh（需用 tl.extra.cuda.libdevice.tanh）, gelu, silu, swish 数学函数: exp, log, sqrt, sin, cos, abs

优化策略

1. 连续内存访问优化

张量在内存中连续存储时，可用一维指针遍历，避免多维索引开销。

方案 1: 转连续 + 一维访问（推荐）

class ModelNew(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input_tensor):
        # 非连续张量转为连续（一次性开销）
        if not input_tensor.is_contiguous():
            input_tensor = input_tensor.contiguous()
        
        output_tensor = torch.empty_like(input_tensor)
        n_elements = input_tensor.numel()
        grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
        
        elementwise_kernel[grid](input_tensor, output_tensor, n_elements, BLOCK_SIZE)
        return output_tensor

@triton.jit
def elementwise_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    data = tl.load(input_ptr + offsets, mask=mask)
    result = compute(data)  # 你的计算逻辑
    tl.store(output_ptr + offsets, result, mask=mask)

优势:

• .contiguous() 一次性开销 vs stride 每次访问都有开销
• 更好的合并访问（coalesced access）
• 编译器优化更容易

方案 2: 使用 stride 访问（不推荐）

仅当无法调用 .contiguous() 时使用。

2. BLOCK_SIZE 选择

• 推荐值: 256, 512, 1024
• 原则: 平衡并行度和资源占用
• GPU 考量:
  • 更大的 BLOCK_SIZE → 更少的 block 启动开销，但可能降低 occupancy
  • 更小的 BLOCK_SIZE → 更细粒度的并行，但启动开销增加
  • 确保 Grid 大小足够大以充分利用 GPU

3. Warp 配置

Element-wise 算子通常使用较少的 warp：

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4),
        triton.Config({'BLOCK_SIZE': 512}, num_warps=2),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8),
    ],
    key=['n_elements'],
    restore_value=['output_ptr'],  # 必须：列出所有输出指针参数名
)
@triton.jit
def optimized_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    data = tl.load(input_ptr + offsets, mask=mask)
    result = compute(data)
    tl.store(output_ptr + offsets, result, mask=mask)