InferLLM 框架中 ARM 优化模块的 optimized.h 分析

optimized.h 文件是 InferLLM 框架中 ARM 优化模块的核心部分，它实现了各种基础向量运算函数，这些函数被 kernel.cpp 中的高级计算函数调用。下面对其代码结构和功能实现进行详细分析。

1. 文件结构概述

#pragma once

#include <assert.h>
#include "arm_neon.h"
#include "kern/kernel_define.h"

namespace inferllm {
namespace opt {
    // 各种优化的向量运算函数
}  // namespace opt
}  // namespace inferllm

该文件包含了必要的头文件，其中 arm_neon.h 是 ARM NEON 指令集的头文件，提供了 NEON 指令集的内联函数。所有函数都定义在 inferllm::opt 命名空间中。

2. 基础元素级运算函数

2.1 向量加法

inline void elemwise_vector_add(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    for (int i = 0; i < n; i++) {
        z[i] = x[i] + y[i];
    }
}

这个函数实现了向量加法，将两个向量 x 和 y 相加，结果存储在向量 z 中。使用 __restrict 关键字告诉编译器这些指针不会重叠，有助于编译器生成更优化的代码。

2.2 向量乘法

inline void elemwise_vector_mul(
        const int n, const float* __restrict x, const float* __restrict y,
        float* __restrict z) {
    for (int i = 0; i < n; i++) {
        z[i] = x[i] * y[i];
    }
}

这个函数实现了向量乘法，将两个向量 x 和 y 对应元素相乘，结果存储在向量 z 中。

2.3 SiLU 激活函数

inline void elemwise_vector_silu(
        const int n, const float* __restrict x, float* __restrict z) {
    for (int i = 0; i < n; i++) {
        z[i] = x[i] / (1 + exp(-x[i]));
    }
}

这个函数实现了 SiLU 激活函数，也称为 Swish 激活函数，公式为 f(x) = x * sigmoid(x)，这里使用等价形式 f(x) = x / (1 + exp(-x))。

2.4 GELU 激活函数

inline void elemwise_vector_gelu(
        const int n, const float* __restrict x, float* __restrict z) {
    for (int i = 0; i < n; i++) {
        float src = x[i];
        z[i] = 0.5 * src * (1 + tanh(sqrt(2.0 / PI) * (src + PGELU * src * src * src)));
    }
}

这个函数实现了 GELU 激活函数，使用近似公式 f(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))，其中 PGELU 是常数 0.044715。

2.5 向量缩放

inline void elemwise_vec_scale(
        const int n, const float* __restrict x, float scale, float* __restrict z) {
    int i = 0;
    for (; i < n; i++) {
        z[i] = x[i] * scale;
    }
}

这个函数实现了向量缩放，将向量 x 的每个元素乘以缩放因子 scale，结果存储在向量 z 中。

3. 归约运算函数

3.1 平方和归约

inline float reduce_square_sum(const int n, const float* __restrict x) {
    float sum = 0.0f;
    for (int i = 0; i < n; i++) {
        sum += x[i] * x[i];
    }
    return sum;
}

这个函数计算向量 x 的平方和，用于 RMS 归一化。

3.2 最大值归约

inline float reduce_max(const int n, const float* __restrict x) {
    float max = -INFINITY;
    for (int i = 0; i < n; i++) {
        max = std::max(max, x[i]);
    }
    return max;
}

这个函数找到向量 x 的最大值，用于 Softmax 计算。

3.3 减去最大值并计算指数和

inline float select_sub_max_and_reduce_sum(
        const int n, const float* __restrict x, float* __restrict y, const float max) {
    float sum = 0.0f;
    for (uint32_t i = 0; i < n; i++) {
        if (x[i] == -INFINITY) {
            y[i] = 0.0f;
        } else {
            float val = exp(x[i] - max);
            sum += val;
            y[i] = val;
        }
    }
    return sum;
}

