陷阱四：控制流与数据流的“冲突”——动态shape下的融合噩梦

说实话，做AI芯片这几年，让我半夜被电话吵醒的bug，十有八九跟动态shape有关。

静态shape时代，算子融合就像搭积木。每个算子的输入输出尺寸都定死了，编译器可以提前规划好一切。但动态shape一来，整个局面就变了——控制流和数据流开始打架，融合策略动不动就崩。

我有个项目，客户要求模型支持可变分辨率输入。一开始觉得没什么，不就是把batch_size和H, W设成变量嘛。结果一跑起来，性能直接腰斩。查了三天，发现是融合后的算子内部，数据依赖关系在运行时全乱了。

4.1 动态shape到底难在哪？

先看一个最简单的例子。假设我们要融合两个算子：

// 伪代码：静态shape融合
// 输入尺寸固定为 [1, 3, 224, 224]
Tensor x = input(1, 3, 224, 224);
Tensor y = relu(x);      // shape不变
Tensor z = conv(y);      // shape已知，可提前分配内存
// 编译器可以放心地把relu和conv合并成一个kernel

静态shape下，编译器能精确知道每一步的内存布局。融合后，relu的输出可以直接作为conv的输入，中间结果甚至不需要写回显存。

但动态shape一来：

// 动态shape：输入尺寸运行时才确定
Tensor x = input(batch, 3, H, W);  // batch, H, W都是变量
Tensor y = relu(x);                // shape = [batch, 3, H, W]
Tensor z = conv(y);                // conv的输出shape依赖于H, W
// 问题：融合后，中间内存该分配多大？

嗯，这里要注意。动态shape下，编译器在编译期根本不知道H和W是多少。它只能生成一段代码，在运行时根据实际输入去计算shape、分配内存。

这就引出了第一个冲突点：控制流（运行时shape计算）和数据流（算子执行）被强行耦合在了一起。

4.2 融合后的“幽灵依赖”

我个人习惯把动态shape下的融合问题分成两类：

1. 显式依赖：算子的输出shape直接依赖于输入shape，比如reshape、slice
2. 隐式依赖：算子的执行路径依赖于数据内容，比如条件判断、循环

显式依赖还好处理，大不了在运行时多算一步shape。真正要命的是隐式依赖。

举个例子，一个动态shape模型里常见的模式：

// 动态shape下的条件分支
if (x.shape[0] > 32) {
    y = big_kernel(x);   // 大batch用大kernel
} else {
    y = small_kernel(x); // 小batch用小kernel
}
z = relu(y);

你想把big_kernel或small_kernel和后面的relu融合。但问题是——编译器在编译期不知道走哪条分支。

你可能会说：「那就两条分支都融合，运行时根据条件选一个呗。」

听起来简单，但实际做起来坑很多。我在项目中遇到过这种情况：两条分支融合后的kernel，对内存的需求完全不同。大分支需要16MB的共享内存，小分支只需要2MB。你按大的分配，小分支浪费；按小的分配，大分支直接崩。

4.3 实战避坑：动态shape融合的三种策略

踩了这么多坑，我总结了三套打法。根据不同场景选：

策略一：保守融合

说白了就是「不融合」。对于动态shape变化剧烈的场景，干脆放弃融合，让每个算子独立执行。虽然性能差一些，但至少不会出运行时错误。

我曾经在一个边缘设备项目上就这么干过。模型输入从240p到4K随机变化，融合后的kernel在低分辨率下性能提升明显，但一到高分辨率就内存溢出。最后我选择了只融合那些shape绝对不变的算子，其他的保持原样。

策略二：自适应融合

在融合后的kernel内部，加入运行时shape检查逻辑。根据实际shape动态调整内存分配和计算策略。

// 自适应融合kernel伪代码
__global__ void fused_kernel(Tensor input, Tensor output) {
    int H = input.shape[2];
    int W = input.shape[3];
    
    // 根据实际shape动态选择计算路径
    if (H * W < 1024 * 1024) {
        // 小尺寸：用寄存器缓存
        compute_small(input, output);
    } else {
        // 大尺寸：用全局内存
        compute_large(input, output);
    }
}

这种策略的难点在于，你要为每一种可能的shape组合都准备好计算路径。我建议只覆盖90%的常见情况，剩下的10%走fallback路径。

策略三：JIT编译融合

这是最彻底的办法。在运行时，拿到实际的shape后，再动态生成并编译融合后的kernel。这样既能保证最优性能，又能处理任意shape。

但代价也很明显——JIT编译本身有开销。我记得有一次，模型每帧输入shape都不同，JIT编译的时间比推理时间还长。后来加了个kernel缓存，才把问题解决。

4.4 控制流与数据流的解耦设计

说到底，动态shape下融合的根因，是控制流和数据流搅在了一起。我个人的做法是：在编译期把控制流和数据流彻底分开。

具体来说：

• 控制流部分（shape计算、条件判断）单独提取出来，生成一段「调度代码」
• 数据流部分（实际的矩阵运算）保持纯数据驱动，不依赖任何运行时变量
• 调度代码在CPU上执行，数据流在AI芯片上执行

这样做的好处是，融合后的kernel可以完全专注于计算，不需要处理动态shape的逻辑。shape变化时，只需要重新生成调度代码，数据流kernel可以复用。

下面是我画的一张示意图，展示这种解耦设计：

4.5 避坑指南

最后，分享几个我亲身踩过的坑，希望能帮你少走弯路：

从那以后，我给自己定了三条铁律：

1. 动态shape的算子，融合前必须做边界测试——用最小、最大、以及中间值各跑一遍
2. 融合后的kernel，内部不要依赖任何「隐式假设」——比如不要假设H是偶数、W是32的倍数
3. 保留fallback路径——一旦融合kernel出错，能自动切换到非融合版本

你想想看，动态shape本身已经够复杂了，再加上融合，简直就是把两个难题叠在一起。我的建议是：能不融合就不融合，非要融合就做最保守的方案。性能差一点可以接受，但运行时崩溃是不能接受的。

嗯，关于动态shape下的融合陷阱，就先聊到这里。记住一句话：控制流和数据流，该分开时就分开，别硬融。