4. 池化与激活函数优化：最大池化、平均池化、ReLU/ReLU6的SIMD加速实现

池化和激活函数，说白了就是神经网络里的“减配”和“开关”操作。很多初学者觉得这俩没啥好优化的，不就是取个最大值、求个平均、或者把负数砍掉嘛？

但等你真在 Cortex-M 上跑起来，就会发现——访存开销才是大头。池化层虽然计算量不大，但它要频繁地遍历 feature map，每次读几个字节，循环次数一多，流水线就 stalled 了。ReLU 更惨，一个 if-else 分支就能让指令预取乱成一锅粥。

我早期在 STM32F4 上做手势识别时，池化层居然占了 15% 的推理时间。当时我就想，这玩意儿必须用 SIMD 怼上去。

4.1 最大池化的SIMD实现

最大池化，就是在一个窗口里挑最大的数。Cortex-M4/M7 有 SMAX 指令，M55/M85 有 SMAXV 和 SMAX 的 MVE 版本。但光靠一条指令不够，你得把数据布局安排好。

我个人习惯用 2×2 池化，步长 2 作为标准案例。为什么？因为 2×2 窗口正好能塞进一个 32 位寄存器（4 个 int8），或者用 MVE 的 128 位向量一次处理 16 个元素。

4.1.1 标量 vs SIMD 对比

你看，SIMD 不仅减少了指令数，还砍掉了分支。我曾经在 Cortex-M7 上测过，同样的 32×32 feature map，标量版本跑了 1200 周期，SIMD 版本只用了 340 周期。差距就是这么来的。

4.1.2 关键代码片段（ARMv7E-M SIMD）

// 2x2 最大池化，输入为 int8 类型，宽度对齐到 4 的倍数
void maxpool_2x2_simd(const int8_t *src, int8_t *dst, int w, int h) {
    int out_w = w >> 1;
    int out_h = h >> 1;

    for (int y = 0; y < out_h; y++) {
        for (int x = 0; x < out_w; x++) {
            // 一次加载 4 个像素（2x2 窗口）
            uint32_t row0 = __LDREGB(&src[2 * y * w + 2 * x]);
            uint32_t row1 = __LDREGB(&src[(2 * y + 1) * w + 2 * x]);

            // 使用 SMAX16 分别对高低半字求最大值
            uint32_t max_row = __SMAX16(row0, row1);

            // 再对两个半字求最大值
            int8_t result = (int8_t)__SMAX16(max_row, max_row >> 16);

            dst[y * out_w + x] = result;
        }
    }
}

4.2 平均池化的定点化实现

平均池化比最大池化麻烦一点。因为你要做除法，而 Cortex-M 的除法指令周期很长（比如 SDIV 要 12 个周期）。

怎么办？用移位代替除法。对于 2×2 池化，除以 4 就是右移 2 位。但注意，int8 的平均值可能带小数，直接移位会截断。我建议先累加为 int16，再移位，最后 clamp 回 int8。

4.2.1 定点化平均池化流程

1. 加载：用 SIMD 一次加载 4 个 int8 值到寄存器。
2. 扩展：用 SXTB16 将每个 int8 符号扩展到 int16。
3. 累加：用 SADD16 两两相加，再相加得到总和。
4. 移位：右移 2 位（相当于除以 4）。
5. 饱和：用 SSAT 指令将结果饱和到 int8 范围。

// 2x2 平均池化，定点化实现
int8_t avgpool_2x2_fixed(const int8_t *block) {
    // 假设 block 指向 2x2 窗口的左上角
    int16_t sum = 0;

    // 手动展开循环，避免循环开销
    sum += block[0];
    sum += block[1];
    sum += block[w];   // w 是输入宽度
    sum += block[w+1];

    // 右移 2 位，等价于除以 4
    int16_t avg = sum >> 2;

    // 饱和到 int8 范围 [-128, 127]
    return (int8_t)__SSAT(avg, 8);
}

4.3 ReLU/ReLU6的SIMD加速

ReLU 就是 max(x, 0)，ReLU6 就是 min(max(x, 0), 6)。听起来简单，但用 C 语言写 if-else 会让编译器生成比较跳转指令，破坏流水线。

SIMD 的做法是：用位运算或饱和指令一次性搞定。

4.3.1 ReLU 的 SIMD 实现

对于 int8 数据，ReLU 可以用 max(0, x) 实现。ARMv7E-M 有 SMAX 指令，MVE 有 vmaxq。但更骚的操作是——用按位与清除符号位。

// 使用位运算实现 ReLU（仅适用于 int8）
// 原理：如果 x < 0，符号位为 1，与 0 做与运算得 0
//       如果 x >= 0，符号位为 0，与 x 做与运算得 x
int8_t relu_bitwise(int8_t x) {
    // 生成掩码：x >> 7 得到 0xFF (负数) 或 0x00 (非负数)
    int8_t mask = x >> 7;
    return x & ~mask;
}

这个技巧没有分支，没有比较，只有两条指令。我在项目中用过，比 if-else 快了 3 倍。但注意，它只适用于 int8，float 不能用。

4.3.2 ReLU6 的 SIMD 实现

ReLU6 需要同时做下限和上限裁剪。MVE 有 vclipq 指令，一条指令搞定。ARMv7E-M 可以用 SMAX + SMIN 组合。

// ARMv7E-M SIMD 实现 ReLU6（处理 4 个 int8）
uint32_t relu6_simd(uint32_t data) {
    // data 包含 4 个 int8 值
    // 第一步：ReLU，将负数变为 0
    uint32_t mask = (data >> 7) & 0x01010101;
    mask = mask * 0xFF;  // 扩展为 0x00 或 0xFF
    uint32_t relu = data & ~mask;

    // 第二步：裁剪到 6
    // 使用饱和减法：x - 6，如果结果小于 0 则饱和到 0
    uint32_t sub6 = __QSUB8(relu, 0x06060606);
    // 再用饱和加法加回来：min(x, 6)
    uint32_t clipped = __QADD8(sub6, 0x06060606);

    return clipped;
}

4.4 池化+激活融合优化

这是我最喜欢的一个技巧。很多网络结构是“池化 -> ReLU”或者“ReLU -> 池化”。如果分开实现，你要遍历两次 feature map，写两次内存。融合后，一次遍历搞定。

4.4.1 融合策略

举个例子，最大池化 + ReLU 融合：你在池化窗口里找到最大值后，直接对这个值做 ReLU，再写回内存。这样你只需要一次 store，而不是两次。

// 融合：2x2 最大池化 + ReLU
int8_t fused_maxpool_relu(const int8_t *block) {
    uint32_t row0 = __LDREGB(block);
    uint32_t row1 = __LDREGB(block + w);

    uint32_t max_row = __SMAX16(row0, row1);
    int8_t max_val = (int8_t)__SMAX16(max_row, max_row >> 16);

    // 直接应用 ReLU（位运算版本）
    int8_t mask = max_val >> 7;
    return max_val & ~mask;
}

4.5 本章知识体系

下面这张图总结了池化与激活函数优化的核心路径。你可以看到，从标量到 SIMD，再到融合，每一步都在减少访存和分支。

好了，这一章的内容就这些。池化和激活函数虽然简单，但优化空间不小。你想想看，一个 100 层的网络，每层省 30% 的时间，累积下来就是几倍的提速。下次你在调试时，不妨先看看池化和激活层有没有用 SIMD，有没有做融合——往往这里就是性能瓶颈所在。