这个函数将向量 x 的每个元素减去最大值 max，然后计算指数，结果存储在向量 y 中，并返回指数和。这是 Softmax 计算的一部分，减去最大值可以提高数值稳定性。

4. 矩阵乘法函数

4.1 带偏移的矩阵乘法

inline void compute_src_offset_embd_matmul(
        const float* __restrict srcq_head, int offsetq,
        const float* __restrict srck_head, int offsetk, float* dst_head, int seqlen,
        int length, int sub_embd) {
    for (uint32_t row = 0; row < seqlen; row++) {
        auto p_srcq = srcq_head + row * offsetq;
        uint32_t len = 0;
        for (; len + 3 < length; len += 4) {
            auto p_dst = dst_head + row * length + len;
            auto p_srck0 = srck_head + len * offsetk;
            auto p_srck1 = srck_head + (len + 1) * offsetk;
            auto p_srck2 = srck_head + (len + 2) * offsetk;
            auto p_srck3 = srck_head + (len + 3) * offsetk;
            float sum0 = 0;
            float sum1 = 0;
            float sum2 = 0;
            float sum3 = 0;
            for (uint32_t k = 0; k < sub_embd; k++) {
                sum0 += p_srck0[k] * p_srcq[k];
                sum1 += p_srck1[k] * p_srcq[k];
                sum2 += p_srck2[k] * p_srcq[k];
                sum3 += p_srck3[k] * p_srcq[k];
            }
            p_dst[0] = sum0;
            p_dst[1] = sum1;
            p_dst[2] = sum2;
            p_dst[3] = sum3;
        }
        for (; len < length; len++) {
            auto p_dst = dst_head + row * length + len;
            auto p_srck = srck_head + len * offsetk;
            float sum = 0;
            for (uint32_t k = 0; k < sub_embd; k++) {
                sum += p_srck[k] * p_srcq[k];
            }
            *p_dst = sum;
        }
    }
}

这个函数实现了带偏移的矩阵乘法，用于多头注意力中的 Q 和 K 的矩阵乘法。它使用分块处理策略，每次处理 4 列，提高计算效率。

4.2 带不连续目标的矩阵乘法

inline void comput_matmul_with_dst_uncontinue(
        float* __restrict dst, int offset_dst, const float* __restrict srcv,
        int offset_v, const float* __restrict srcqk, int seqlen, int length, int K) {
    for (uint32_t row = 0; row < seqlen; row++) {
        auto p_qk = srcqk + row * length;
        for (uint32_t len = 0; len < K; len++) {
            auto p_dst = dst + row * offset_dst + len;
            auto p_v = srcv + len;
            float sum = 0;
            for (uint32_t k = 0; k < length; k++) {
                sum += p_v[k * offset_v] * p_qk[k];
            }
            *p_dst = sum;
        }
    }
}

这个函数实现了带不连续目标的矩阵乘法，用于多头注意力中的 QK 和 V 的矩阵乘法。目标矩阵 dst 的元素不是连续存储的，而是按照 offset_dst 的偏移存储。

5. 量化点积计算

inline float vec_vec_dot_q40_with_q80(
        const int n, const void* __restrict vx, const void* __restrict vy) {
    // ... 前面部分已分析 ...

    // 根据是否支持 ARM 点积指令选择不同的实现
#if defined(__ARM_FEATURE_DOTPROD)
    // 使用 SDOT 指令计算点积
    const int32x4_t p_0 =
            vdotq_s32(vdotq_s32(vdupq_n_s32(0), v0_0ls, v1_0ls), v0_0hs, v1_0hs);
    const int32x4_t p_1 =
            vdotq_s32(vdotq_s32(vdupq_n_s32(0), v0_1ls, v1_1ls), v0_1hs, v1_1hs);

    // 乘以量化参数并累加
    sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(p_0), x0->d * y0->d);
    sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(p_1), x1->d * y1->d);
