如何利用 ARM i8mm 指令集加速端侧量化模型推理

在端侧 AI 推理（如手机、嵌入式设备）中，Int8 量化是提升性能、降低能耗的核心技术。传统的 ARM NEON 指令集虽有 SDOT (点乘) 指令，但在处理大规模矩阵乘法时仍显吃力。ARMv8.2-A 引入的可选扩展 i8mm (Int8 Matrix Multiplication) 指令集，通过在单周期内完成更高效的矩阵计算，为端侧推理带来了质的飞跃。

什么是 i8mm？

i8mm 是一组专门为 8 位整数矩阵乘法设计的扩展指令。它最核心的指令是 SMMLA (Signed Matrix Multiply-Accumulate)。与 SDOT 每次处理 4 个元素并累加到一个结果不同，SMMLA 可以直接处理 2×8 和 8×2 的小矩阵块，生成 2×2 的 32 位累加结果。这种二维维度的并行度极大减少了寄存器存取压力。

环境准备

要使用 i8mm，你需要满足以下条件：
1. 硬件支持：CPU 核心需基于 ARMv8.2-A 或更高架构（如 Cortex-X1, A78, A510 等）。
2. 编译器支持：GCC 10+ 或 LLVM 11+。
3. 编译标志：在编译参数中添加 -march=armv8.2-a+i8mm。

实操：如何检测 i8mm 支持

在代码中，我们首先需要判断当前运行的 CPU 是否支持 i8mm：

#include <sys/auxv.h>
#include <asm/hwcap.h>
#include <iostream>

bool check_i8mm_support() {
    unsigned long hwcaps2 = getauxval(AT_HWCAP2);
    if (hwcaps2 & HWCAP2_I8MM) {
        return true;
    }
    return false;
}

int main() {
    if (check_i8mm_support()) {
        std::cout << "i8mm is supported!" << std::endl;
    } else {
        std::cout << "i8mm not supported, fallback to NEON." << std::endl;
    }
    return 0;
}

核心指令应用：vmmlaq_s32

在 C++ 中，我们通常使用 Intrinsics（内联函数）来调用 i8mm 指令。核心函数是 vmmlaq_s32，它执行以下操作：

C = A (2×8 signed int8) * B (8×2 signed int8) + C (2×2 signed int32)

以下是一个简单的 2×8 与 8×2 矩阵乘法的实现示例：

#include <arm_neon.h>
#include <vector>

void gemm_i8mm_kernel(int8_t* a_ptr, int8_t* b_ptr, int32_t* c_ptr) {
    // 加载 A 矩阵 (2x8 int8 -> 128-bit vector)
    int8x16_t va = vld1q_s8(a_ptr);
    // 加载 B 矩阵 (8x2 int8 -> 128-bit vector)
    int8x16_t vb = vld1q_s8(b_ptr);
    // 加载 C 矩阵初始值 (2x2 int32 -> 128-bit vector)
    int32x4_t vc = vld1q_s32(c_ptr);

    // 调用 i8mm 核心指令
    // res[0] = A[0..7] * B[col0] + C[0]
    // res[1] = A[0..7] * B[col1] + C[1]
    // res[2] = A[8..15] * B[col0] + C[2]
    // res[3] = A[8..15] * B[col1] + C[3]
    int32x4_t res = vmmlaq_s32(vc, va, vb);

    // 存储结果
    vst1q_s32(c_ptr, res);
}

性能优势分析

1. 更高吞吐量：相比 SDOT 指令，SMMLA 在相同周期内处理的数据量翻倍，吞吐量理论提升约 2 倍。
2. 降低访存压力：由于 SMMLA 内部完成了更多维度的累加，减少了中间变量在寄存器与缓存之间的交换。
3. 适配 NCNN/MNN：目前主流端侧推理框架如 NCNN 和 MNN 都已经针对 i8mm 进行了汇编级优化。在开启 i8mm 后，典型的卷积神经网络（如 MobileNetV2）在量化模式下的速度通常能提升 30%-50%。

总结

ARM i8mm 是目前端侧 AI 推理性能优化的“秘密武器”。如果你的应用场景涉及实时视频处理、手势识别或语音识别，且运行在新款 ARM 处理器上，手动或通过框架接入 i8mm 能够显著降低延迟。开发者应优先检查 SMMLA 的可用性，并结合内存布局（如 Data Packing）充分发挥其计算威力。