Arm KleidiAI软件库的功能解析

作者：Arm 工程部首席软件工程师 Gian Marco Iodice

在持续快速发展的人工智能 (AI) 时代，Arm 坚定地支持全球数百万开发者，确保他们能够获得 AI 创新开发所需的性能、工具和软件库，从而顺利打造下一波令人惊叹的 AI 体验。为此，Arm 于近日推出了 Arm Kleidi，这是一项广泛的软件和软件社区参与计划，旨在加速 AI 发展。其中的第一个举措是推出面向热门 AI 框架的 Arm Kleidi 软件库。这使开发者可以直接取得 Arm CPU 的出色 AI 功能，而如今全球从云端到边缘侧的大多数 AI 推理工作负载都在这些 Arm CPU 上运行。

Arm KleidiAI

我们所推出的开创性软件库 Arm KleidiAI，旨在提升 AI 在 Arm CPU 上的性能。KleidiAI 的命名来自于希腊语“kleidi”，意为“钥匙”，象征其在提升 Arm CPU 上 AI 性能方面发挥着关键作用。在开发该项目的过程中，我们认真考虑了框架开发者的需求，致力于提供一个紧凑、有影响力且可轻松适用于各类 AI 框架的开源库。

尽管 KleidiAI 仍处于早期阶段，但已帮助 Google MediaPipe 和 XNNPACK 团队将开源大语言模型 (LLM) Gemma 的性能提高了 25%。

我们的目标不仅仅是将 KleidiAI 打造成 AI 优化例程的集合，更是希望该项目能成为学习 Arm CPU 上软件优化最佳实践的知识库。因此，我们诚挚邀请开发者立即加入这激动人心的学习之旅，并提供反馈意见，共同改进我们的产品。

接下来，我将详细探讨 KleidiAI 的初始功能。点击阅读原文，可获取一个关键函数的分步运行指南，该函数用于加速 Gemma LLM 的 4 位整数矩阵乘法例程。

微内核

首先介绍 KleidiAI 库中提供的微内核。KleidiAI 是面向 AI 框架开发者的开源库，可以为 Arm CPU 提供经过优化的性能关键型例程。这些例程是以高性能加速给定算子所需的近乎最小规模的软件，通常称为微内核（或 uKernel）。

以通过 Winograd 算法执行的 2D 卷积算子为例，该计算涉及四个主要运算：

Winograd 输入转换

Winograd 滤波转换

矩阵乘法

Winograd 输出转换

上述每个运算都是一个微内核。但为什么上述运算不叫“内核”或“函数”呢？如下图所示，微内核一词强调了其处理输出张量部分的能力：

图：微内核案例

如上图中的两个案例所示，微内核仅能处理整个输出的一部分。这一设计决策支持细粒度优化，例如，提供高效串联多个微内核的灵活性，从而进一步提升 AI 框架的性能。

微内核可用于各种 Arm 架构、技术和计算参数。例如，在 matmul_clamp_f32_qai8dxp_qsi4cxp 文件夹中，来看下具有按通道量化功能的 Int4 矩阵乘法例程：

该文件夹包含用于 4 位整数矩阵乘法的关键 LLM 微内核。在该文件夹中可以看到，使用 Arm 点积或 i8mm 扩展优化的微内核使用 Neon 汇编来尽可能提高效率，以及在处理的最小输出块方面有不同的多种变体。

文件夹中的每个 .c 和 .h 文件对均代表一种微内核变体。这些变体的计算参数（如块大小）、使用的 Arm 技术（如 Arm Neon）和特定 Arm 架构特性（如 FEAT_DotProd）各有不同。

所有微内核变体共享相同的功能和界面，从而保持一致性。一致性是易于采用的关键因素，接下来我将谈谈我们还采取了哪些其他措施来帮助框架开发者轻松集成微内核。

易于采用

我们深知，在 AI 框架中集成新库时面临着诸多挑战，如库的大小、外部依赖关系和文档。因此，我们努力收集合作伙伴的反馈意见，并将其纳入我们的未来计划之中，尽可能帮助开发者顺畅完成集成。

为了实现这一点，KleidiAI 的设计原则是让框架开发者能够轻松集成所需的微内核。只需拉取相应的 .c 和 .h 文件，以及所有微内核共享的通用头文件 (kai_common.h)，即可集成每个微内核。我们称之为“三文件微内核依赖” (Three-file Micro-kernel Dependency)。这样我们便实现了这一目标。

您可以参考我们准备的相关指南，学习如何使用微内核。该指南探讨了 4 位整数矩阵乘法微内核的使用，这有助于提高 Gemma LLM 的性能。