AVX SIMD in matrix multiplication

Question

I have coded the following C function for multiplying two NxN matrices and using AVX vectors to speed up the calculation. It works but the speedup is not what is to be expected(some scalar code is faster).

What are my major mistakes when trying to use AVX? (I'm new to programming with it so I don't fully understand how the pitfalls when programming for performance)

void version5(int mat1[N][N], int mat2[N][N], int result[N][N])
{
__m256i vec_multi = _mm256_setzero_si256(); //Initialize vector to zero
__m256i vec_mat1 = _mm256_setzero_si256(); //Initialize vector to zero
__m256i vec_mat2 = _mm256_setzero_si256(); //Initialize vector to zero

int i, j, k;
for (i = 0; i < N; i += 8)
{
    for (j = 0; j < N; ++j)
    {
        //The cost heavy storing is done once for each iteration and then use them at all places needed, stores column of matrix2 in chunks of eight, uses setr to have them in the right order
        vec_mat2 = _mm256_setr_epi32(mat2[i][j], mat2[i+1][j],mat2[i+2][j],mat2[i+3][j], mat2[i+4][j],mat2[i+5][j],mat2[i+6][j],mat2[i+7][j]);

        for (k = 0; k < N; ++k) 
        {
            vec_mat1 = _mm256_loadu_si256(&mat1[k][i]); //Stores row of first matrix (eight in each iteration)
            vec_multi = _mm256_add_epi32(vec_multi ,_mm256_mullo_epi32(vec_mat1, vec_mat2));//Multiplies the vectors

            result[k][j] += _mm256_extract_epi32(vec_multi, 0) + _mm256_extract_epi32(vec_multi, 1) +_mm256_extract_epi32(vec_multi, 2) +_mm256_extract_epi32(vec_multi, 3) +_mm256_extract_epi32(vec_multi, 4) +_mm256_extract_epi32(vec_multi, 5) +_mm256_extract_epi32(vec_multi, 6) +_mm256_extract_epi32(vec_multi, 7);

            vec_multi = _mm256_setzero_si256();
        }
    }
}
}

Answer 1

The biggest performance problem (for small matrices), and a common mistake, is doing this:

_mm256_extract_epi32(vec_multi, 0) + _mm256_extract_epi32(vec_multi, 1) +_mm256_extract_epi32(vec_multi, 2) +_mm256_extract_epi32(vec_multi, 3) +_mm256_extract_epi32(vec_multi, 4) +_mm256_extract_epi32(vec_multi, 5) +_mm256_extract_epi32(vec_multi, 6) +_mm256_extract_epi32(vec_multi, 7);

This dominates not just horizontal screen space, but also the performance (or lack thereof) of the loop.

Actual extracts (an extract with an index of 0 is turned into vmovd by a reasonable compiler, which is less of a problem) have a throughput of only 1 per cycle on typical CPUs. There are 7 here, so even with just this horizontal addition the loop body could only execute once every 7 cycles (but there is other stuff in there too so it's worse).

There are slightly faster ways to do horizontal sums, for example (not really tuned or anything, just some simple "better than totally naive" hsum)

int hsum_epi32(__m256i x)
{
    __m128i l = _mm256_extracti128_si256(x, 0); //implicit
    __m128i h = _mm256_extracti128_si256(x, 1);
    l = _mm_add_epi32(l, h);
    l = _mm_hadd_epi32(l, l);
    return _mm_extract_epi32(l, 0) + _mm_extract_epi32(l, 1);
}

But that is not great either. Ideally there should just be no horizontal operation in the inner loop. Even more ideally, not anywhere. And that can be arranged: a block of 8 results can be computed by rearranging the computation so that horizontal operations turn into broadcasts. Broadcasting from memory is pretty cheap (naturally it sort of "wastes" the load by loading only one thing, but it's not a slow operation), and there are much fewer of them, so that's probably better.

Sort of like this (to show the general idea, not tested)

for (int i = 0; i < N; i++) {
  for (int j = 0; j < N; j += 8) {
    auto sum = _mm256_setzero_si256();
    for (int k = 0; k < N; k++) {
      auto bc_mat1 = _mm256_set1_epi32(mat1[i][k]);
      auto vec_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j]);
      auto prod = _mm256_mullo_epi32(bc_mat1, vec_mat2);
      sum = _mm256_add_epi32(sum, prod);
    }
    _mm256_storeu_si256((__m256i*)&result[i][j], sum);
  }
}

