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

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));
     }
});

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

Efficient gather (of whole rows) from a large matrix

c

performance

caching

linear-algebra

avx2