静态形状与动态形状

机器学习编译中的一个主要区别在于张量的维度是在编译时已知还是在运行时才知晓。这个区别决定了内存管理策略以及编译器能对生成代码施加的具体优化。

静态形状的优点

在静态形状体系下，计算图中的每个张量的每个维度都是一个固定的整数，在程序运行前就已得知。这在计算机视觉模型中很常见。例如，标准的ResNet50实现通常需要精确的 $(1, 3, 224, 224)$ 输入。

当编译器遇到静态形状时，它拥有完整的信息。它可以计算整个推理过程所需的精确内存占用。编译器无需在执行期间动态分配和释放内存，而是可以创建静态内存规划。它在一个单一连续的内存块中预先计算每个张量的偏移量，该内存块常被称为内存竞技场。

此外，静态形状使得激进的循环优化成为可能。如果一个循环迭代一个大小为64的维度，编译器可以展开该循环，或将其精确映射到宽度为8或16的向量单元。它不需要插入边界检查或处理数据大小与硬件向量宽度不匹配情况下的剩余循环。

编译流程比较。静态形状允许直接生成具有固定边界的代码，而动态形状则需要中间的形状解析步骤。

动态形状的灵活性

实际应用常常不遵循固定维度。自然语言处理模型处理长度不同的句子。目标检测模型根据图像内容输出可变数量的边界框。在这些情况下，一个或多个维度是符号化的。

A动态形状在IR中不表示为像 $64$ 这样的字面整数，而是表示为一个变量，通常记作 $N$、$M$ 或 $?$。当编译器处理动态形状时，它无法预先计算精确的内存偏移量。相反，它必须生成在运行时执行“形状推断”的代码。

考虑张量 $A$ (形状为 $(M, K)$) 与张量 $B$ (形状为 $(K, N)$) 之间的矩阵乘法。在动态设置中，编译器生成以下指令：

1. 读取 $M$、$K$ 和 $N$ 的运行时值。
2. 验证内部维度是否匹配（运行时断言）。
3. 计算输出大小 $(M, N)$。
4. 根据此计算为结果分配内存。

这种灵活性也伴随着代价。生成的二进制文件必须包含管理这些形状的额外逻辑。此外，算术核心（kernel）变得通用。为“大小 $N$”编译的核心通常不如为“大小 1024”编译的核心高效，因为编译器无法硬编码优化常数或假定数据对齐。

符号形状传播

为了弥合高级框架与低级代码之间的鸿沟，机器学习编译器使用符号形状传播。即使确切值未知，形状之间的关系也是确定性的。

如果一个输入张量的形状是 $(Batch, 128)$，并通过一个有64个单元的密集层，输出形状将是 $(Batch, 64)$。编译器在图中跟踪符号 $Batch$。如果后续操作尝试将此张量重塑为 $(Batch, 32, 2)$，编译器可以静态验证这是有效的，因为 $32 \times 2 = 64$。然而，如果操作尝试将其重塑为 $(Batch, 30, 2)$，编译器可以在编译时抛出错误，即使不知道 $Batch$ 的值。

符号维度在计算图中的流向。编译器跟踪变量 'N' 以确保图的有效性，而无需知晓其运行时值。

混合方法与桶化

由于动态形状会抑制优化，工程师在部署模型时常采用一种称为桶化（或填充）的混合策略。系统不支持任意输入大小，而是支持一组离散的“桶”。

例如，一个服务系统可能会为序列长度为32、64、128和256编译特定的核心。如果用户发送一个长度为50的请求，运行时会将数据填充到长度64，并执行针对64优化的核心。这种方法以少量计算（处理填充）换取静态编译核心的高效率。

这种权衡在性能特点中可见。静态编译能带来峰值性能，但每遇到新形状都需要重新编译。动态编译能处理所有情况，但会带有一个固定的开销。

跨序列长度的延迟比较。静态优化点（红色）通常在特定设计点上提供比通用动态执行（蓝线）更低的延迟。

即时 (JIT) 专门化

现代机器学习编译器常使用JIT编译来处理动态形状，同时保留优化优势。当模型接收到具有特定形状的输入（例如，批大小为1）时，JIT编译器会即时为该形状生成一个专门的核心并进行缓存。如果下一个输入具有相同的形状，则重用缓存的核心。如果形状明显改变，则会编译一个新的核心。

此方法假设输入形状的分布不是随机的。在许多生产环境中，输入大小遵循幂律分布，这意味着少数独特的形状占据了大部分流量。这使得编译器可以针对常见情况进行专门化，同时为稀有形状回退到通用的动态核心。