打破库依赖与93%峰值效率！Intel提出MLIR驱动的编译器自动生成NanoKernel实现高性能矩阵乘法内核

关键词： MLIR 、Nanokernels 、 Microkernels 、Matmul、Vectorization、Compiler

超微内核（Nanokernel） 指寄存器级别的最小计算单元，专为特定硬件指令集优化，可作为可组合的、目标无关的编译器 IR 到目标特定指令的 kernel。

论文标题：Library Liberation: Competitive Performance Matmul Through Compiler-composed Nanokernels
论文链接：https://arxiv.org/pdf/2511.13764

在 AI 和机器学习 workload 飞速发展的今天，矩阵乘法（GEMM）作为核心计算原语，其性能直接决定了模型训练和推理的效率。但长期以来，想要让矩阵乘法逼近硬件峰值性能，要么需要专家手工编写目标硬件专属内核，要么依赖 CUTLASS、libxsmm 等厂商优化库——前者门槛极高，后者则存在兼容性滞后、依赖复杂、难以适配新兴硬件的问题。

最近 Intel 团队在 arXiv 上发表的论文《Library Liberation: Competitive Performance Matmul Through Compiler-composed Nanokernels》，给出了一套颠覆性解决方案：让编译器直接自动生成高性能超微内核（Nanokernels） ，彻底摆脱对外部库和运行时 JIT 编译的依赖，同时性能媲美甚至超越主流优化库。这篇文章我们就来深入拆解这项工作的核心创新和技术细节。

一、核心问题：为什么传统方案难以兼顾性能与灵活性？

在深入创新点前，我们先搞懂传统矩阵乘法优化的痛点：

性能与硬件强绑定：不同 CPU 架构（如 AMX、AVX512、AVX2）的向量寄存器宽度、指令集特性不同，手工编写的内核无法跨架构复用，优化库也需针对性适配。
库依赖的“枷锁”：主流方案依赖 libxsmm 等外部库，这些库不仅会引入 runtime 开销，还难以快速支持新硬件特性（如 VNNI 指令）和新型算子模式。注，VNNI 是向量神经网络指令，x86 架构中针对 AI workload 优化的指令集，支持打包数据的融合乘加运算；
内存墙瓶颈：矩阵乘法的计算强度高，但连续操作间的数据搬运容易成为瓶颈，传统方案的内存布局优化难以与编译流程深度融合。

而这篇论文的核心目标，就是通过编译器原生能力， 解决“性能、灵活性、可扩展性”三者不可兼得的难题。

二、核心创新点：三大突破打破传统局限

这篇工作的创新并非单点优化，而是一套端到端的编译方案重构 ，核心包含三大突破：

创新 1：编译器原生超微内核生成，彻底去库依赖

传统方案中，微内核（Microkernel）的优化逻辑被封装在外部库中，编译器仅负责调用。而本文提出：将库的专家优化知识直接嵌入编译器，通过 MLIR 的多级下转换（Multi-level Lowering）自动生成纳米内核 。

超微内核（Nanokernel）： 是比微内核更精细的计算单元，专为寄存器级计算设计 ，可作为可组合的、目标无关的编译器 IR 到目标特定指令的降低核心，进而形成适配不同硬件的微内核。
关键优势： 无需外部库调用、无 JIT runtime 开销 ，实现了从高级算子到机器码的纯编译器 AOT（提前编译）统一路径，且能天然适配新硬件（只需扩展 MLIR dialect 支持）。

创新 2：寄存器级智能分块+向量化，实现近最优寄存器利用率

寄存器溢出是内核性能的“隐形杀手”，本文通过精细化的寄存器分块策略，确保计算全程无溢出：

分块维度：覆盖 M（矩阵 A 的行）、N（矩阵 B 的列）、K（公共维度）和 batch 维度，batch 固定为 1，K 根据指令集适配（1/VNNI/32）。
动态适配：根据硬件寄存器资源精确计算分块大小，以实现“用满但不超”的资源利用。例如 AVX512（32 个向量寄存器）采用 M=4、N=64，恰好使用 29 个寄存器（24 累加器+4 广播+1 加载）；AVX2（16 个向量寄存器）采用 M=3、N=32，用满 16 个寄存器。
兜底策略：当分块可能导致溢出时，自动交换 A/B 的广播/加载角色，确保寄存器充分利用（如图 2(d)）。

