使用 torch.compile 加速 PyTorch Inductor 的 CPU 推理

作者：英特尔

内容提要

• 尽管 PyTorch* Inductor C++/OpenMP* 后端使用户能够利用现代 CPU 架构和并行处理，但它缺乏优化，导致该后端在端到端性能方面表现不如 eager 模式。
• 英特尔采用混合策略优化了 Inductor 后端，该策略将操作分为两类：Conv/GEMM 以及非 Conv/GEMM 的逐元素和归约操作。
• 对于流行的深度学习模型，这种混合策略与 eager 模式相比表现出可喜的性能提升，并提高了 C++/OpenMP 后端对 PyTorch 模型的效率和可靠性。

Inductor 后端面临的挑战

PyTorch Inductor C++/OpenMP 后端使用户能够利用现代 CPU 架构和并行处理来加速计算。

然而，在其开发的早期阶段，该后端缺乏一些优化，使其无法充分利用 CPU 计算能力。结果是，对于大多数模型，C++/OpenMP 后端在端到端性能方面表现不如 eager 模式，其中 45% 的 TorchBench 模型、100% 的 Hugging Face 模型和 75% 的 TIMM 模型表现不如 eager 模式。

在本文中，我们将重点介绍英特尔对 Inductor CPU 后端的优化，包括相关技术和结果。

我们采用混合策略优化了该后端，该策略将操作分为两类：Conv/GEMM 以及非 Conv/GEMM 的逐元素和归约操作。前者利用 oneDNN 性能库进行后置操作融合和权重预打包进行优化，而后者则在 C++ 代码生成中采用显式向量化进行优化。

这种混合策略与 eager 模式相比表现出可喜的性能提升，特别是在 Inductor Hugging Face、Inductor TorchBench 和 Inductor TIMM 等流行的深度学习模型上。总的来说，英特尔的优化提高了 C++/OpenMP 后端对 PyTorch 模型的效率和可靠性。

图 1：性能加速比趋势

英特尔混合优化性能状态

与采用混合优化的 eager 模式相比，C++/OpenMP 后端显示出可喜的性能提升。我们测量了 TorchBench、Hugging Face 和 TIMM 这三个 Inductor 基准测试套件的性能，结果如下。（注意：我们每周在 GitHub 上发布两次性能数据。）

总的来说，这些优化有助于确保 C++/OpenMP 后端为 PyTorch 模型提供高效可靠的支持。

通过率

+----------+------------+-------------+-------------+
| Compiler | torchbench | huggingface | timm_models |
+----------+------------+-------------+-------------+
| inductor | 93%, 56/60 | 96%, 44/46  | 100%, 61/61 |
+----------+------------+-------------+-------------+

几何平均加速比（单路多线程）

+----------+------------+-------------+-------------+
| Compiler | torchbench | huggingface | timm_models |
+----------+------------+-------------+-------------+
| inductor |   1.39x    |    1.20x    |    1.73x    |
+----------+------------+-------------+-------------+

单个模型性能

图 2：TorchBench FP32 性能（单路多线程）

图 3：Hugging Face FP32 性能（单路多线程）

图 4：TIMM FP32 性能（单路多线程）

几何平均加速比（单核单线程）

+----------+------------+-------------+-------------+
| Compiler | torchbench | huggingface | timm_models |
+----------+------------+-------------+-------------+
| inductor |    1.29x   |    1.15x    |    1.37x    |
+----------+------------+-------------+-------------+

图 5：TorchBench FP32 性能（单路单线程）

图 6：Hugging Face FP32 性能（单路单线程）

图 7：TIMM FP32 性能（单路单线程）

技术深入探讨

现在，让我们仔细看看 Inductor C++/OpenMP 后端中使用的两个主要优化：

1. 通过 oneDNN 库进行权重预打包和后置操作融合
2. 在 Inductor C++ 代码生成中进行显式向量化

通过 oneDNN 进行权重预打包和后置操作融合

oneDNN 库是 Intel® oneAPI 深度神经网络库的简称，它提供了一系列后置操作融合（即将卷积和矩阵乘与后续操作融合），这些融合可以使流行模型受益。 Intel® Extension for PyTorch 已经实现了这些融合的大部分，并取得了显著的性能提升。因此，我们将英特尔 PyTorch 扩展中应用的所有这些融合上游贡献到了 Inductor，使更广泛的模型能够从这些优化中受益。我们将这些融合定义为 mkldnn 命名空间下的运算符。这使得 Python 模块可以直接调用这些 mkldnn 操作。

