Arm KleidiAI軟件庫(kù)的功能解析

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

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

Arm KleidiAI

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

盡管 KleidiAI 仍處于早期階段，但已幫助 Google MediaPipe 和 XNNPACK 團(tuán)隊(duì)將開源大語(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)開發(fā)者立即加入這激動(dòng)人心的學(xué)習(xí)之旅，并提供反饋意見，共同改進(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 框架開發(fā)者的開源庫(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)幫助框架開發(fā)者輕松集成微內(nèi)核。

易于采用

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

為了實(shí)現(xiàn)這一點(diǎn)，KleidiAI 的設(shè)計(jì)原則是讓框架開發(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 的性能。