基于ARM NEON与SSE的外层循环向量化实现方案咨询

原始循环代码 #

for (int i = 0; i < n; ++i) {
   b[i][0] = 0.0;
   for (int j = 1; j < n; ++j) {
      b[i][j] = b[i][j - 1] + a[i][j]; 
   }
}

问题分析 #

这段代码的内层循环存在循环携带依赖（b[i][j]依赖前一位置的b[i][j-1]），无法直接做内层向量化。计划采用外层循环向量化：并行处理4个不同i的内层循环实例，按j的顺序逐列迭代（比如先处理所有i的j=1，再处理j=2等）。但NEON和SSE不支持gather/scatter指令，无法直接高效访问跨行的非连续内存（不同i的b[i][j]在内存中跨行分布）。

高效实现思路 #

1. 内存布局转置（核心优化） #

既然跨行访问内存效率低且无法直接向量化，最直接的解决方法是先转置数组a和b：

• 转置后，原a[i][j]变为a_trans[j][i]，b[i][j]变为b_trans[j][i]
• 原始逻辑转置后变为：对每一行j，b_trans[j][0] = 0.0，然后b_trans[j][i] = b_trans[j][i-1] + a_trans[j][i]（i从1到n-1）
• 此时内层循环操作的是连续内存的元素，完全可以直接用NEON/SSE做向量化，无需处理跨行访问问题

转置的开销会被后续向量化带来的性能收益抵消，尤其当n较大时。转置本身也可以用向量指令加速实现。

2. 分块处理（避免全转置的折中方案） #

如果不想做全数组转置，可采用分块策略：

• 将数组按i维度分成大小为4的块（对应NEON/SSE的4个float32向量宽度），比如处理i=0-3、i=4-7等块
• 对每个块，提取连续4行的数据，将每一列的4个元素打包成一个向量（相当于对块内数据做局部转置）
• 按j迭代，对每个列向量做累加：初始向量为{0.0, 0.0, 0.0, 0.0}，每一步加上当前列的a向量，结果存入b的列向量
• 处理完所有j后，再将块内的列向量数据转回行布局存回原数组

这种方法避免了全数组转置，仅在块内做局部转置，适合内存受限的场景。

3. NEON/SSE具体实现示例 #

NEON分块实现

// 假设n是4的倍数，处理i维度的块
for (int i_block = 0; i_block < n; i_block += 4) {
    // 初始化b的第0列向量
    float32x4_t b_prev = vdupq_n_f32(0.0f);
    vst1q_f32(&b[i_block][0], b_prev);
    
    for (int j = 1; j < n; ++j) {
        // 加载当前列的4个元素（需提前完成块内列的连续存储）
        float32x4_t a_col = vld1q_f32(&a[i_block][j]);
        // 累加操作
        b_prev = vaddq_f32(b_prev, a_col);
        // 存储回b的当前列
        vst1q_f32(&b[i_block][j], b_prev);
    }
}

注：如果原始数组是行优先布局，需要先对块内的4行做列提取打包，确保&a[i_block][j]指向连续的列元素内存。

SSE分块实现

// 假设n是4的倍数
for (int i_block = 0; i_block < n; i_block += 4) {
    __m128 b_prev = _mm_set_ps(0.0f, 0.0f, 0.0f, 0.0f);
    _mm_store_ps(&b[i_block][0], b_prev);
    
    for (int j = 1; j < n; ++j) {
        // 加载当前列的4个元素（需提前完成块内列的连续存储）
        __m128 a_col = _mm_load_ps(&a[i_block][j]);
        b_prev = _mm_add_ps(b_prev, a_col);
        _mm_store_ps(&b[i_block][j], b_prev);
    }
}

4. 边界处理 #

如果n不是4的倍数，剩余不足4个i的部分用标量循环处理即可，避免复杂的向量边界逻辑。

内容的提问来源于stack exchange，提问作者Bogi