当前定义的融合操作如下。您可以在 RegisterMkldnnOpContextClass.cpp 中找到这些定义的融合操作。

• _linear_pointwise：融合 Linear 及其后置的一元逐元素操作
• _linear_pointwise.binary：融合 Linear 及其后置的二元逐元素操作
• _convolution_pointwise：融合 Convolution 及其后置的一元逐元素操作
• _convolution_pointwise.binary：融合 Convolution 及其后置的二元逐元素操作

详细的融合模式定义在 mkldnn.py 文件中：convolution/linear + sigmoid/hardsigmoid/tanh/hardtanh/hardswish/leaky_relu/gelu/relu/relu6/siluconvolution/linear + add/add_/iadd/sub/sub_

在 Inductor 端，我们将这些融合应用于已降低的 FX 图。我们将 mkldnn_fuse_fx 定义为应用所有融合的入口点。相关代码片段如下

def mkldnn_fuse_fx(gm: torch.fx.GraphModule, example_inputs):
    ...
    gm = fuse_unary(gm)
    gm = fuse_binary(gm)
    ...
    if config.cpp.weight_prepack:
        gm = pack_module(gm)
    return gm

在 mkldnn_fuse_fx 函数中，我们对尚未降低的 FX 图应用融合。要融合卷积/线性及其后续的逐元素操作，我们按如下方式调用 fuse_unary 和 fuse_binary

   gm = fuse_unary(gm)
   gm = fuse_binary(gm)

除了后置操作融合外，我们还应用权重预打包来进一步提高 Conv/GEMM 性能

   gm = pack_module(gm)

权重预打包涉及将权重张量重新排列为分块布局，这

• 与 NCHW 或 NHWC 等普通格式相比，可以改善向量化和缓存复用；并且
• 可以帮助避免运行时权重重新排序，从而减少开销并提高性能；并且
• 作为权衡，会增加内存使用。

由于这些原因，我们在 Inductor 中提供了 config.cpp.weight_prepack 标志，以便为用户提供对该优化的更多控制，允许他们根据自己的具体需求启用它。

Inductor C++ 代码生成中的显式向量化

向量化是一种关键的优化技术，可以显著提高数值计算的性能。通过利用 SIMD（单指令多数据）指令，向量化可以在单个处理器核心上同时执行多个计算，从而显著提高性能。

在 Inductor C++/OpenMP 后端中，我们利用 aten 向量化库来促进实现，使用 Intel® AVX2 和 Intel® AVX-512 ISA（指令集架构）选项进行向量化。Aten 向量化支持多种平台，包括 x86 和 Arm，以及多种数据类型。通过添加更多 VecISA 子类，可以轻松扩展以支持其他 ISA。这使得 Inductor 在未来可以轻松支持其他平台和数据类型。

由于平台差异，Inductor 的 C++/OpenMP 后端在代码生成开始时会首先检测 CPU 特性以确定向量化位宽。默认情况下，如果机器同时支持 AVX-512 和 AVX2，后端将选择 512 位向量化。

如果硬件支持向量化，C++/OpenMP 后端首先会检测循环体是否可以向量化。主要有三种情况下我们无法生成带向量化的内核：

1. 循环体缺乏向量内在函数支持，例如 rand 和 atomic_add。
2. 循环体缺乏高效的向量内在函数支持，例如非连续的 load/store。
3. 向量化尚未支持但正在进行中的数据类型，例如整数、double、half 和 bfloat16。

为了解决这个问题，C++/OpenMP 后端使用 CppVecKernelChecker 来检测特定循环体中的所有操作是否可以向量化。一般来说，我们通过识别操作是否依赖于上下文将其分为两类。

对于大多数逐元素操作，例如 add、sub、relu，向量化是直接的，并且它们的执行不依赖于上下文。

然而，对于某些其他操作，它们的语义更复杂，并且它们的执行通过静态分析依赖于上下文。

例如，考虑 where 操作，它接受 mask、true_value 和 false_value 作为输入，其中 mask 值从 uint8 张量加载。fx 图可能如下所示：

