CANN算子融合原理与实践：为什么融合能让延迟打三折

有个团队做过一个实验：同一个ResNet-50模型，分别用融合前和融合后的算子跑。融合前跑出来45ms，融合后14ms——快了3.2倍。他们问我为什么会快这么多。我说，本质上就是省掉了HBM读写中间结果的次数。

这个概念听起来简单，但做起来有很多细节。融合不只是“把两个算子拼在一起”这么简单——你要处理内存布局兼容、tiling策略协调、梯度反向传播等等问题。这篇把算子融合的原理、约束和实践讲透。

融合为什么快：HBM访问次数的本质

一个普通卷积层（Conv → BatchNorm → ReLU）的执行过程：

输入 → Conv算子 → 写HBM(中间结果) → BatchNorm算子 → 写HBM(中间结果) → ReLU算子 → 写HBM(输出)

总HBM访问次数：
  读输入: 1次
  Conv写中间: 1次
  BatchNorm读中间: 1次
  BatchNorm写中间: 1次
  ReLU读中间: 1次
  ReLU写输出: 1次
  ─────────────────
  总计: 6次HBM读写

融合成 FusedConvBnRelu 之后：

输入 → FusedConvBnRelu算子 → 写HBM(输出)

总HBM访问次数：
  读输入: 1次
  写输出: 1次
  ─────────────────
  总计: 2次HBM读写

省了4次HBM访问，理论上快3倍。实测快2.5倍（因为融合后的kernel内部还是有寄存器/UB的读写）。

融合的数学基础：为什么有些算子能融合，有些不能

不是所有相邻的算子都能融合。能融合的条件：前一个算子的输出恰好是后一个算子的输入，不需要额外处理。

# 可以融合：数学上等价，内存布局兼容
Conv(x) + BatchNorm(Conv(x)) + ReLU(BatchNorm(Conv(x)))
= FusedConvBnRelu(x)
# 融合条件：
#   1. BatchNorm的参数可以吸收到Conv的weight/bias里（数学推导可得）
#   2. ReLU不改变tensor的shape和dtype
#   3. 三个算子都在同一个计算核上支持

# 不能融合：布局不兼容
Conv(x) → NC1HWC0格式
Reshape(x) → NHWC格式  ← 格式变了
Conv(Reshape(x)) → 需要格式转换，不能直接融合

# 不能融合：数据依赖
Conv(x) + Concat([Conv(x), Pool(x)])  ← Pool(x)是独立分支
Concat([Conv(x), Pool(x)])  ← Conv和Pool没有依赖，可以各自独立融合

BatchNorm为什么能被融合：BatchNorm的本质是 $y=\frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta$。这可以转化成 $y=x\cdot \alpha +{\beta }^{\prime }$，其中 $\alpha =\gamma /\sqrt{{\sigma }^2+ϵ}$ 和 ${\beta }^{\prime }=\beta -\alpha \cdot \mu$ 都是常数。所以BatchNorm就是一个乘加操作，可以直接写进Conv的weight和bias里。

融合的三种类型

CANN里的算子融合分三种类型，收益不同：

• 类型1：垂直融合（Vertical Fusion）—— 收益最大
  沿计算图纵向融合，前后依赖的算子合并成一个。

  # Conv → BN → ReLU → Conv → BN → ReLU
  # 融合成：
  #   FusedConvBnRelu → FusedConvBnRelu
  #
  # 收益：省掉中间结果的HBM读写
  # 典型收益：1.5x ~ 2.5x
  

• 类型2：水平融合（Horizontal Fusion）—— 收益中等
  沿计算图横向融合，并行的算子合并成一个。

  # Conv1(x) ─┬─→ Concat
  # Conv2(x) ─┘
  #
  # 融合成：
  #   FusedConv1Conv2Concat(x)
  #
  # 收益：一次kernel launch替代两次，减少kernel调度开销
  # 典型收益：1.1x ~ 1.3x
  

• 类型3：替换融合（Replacement Fusion）—— 收益不确定
  用更优的算子替换原有算子组合。

  # x @ W1 @ W2
  # 替换成：
  #   x @ FusedLinear(W1, W2)  ← 一次MatMul替代两次
  #
  # 收益：取决于具体shape
  # 典型收益：1.2x ~ 1.8x
  

