关注了就能看到更多这么棒的文章哦～

Intel AMX support in 5.16

By Jonathan Corbet
November 8, 2021
DeepL assisted translation
https://lwn.net/Articles/874846/

x86 指令集中有非常多的指令，但这并没有阻碍它继续扩充。即将推出的英特尔处理器中将推出以 "高级矩阵扩展"（AMX，Advanced Matrix Extensions）为名的一套新指令，用于对矩阵数据进行操作。经过一个有些波折的开发过程之后，对 AMX 的支持可以进入即将到来的 5.16 内核了。要使用这套指令集的话，应用程序开发人员就必须改变一些做法。

AMX（在 https://software.intel.com/content/dam/develop/external/us/en/documents-tps/architecture-instruction-set-extensions-programming-reference.pdf 中有介绍）是一个通用架构，用来加速 x86 处理器上的矩阵操作。其最初形式中实现了一组多达 8 个 "tiles" （也就是 16 行 64 字节的阵列）。程序员可以在这些 tiles 中存储任意尺寸的矩阵。16x16 的 32 位浮点值的矩阵可以支持，此外也支持其他规格的矩阵。目前支持的一个操作是将存储在两个 tile 中的矩阵相乘，然后将结果写入到第三个 tile 中。通过串联这类操作，就可以实现任意大小的矩阵乘法。显然，未来还会支持其他种类的矩阵操作。

虽然看起来 AMX 可能是一个用于数值分析场合的功能，但其真正的目标用例似乎是机器学习这种应用。这可以解释为什么是支持 16 位浮点运算，而不是 64 位。

AMX 在设计上就允许内核来控制这些功能是否可以被某个特定进程来使用。设计成这样是有几个原因，其中之一是正如人们能预料到的，AMX 指令使用了大量的处理器资源。一个进程如果在共享的计算机（shared computer）上做大量 AMX 计算的进程可能会对其他进程产生很大影响。但是，除非内核和用户空间进程都为此准备好了，否则 AMX 也无法得到正确支持。

Development process

对 AMX 的支持是由 Chang Bae 在 2020 年 10 月首次发布出来的，但得到的 review 很少。到了 2 月份第 4 版出来的时候，更多的开发者开始关注它，他们对这个功能当前被整合到内核现有的浮点单元（FPU）代码中的这种做法不是很满意。接下来各种版本就接踵而至，作者的挫折感似乎也在不断累积。9月底的时候，Len Brown 发布了一份对话记录，其中似乎列出了接下来前进的方向。

不幸的是，第二天发布的第 11 版似乎忽略了许多已经做出的决定。这一版就引起了 Thomas Gleixner 的强烈指责，他认为这个功能是在强行要塞进内核，而完全没有听取人们的抱怨。AMX 的情况并不乐观，但是幕后人们仍在继续坐着工作。10 月中旬的时候，Gleixner 发表了他对 FPU 代码的大规模 rework，希望能减轻在 kernel 里支持 AMX 的代价。不久之后，一个新的 AMX patchset 就出现了，这就是 5.16 中得到的最终成果。

Gleixner 对这部分代码的 pull request 中也承认，这些代码还是相对不成熟的：

注意，这部分代码还是比较新的，尽管对 AMX 的支持已经进行了一年多了。

对 FPU 代码的重大重构，使其能够进行适当的整合，这个工作是在 3 周前开始的。重构现有的 FPU 代码和原始的 AMX patch 花了一周时间，并进行了广泛的 review 和 test。唯一没有经历 review 和 test 的内容也仅限于支持 AMX 的系统，这对那些不使用英特尔和他们的 early access program （新功能的早期使用项目）的人员来说完全没有影响。像往常一样，可能会有潜伏的龙（也就是指潜在的 bug），但到目前为止，细化之后的重构已经经受住了考验，最终尚未发现的那些影响是可以使用 bisect 方式来确认出来源的，在 5.16 版本发布之前应该可以轻松解决。著名的最后补充…

