AVX SIMD in matrix multiplication

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

• You'll receive better reviews if you show a complete example. For example, I recommend that you edit to show the necessary #include lines, and a main() that shows how to call your function. It's not mandatory, but it really helps! – Toby Speight Oct 10 '17 at 12:26
• Right now you have your k loop incrementing by one, but you're commenting that you increment by 8 - what is intended here? – Dannnno Oct 10 '17 at 12:42
• @Dannnno Thanks for pointing it out. I moved to 8 increment to i instead since I set vec_mat2 and then use it until it is no longer needed. The vector operates on 8 integers at a time. The vec_mat2 holds eight element from one column and is because of row-major storing in memory the slowest to retrive from memory. – Henrik Ståhlberg Oct 10 '17 at 12:45
• You say some scalar code is faster - can you include which scalar code, and at what problem sizes? If its just the normal naive mmult, and/or at all problem sizes, feel free to just say that. – Dannnno Oct 11 '17 at 20:41
• I might be totally off, but if you new to programming, I personally would suggest, write the cleanest most readable version of your multiplication and let your compiler do the job for you. Maybe help him with an compiler specific pragma or similar to indicate hey that's an operation for parallel execution and usually the compiler will do a much better job than you would have done. – ExOfDe Oct 13 '17 at 22:34

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);
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 prod = _mm256_mullo_epi32(bc_mat1, vec_mat2);
}
_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 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);
}
_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 {
}
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 vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j + 8]);
}
_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 {
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 vecB_mat2 = _mm256_loadu_si256((__m256i*)&mat2[k][j + 8]);
auto bc_mat1_2 = _mm256_set1_epi32(mat1[i + 1][k]);
}
_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 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 {
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));
auto bc_mat1_2 = _mm256_set1_epi32(mat1[i + 1][k]);
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.

• Thanks! Great diskussion now I have some ideas how to improve it and where I been going wrong. (I suspected that horizontal sum was a big problem). Since it's about learning AVX I'm gonne take your points and try to improve it myself. – Henrik Ståhlberg Oct 15 '17 at 6:08

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)

• For the same size progression, 23 51 110 230 570 1600 (time to zero out the result was not counted) – harold Oct 15 '17 at 7:50
• If I add tiling and unroll the j loop by 4 I can make it 18 35 70 140 330 825 which is very good, the last one mysteriously breaks the pattern, perhaps I could not find good tiling parameters for it – harold Oct 15 '17 at 8:10