CANN的融合规则引擎

CANN内置了一套融合规则引擎（opscene的核心），自动识别可融合的算子模式。

# 融合规则以 JSON 的形式配置
# 文件路径：$CAN_PATH/opp/op_impl/built-in/ops/fusion_rule/
# 可以查看/修改这些规则

# 示例：conv_bn_relu 的融合规则
{
    "name": "ConvBatchNormReLU",
    "pattern": [
        {
            "op_type": "Conv",
            "next": {
                "op_type": "BatchNorm",
                "next": {
                    "op_type": "ReLU"
                }
            }
        }
    ],
    "fused_op": "FusedConvBnRelu",
    "conditions": [
        "input.shape[1] == bn.num_features",
        "conv.stride == (1, 1) or conv.stride == (1, 1, 1)",
        "conv.padding in ['same', (0,0)]"
    ]
}

# 查看当前生效的融合规则
from cann import FusionRuleInspector

inspector = FusionRuleInspector()

# 列出所有融合规则
rules = inspector.list_rules()
print(f"共有 {len(rules)} 条融合规则：\n")

for rule in rules[:10]:  # 打印前10条
    print(f"  {rule.name}")
    print(f"    模式: {rule.pattern}")
    print(f"    融合后: {rule.fused_op}")
    print()

# 统计融合情况
stats = inspector.get_fusion_stats(model_path="resnet50.om")
print(f"融合前算子数: {stats.original_op_count}")
print(f"融合后算子数: {stats.fused_op_count}")
print(f"融合减少: {stats.original_op_count - stats.fused_op_count} 个算子")
print(f"融合率: {(1 - stats.fused_op_count/stats.original_op_count)*100:.1f}%")

手动指定融合策略

自动融合有时候会失败（某些边界情况判断不了），这时候可以手动指定。

import torch
from cann import FusionConfig

# 配置融合策略
config = FusionConfig(
    # 强制开启某些融合
    enable_fusions=[
        "ConvBnRelu",
        "ConvAdd",
        "MatMulBias",
        "LayerNormProjection"  # 自定义融合
    ],
    # 禁用某些融合（debug用）
    disable_fusions=[
        "FusedAttentionMask"  # 禁用这个融合，看看是不是它的bug
    ]
)

# 编译模型时传入配置
from cann import ATCCompiler

compiler = ATCCompiler(
    model_path="resnet50.onnx",
    output_path="resnet50_fused.om",
    fusion_config=config  # ← 传入融合配置
)

# 检查某个融合是否生效
inspector = FusionRuleInspector()
result = inspector.check_fusion(model_path="resnet50_fused.om")

for fusion in result.fusions_applied:
    print(f"✅ {fusion.name}: 融合成功")
    
for fusion in result.fusions_failed:
    print(f"❌ {fusion.name}: 融合失败")
    print(f"   原因: {fusion.failure_reason}")

自定义融合：把多个算子绑在一起

如果内置的融合规则不够用，可以写自定义融合。这需要你用Ascend C写一个融合算子，然后注册到融合规则引擎。

// 自定义融合算子：Conv + BatchNorm + ReLU + AddResidual
// 文件：my_fusion_op.cpp

extern "C" __global__ __aicore__ void fused_conv_bn_relu_add(
    half* output,
    const half* input,
    const half* conv_weight,
    const half* bn_gamma,
    const half* bn_beta,
    const half* bn_mean,
    const half* bn_var,
    const half* residual_input,
    const float epsilon = 1e-5
) {
    // 1. 读取输入
    // 2. Conv 计算
    // 3. BatchNorm 计算（融合进 Conv 的计算里）
    // 4. ReLU 计算
    // 5. 加上残差 (Add)
    // 6. 写输出
    
    // 一次性完成 4 个操作，HBM 只需要 2 次访问
}

# 注册自定义融合规则
from cann import FusionRuleRegistry

registry = FusionRuleRegistry()

# 添加自定义融合规则
registry.add_rule({
    "name": "ConvBnReluAddResidual",
    "pattern": [
        {"op_type": "Conv"},
        {"op_type": "BatchNorm"},
        {"op_type": "ReLU"},
        {"op_type": "Add"},  # 残差连接
    ],
    "fused_op": "FusedConvBnReluAddResidual",
    "kernel_path": "/path/to/my_fusion_op.cpp",  # Ascend C 实现
    "conditions": [
        "conv.output_channels == bn.num_features",
        "residual_input.shape == output.shape"
    ]
})

