融合策略与硬件映射：如何根据硬件特性选择融合策略

做算子融合这么多年，我最大的体会就是：没有银弹。同一个融合策略，在A卡上跑得飞起，换到B卡上可能反而变慢。说白了，融合策略的选择，本质上是在跟硬件的“脾气”打交道。

我个人习惯把硬件映射分成三个维度来看：计算密度、内存带宽、寄存器与共享内存。这三个东西互相牵制，你压了这头，那头可能就翘起来了。

计算密度 vs 内存带宽：融合的“跷跷板”

先问一个问题：什么时候该融合？什么时候不该？

我遇到过不少团队，一上来就把所有小算子全揉在一起，结果性能反而下降。为什么？因为融合的本质是用计算换访存。你把两个kernel合并成一个，减少了全局内存的读写，但中间结果可能要从寄存器或共享内存里走，计算量也可能翻倍。

这里有个关键指标叫计算密度（Arithmetic Intensity），就是“每次访存能做多少次计算”。

场景	计算密度	融合建议
Element-wise 操作链（ReLU + Add + Mul）	极低（< 1 FLOP/byte）	强烈推荐融合，带宽瓶颈明显
卷积 + 激活函数	中等（1-10 FLOP/byte）	视硬件带宽而定，HBM2/3 上收益大
大矩阵乘法 + 后续操作	高（> 10 FLOP/byte）	谨慎融合，可能增加寄存器压力

我的经验法则：当计算密度低于 2 FLOP/byte 时，融合几乎总是赢家。高于 10 时，先做 profiling，别盲目融合。

我记得有一次优化一个 Transformer 的 LayerNorm + Add + Residual 链。原始实现是三个分开的 kernel，每次都要读写一遍全局内存。融合之后，中间结果全留在寄存器里，延迟直接降了 40%。但同样的策略用在卷积 + BatchNorm 上，因为 BatchNorm 需要统计全局的均值和方差，融合反而让代码变得复杂，收益也不明显。

寄存器压力：那个看不见的瓶颈

嗯，这里要重点说一下寄存器压力。很多做融合的人容易忽略这个。

GPU 的寄存器是有限的。每个 SM 就那么几十 KB，分给每个线程后，你能用的寄存器数量是固定的。你融合的算子越多，中间变量就越多，寄存器占用就越高。

一旦寄存器溢出（spill），数据就会被踢到本地内存（Local Memory），而本地内存其实是在全局内存里划了一块。这一 spill，性能直接崩盘。

我曾经踩过的坑：有一次融合了 5 个 element-wise 操作，看起来很美。结果编译后发现寄存器占用从 32 飙到了 64，occupancy 从 50% 掉到了 25%。性能反而比分开跑还慢。后来我拆成了两组融合，每组 2-3 个操作，occupancy 保住了，性能才回来。

怎么判断寄存器压力？我一般看两个指标：

编译报告中的 register count：超过硬件阈值就要警惕
occupancy 变化：如果融合后 occupancy 下降超过 20%，大概率是寄存器溢出

Shared Memory 优化：融合的“秘密武器”

Shared Memory 是 GPU 上的一块可编程缓存，速度比全局内存快两个数量级。融合策略里，Shared Memory 用得好的话，能解决很多寄存器放不下的问题。

我个人的做法是：把融合后的中间结果，按 tile 分块放到 Shared Memory 里。这样每个线程块内的线程可以共享数据，不用反复读全局内存。

举个例子，假设你要融合一个卷积 + 一个池化操作。传统做法是卷积写全局内存，池化再读。用 Shared Memory 的话，卷积的输出直接写到 Shared Memory，池化线程直接从 Shared Memory 里读，省掉了一次全局内存的读写。

// 伪代码：融合卷积 + 池化
__global__ void fused_conv_pool(...) {
    __shared__ float smem[TILE_SIZE][TILE_SIZE];
    
    // 阶段1：卷积计算，结果写到 smem
    float conv_out = convolution_pixel(...);
    smem[tx][ty] = conv_out;
    __syncthreads();
    
    // 阶段2：池化直接从 smem 读
    if (tx % 2 == 0 && ty % 2 == 0) {
        float max_val = max(smem[tx][ty], smem[tx+1][ty],
                           smem[tx][ty+1], smem[tx+1][ty+1]);
        output[out_idx] = max_val;
    }
}

小技巧：Shared Memory 的大小是有限的（通常 48KB-164KB）。如果融合的算子太多，Shared Memory 不够用，可以考虑用“分阶段融合”——先融合一部分，写回全局内存，再融合下一部分。虽然多了一次全局内存读写，但比寄存器溢出强。

硬件特性决定融合策略

不同的硬件，融合策略的侧重点完全不同。我整理了一张对照表：

硬件特性	典型设备	融合策略倾向
高带宽（HBM2e/HBM3）	NVIDIA A100/H100	计算密集型融合，带宽不是瓶颈
低带宽（GDDR6）	消费级显卡	优先融合访存密集型操作
大 Shared Memory（>96KB）	AMD MI250	大胆使用 Shared Memory 做中间缓存
小寄存器文件（256KB/SM）	嵌入式 GPU	控制融合深度，避免寄存器溢出

说白了，你得先搞清楚你的目标硬件“缺什么”。缺带宽，就融合访存；缺计算，就别乱加计算量；缺寄存器，就拆小融合粒度。

一个实战案例：MobileNet 的融合优化

我记得优化 MobileNetV2 时，它的 bottleneck 结构里有一个 Depthwise Conv + Pointwise Conv + ReLU6 的链条。原始实现是三个 kernel，每次都要读写特征图。

我当时的做法是：

先分析计算密度：Depthwise Conv 的计算密度很低（约 0.5 FLOP/byte），Pointwise Conv 稍高（约 4 FLOP/byte）。整体链条属于访存瓶颈。
决定融合策略：把 Depthwise Conv 和 ReLU6 融合，中间结果留在寄存器里。Pointwise Conv 单独保留，因为它的计算量较大，融合后寄存器压力会超标。
Shared Memory 辅助：在 Pointwise Conv 里，用 Shared Memory 做输入数据的 tile 缓存，减少全局内存的重复读取。

最终效果：延迟降低了 25%，occupancy 保持在 40% 以上。如果强行三合一，occupancy 掉到 20%，性能反而倒退。

核心结论：融合不是越多越好，而是“够用就好”。硬件特性决定了融合的上限，寄存器压力和 Shared Memory 大小决定了融合的边界。

下面这张图总结了本章的核心逻辑：从硬件特性出发，经过计算密度评估和资源约束检查，最终确定融合策略。

最后说一句：融合策略的选择，永远是在“减少访存”和“控制资源占用”之间找平衡。没有哪个策略是万能的，但只要你把硬件的脾气摸透了，融合就能成为你最趁手的优化工具。

📚 课程目录

融合策略与硬件映射：如何根据硬件特性选择融合策略

计算密度 vs 内存带宽：融合的“跷跷板”

寄存器压力：那个看不见的瓶颈

Shared Memory 优化：融合的“秘密武器”

硬件特性决定融合策略

一个实战案例：MobileNet 的融合优化