oneDNN Graph 在 PyTorch 2.0 中作为 Beta 功能受支持，它利用积极的融合模式来加速 x86-64 机器（尤其是 Intel® Xeon® 可扩展处理器）上的推理。
oneDNN Graph API 通过灵活的图 API 扩展了 oneDNN，以最大限度地优化 AI 硬件上高效代码的生成。它自动识别要通过融合加速的图分区。融合模式 专注于融合计算密集型操作，例如卷积、矩阵乘法及其相邻操作，适用于推理和训练用例。
在 PyTorch 2.0 及更高版本中，oneDNN Graph 可以通过 Float32 和 BFloat16（支持 PyTorch 的自动混合精度）数据类型帮助加速 x86-64 CPU（主要是基于 Intel Xeon 处理器的机器）上的推理。对于 BFloat16，加速仅限于支持 AVX512_BF16 ISA（指令集架构）的机器以及也支持 AMX_BF16 ISA 的机器。
oneDNN Graph 用法
从用户的角度来看，用法非常简单直观，代码中唯一的变化是 API 调用。为了利用 oneDNN Graph 和 JIT 追踪，模型使用示例输入进行分析，如下图 1 所示。
图 1：演示使用 oneDNN Graph 的代码片段
oneDNN Graph 接收模型的图并根据示例输入的输入形状识别运算符融合的候选对象。目前，仅支持静态形状。这意味着任何其他输入形状既不受支持也不会获得任何性能优势。
测量
为了确保结果的可重现性，我们使用了 TorchBench 的 一个分支 来测量在 AWS m7i.16xlarge 实例（使用第 4 代 Intel® Xeon® 可扩展处理器）上某些视觉模型的推理加速。
比较的基线是 torch.jit.optimize_for_inference，它只支持 Float32 数据类型。每个模型的批次大小都基于它们在 TorchBench 中使用的相应批次大小。
在图 2 中，我们展示了使用 oneDNN Graph 相对于仅使用 PyTorch 的推理加速。使用 oneDNN Graph 对于 Float32 数据类型的几何平均加速是 1.24 倍，对于 BFloat16 数据类型的几何平均加速是 3.31 倍1。
图 2：使用 oneDNN Graph 相对于默认 CPU JIT Fuser（仅使用 Float32 数据类型）的推理加速
未来工作
oneDNN Graph 目前通过 TorchScript 在 PyTorch 中受支持，但 Intel 已经在进行工作，将其与 Inductor-CPU 后端集成，作为未来 PyTorch 版本中的原型功能，并且 Dynamo 使支持 PyTorch 的动态形状更容易，我们希望通过 Inductor-CPU 引入动态形状支持。我们还计划添加 int8 量化支持。
致谢
本博客中展示的结果是 Meta 和 Intel PyTorch 团队的共同努力。特别感谢来自 Meta 的 Elias Ellison，他花费宝贵时间彻底审查了 PR 并提供了有益的反馈。