FPU 代码是相对比较棘手、底层性的工作，所以在新的工作中发现一两个潜伏的 bug 确实是不奇怪的。

Using AMX

如上所述，内核能够控制哪些进程可以使用 AMX 指令。用户空间进程要想使用它，第一步就是利用一个新的 arch_prctl()命令（名为 ARCH_GET_XCOMP_SUPP）来获得系统能支持的 feature 列表。如果返回结果中相应 bit 被置位了，那么就可以直接使用 AMX 了。此外还有另一个 arch_prctl()命令（名为 ARCH_REQ_XCOMP_PERM）可以用来请求获取权限来使用 AMX。这里将会进行一些检查（其中一个检查会在本文下面介绍），security modules 也有机会表达意见。不过，通常情况下这些 request 都会被批准。所获取的权限适用于此进程中的所有线程，并在 fork 后也会继续拥有这个权限，不过，调用 execve() 的话就会丧失原进程的特有权限。

AMX 带来的一个挑战是处理器在运行 AMX 指令时可能创建大量的内部状态（internal state）。如果 CPU 在运行这部分指令过程中被中断打断，这些状态就必须保存在某个地方，否则会丢失许多进展。因此，如果一个进程使用了 AMX，内核就必须在做其他事情之前先要在其中断处理程序中保存 10KB 左右的数据。这种保存工作是通过 XSAVE 指令完成的。

内核针对这个目的，为每个进程分配了相应的内存空间。不过，为系统中的每个进程都分配 10KB 的内存空间会浪费很多内存。而且大多数进程不会使用 AMX 指令。不过好消息是，处理器可以配置成在进程第一次执行 AMX 指令时 trap 进入 kernel，然后内核可以检查是否被允许使用这些指令，允许的话就分配一个适当大小的 buffer 来保存 FPU 状态并允许此次操作继续进行。

这里有一个与 sigaltstack() 系统调用有关的潜在问题，这个调用允许线程来建立一个新的堆栈用来处理信号。如果此进程使用了 AMX，那么这个堆栈也必须要足够大，才够用于保存 FPU 的状态。多年来，开发者一直被告知要用 MINSIGSTKSZ 来作为这个堆栈的最小 size，也就是 2KB，但对于使用 AMX 的进程来说这就远远不够了。事实上，这个 size 甚至都不够用来使用 AVX-512 extension，过去曾经因为这个原因而引发了一些 stack 被破坏的问题。

为了在 AMX 中避免这个问题，内核将仔细检查，确保所有的 signal stack 都足够大。这个检查是在每次调用 sigaltstack()时进行的，但是当一个进程首次请求获取 AMX 权限时也会对当前堆栈进行检查。不使用 AMX 的进程将不需要那么大的堆栈，因此，不会在这些检查的时候被拒绝。而那些想要使用 AMX 的进程则只有拥有足够大的 signal stack 时才允许继续执行。

等这些检查相关的基础设施都到位了之后，内核就可以确保那些使用 AVX-512 的进程有足够大小的 signal stack 了。但是，如果强制确保这个条件的话，有可能破坏现在看起来还能正常工作的应用程序，比如也许它们的 signal handler 处理程序从未被实际调用过。为了避免这种问题，准备了一个内核配置选项（STRICT_SIGALTSTACK_SIZE）和一个命令行选项（strict_sas_size=），它们都可以被用来控制在使用 AVX-512 时是否进行严格检查。

假设所有这些都能配合好，那么这也就是 AMX 在 5.16 中得以支持所采取的形式了。想了解更多信息的人可以看看包含 AMX test case 的 commit 以及一些关于 arch_prctl()命令的文档。同时，在接下来的九周左右的时间里，请留意是有有 bug 的出现。

全文完
LWN 文章遵循 CC BY-SA 4.0 许可协议。

欢迎分享、转载及基于现有协议再创作～

长按下面二维码关注，关注 LWN 深度文章以及开源社区的各种新近言论～