Efficient gather (of whole rows) from a large matrix - c

Question

I am trying to perform a simple operation. I have a matrix that is A x B by size. I have a list of indices of length C, and I want to make a C x B matrix by collecting rows from the first matrix according to the indices. i.e. index i tells me which row from the first matrix I put into row i in the second matrix.
I presorted the indices so the algorithm is input stationary: I load in the row from the A x B matrix and write that row to all the rows in the C x B matrix.
The code looks something like this:
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);
}
}
}
The code is entirely bottlenecked by write to memory.
This code is efficient when C is small. However it scales very poorly when C increases, which I surmise is due to cache behavior. (It takes 10x time when C = 1024 compared to C = 256).
I tried blocking in the C dimension:
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);
}
}
}
}
This actually slows down the code more.
Any suggestions?

Answer 1

It seems the inner loop is a mere streamed copy operation. Cache wouldn't matter in such a case. Rather try using simple memcpy() instead so the compiler can yield better execution code, hopefully.
//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));
Appendix
If satisfiable results won't be obtained, in the last resort, take C++ and replace the outer loop with parllel_for(). Then it may be possible to make the cache(or otherwise pipeline?) work a little bit better.
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));
}
});