There is some room for improvement. The memory access pattern through mat2 is not ideal, the inner loop iterates over the rows and only half of a cache line is used every time, for a large enough matrix that means half of every cache miss is wasted. It's probably better to unroll the middle loop by 2 (or 4) again. As a bonus, we get to re-use the broadcast from mat1. Sort of like this (not tested, showing 2x)

for (int i = 0; i < N; i++) {
  for (int j = 0; j < N; j += 16) {
    auto sumA = _mm256_setzero_si256();
    auto sumB = _mm256_setzero_si256();
    for (int k = 0; k < N; k++) {
      auto bc_mat1 = _mm256_set1_epi32(mat1[i][k]);
      auto vecA_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j]);
      auto vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j + 8]);
      auto prodA = _mm256_mullo_epi32(bc_mat1, vecA_mat2);
      auto prodB = _mm256_mullo_epi32(bc_mat1, vecB_mat2);
      sumA = _mm256_add_epi32(sumA, prodA);
      sumB = _mm256_add_epi32(sumB, prodB);
    }
    _mm256_storeu_si256((__m256i*)&result[i][j], sumA);
    _mm256_storeu_si256((__m256i*)&result[i][j + 8], sumB);
  }
}

But, it's probably still possible to improve it, and I didn't test anything.

---

I've done some testing now, and made some improvement specifically for larger matrices (below N=64 it doesn't really help).

Some results on Haswell, compiled with MSVC^[1] 2017, measuring the time (in cycles) per element of the result matrix (so you can mentally compare it to how much time it should take). Time results were eyeballed and rounded to a "typical" value.

       64 128  256  512  1024  2048
naive  90 305 1360 2700  9500   TLE
v1     19  45  170  340  1460  7900
v2     18  40  170  245  1030  4300
v3     22  44   85  150   380   950
v4     17  35   70  150   310   750
v5     18  35   70  140   275   550
ideal  16  32   64  128   256   512

ideal is based on the throughput of vpmulld on Haswell, which is one every two cycles. So per cycle there can be 4 multiplications, we need N of them, so N/4 is the ideal time per element. For small sizes that's not so hard to get near, but for bigger matrices the memory access pattern messes up everything. v1 and v2 are the versions from above.

v3 adds loop tiling to significantly improve the performance for bigger matrices. Understandably this causes some overhead, noticeable for smaller matrices. v4 unrolls the i loop by 2x. Annoyingly, the best choice for the block size depends not only on cache size, but also on the size of the matrix. The times above are not all with the same parameters, but tuned a bit per N.

I'm not sure where to go from here but it seems as though some improvement should still be possible, for bigger sizes it's still a decent factor away from the ideal.

v3:

size_t jb = std::min(512u, N);
size_t kb = std::min(24u, N);

for (size_t jj = 0; jj < N; jj += jb)
{
    for (size_t kk = 0; kk < N; kk += kb)
    {
        for (size_t i = 0; i < N; i += 1) {
            for (size_t j = jj; j < jj + jb; j += 16) {
                __m256i sumA_1, sumB_1;
                if (kk == 0) {
                    sumA_1 = sumB_1 = _mm256_setzero_si256();
                }
                else {
                    sumA_1 = _mm256_load_si256((__m256i*)&result[i][j]);
                    sumB_1 = _mm256_load_si256((__m256i*)&result[i][j + 8]);
                }
                size_t limit = std::min(N, kk + kb);
                for (size_t k = kk; k < limit; k++) {
                    auto bc_mat1_1 = _mm256_set1_epi32(mat1[i][k]);
                    auto vecA_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j]);
                    auto vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j + 8]);
                    sumA_1 = _mm256_add_epi32(sumA_1, _mm256_mullo_epi32(bc_mat1_1, vecA_mat2));
                    sumB_1 = _mm256_add_epi32(sumB_1, _mm256_mullo_epi32(bc_mat1_1, vecB_mat2));
                }
                _mm256_storeu_si256((__m256i*)&result[i][j], sumA_1);
                _mm256_storeu_si256((__m256i*)&result[i][j + 8], sumB_1);
            }
        }
    }
}