print("自定义融合规则注册成功！")

融合的性能边界：什么时候融合反而变慢

融合不是万能的。某些情况下，融合会变慢：

# 情况1：融合后kernel太大，放不进Unified Buffer
# 不好：把8个算子融合成一个
Conv + BN + ReLU + Conv + BN + ReLU + Conv + BN
→ FusedConv8  ← UB放不下，只能分块执行，开销巨大

# 好：分成2组融合
[Conv + BN + ReLU] → [Conv + BN + ReLU] → [Conv + BN]
→ FusedConvBnRelu × 3  ← 每个UB都能放得下

# 情况2：融合后kernel的tiling选择变差
# 小tensor融合后，Cube利用率反而下降

# 情况3：有条件分支的算子融合
# if x > 0: y = relu(x) else: y = x
# 这个没法融合成简单的数学公式

# 判断方法：用profiler对比融合前后的性能
from cann_colt_profiler import compare_fusion

result = compare_fusion(
    model_path="resnet50.onnx",
    fusion_name="ConvBnRelu",
    before=True,   # 融合前
    after=False    # 融合后
)

print(f"融合前延迟: {result.before_latency_ms:.3f}ms")
print(f"融合后延迟: {result.after_latency_ms:.3f}ms")
print(f"加速比: {result.before_latency_ms / result.after_latency_ms:.2f}x")

梯度反向传播中的融合问题

训练阶段，融合的另一个问题是反向传播。融合的前向算子，它的梯度也需要对应融合。

# 前向：Conv + BN + ReLU = FusedConvBnRelu
# 反向：dInput = FusedConvBnRelu_grad(dOutput)
#
# 问题：FusedConvBnRelu 的梯度算子必须跟前向算子完全对应
# 如果你融合了前向但没融合反向，或者融合了错误的反向，
# 会导致梯度计算错误

# 验证梯度正确性（融合前后应该一致）
import torch

model_fused = FusedConvBnRelu()
model_unfused = nn.Sequential(Conv(), BatchNorm(), ReLU())

x = torch.randn(1, 64, 56, 56, requires_grad=True).npu()

# 前向对比
y1 = model_fused(x)
y2 = model_unfused(x)
print(f"前向输出差异: {(y1 - y2).abs().max().item():.6f}")

# 反向对比
y1.backward(gradient=torch.ones_like(y1))
x_grad_fused = x.grad.clone()
model_unfused.zero_grad()
x.grad.zero_()
y2.backward(gradient=torch.ones_like(y2))
x_grad_unfused = x.grad

print(f"反向梯度差异: {(x_grad_fused - x_grad_unfused).abs().max().item():.6f}")

# 如果梯度差异 > 1e-5，说明融合的反向算子写错了

融合的调试清单

# 当融合没有按预期工作时，按这个清单排查

# 1. 确认融合规则是否生效
inspector = FusionRuleInspector()
result = inspector.check_fusion(model_path="model.om")

print("已应用的融合：")
for fusion in result.fusions_applied:
    print(f"  ✅ {fusion.name}")
    
print("\n融合失败：")
for fusion in result.fusions_failed:
    print(f"  ❌ {fusion.name}: {fusion.failure_reason}")

# 2. 确认数据格式兼容
# 融合要求输入输出格式一致
fmt_in = get_tensor_format(tensors[0])
fmt_out = get_tensor_format(tensors[-1])
if fmt_in != fmt_out:
    print(f"⚠ 格式不兼容: {fmt_in} → {fmt_out}")
    print("  需要在融合前后加 FormatConvert")

# 3. 确认shape兼容
# 融合要求输入输出shape对应
print(f"输入shape: {input.shape}")
print(f"输出shape: {output.shape}")
# 如果中间有reshape/slice，融合会失败

# 4. 确认tiling能放下
ub_size = 2 * 1024 * 1024  # 2MB
estimated_ub_usage = estimate_fused_op_ub(model)
print(f"UB使用量: {estimated_ub_usage/1024:.1f}KB / 2048KB")
if estimated_ub_usage > ub_size:
    print("⚠ UB放不下，需要分块或减少融合层数")