graph():
    %ops : [#users=9] = placeholder[target=ops]
    %get_index : [#users=1] = call_module[target=get_index](args = (index0,), kwargs = {})
    %load : [#users=1] = call_method[target=load](args = (%ops, arg1_1, %get_index), kwargs = {})
    %to_dtype : [#users=1] = call_method[target=to_dtype](args = (%ops, %load, torch.bool), kwargs = {})
    ...
    %where : [#users=1] = call_method[target=where](args = (%ops, %to_dtype, %to_dtype_2, %to_dtype_3), kwargs = {})

关于 uint8，它是一种通用数据类型，可以用于计算，但不仅限于用作 mask 的布尔值。因此，我们需要对其上下文进行静态分析。特别是，CppVecKernelChecker 会检查 uint8 张量是否仅由 to_dtype 使用，并且 to_dtype 是否仅由 where 使用。如果满足，则可以向量化。否则，将回退到标量版本。生成的代码可能如下所示：

标量版本

auto tmp0 = in_ptr0[i1 + (17*i0)];
auto tmp3 = in_ptr1[i1 + (17*i0)];
auto tmp1 = static_cast<bool>(tmp0);
auto tmp2 = static_cast<float>(-33.0);
auto tmp4 = tmp1 ? tmp2 : tmp3;
tmp5 = std::max(tmp5, tmp4);

向量化版本

float g_tmp_buffer_in_ptr0[16] = {0};
// Convert the flag to float for vectorization. 
flag_to_float(in_ptr0 + (16*i1) + (17*i0), g_tmp_buffer_in_ptr0, 16);
auto tmp0 = at::vec::Vectorized<float>::loadu(g_tmp_buffer_in_ptr0);
auto tmp3 = at::vec::Vectorized<float>::loadu(in_ptr1 + (16*i1) + (17*i0));
auto tmp1 = (tmp0);
auto tmp2 = at::vec::Vectorized<float>(static_cast<float>(-33.0));
auto tmp4 = decltype(tmp2)::blendv(tmp3, tmp2, tmp1);

除了上下文分析之外，C++/OpenMP 后端还包含其他几种与向量化相关的优化。其中包括：

• 支持转置加载的平铺内核实现 - cpp.py
• 基于值范围的数据类型降级 - cpp.py
• 将 sleef 实现替换为 oneDNN/oneMKL 实现以优化 aten 向量化 - #94577, #92289, #91613

总而言之，我们研究了 Inductor C++ 后端对 150 个基准模型的 FP32 训练和推理的向量化优化，其中 90% 的推理内核和 71% 的训练内核实现了向量化。

在推理方面，共生成了 28,185 个 CPP 内核，其中 25,579 个（90%）是向量化的，而其余 10% 是标量。在训练方面，共生成了 103,084 个内核，其中 73,909 个（71%）是向量化的，29% 未向量化。

结果表明，推理内核的向量化相当令人印象深刻（由于我们刚开始着手训练方面的工作，训练内核还有一些工作要做）。对剩余未向量化的内核进行了分类分析，突出了提高向量化覆盖率的后续步骤：索引相关操作、int64 支持、垂直归约、带回退的向量化等等。

此外，我们还通过缓冲区重用和 CppWrapper 等其他优化对 C++/OpenMP 后端进行了优化。

未来工作

下一步，我们将继续优化 C++/OpenMP 后端，并将其扩展以支持更多数据类型。这包括：

1. 提高向量化覆盖率
2. 支持和优化包括 BF16、FP16、量化在内的低精度内核
3. 训练优化
4. 循环分块
5. 自动调优
6. Conv/GEMM 内核的进一步融合优化。
7. 探索其他代码生成路径：clang/llvm/triton

总结

Inductor C++/OpenMP 后端是一个灵活高效的 CPU 后端。本博客描述了 Inductor C++/OpenMP 后端用于 TorchBench、Hugging

Face 和 TIMM 这三个基准测试套件的推理和训练所使用的优化。主要优化包括通过 oneDNN 库进行的权重预打包和后置操作融合，以及在 Inductor C++ 代码生成中使用 AVX2 和 AVX-512 指令进行的显式向量化。

结果显示，90% 的推理内核和 71% 的训练内核实现了向量化，这表明推理的向量化令人印象深刻，而训练方面还有改进空间。此外，我们还应用了缓冲区重用和 CppWrapper 等其他优化。我们将持续关注上述未来工作，以进一步提高性能。

致谢

本博客文章中呈现的结果是英特尔 PyTorch 团队与 Meta 之间协作努力的成果。我们衷心感谢 @jansel、@desertfire 和 @Chillee 在整个开发过程中做出的宝贵贡献和坚定支持。他们的专业知识和奉献精神对于实现本文讨论的优化和性能提升至关重要。

配置详情

硬件详情

软件详情

配置

• Intel OpenMP
• Jemalloc - oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1
• 单路多线程：实例数：1；每实例核心数：32
• 单核单线程：实例数：1；每实例核心数：1