v4:

size_t jb = std::min(512u, N);
size_t kb = std::min(24u, N);

for (size_t jj = 0; jj < N; jj += jb)
{
    for (size_t kk = 0; kk < N; kk += kb)
    {
        for (size_t i = 0; i < N; i += 2) {
            for (size_t j = jj; j < jj + jb; j += 16) {
                __m256i sumA_1, sumB_1, sumA_2, sumB_2;
                if (kk == 0) {
                    sumA_1 = sumB_1 = sumA_2 = sumB_2 = _mm256_setzero_si256();
                }
                else {
                    sumA_1 = _mm256_load_si256((__m256i*)&result[i][j]);
                    sumB_1 = _mm256_load_si256((__m256i*)&result[i][j + 8]);
                    sumA_2 = _mm256_load_si256((__m256i*)&result[i + 1][j]);
                    sumB_2 = _mm256_load_si256((__m256i*)&result[i + 1][j + 8]);
                }
                size_t limit = std::min(N, kk + kb);
                for (size_t k = kk; k < limit; k++) {
                    auto bc_mat1_1 = _mm256_set1_epi32(mat1[i][k]);
                    auto vecA_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j]);
                    auto vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j + 8]);
                    sumA_1 = _mm256_add_epi32(sumA_1, _mm256_mullo_epi32(bc_mat1_1, vecA_mat2));
                    sumB_1 = _mm256_add_epi32(sumB_1, _mm256_mullo_epi32(bc_mat1_1, vecB_mat2));
                    auto bc_mat1_2 = _mm256_set1_epi32(mat1[i + 1][k]);
                    sumA_2 = _mm256_add_epi32(sumA_2, _mm256_mullo_epi32(bc_mat1_2, vecA_mat2));
                    sumB_2 = _mm256_add_epi32(sumB_2, _mm256_mullo_epi32(bc_mat1_2, vecB_mat2));
                }
                _mm256_storeu_si256((__m256i*)&result[i][j], sumA_1);
                _mm256_storeu_si256((__m256i*)&result[i][j + 8], sumB_1);
                _mm256_storeu_si256((__m256i*)&result[i + 1][j], sumA_2);
                _mm256_storeu_si256((__m256i*)&result[i + 1][j + 8], sumB_2);
            }
        }
    }
}

---

Yet an other version, with even more tiling and with rearranging matrix 2. Added to the table of times above. Of course, some time can be saved if that matrix can be assumed to already be in that order, but I counted the rearranging in the benchmarks. That overhead scales as O(N²) while the meat of the algorithm scales as O(N³) so for a large matrix it does not represent a significant cost anyway.

It seems to behave well now, staying around 110% of the theoretical optimum for any size I test. Perhaps some small tweaks are still possible. For example, unrolling the i loop by 4 instead of 2 improved it slightly in my tests, but the difference is kind of small.

v5:

size_t ib = std::min(256, (int)N);
size_t jb = std::min(512, (int)N);
size_t kb = std::min(16, (int)N);

int *mat2 = (int*)_aligned_malloc(N * N * sizeof(int), 32);
size_t m2idx = 0;
for (size_t jj = 0; jj < N; jj += jb)
{
    for (size_t kk = 0; kk < N; kk += kb)
    {           
        for (size_t j = jj; j < jj + jb; j += 16)
        {
            for (size_t k = kk; k < kk + kb; k++)
            {
                auto vecA_mat2 = _mm256_loadu_si256((__m256i*)&mat2in[k][j]);
                auto vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2in[k][j + 8]);
                _mm256_storeu_si256((__m256i*)&mat2[m2idx], vecA_mat2);
                _mm256_storeu_si256((__m256i*)&mat2[m2idx + 8], vecB_mat2);
                m2idx += 16;
            }
        }
    }
}


