作者：Arm 工程部首席軟件工程師 Gian Marco Iodice

在持續(xù)快速發(fā)展的人工智能 (AI) 時(shí)代，Arm 堅(jiān)定地支持全球數(shù)百萬(wàn)開(kāi)發(fā)者，確保他們能夠獲得 AI 創(chuàng)新開(kāi)發(fā)所需的性能、工具和軟件庫(kù)，從而順利打造下一波令人驚嘆的 AI 體驗(yàn)。為此，Arm 于近日推出了 Arm Kleidi，這是一項(xiàng)廣泛的軟件和軟件社區(qū)參與計(jì)劃，旨在加速 AI 發(fā)展。其中的第一個(gè)舉措是推出面向熱門(mén) AI 框架的 Arm Kleidi 軟件庫(kù)。這使開(kāi)發(fā)者可以直接取得 Arm CPU 的出色 AI 功能，而如今全球從云端到邊緣側(cè)的大多數(shù) AI 推理工作負(fù)載都在這些 Arm CPU 上運(yùn)行。

Arm KleidiAI

我們所推出的開(kāi)創(chuàng)性軟件庫(kù) Arm KleidiAI，旨在提升 AI 在 Arm CPU 上的性能。KleidiAI 的命名來(lái)自于希臘語(yǔ)“kleidi”，意為“鑰匙”，象征其在提升 Arm CPU 上 AI 性能方面發(fā)揮著關(guān)鍵作用。在開(kāi)發(fā)該項(xiàng)目的過(guò)程中，我們認(rèn)真考慮了框架開(kāi)發(fā)者的需求，致力于提供一個(gè)緊湊、有影響力且可輕松適用于各類 AI 框架的開(kāi)源庫(kù)。

盡管 KleidiAI 仍處于早期階段，但已幫助 Google MediaPipe 和 XNNPACK 團(tuán)隊(duì)將開(kāi)源大語(yǔ)言模型 (LLM) Gemma 的性能提高了 25%。

我們的目標(biāo)不僅僅是將 KleidiAI 打造成 AI 優(yōu)化例程的集合，更是希望該項(xiàng)目能成為學(xué)習(xí) Arm CPU 上軟件優(yōu)化最佳實(shí)踐的知識(shí)庫(kù)。因此，我們誠(chéng)摯邀請(qǐng)開(kāi)發(fā)者立即加入這激動(dòng)人心的學(xué)習(xí)之旅，并提供反饋意見(jiàn)，共同改進(jìn)我們的產(chǎn)品。

接下來(lái)，我將詳細(xì)探討 KleidiAI 的初始功能。點(diǎn)擊閱讀原文，可獲取一個(gè)關(guān)鍵函數(shù)的分步運(yùn)行指南，該函數(shù)用于加速 Gemma LLM 的 4 位整數(shù)矩陣乘法例程。

微內(nèi)核

首先介紹 KleidiAI 庫(kù)中提供的微內(nèi)核。KleidiAI 是面向 AI 框架開(kāi)發(fā)者的開(kāi)源庫(kù)，可以為 Arm CPU 提供經(jīng)過(guò)優(yōu)化的性能關(guān)鍵型例程。這些例程是以高性能加速給定算子所需的近乎最小規(guī)模的軟件，通常稱為微內(nèi)核（或 uKernel）。

以通過(guò) Winograd 算法執(zhí)行的 2D 卷積算子為例，該計(jì)算涉及四個(gè)主要運(yùn)算：

Winograd 輸入轉(zhuǎn)換

Winograd 濾波轉(zhuǎn)換

矩陣乘法

Winograd 輸出轉(zhuǎn)換

上述每個(gè)運(yùn)算都是一個(gè)微內(nèi)核。但為什么上述運(yùn)算不叫“內(nèi)核”或“函數(shù)”呢？如下圖所示，微內(nèi)核一詞強(qiáng)調(diào)了其處理輸出張量部分的能力：

圖：微內(nèi)核案例

如上圖中的兩個(gè)案例所示，微內(nèi)核僅能處理整個(gè)輸出的一部分。這一設(shè)計(jì)決策支持細(xì)粒度優(yōu)化，例如，提供高效串聯(lián)多個(gè)微內(nèi)核的靈活性，從而進(jìn)一步提升 AI 框架的性能。

微內(nèi)核可用于各種 Arm 架構(gòu)、技術(shù)和計(jì)算參數(shù)。例如，在 matmul_clamp_f32_qai8dxp_qsi4cxp 文件夾中，來(lái)看下具有按通道量化功能的 Int4 矩陣乘法例程：

該文件夾包含用于 4 位整數(shù)矩陣乘法的關(guān)鍵 LLM 微內(nèi)核。在該文件夾中可以看到，使用 Arm 點(diǎn)積或 i8mm 擴(kuò)展優(yōu)化的微內(nèi)核使用 Neon 匯編來(lái)盡可能提高效率，以及在處理的最小輸出塊方面有不同的多種變體。

文件夾中的每個(gè) .c 和 .h 文件對(duì)均代表一種微內(nèi)核變體。這些變體的計(jì)算參數(shù)（如塊大?。?、使用的 Arm 技術(shù)（如 Arm Neon）和特定 Arm 架構(gòu)特性（如 FEAT_DotProd）各有不同。

所有微內(nèi)核變體共享相同的功能和界面，從而保持一致性。一致性是易于采用的關(guān)鍵因素，接下來(lái)我將談?wù)勎覀冞€采取了哪些其他措施來(lái)幫助框架開(kāi)發(fā)者輕松集成微內(nèi)核。

易于采用

我們深知，在 AI 框架中集成新庫(kù)時(shí)面臨著諸多挑戰(zhàn)，如庫(kù)的大小、外部依賴關(guān)系和文檔。因此，我們努力收集合作伙伴的反饋意見(jiàn)，并將其納入我們的未來(lái)計(jì)劃之中，盡可能幫助開(kāi)發(fā)者順暢完成集成。

為了實(shí)現(xiàn)這一點(diǎn)，KleidiAI 的設(shè)計(jì)原則是讓框架開(kāi)發(fā)者能夠輕松集成所需的微內(nèi)核。只需拉取相應(yīng)的 .c 和 .h 文件，以及所有微內(nèi)核共享的通用頭文件 (kai_common.h)，即可集成每個(gè)微內(nèi)核。我們稱之為“三文件微內(nèi)核依賴” (Three-file Micro-kernel Dependency)。這樣我們便實(shí)現(xiàn)了這一目標(biāo)。

您可以參考我們準(zhǔn)備的相關(guān)指南，學(xué)習(xí)如何使用微內(nèi)核。該指南探討了 4 位整數(shù)矩陣乘法微內(nèi)核的使用，這有助于提高 Gemma LLM 的性能。