triton-ascend-reduce

Reduce 算子优化

适用于需要聚合多个值的归约操作

适用算子

基础归约: sum, mean, max, min, prod 归一化: softmax, logsoftmax, layernorm, rmsnorm, groupnorm, batchnorm 统计: variance, std

关键性能优化：计算重组（延迟归约）

Ascend 上 tl.sum/tl.max/tl.min 等归约指令开销较大，循环内每次迭代都调用归约会成为性能瓶颈。核心思路：循环内只做逐元素累加（+=），循环结束后再执行一次归约。

反模式 vs 正确范式

# ❌ 反模式：循环内每次都调 tl.sum，产生 N/BLOCK_SIZE 次归约
total = 0.0
for offset in range(0, N, BLOCK_SIZE):
    block = tl.load(ptr + offset + tl.arange(0, BLOCK_SIZE), ...)
    total += tl.sum(block)  # 每次迭代都归约 → 开销大

# ✅ 正确：循环内只做逐元素累加，最后一次性归约
acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
for offset in range(0, N, BLOCK_SIZE):
    block = tl.load(ptr + offset + tl.arange(0, BLOCK_SIZE), ...)
    acc += block               # 逐元素 add，无归约开销
total = tl.sum(acc)            # 仅此一次归约

2D 场景（沿某一轴归约）

# ❌ 反模式：循环内每次沿 axis=0 归约
acc_1d = tl.zeros((BLOCK_N,), dtype=tl.float32)
for m_start in range(0, M, BLOCK_M):
    tile = tl.load(...)  # [BLOCK_M, BLOCK_N]
    acc_1d += tl.sum(tile, axis=0)  # 每次迭代都归约

# ✅ 正确：保持 2D 累加器，最后一次性归约
acc_2d = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for m_start in range(0, M, BLOCK_M):
    tile = tl.load(...)  # [BLOCK_M, BLOCK_N]
    acc_2d += tile                      # 保持 2D，无归约
result = tl.sum(acc_2d, axis=0)         # 最后一次归约 → [BLOCK_N]

适用条件

• 可结合律操作：sum（+=）、prod（*=）等满足结合律的操作均可使用此范式
• 非 sum 归约（max/min）也适用：循环内用 tl.maximum/tl.minimum 逐元素取极值，最后一次 tl.max/tl.min
• UB 容量权衡：2D 累加器占用更多 UB（统一缓冲区），需确保 BLOCK_M × BLOCK_N × dtype_size 不超出 UB 容量。当 UB 不够时可适当减小 BLOCK_SIZE
• 掩码处理：累加器初始化为归约的幺元（sum → 0、prod → 1、max → -inf、min → inf），用 other=幺元 处理边界

完整示例：Sum reduction over a dimension

@triton.jit
def sum_reduce_kernel(
    x_ptr, y_ptr,
    B: tl.constexpr, M: tl.constexpr, N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    NUM_CORES: tl.constexpr = 20,
):
    """Input X[B, M, N] → Output Y[B, N]，沿 M 轴求和"""
    pid = tl.program_id(0)
    num_blocks_n = tl.cdiv(N, BLOCK_SIZE_N)
    total_blocks = B * num_blocks_n

    for block_idx in range(pid, total_blocks, NUM_CORES):
        b_idx = block_idx // num_blocks_n
        n_start = (block_idx % num_blocks_n) * BLOCK_SIZE_N
        n_offsets = n_start + tl.arange(0, BLOCK_SIZE_N)
        n_mask = n_offsets < N

        # 2D 累加器，延迟归约
        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

        for m_start in range(0, M, BLOCK_SIZE_M):
            m_offsets = m_start + tl.arange(0, BLOCK_SIZE_M)
            m_mask = m_offsets < M
            x_offset = b_idx * M * N + m_offsets[:, None] * N + n_offsets[None, :]
            x_block = tl.load(x_ptr + x_offset, mask=m_mask[:, None] & n_mask[None, :], other=0.0)
            acc += x_block  # 逐元素累加，不归约

        result = tl.sum(acc, axis=0)  # 循环结束后一次性归约
        tl.store(y_ptr + b_idx * N + n_offsets, result, mask=n_mask)