利用 Intel® 高階矩陣擴充套件¶

建立日期：2023 年 6 月 14 日 | 最後更新日期：2023 年 6 月 14 日 | 最後驗證日期：未驗證

引言¶

高階矩陣擴充套件 (AMX)，也稱為 Intel® 高階矩陣擴充套件 (Intel® AMX)，是一種 x86 擴充套件，它引入了兩個新元件：一個稱為“tiles”的二維暫存器檔案和一個能夠操作這些 tiles 的瓦片矩陣乘法 (TMUL) 加速器。AMX 旨在處理矩陣，以加速 CPU 上的深度學習訓練和推理，非常適合自然語言處理、推薦系統和影像識別等工作負載。

英特爾透過第四代 Intel® Xeon® 可擴充套件處理器和 Intel® AMX 提升 AI 能力，提供比上一代高 3 到 10 倍的推理和訓練效能，詳見使用 Intel® AMX 加速 AI 工作負載。與執行 Intel® Advanced Vector Extensions 512 Neural Network Instructions (Intel® AVX-512 VNNI) 的第三代 Intel Xeon 可擴充套件處理器相比，執行 Intel AMX 的第四代 Intel Xeon 可擴充套件處理器每週期可執行 2,048 次 INT8 操作，而前者每週期為 256 次 INT8 操作。它們每週期還可以執行 1,024 次 BF16 操作，而每週期為 64 次 FP32 操作，詳見使用 Intel® AMX 加速 AI 工作負載第 4 頁。關於 AMX 的更多詳細資訊，請參閱Intel® AMX 概述。

PyTorch 中的 AMX¶

PyTorch 利用其後端 oneDNN 對 BFloat16 和 INT8 量化的計算密集型運算元使用 AMX 加速，以便在支援 AMX 的 x86 CPU 上獲得更高的開箱即用效能。有關 oneDNN 的更多詳細資訊，請參閱oneDNN。

該操作完全由 oneDNN 根據生成的執行程式碼路徑處理。例如，當支援的操作在支援 AMX 的硬體平臺上執行到 oneDNN 實現時，AMX 指令將自動在 oneDNN 內部呼叫。由於 oneDNN 是 PyTorch CPU 的預設加速庫，因此無需手動操作即可啟用 AMX 支援。

利用 AMX 處理工作負載的指南¶

本節提供了關於如何利用 AMX 處理各種工作負載的指南。

BFloat16 資料型別

使用 torch.cpu.amp 或 torch.autocast("cpu") 將為支援的運算元利用 AMX 加速。

model = model.to(memory_format=torch.channels_last)
with torch.cpu.amp.autocast():
   output = model(input)

注意

使用 torch.channels_last 記憶體格式以獲得更好的效能。

量化
- 應用量化將為支援的運算元利用 AMX 加速。
torch.compile
- 當生成的圖模型在支援的運算元中遇到 oneDNN 實現時，AMX 加速將被啟用。

注意

在支援 AMX 的 CPU 上使用 PyTorch 時，框架預設會自動啟用 AMX 使用。這意味著 PyTorch 將盡可能嘗試利用 AMX 功能來加速矩陣乘法運算。然而，需要注意的是，是否分派到 AMX 核心最終取決於 oneDNN 庫和量化後端的內部最佳化策略，PyTorch 依賴這些策略進行效能增強。PyTorch 和 oneDNN 庫內部處理 AMX 利用的具體細節可能會隨著框架的更新和改進而發生變化。

可利用 AMX 的 CPU 運算元：¶

可利用 AMX 的 BF16 CPU 運算元

conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
bmm
mm
baddbmm
addmm
addbmm
linear
matmul

可利用 AMX 的量化 CPU 運算元

conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear

確認 AMX 是否被利用¶

設定環境變數 export ONEDNN_VERBOSE=1，或使用 torch.backends.mkldnn.verbose 來啟用 oneDNN 輸出詳細訊息。

with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
    with torch.cpu.amp.autocast():
        model(input)

例如，獲取 oneDNN 詳細輸出

onednn_verbose,info,oneDNN v2.7.3 (commit 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
onednn_verbose,info,cpu,runtime:OpenMP,nthr:128
onednn_verbose,info,cpu,isa:Intel AVX-512 with float16, Intel DL Boost and bfloat16 support and Intel AMX with bfloat16 and 8-bit integer support
onednn_verbose,info,gpu,runtime:none
onednn_verbose,info,prim_template:operation,engine,primitive,implementation,prop_kind,memory_descriptors,attributes,auxiliary,problem_desc,exec_time
onednn_verbose,exec,cpu,reorder,simple:any,undef,src_f32::blocked:a:f0 dst_f32::blocked:a:f0,attr-scratchpad:user ,,2,5.2561
...
onednn_verbose,exec,cpu,convolution,jit:avx512_core_amx_bf16,forward_training,src_bf16::blocked:acdb:f0 wei_bf16:p:blocked:ABcd16b16a2b:f0 bia_f32::blocked:a:f0 dst_bf16::blocked:acdb:f0,attr-scratchpad:user ,alg:convolution_direct,mb7_ic2oc1_ih224oh111kh3sh2dh1ph1_iw224ow111kw3sw2dw1pw1,0.628906
...
onednn_verbose,exec,cpu,matmul,brg:avx512_core_amx_int8,undef,src_s8::blocked:ab:f0 wei_s8:p:blocked:BA16a64b4a:f0 dst_s8::blocked:ab:f0,attr-scratchpad:user ,,1x30522:30522x768:1x768,7.66382
...

如果你看到針對 BFloat16 的 avx512_core_amx_bf16 或針對 INT8 量化的 avx512_core_amx_int8 的詳細輸出，則表示 AMX 已啟用。

結論¶

在本教程中，我們簡要介紹了 AMX，如何在 PyTorch 中利用 AMX 加速工作負載，以及如何確認 AMX 正在被利用。

隨著 PyTorch 和 oneDNN 的改進和更新，AMX 的利用方式可能會相應發生變化。

一如既往，如果您遇到任何問題或有任何疑問，可以使用論壇或GitHub Issues與我們聯絡。