大矩阵按索引整行Efficient gather操作性能优化问题咨询

根因分析 #

你的代码性能瓶颈和分块优化失效的核心原因如下：

1. 原逻辑写入output时目标行dest随机跳转，当行宽等于C时，C=1024对应单行长4KB，单次写入整行会占用多组缓存线，随机跳转时缓存命中率极低，写入带宽被严重浪费。
2. 你尝试的C维度分块反而恶化性能：该方案将原矩阵行的读取次数从1次放大到了K倍，input侧的读取带宽占用直接翻K倍，完全抵消了分块可能带来的写入缓存收益。

可行优化建议 #

• 调整循环顺序，复用input加载结果
  把列遍历循环放到目标索引遍历循环外层，同一个input行的AVX向量只需要加载一次，就可以写入所有关联的dest行对应位置，大幅减少input侧的重复读取开销，修改后的核心逻辑如下：

  for(int i = 0;i < A; i ++)
  {
      int k_start = offsets[i], k_end = offsets[i+1];
      // 外层遍历列维度，同一向量只加载一次
      for(int j = 0;j < C/8; j++)
      {
          __m256 a = _mm256_load_ps(&input[i * C + j * 8]);
          // 内层遍历所有目标行，写入同一个向量
          for(int k = k_start; k < k_end; k ++)
          {
              int dest = index1[k];
              _mm256_store_ps(&output[dest * C + j * 8] ,a);
          }
      }
  }
  

• 使用非临时写入指令绕过缓存
  你的output数据写入后不会立即复用，把_mm256_store_ps替换为_mm256_stream_ps，可以绕过CPU缓存写入，避免output数据挤占input需要的缓存空间，同时减少缓存回写的额外开销。注意该指令要求写入地址是32字节对齐，只要你定义output数组时保证了AVX对齐就可以直接使用。
• 优化索引排序规则
  对同一个原行i对应的所有dest索引按升序排序，保证写入output的行号尽可能连续，大幅提升写入时的缓存命中率。
• 调整分块策略
  不要在列维度C做分块，改为对原行维度A分块：每次取S个原行处理，S的大小设置为刚好让S行input数据能完全放入L2缓存（比如L2缓存为256KB、单行长4KB时，S设为64），既可以保证input侧的缓存命中率，又不会额外增加input读取次数。
• 开启编译器优化参数
  编译时添加-O3 -march=native -fstrict-aliasing参数，让编译器自动做循环展开、指令调度优化，进一步提升运行效率。

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