作者：Intel

oneDNN Graph 在 PyTorch 2.0 中作为测试版功能提供支持，它利用激进的融合模式来加速 x86-64 机器上的推理，尤其是在英特尔® 至强® 可扩展处理器上。

oneDNN Graph API 通过灵活的图 API 扩展了 oneDNN，最大限度地提高了在 AI 硬件上生成高效代码的优化机会。它自动识别需要通过融合加速的图分区。融合模式专注于融合计算密集型操作，例如卷积、matmul 及其邻近操作，适用于推理和训练两种用例。

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

oneDNN Graph 用法

从用户的角度来看，用法非常简单直观，代码中唯一的变化是一个 API 调用。要结合 JIT-tracing 利用 oneDNN Graph，需要使用一个示例子输入对模型进行性能分析，如以下图 1 所示。

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

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

oneDNN Graph 接收模型的计算图，并根据示例子输入的形状识别算子融合的候选。目前仅支持静态形状。这意味着任何其他输入形状既不受支持，也不会获得任何性能优势。

测量结果

为了确保结果的可重复性，我们使用了 TorchBench 的一个 fork，在使用了第 4 代英特尔® 至强® 可扩展处理器的 AWS m7i.16xlarge 实例上测量了一些 Vision 模型的推理加速。

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

在图 2 中，我们展示了使用 oneDNN Graph 相对于单独使用 PyTorch 的推理加速。使用 oneDNN Graph 进行 Float32 数据类型的几何平均加速为 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 中得到支持，但 Intel 已经在进行工作，计划在未来的 PyTorch 版本中将其作为原型功能与 Inductor-CPU 后端集成，并且 Dynamo 使 PyTorch 更容易支持动态形状，我们希望在 Inductor-CPU 中引入动态形状支持。我们还计划增加 int8 量化支持。

致谢

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