由英特尔提供

支持在 PyTorch 2.0 中作为测试功能使用，oneDNN 图利用激进的融合模式来加速 x86-64 机器上的推理，尤其是在 Intel® Xeon® 可扩展处理器上。

oneDNN 图 API 通过灵活的图 API 扩展 oneDNN，以最大化生成高效代码的优化机会，用于 AI 硬件。它自动识别需要通过融合加速的图分区。融合模式专注于融合计算密集型操作，如卷积、矩阵乘法及其相邻操作，适用于推理和训练用例。

在 PyTorch 2.0 及更高版本中，oneDNN 图可以帮助加速 x86-64 CPU（主要是基于 Intel Xeon 处理器的机器）上的推理，支持 Float32 和 BFloat16（使用 PyTorch 的自动混合精度支持）数据类型。使用 BFloat16 时，加速速度仅限于支持 AVX512_BF16 ISA（指令集架构）的机器，以及也支持 AMX_BF16 ISA 的机器。

oneDNN 图使用说明

从用户的角度来看，使用方法非常简单直观，唯一的代码更改是 API 调用。要利用 oneDNN Graph 和 JIT-tracing，可以使用以下示例输入对模型进行性能分析，如图 1 所示。

Figure 1. A code-snippet that demonstrates using oneDNN Graph

图 1：演示使用 oneDNN Graph 的代码片段

oneDNN Graph 接收模型的图，并根据示例输入的输入形状识别操作融合的候选者。目前仅支持静态形状。这意味着任何其他输入形状都不会得到支持，也不会获得任何性能提升。

测量结果

为确保结果的再现性，我们使用了一个 TorchBench 的分支来测量某些视觉模型在 AWS m7i.16xlarge 实例上的推理速度提升，该实例使用第 4 代英特尔® 至强® 可扩展处理器。

比较的基线是 torch.jit.optimize_for_inference，它仅支持 Float32 数据类型。每个模型的批量大小基于在 TorchBench 中使用的相应批量大小。

在图 2 中，我们展示了使用 oneDNN Graph 与仅使用 PyTorch 相比的推理速度提升。对于 Float32 数据类型，使用 oneDNN Graph 的几何平均速度提升为 1.24 倍，对于 BFloat16 数据类型，几何平均速度提升为 3.31 倍。

Figure 2. Inference speedup with oneDNN Graph over default CPU JIT Fuser (which only uses Float32 datatype)

图 2：使用 oneDNN Graph 相对于默认 CPU JIT Fuser（仅使用 Float32 数据类型）的推理速度提升

未来工作

目前，oneDNN Graph 通过 TorchScript 支持在 PyTorch 中使用，但英特尔已经开始着手将其与 Inductor-CPU 后端集成，作为未来 PyTorch 版本中的原型功能。同时，Dynamo 使 PyTorch 支持动态形状变得更加容易，我们希望介绍 Inductor-CPU 的动态形状支持。我们还计划添加 int8 量化支持。

致谢

本博客中展示的结果是 Meta 和英特尔 PyTorch 团队共同努力的结果。特别感谢 Meta 的 Elias Ellison，他花费宝贵时间仔细审查了 PR 并提供了有益的反馈。