图 2：文献《Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning Workloads》中提出的 BRGEMM 寄存器分块方案，适用于 (a) 32 个向量寄存器、(b) 16 个向量寄存器、(c) 8 个二维寄存器瓦片，以及 (d) 16 个向量寄存器（A 和 B 的加载顺序互换）的场景。该图展示了针对不同 CPU 架构向量资源的寄存器分块策略，核心目标是避免寄存器溢出并最大化寄存器利用率。例如，AVX512 架构（32 个向量寄存器，图 a）采用 M:4、N:64 分块，仅使用 29 个寄存器（24 个累加器、4 个 A 广播寄存器、1 个 B 加载寄存器）；AVX2 架构（16 个向量寄存器，图 b）用 M:3、N:32 分块，耗尽所有寄存器；Intel AMX 架构（图 c）则用 8 个 16×16 FP32 二维瓦片存储数据。当分块导致寄存器不足时（如图 d 的 16 寄存器场景），通过互换 A 广播与 B 加载角色，仍能实现无溢出计算，为跨架构超微核生成提供了分块设计依据。

创新 3：多架构与多精度统一适配，基于 MLIR Dialect 打通底层

本文的技术核心在于利用 MLIR 的 dialect 系统，通过抽象不同硬件的指令特性，实现“一套编译逻辑，多架构适配”的目标。

核心 Dialect 组合：amx dialect（适配 AMX 架构的瓦片操作）、x86vector dialect（适配 AVX 系列指令）、vector dialect（通用向量操作）。
多精度支持：方案天然支持 FP32、BF16（重点优化）和 INT8（未来扩展）。其中，BF16 通过 VNNI 打包布局来提升 SIMD 计算效率。

三、工作方法详解：从矩阵乘法到高性能超微内核的完整流程

该方案的核心流程可概括为三个步骤：寄存器分块 → 向量化 → 超微内核生成。整个过程完全基于 MLIR 上游变换，无需自定义黑盒逻辑。

第一步：寄存器分块 —— 让计算“住进”寄存器

目标是将 Batch-Reduce GEMM（BRGEMM，支持批量矩阵乘法累积，适用于 Transformer 等模型中的注意力机制计算）分解为寄存器可容纳的子问题，避免数据频繁进出内存。

BRGEMM 是本方案的核心计算模型，相比传统 GEMM 具有两大优势：
1. 累积的瓦片数据全程驻留寄存器，大幅减少内存访问。
2. 支持与元素级操作（如偏置加法、ReLU激活）进行融合，提升数据局部性。

分块策略必须严格匹配硬件特性。论文给出了针对不同架构的最优分块配置：

如上图所示，分块设计的核心原则是“避免溢出”与“最大化利用率”。例如，在 (a) 所示的 AVX512 架构（32个向量寄存器）场景中，通过选择 M=4、N=64 的分块，精确使用了 29 个寄存器（24 个累加器 + 4 个广播寄存器 + 1 个加载寄存器），在寄存器数量限制内实现了无溢出的高效计算。

第二步：向量化 —— 激活 SIMD 并行

将分块后的 BRGEMM 计算转换为 vector.contract 操作，这是生成纳米内核的入口：
* 数据移动：通过 vector.transfer_read/write 操作实现向量化的数据加载与存储。
* 计算表达：通过 vector.contract 定义矩阵乘法的收缩运算，明确并行维度和归约维度。

第三步：超微内核生成 —— 适配硬件的指令级优化

这是最核心的环节，根据目标架构、数据类型（FP32/BF16）和数据布局（VNNI-packed/Flat），生成针对性的超微内核。

关键技术 1：数据布局优化（VNNI 打包）

对于 BF16 数据类型，x86 架构的 VNNI 指令要求数据按特定格式打包（每 2 个 BF16 元素为一组），否则会导致内存访问不连续。本方案通过 MLIR 变换自动完成从 Flat 布局到 VNNI-packed 布局的转换。在生成的代码中，对于 AVX 目标，该转换映射到底层的 vpunpcklwd/hwd 指令；而在 AMX 等目标的 IR 层面，则通过 vector.shuffle 操作来表达。

关键技术 2：架构感知的 Dialect 映射

不同架构的超微内核通过组合相应的 MLIR dialect 操作生成。论文清晰地展示了 dialect 与目标架构的映射关系。

关键技术 3：超微内核生成算法

论文提出了统一的超微内核生成框架（基于原论文 Algorithm 1 的核心流程）：
1. 加载 C 矩阵的子块到寄存器作为累加器。
2. 对 batch 维度和 K 维度进行循环，加载 A/B 矩阵的相应子块。
3. 根据数据类型和目标架构，执行乘法-累加运算（FMA 或点积指令）。
4. 若为 Flat 布局，则在结果写回 C 矩阵前完成数据重排。

