PyTorchのCPU最適化をoneDNNに集約
AMDとIntelは、PyTorchでAI推論を実行する際のx86 CPU向け最適化を、オープンソースライブラリ「oneDNN」に統合します。AMD独自の最適化ライブラリ「ZenDNN」の機能をoneDNNに取り込み、両社のCPUで共通のソフトウェア経路を用意する計画です。
PyTorchは機械学習モデルの開発や実行に使われるソフトウェア基盤です。推論では、学習済みモデルから結果を得ます。CPUごとに異なる最適化の利用方法をoneDNNにまとめ、開発者がAMD製とIntel製のx86 CPUを扱う際のソフトウェア上の分断を減らすことを目指します。
今回の統合は、x86 Ecosystem Advisory Group（EAG）の提言を受けたものです。EAGにはAMDとIntelが参加しており、両社は2024年に設立したと発表しています。CPUの命令セットを共有するx86環境に向け、最適化でも協力します。
ZenDNNの機能を共通ライブラリへ
ZenDNNは、AMDのZenアーキテクチャを採用したCPU上で、深層学習の推論を高速化するライブラリです。コンピュータービジョン、自然言語処理、推薦システムなどを対象とし、プラグインを通じてPyTorchやTensorFlowと連携します。
AMDはZenDNNの最適化をoneDNNへ統合しながら、AMD CPU向けの最適化も続けます。oneDNNはIntelが開発したライブラリで、現在はUnified Acceleration Foundation（UXL Foundation）の下で運営されています。
UXL Foundationは複数の企業が参加する団体で、異なるベンダーや世代のハードウェアで使えるオープンなソフトウェア標準とプロジェクトを進めています。ワーキンググループなどを通じて、仕様やオープンソース実装への貢献をまとめています。AMDは同団体の運営に関わるステアリングメンバーにも加わりました。
開発者向けの共通経路と今後の拡張
oneDNNは、CPUの機能を実行時に検出し、ハードウェアに合わせたカーネルを生成するライブラリです。行列積や畳み込みなど、AIや機械学習で使われる演算向けの最適化命令も備えています。
AMDとIntelは、oneDNNをPyTorchにおけるx86向けアクセラレーションの標準的なオープンソースライブラリとして位置づけています。両社のCPU向け最適化を共通の場所で利用できるようにし、AIアプリケーションの開発、保守、展開を簡素化するのが狙いです。
両社は別途、x86向けの次世代命令セット「AI Compute Extensions（ACE）」でも協力しています。ACEはAI処理の中核となる行列積を高速化する設計で、将来のoneDNN更新を通じてPyTorchに組み込む計画です。
oneDNN Graphが示すCPU推論の処理
oneDNN Graphは、oneDNNのグラフAPIを使って処理をまとめる仕組みです。畳み込みや行列積とその周辺演算を融合し、推論時の処理を最適化します。
oneDNN GraphはPyTorch 2.0でベータ機能として提供されました。例示された利用方法では、モデルのグラフから融合候補を選びます。この方法は入力形状が固定されたモデルを対象とし、異なる入力形状には対応しません。
この推論方法が扱うデータ型はFloat32とBFloat16です。BFloat16の高速化には、対応するAVX512-BF16またはAMX-BF16命令セットが必要です。