#else
    // 不支持 SDOT 指令时的模拟实现
    // 使用 vmull_s8 计算 8 位整数的乘法，结果为 16 位整数
    const int16x8_t pl0l = vmull_s8(vget_low_s8(v0_0ls), vget_low_s8(v1_0ls));
    const int16x8_t pl0h = vmull_s8(vget_high_s8(v0_0ls), vget_high_s8(v1_0ls));
    const int16x8_t ph0l = vmull_s8(vget_low_s8(v0_0hs), vget_low_s8(v1_0hs));
    const int16x8_t ph0h = vmull_s8(vget_high_s8(v0_0hs), vget_high_s8(v1_0hs));

    const int16x8_t pl1l = vmull_s8(vget_low_s8(v0_1ls), vget_low_s8(v1_1ls));
    const int16x8_t pl1h = vmull_s8(vget_high_s8(v0_1ls), vget_high_s8(v1_1ls));
    const int16x8_t ph1l = vmull_s8(vget_low_s8(v0_1hs), vget_low_s8(v1_1hs));
    const int16x8_t ph1h = vmull_s8(vget_high_s8(v0_1hs), vget_high_s8(v1_1hs));

    // 使用 vpaddlq_s16 计算相邻元素的和，结果为 32 位整数
    const int32x4_t pl0 = vaddq_s32(vpaddlq_s16(pl0l), vpaddlq_s16(pl0h));
    const int32x4_t ph0 = vaddq_s32(vpaddlq_s16(ph0l), vpaddlq_s16(ph0h));
    const int32x4_t pl1 = vaddq_s32(vpaddlq_s16(pl1l), vpaddlq_s16(pl1h));
    const int32x4_t ph1 = vaddq_s32(vpaddlq_s16(ph1l), vpaddlq_s16(ph1h));

    // 乘以量化参数并累加
    sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(vaddq_s32(pl0, ph0)), x0->d * y0->d);
    sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(vaddq_s32(pl1, ph1)), x1->d * y1->d);
#endif
    }
    // 水平求和并返回结果
    return vaddvq_f32(sumv0) + vaddvq_f32(sumv1);
}

这个函数实现了 4 位整数量化与 8 位整数量化的点积计算，是矩阵乘法中的核心计算部分。它使用 NEON 指令集进行向量化计算，并根据处理器是否支持 SDOT 指令选择不同的实现。

主要优化点包括：

1. 条件编译：根据 __ARM_FEATURE_DOTPROD 宏判断处理器是否支持 SDOT 指令，选择不同的实现
2. NEON 指令集优化：使用 NEON 指令集进行向量化计算，每次处理多个数据元素
3. 分块处理：每次处理两个块，减少循环开销
4. 并行计算：使用多个向量寄存器并行计算，提高指令级并行性
5. 内存访问优化：使用连续的内存访问模式，提高缓存命中率

6. NEON 指令集优化分析

ARM 优化模块使用 NEON 指令集进行向量化计算，提高计算效率。NEON 指令集是 ARM 处理器的 SIMD（单指令多数据）扩展，可以同时处理多个数据元素。

6.1 NEON 指令集的基本操作

1. 数据加载：使用 vld1q_u8、vld1q_s8 等指令加载数据到向量寄存器
2. 数据存储：使用 vst1q_f32 等指令将向量寄存器中的数据存储到内存
3. 算术运算：使用 vaddq_f32、vmulq_f32 等指令进行向量加法、乘法等运算
4. 位运算：使用 vandq_u8、vshrq_n_u8 等指令进行位运算
5. 类型转换：使用 vreinterpretq_s8_u8、vcvtq_f32_s32 等指令进行类型转换
6. 归约运算：使用 vaddvq_f32 等指令进行水平求和

6.2 NEON 指令集在 optimized.h 中的应用

1. 向量加法：可以使用 vaddq_f32 指令同时处理 4 个浮点数
2. 向量乘法：可以使用 vmulq_f32 指令同时处理 4 个浮点数
3. 点积计算：可以使用 vdotq_s32 指令（如果支持）或模拟实现计算点积
4. 归约运算：可以使用 vaddvq_f32 指令计算向量的水平求和

