5
\$\begingroup\$

The goal of the code review by order of importance (i.e. What I hope to hear from you):

  1. I've verified correctness using a straightforward matrix multiply function though I am open to those who want to verify correctness.

  2. I'm just starting to learn about the benefits and challenges posed by coding things in SSE, while at the same time, keeping in mind hardware level caching. As such, input on room for efficiency gains would be greatly appreciated.

  3. Coding style. Obviously there is always some room for improvement here so any thoughts on this is greatly appreciated.

Some notes on SSE:

  1. Writing out from SSE is always expensive, especially when trying to extract some but not all of the values in the vector lanes (using for example _mm_extract_epi32
  2. Some SSE instructions are quite expensive. For example _mm_hadd_epi32 is quite expensive, requiring that I find ways to minimize its use.

The goal of the code:

The function below is me trying to learn caching and SSE. I've done both to a limited extent, separately. And this is just me trying to combine the two.

The function performs a matrix multiply on matrices a and b, both of size n x n and saves the results in c. All of these matrices are 16 byte aligned in memory.

It first transposes b into t so the data in matrix b, now t, can be accessed in row-major order.

It then reads block number rows of a, and iterates those rows over block number of rows in t using SSE to perform the multiplication and summation. This "blocking" is the bit that helps with caching and is a typical cache optimization technique. These are implemented using the ii and jj loops.

The main inner loop is the i and j loops. The i loop iterates over each row in each block in a and the j loop iterates over each row in each block in t.

The j loop accumulates 4 rows worth of products in the correct order into a single SSE register before writing out the results to c.

Thanks in advance.

Code:

   int fast_matrix_multiply_simd(uint32_t * a, 
                                 uint32_t * b, 
                                 uint32_t * c, 
                                 size_t n, 
                                 size_t block){

        // for simplicity, the following limitations are enfoced
        // block should be a multiple of 4
        if (block % 4 != 0){
            return -1;
        }

        // n should be a multiple of block
        if (n % block != 0){
            return -1;
        }

        // create a temporary transpose matrix -- created in dynamic memory
        uint32_t * t = transpose_matrix_block(b, n);
        if (!t){
            return -1;
        }

        __m128i * m128_a, * m128_t, * m128_c;
        size_t en = n;

        // number of iterations of 128bit loads per row
        size_t m128_count = (n * sizeof(uint32_t) * 8)/128;

        // need a temporary place to store result vectors
        // which will be summed
        __m128i loaded_temp[4];

        // values that update at each loop
        size_t at_i, at_j, i, j, k, ii_block, jj_block;

        // blocking for cache concious purposes using the outer loop
        for (size_t ii = 0; ii < en; ii += block){
            ii_block = ii + block;
            for (size_t jj = 0; jj < en; jj += block){
                jj_block = jj + block;

                // main inner loop, i and j loop over rows of a and t respectively
                for (i = ii; i < ii_block; ++i){

                    at_i = n * i;
                    m128_a = (__m128i*)&(a[at_i]);
                    m128_c = (__m128i*)&(c[at_i + jj]);

                    for (j = jj; j < jj_block; j+=4){
                        at_j = n * j;
                        m128_t = (__m128i*)&(t[at_j]);

                        // essentially, unroll an inner loop that multiplies a row
                        // in the "a" matrix by 4 rows in the "t" matrix, to take 
                        // advantage of 4 SSE lanes, Accumulate the sums of each
                        // in a single SSE vector, write into "c" matrix, then
                        // move on to the next block in "t"

                        loaded_temp[0] = _mm_set_epi32(0,0,0,0);
                        for (k = 0; k < m128_count; ++k){
                            loaded_temp[0] = _mm_add_epi32(loaded_temp[0],_mm_mullo_epi32(m128_a[k], m128_t[k]));
                        }

                        m128_t = (__m128i*)(t + at_j + n);
                        loaded_temp[1] = _mm_set_epi32(0,0,0,0);
                        for (k = 0; k < m128_count; ++k){
                            loaded_temp[1] = _mm_add_epi32(loaded_temp[1],_mm_mullo_epi32(m128_a[k], m128_t[k]));
                        }

                        m128_t = (__m128i*)(t + at_j + n + n);
                        loaded_temp[2] = _mm_set_epi32(0,0,0,0);
                        for (k = 0; k < m128_count; ++k){
                            loaded_temp[2] = _mm_add_epi32(loaded_temp[2],_mm_mullo_epi32(m128_a[k], m128_t[k]));
                        }

                        m128_t = (__m128i*)(t + at_j + n + n + n);
                        loaded_temp[3] = _mm_set_epi32(0,0,0,0);
                        for (k = 0; k < m128_count; ++k){
                            loaded_temp[3] = _mm_add_epi32(loaded_temp[3],_mm_mullo_epi32(m128_a[k], m128_t[k]));
                        }

                        // combine first two
                        loaded_temp[0] = _mm_hadd_epi32(loaded_temp[0], loaded_temp[1]);
                        loaded_temp[0] = _mm_hadd_epi32(loaded_temp[0], loaded_temp[0]);
                        loaded_temp[0] = _mm_and_si128(_mm_set_epi32(-1,0,0,-1),loaded_temp[0]);

                        // combine second 2
                        loaded_temp[2] = _mm_hadd_epi32(loaded_temp[2], loaded_temp[3]);
                        loaded_temp[2] = _mm_hadd_epi32(loaded_temp[2], loaded_temp[2]);
                        loaded_temp[2] = _mm_and_si128(_mm_set_epi32(-1,0,0,-1),loaded_temp[2]);

                        // combine totals
                        loaded_temp[0] = _mm_hadd_epi32(loaded_temp[0], loaded_temp[2]);

                        *m128_c = loaded_temp[0];
                        ++m128_c;
                    } // ends j loop
                } // ends i loop
                // end main inner loop

            } // ends jj block loop
        } // ends ii block loop

        free(t); // Free the transposed matrix since this is no longer needed
        return 0;
    }
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.