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

由英特尔提供

概述

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

电感器后端挑战

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

然而，在其开发的早期阶段，后端缺乏一些优化，这阻碍了它充分利用 CPU 计算能力。因此，对于大多数模型，C++/OpenMP 后端在端到端性能方面表现不如急切模式，45%的 TorchBench、100%的 Hugging Face 和 75%的 TIMM 模型的表现不如急切模式。

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

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

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

图 1：性能加速比趋势

英特尔混合优化性能状态

与具有混合优化的急切模式相比，C++/OpenMP 后端显示出有希望的性能提升。我们测量了三个电感器基准测试套件——火炬基准、Hugging Face 和 TIMM——的结果如下。（注意：我们每周在 GitHub 上发布两次性能数据。）

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

通过率

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

几何平均速度提升（单 Socket 多线程）

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

单个模型性能

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

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

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

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

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

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

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

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

技术深度解析

现在，让我们更深入地看看 Inductor C++/OpenMP 后端中使用的两种主要优化：

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

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

Intel® oneAPI 深度神经网络库的简称，oneDNN 库提供了一系列后操作融合（例如，融合卷积和矩阵乘法及其后续操作），这可以惠及流行的模型。Intel®对 PyTorch 的扩展已实现这些融合中的大多数，并取得了显著的性能提升。因此，我们将已在 Intel 的 PyTorch 扩展中应用的这些融合全部上推到 Inductor，使更广泛的模型能够受益于这些优化。我们已将这些融合定义为 mkldnn 命名空间下的操作符。这允许 Python 模块直接调用这些 mkldnn 操作。

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

• _linear_pointwise : 将线性操作及其后一元元素级操作融合
• _linear_pointwise.binary : 将线性操作及其后二元元素级操作融合
• _convolution_pointwise : 将卷积操作及其后一元元素级操作融合
• _convolution_pointwise.binary : 将卷积操作及其后二元元素级操作融合

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

在电感器一侧，我们在已经降低的 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. 数据类型中尚未支持向量化的有整数、双精度、半精度和 bf16，但相关工作正在进行中。

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

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

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

例如，让我们考虑一个 where 操作，它接受掩码 true_value 和 false_value ，而掩码值是从 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 ，它是一种通用数据类型，可用于计算，但不仅限于用作掩码的布尔值。因此，我们需要静态地分析其上下文。特别是，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
• 用 oneDNN/oneMKL 实现替换 sleef 实现以优化 aten 向量化 - #94577, #92289, #91613

概括来说，我们研究了 Inductor C++后端在 FP32 训练和推理 150 个基准模型中的向量化优化，其中 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 在三个基准测试套件（TorchBench、Hugging Face 和 TIMM）的推理和训练中使用的优化。主要优化包括通过 oneDNN 库进行权重预打包和后操作融合，以及在 Inductor C++ 代码生成中使用 AVX2 和 AVX-512 指令进行显式向量化。

面部识别和 TIMM。主要优化包括通过 oneDNN 库进行权重预打包和后操作融合，以及在 Inductor C++ 代码生成中使用 AVX2 和 AVX-512 指令进行显式向量化。

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

致谢

本博客文章中展示的结果是英特尔 PyTorch 团队与 Meta 合作的结晶。我们想对@jansel、@desertfire 和@Chillee 表示衷心的感谢，他们在整个开发过程中做出了宝贵的贡献，并给予了坚定不移的支持。他们的专业知识和奉献精神对于实现此处讨论的优化和性能提升至关重要。

配置详情

硬件详情

软件详细信息

配置

• 英特尔 OpenMP
• Jemalloc - oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1
• 单核多线程：实例数量：1；每个实例的核数：32
• 单核单线程：实例数量：1；每个实例的核数：1