6.3 NEON 指令集的优化效果

使用 NEON 指令集可以显著提高计算效率，特别是对于大型矩阵乘法和向量运算。在支持 SDOT 指令的处理器上，点积计算的性能更高。

7. 优化策略分析

7.1 向量化计算

ARM 优化模块使用 NEON 指令集进行向量化计算，每次处理多个数据元素，提高计算效率。例如，在点积计算中，每次处理 32 个 4 位整数和 32 个 8 位整数。

7.2 分块处理

ARM 优化模块使用分块处理策略，将数据分成多个块进行处理，减少循环开销。例如，在矩阵乘法中，每次处理 4 列，在点积计算中，每次处理两个块。

7.3 条件编译

ARM 优化模块使用条件编译，根据处理器支持的指令集选择不同的实现。例如，在点积计算中，根据处理器是否支持 SDOT 指令选择不同的实现。

7.4 内存访问优化

ARM 优化模块优化了内存访问模式，使用连续的内存访问模式，提高缓存命中率。例如，在矩阵乘法中，一次加载输入数据，然后计算多个输出元素。

8. 与 naive 实现的比较

ARM 优化模块中的函数与 naive 模块中的函数相比，主要区别在于：

1. 向量化实现：ARM 优化模块使用 NEON 指令集进行向量化计算，naive 模块使用标量实现
2. 分块处理：ARM 优化模块使用分块处理策略，naive 模块使用简单的循环
3. 条件编译：ARM 优化模块使用条件编译，根据处理器支持的指令集选择不同的实现，naive 模块没有这种优化
4. 内存访问优化：ARM 优化模块优化了内存访问模式，naive 模块没有这种优化

在实际应用中，ARM 优化模块的性能明显优于 naive 模块，特别是在支持 NEON 指令集的 ARM 处理器上。

9. 未来优化方向

基于当前实现，可以考虑以下优化方向：

9.1 更多 NEON 指令集优化

1. 使用 NEON 指令集优化基础向量运算函数：目前 elemwise_vector_add、elemwise_vector_mul 等函数没有使用 NEON 指令集优化，可以添加 NEON 实现
2. 使用 NEON 指令集优化激活函数：目前 elemwise_vector_silu、elemwise_vector_gelu 等函数没有使用 NEON 指令集优化，可以添加 NEON 实现
3. 使用 NEON 指令集优化归约运算：目前 reduce_square_sum、reduce_max 等函数没有使用 NEON 指令集优化，可以添加 NEON 实现

9.2 更多 ARM 指令集支持

1. 支持 ARMv8.2-A 的 FP16 指令：使用半精度浮点数进行计算，减少内存占用和计算量
2. 支持 ARMv8.2-A 的 DotProd 指令：加速点积计算，提高矩阵乘法性能
3. 支持 ARMv8.6-A 的 BFloat16 指令：使用 BFloat16 进行计算，减少内存占用和计算量

9.3 更高效的算法

1. 使用 Winograd 算法优化矩阵乘法：减少乘法次数，提高计算效率
2. 使用 Flash Attention 算法优化注意力计算：减少内存占用和计算量
3. 使用混合精度计算提高性能：在不同的计算阶段使用不同的精度，平衡精度和性能

总结

InferLLM 框架中的 ARM 优化模块通过 optimized.h 文件实现了各种基础向量运算函数，这些函数被 kernel.cpp 中的高级计算函数调用。optimized.h 文件中的函数使用 NEON 指令集进行向量化计算，使用分块处理策略和条件编译等优化技术，提高大语言模型推理的性能。

与 naive 模块相比，ARM 优化模块的性能明显更高，特别是在支持 NEON 指令集的 ARM 处理器上。未来可以考虑添加更多 NEON 指令集优化、支持更多 ARM 指令集和使用更高效的算法等方向进行优化，进一步提高大语言模型在 ARM 平台上的推理性能。