triton-ascend-fused-operator-optimization

融合算子深度优化方法论

优化前：性能天花板分析框架

优化前先分析算子的瓶颈类型，选择正确的优化方向，避免在物理极限附近浪费时间：

瓶颈类型	判断方法	优化方向	典型天花板
内存带宽受限	计算量少、数据搬运多	减少 HBM 读写次数	~1.5-2x
多次遍历	同一数据被读取 3+ 次	多 pass 合并为单 pass	~3-4x
数据访问模式	非连续/strided 访问	重构为连续访问	~5-20x
计算主导	matmul weight 矩阵大	融合几乎无效	~1.0x

天花板计算方法

理论加速比 = baseline 总 HBM 访问量 / 优化后总 HBM 访问量

示例：对于 y = f(x) * z 类融合：

• Baseline（2 个 PyTorch op）：读 x → 写 f(x) → 读 f(x) + z → 写 y = 4 次
• Triton 融合：读 x + z → 写 y = 2 次
• 理论上限 = 4/2 = 2x

实际天花板更低的原因：baseline 中间 tensor 常命中 L2 cache，等效减少了 HBM 访问次数。

Matmul 主导型融合的判断

优化方法 1：多 Pass 合并

适用条件

算子对同一数据进行多次独立遍历（如 softmax 的 max→exp_sum→normalize，或 topk 的多次扫描）。

方法

将所有 pass 合并为单次遍历，在寄存器内完成全部计算：

# 次优：多次遍历
max_val = pass_find_max(data)          # 遍历 1
exp_sum = pass_compute_exp(data)       # 遍历 2
topk = pass_find_topk(data)            # 遍历 3+

# 推荐：单次遍历
data = tl.load(...)
max_val = tl.max(data, axis=0)
exp_vals = tl.math.exp(data - max_val)
exp_sum = tl.sum(exp_vals, axis=0)
probs = exp_vals / exp_sum
# topk 直接在同一 block 内完成
first_val = tl.max(probs, axis=0)

关键约束

当归约维度能放入单个 BLOCK 时效果最佳；维度过大则需分块归约，收益递减。

优化方法 2：数据访问模式重构

适用条件

算子涉及 strided / 非连续访问模式（如需要访问相邻元素的配对计算）。

方法

从按元素展平处理，改为按语义分组处理，使相关元素落在同一 block 内：

# 次优：展平后按 flat_idx 处理，需跨 stride 访问配对元素
for block_id in range(pid, total_elements // BLOCK_SIZE, CORE_NUM):
    flat_idx = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    d = flat_idx % D
    d_pair = d ^ 1  # 跨 stride 随机访问

# 推荐：按语义维度分组，内层连续加载
for group_idx in range(pid, total_groups, CORE_NUM):
    # 计算分组坐标
    for d_start in range(0, D, BLOCK_D):
        d_offsets = d_start + tl.arange(0, BLOCK_D)
        # 配对元素天然在同一 block 内
        vals = tl.load(ptr + base + d_offsets)

为什么有效

Ascend 硬件对非连续访问有显著性能惩罚。重构后连续加载减少 gather 操作，性能差距可达数倍至数十倍。

优化方法 3：Normalization 两阶段决策

适用条件

LayerNorm / RMSNorm / GroupNorm 等需要先统计再归一化的算子。

结论：两阶段（2-pass）优于单 Pass