for (size_t ii = 0; ii < N; ii += ib) {
    for (size_t jj = 0; jj < N; jj += jb) {
        for (size_t kk = 0; kk < N; kk += kb) {
            for (size_t i = ii; i < ii + ib; i += 2) {
                for (size_t j = jj; j < jj + jb; j += 16) {
                    size_t m2idx = (j - jj) * kb + kk * jb + jj * N;
                    __m256i sumA_1, sumB_1, sumA_2, sumB_2;
                    if (kk == 0) {
                        sumA_1 = sumB_1 = sumA_2 = sumB_2 = _mm256_setzero_si256();
                    }
                    else {
                        sumA_1 = _mm256_load_si256((__m256i*)&result[i][j]);
                        sumB_1 = _mm256_load_si256((__m256i*)&result[i][j + 8]);
                        sumA_2 = _mm256_load_si256((__m256i*)&result[i + 1][j]);
                        sumB_2 = _mm256_load_si256((__m256i*)&result[i + 1][j + 8]);
                    }

                    for (size_t k = kk; k < kk + kb; k++) {
                        auto bc_mat1_1 = _mm256_set1_epi32(mat1[i][k]);
                        auto vecA_mat2 = _mm256_loadu_si256((__m256i*)(mat2 + m2idx));
                        auto vecB_mat2 = _mm256_loadu_si256((__m256i*)(mat2 + m2idx + 8));
                        sumA_1 = _mm256_add_epi32(sumA_1, _mm256_mullo_epi32(bc_mat1_1, vecA_mat2));
                        sumB_1 = _mm256_add_epi32(sumB_1, _mm256_mullo_epi32(bc_mat1_1, vecB_mat2));
                        auto bc_mat1_2 = _mm256_set1_epi32(mat1[i + 1][k]);
                        sumA_2 = _mm256_add_epi32(sumA_2, _mm256_mullo_epi32(bc_mat1_2, vecA_mat2));
                        sumB_2 = _mm256_add_epi32(sumB_2, _mm256_mullo_epi32(bc_mat1_2, vecB_mat2));
                        m2idx += 16;
                    }
                    _mm256_storeu_si256((__m256i*)&result[i][j], sumA_1);
                    _mm256_storeu_si256((__m256i*)&result[i][j + 8], sumB_1);
                    _mm256_storeu_si256((__m256i*)&result[i + 1][j], sumA_2);
                    _mm256_storeu_si256((__m256i*)&result[i + 1][j + 8], sumB_2);
                }
            }
        }
    }
}

_aligned_free(mat2);

---

[1] MSVC is a C++ compiler so my examples are technically not C, but it's really about the approach anyway and it's trivial to convert.

Answer 2

Thanks for all comments and tips! Especially @harold This is my final version with the main goal to eliminate cache misses by never loading columns into a vector.

It removes the costly set and load function by working with pointer loading and storing instead of element. By working row by row the cachemisses are few and memory blocks are used more fully.

@harold if you want to test it against your version I would be interested to see how it stacks up against them.

void version5(int mat1[N][N], int mat2[N][N], int result[N][N])
{
__m256i vec_multi_res = _mm256_setzero_si256(); //Initialize vector to zero
__m256i vec_mat1 = _mm256_setzero_si256(); //Initialize vector to zero
__m256i vec_mat2 = _mm256_setzero_si256(); //Initialize vector to zero

int i, j, k;
for (i = 0; i < N; i++)
{
    for (j = 0; j < N; ++j)
    {
        //Stores one element in mat1 and use it in all computations needed before proceeding
        //Stores as vector to increase computations per cycle
        vec_mat1 = _mm256_set1_epi32(mat1[i][j]);

        for (k = 0; k < N; k += 8)
        {
            vec_mat2 = _mm256_loadu_si256((__m256i*)&mat2[j][k]); //Stores row of second matrix (eight in each iteration)
            vec_multi_res = _mm256_loadu_si256((__m256i*)&result[i][k]); //Loads the result matrix row as a vector
            vec_multi_res = _mm256_add_epi32(vec_multi_res ,_mm256_mullo_epi32(vec_mat1, vec_mat2));//Multiplies the vectors and adds to th the result vector

            _mm256_storeu_si256((__m256i*)&result[i][k], vec_multi_res); //Stores the result vector into the result array
        }
    }
    }
}