从大型矩阵中高效收集（整行）

Question

我正在尝试执行一个简单的操作。我有一个大小为 A x B 的矩阵。我有一个长度为 C 的索引列表，我想通过根据索引从第一个矩阵中收集行来制作一个 C x B 矩阵。即索引 i 告诉我将第一个矩阵中的哪一行放入第二个矩阵的第 i 行。

我对索引进行了预排序，因此算法输入是固定的：我从 A x B 矩阵加载行并将该行写入 C x B 矩阵中的所有行。

代码如下所示：

for(int i = 0;i < A; i ++)
{
    for(int k = offsets[i]; k < offsets[i+1]; k ++)
    {
           int dest = index1[k];
                
           for(int j = 0;j < C/ 8; j++)
           {
                __m256 a = _mm256_load_ps(&input[i * C + j * 8]);
                _mm256_store_ps(&output[dest * C + j * 8] ,a);
           }
     }
 }

代码完全被写入内存所限制。 当 C 较小时，此代码是有效的。然而，当 C 增加时，它的扩展性很差，我推测这是由于缓存行为。 （与 C = 256 相比，C = 1024 需要 10 倍的时间）。

我尝试在 C 维度进行阻塞：

for(int c = 0; c < C; c+= K){
for(int i = 0;i < A; i ++)
{
    for(int k = offsets[i]; k < offsets[i+1]; k ++)
    {
           int dest = index1[k];
                
           for(int j = 0;j < C/ 8 / K; j++)
           {
                __m256 a = _mm256_load_ps(&input[i * C + c + j * 8]);
                _mm256_store_ps(&output[dest * C + c + j * 8] ,a);
           }
     }
 }
}

这实际上会进一步减慢代码速度。

有什么建议吗？

Answer 1

似乎内部循环只是一个流式复制操作。在这种情况下，缓存无关紧要。而是尝试使用简单的 memcpy() 代替，以便编译器可以产生更好的执行代码，希望如此。

//for(int j = 0;j < C/ 8; j++)
//{
//     __m256 a = _mm256_load_ps(&input[i * C + j * 8]);
//     _mm256_store_ps(&output[dest * C + j * 8] ,a);
//}

memcpy(&output[dest * C], &input[i * C], C * sizeof(float));

附录

如果不能得到满意的结果，最后还是选择 C++ 并用 parllel_for() 替换外循环。那么就有可能使缓存（或其他管道？）工作得更好一点。

parallel_for(0, A, [&](const int i) {

    for(int k = offsets[i]; k < offsets[i+1]; k++)
    {
       int dest = index1[k];
       memcpy(&output[dest * C], &input[i * C], C * sizeof(float));
     }
});