利用英特尔®高级矩阵扩展¶

Created On: Jun 14, 2023 | Last Updated: Jun 14, 2023 | Last Verified: Not Verified

简介¶

高级矩阵扩展（AMX），也称为英特尔®高级矩阵扩展（Intel® AMX），是一种x86扩展，它引入了两个新组件：一个称为“tiles”的二维寄存器文件和一个可以操作这些tiles的矩阵乘加加速器（TMUL）。AMX旨在处理矩阵操作，以加速CPU上的深度学习训练和推理，非常适合自然语言处理、推荐系统和图像识别等工作负载。

英特尔通过第4代英特尔®至强®可扩展处理器和英特尔®AMX，提供比前一代高3到10倍的推理和训练性能，详见`通过英特尔®AMX加速AI工作负载`_。与运行英特尔®高级矢量扩展512神经网络指令（Intel® AVX-512 VNNI）的第3代英特尔至强可扩展处理器相比，运行英特尔AMX的第4代至强可扩展处理器每周期可执行2,048次INT8操作，而非256次INT8操作。它们每周期还可执行1,024次BF16操作，而非64次FP32操作，详见`通过英特尔®AMX加速AI工作负载`_的第4页。有关AMX的更多详细信息，参见`英特尔®AMX概述`_。

PyTorch中的AMX¶

PyTorch利用AMX通过其后端oneDNN对使用BFloat16和INT8量化的计算密集型操作符实现开箱即用的高性能支持x86 CPU的AMX功能。有关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
...

如果您获取到``avx512_core_amx_bf16``的BFloat16或``avx512_core_amx_int8``的INT8量化的详细信息，这表明AMX已激活。

结论¶

在本教程中，我们简要介绍了AMX，如何在PyTorch中使用AMX加速工作负载，以及如何确认AMX是否被利用。

随着PyTorch和oneDNN的改进和更新，AMX的利用可能会相应发生变化。

一如既往，如果您遇到任何问题或有任何疑问，可以通过`论坛 <https://discuss.pytorch.org/>`_或`GitHub问题 <https://github.com/pytorch/pytorch/issues>`_联系我们。