以 BF16-AMX 架构的生成为例：
* 使用 amx.tile_load 加载 A/B 的 16×32 子块。
* 使用 amx.tile_mulf 执行瓦片级点积运算，结果累积在 FP32 寄存器中。
* 使用 amx.tile_store 写回结果，若需 BF16 输出则进行类型转换。

四、相关工作对比：本文方案的创新之处

为更清晰地凸显创新，下表对比了本方案与主流方案的核心差异：

上表对比了主流的矩阵乘法优化方案。本文方案的核心优势在于“去库依赖”和“编译器原生适配”，在保留手工优化库高性能的同时，具备了跨架构和快速支持新硬件的灵活性。

五、实验结果：性能媲美甚至超越手工优化库

论文在 Intel 最新三代处理器（Emerald Rapids、Sierra Forest、Arrow Lake）上进行了全面验证，核心结论如下：

1. 性能表现

以 libxsmm 为基准（100%），不同平台和配置下的性能对比如下：

表 2：在 EMR 和 SRF 两台机器上，不同线程数与矩阵大小下，编译器生成超微核与 libxsmm 的几何平均性能百分比对比。实验平台中，EMR 为双路 Xeon Platinum 8592+（支持 AMX/AVX512，2×64 核），SRF 为双路 Xeon 6780E（能效核架构，2×144 E 核，支持 AVX2）。

表格数据显示，在 FP32 场景下，超微核性能接近 libxsmm（EMR 达 91%-94%、SRF 达 91%-93%）；在 BF16 场景中，VNNI 布局性能达 73%-93%，而 Flat 布局在 SRF 上反超 libxsmm（112%-113%）。这表明编译器生成的纳米核在多线程扩展（1-64 线程）和不同矩阵规模（1024-4096）下均表现稳定，尤其在 SRF 的 E 核架构上，BF16-Flat 布局因避免了格式转换开销而展现出显著优势。

表 3：在 ARL 机器的性能核（P 核）和能效核（E 核）上，不同线程数与矩阵大小下，编译器生成超微核与 libxsmm 的几何平均性能百分比对比。ARL 平台为单路 Intel Core Ultra 9 285K（8 个性能核 P+16 个能效核 E，支持 AVX2）。

表格数据凸显了超微核对异构核心的适配能力。在 FP32 场景下，P 核性能达 86%-91%、E 核达 85%-93%，接近 libxsmm；在 BF16 场景下优势显著，VNNI 布局达 89%-98%，Flat 布局更是超越 libxsmm（P 核 131%-132%、E 核 119%-126%）。这是因为 Flat 布局直接利用 AVX2 的打包操作，规避了 libxsmm 将 Flat 转换为 VNNI 的额外开销，在 ARL 的异构核心上更能发挥硬件能效，验证了编译方案对异构核心的优化效果。

2. 关键洞察

内存带宽是 AMX 架构的瓶颈：虽然 AMX 的计算吞吐量更高，但未对齐的加载/存储操作会导致吞吐量降至峰值的 1/4，因此需要优化数据布局。
AVX2 架构的 Flat 布局优势明显：由于本文方案直接利用 AVX2 的打包操作，无需进行布局转换，性能反而超越了依赖 VNNI 的方案。
多线程扩展性良好：在部分高线程配置下（如 64 线程），AMX Flat 布局的性能甚至能反超 libxsmm 约 5%。整体上，AMX Flat 布局的几何平均性能达到 libxsmm 的 93%（VNNI 布局为 91%）。

六、总结与展望

这项工作的核心价值，是为矩阵乘法等核心计算原语提供了一套“编译器原生、去库依赖、跨架构”的优化方案——通过 MLIR 的 dialect 抽象和寄存器级精细优化，使编译器具备了专家级的内核生成能力。

未来的工作方向明确：
1. 上游化到 LLVM/MLIR 主干，让社区直接复用。
2. 扩展支持 INT8 等更多数据类型，适配量化模型。
3. 自动化推导最优分块大小，无需人工调参。

对于 AI 框架开发者和硬件厂商而言，这套方案意味着无需再为不同架构维护大量手工优化内核，只需扩展 MLIR dialect 即可快速支持新硬件，同时保持高性能。对于终端用户，这意味着模型编译将更高效、更灵活，无需依赖复杂的第三方库。

编译器正在成为释放硬件性能的核心枢纽，这项工作无疑为“编译驱动的高性能计算”提供了一个绝佳范例。

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/17466