19. Operator Fusion Introduction | 算子融合导论

难度： Medium | 环境： CPU-first | 标签： 编译系统, 算子融合, 性能优化 | 目标人群： 编译与系统入门者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

(国内推荐：魔搭社区免费实例)

---

本节导读

矩阵乘之后再加偏置、执行激活时，分开的计算内核可能反复写出并读取中间结果。融合把其中一些步骤放在同一计算内核中，让中间值直接被下一步消费，同时也可能增加寄存器需求和布局转换。

本节比较融合前后的中间读写与启动次数，再检查语义、形状和片上资源条件。完成后，你应能算出省掉了哪些读写，并指出还需测量的代价；真实融合计算内核的实现与计时进入算子优化专题。

关键词： 算子融合、中间张量、内存流量、布局转换

---

前置阅读

导语： 先理解图级优化对象，再检查融合怎样改变中间张量和片上资源。

• 09. AI Compilers and Graph Optimization | AI 编译器与计算图优化
• 18. Triton Block Model | Triton Block 模型

Q1：中间张量怎样增加逻辑读写？

串行算子链的每个组间边界都可能写出中间张量，再由下一组读取。在边界张量等大、每组对应一个核函数的假设下，可以用边界数量计算读写字节和启动次数。融合省去的是内部边界，不包含模型输入和最终输出；逻辑读写减少也不等于实际 HBM 流量或耗时按相同比例下降。

Q1 机制计算器：比较分组边界的读写字节

def fusion_traffic(num_ops, tensor_bytes, fused_groups):
    """估算算子组边界造成的中间张量写回和重读字节数。"""
    if (type(num_ops) is not int or type(fused_groups) is not int
            or num_ops < 1 or not 1 <= fused_groups <= num_ops):
        raise ValueError('需要满足 1 <= fused_groups <= num_ops')
    if type(tensor_bytes) is not int or tensor_bytes <= 0:
        raise ValueError('tensor_bytes 必须大于 0')
    materialized_boundaries = fused_groups - 1
    intermediate_bytes = 2 * materialized_boundaries * tensor_bytes
    return {
        'launches': fused_groups,
        'materialized_boundaries': materialized_boundaries,
        'intermediate_bytes': intermediate_bytes,
    }

tensor_bytes = 64 * 1024**2
unfused = fusion_traffic(num_ops=4, tensor_bytes=tensor_bytes, fused_groups=4)
partially_fused = fusion_traffic(num_ops=4, tensor_bytes=tensor_bytes, fused_groups=2)
fully_fused = fusion_traffic(num_ops=4, tensor_bytes=tensor_bytes, fused_groups=1)
print({'unfused': unfused, 'partial': partially_fused, 'full': fully_fused})

assert unfused['intermediate_bytes'] == 6 * tensor_bytes
assert fully_fused['intermediate_bytes'] == 0

Q2：算子融合为什么能改善数据局部性？

融合把原本分开的多个步骤合成一条更连续的执行链。

这样做的结果是：

• 中间结果更容易留在片上；
• 同一批数据更容易被连续复用；
• 数据搬运和计算之间的间隔更短。

数据局部性提升后，硬件就更容易把已有数据重复利用起来，而不是每一步都重新去 HBM 找一遍。这就是为什么融合不只是“少几个函数”，而是直接改变了数据流路径。

Q2 机制计算器：检查片上预算和逻辑复用量

假设数据块可在整个消费阶段驻留片上，计算重复读取减少量。容量检查不包含其他临时值、分配粒度与同步；这只是局部预算，不是实际融合成功的证明。

def onchip_reuse(tile_bytes, reuse_count, onchip_budget_bytes):
    """比较给定局部预算，并估算完整驻留假设下减少的逻辑读取量。"""
    if any(type(value) is not int or value <= 0
           for value in (tile_bytes, reuse_count, onchip_budget_bytes)):
        raise ValueError('tile、复用次数和片上预算必须为正数')
    fits = tile_bytes <= onchip_budget_bytes
    avoided_read_bytes = tile_bytes * (reuse_count - 1) if fits else 0
    return {
        'fits_onchip': fits,
        'avoided_read_bytes': avoided_read_bytes,
        'budget_utilization': tile_bytes / onchip_budget_bytes,
    }

plans = {
    'small_tile': onchip_reuse(32 * 1024, 3, 96 * 1024),
    'large_tile': onchip_reuse(128 * 1024, 3, 96 * 1024),
}
for name, result in plans.items():
    print(name, '->', result)

assert plans['small_tile']['avoided_read_bytes'] == 64 * 1024
assert plans['large_tile']['avoided_read_bytes'] == 0

Q3：哪些条件要先核对，哪些代价需要实测？

融合前先核对语义、形状、片上资源与后端支持，再分别记录布局转换、编译和执行成本。额外布局转换可能抵消一部分收益，但不是不可执行的充分理由。

下面的布尔值表示给定的核对结果，并不替代数值对照、编译或运行。只要前提满足，就保留布局字节数供后续比较；不因它大于零直接排除候选。

Q3 对照验证：布局转换有代价，但不自动阻止融合

def fusion_decision(semantics_match, shape_supported, onchip_fit,
                    backend_lowered, reorder_bytes=0):
    """分开记录给定检查结果与布局成本，不执行编译或数值验证。"""
    checks = {
        'semantics_match': semantics_match,
        'shape_supported': shape_supported,
        'onchip_fit': onchip_fit,
        'backend_lowered': backend_lowered,
    }
    if any(type(value) is not bool for value in checks.values()):
        raise ValueError('检查结果必须是布尔值')
    if type(reorder_bytes) is not int or reorder_bytes < 0:
        raise ValueError('布局转换字节数必须是非负整数')
    failures = [name for name, passed in checks.items() if not passed]
    return {
        'stated_conditions_met': not failures,
        'failures': failures,
        'extra_reorder_bytes': reorder_bytes,
    }


cases = {
    'ready': (True, True, True, True, 0),
    'shape_miss': (True, False, True, True, 0),
    'reorder_cost': (True, True, True, True, 8 * 1024**2),
}
for name, config in cases.items():
    print(name, '->', fusion_decision(*config))
assert fusion_decision(*cases['ready'])['stated_conditions_met']
assert fusion_decision(*cases['reorder_cost'])['stated_conditions_met']
assert fusion_decision(*cases['shape_miss'])['failures'] == ['shape_supported']

相关阅读

导语： 把算子融合继续接到性能调优、归约和 attention kernel，会更直观地看到中间张量是怎么被省掉的。

• 05. Triton 性能调优与基准测试 (Autotune & Profiling)
• 06. Triton 进阶：跨线程归约与数值稳定 (Safe Softmax)
• 08. Triton Flash Attention | 真正的 Flash Attention 前向算子

---