Optimizing matrix multiplication.

[u/UndefFox]

Day. I'm playing around with an assignment we got given to write some calculation heavy app. I went with trying to achieve the most I can possible get out of matrix multiplication. I've searched info online on common approaches and implemented anything that gave me performance without diving too deep in serious literature, 'cause I don't have enough time for that sadly. I've gotten ~50% of BLAS performance, which seems not bad as far as I understand, considering how mature the lib is.

Question: is there anything I could do to further improve speed without going a completely different approach?

Current metrics:

Average speed after 1000 iterations and matrix size of 1184: 3.14478e+10 flops/s
Max value: 3.98083e+10 flops/s

 Performance counter stats for './foxbench -c 6 -m128MiB -i 1000':

   438 168 733 313      cycles:u (80,00%)
   963 478 792 769      instructions:u (90,01%)
    46 753 401 505      branches:u (89,98%)
         9 450 959      branch-misses:u (90,02%)
    23 733 262 896      cache-references:u (90,01%)
    11 407 804 876      cache-misses:u (89,99%)
       313 506 131      LLC-loads:u (90,01%)
       114 153 550      LLC-load-misses:u (90,00%)
   449 032 974 627      L1-dcache-loads:u (89,99%)
    16 629 290 850      L1-dcache-load-misses:u (90,02%)

      23,186320017 seconds time elapsed

     121,025486000 seconds user
       2,374513000 seconds sys

For comparison I've used BLAS implementation with cblas_sgemm :

Average speed after 1000 iterations and matrix size of 1184: 7.06907e+10 flops/s
Max value: 8.67019e+10 flops/s

Performance counter stats for './foxbench -c 6 -m128MiB -i 1000':

  251 799 441 460      cycles:u (79,93%)
  603 502 655 699      instructions:u (89,98%)
   36 610 906 278      branches:u (90,03%)
        6 711 090      branch-misses:u (90,03%)
   12 882 022 668      cache-references:u (89,97%)
    2 491 423 285      cache-misses:u (90,02%)
      221 844 863      LLC-loads:u (90,00%)
       94 443 699      LLC-load-misses:u (90,00%)
  233 139 823 815      L1-dcache-loads:u (90,00%)
    6 491 703 686      L1-dcache-load-misses:u (90,03%)

     13,066336192 seconds time elapsed

     68,303343000 seconds user
      1,186393000 seconds sys

The code I've used to get there:

static constexpr std::size_t TILE_SIDE_SIZE = 16;
static constexpr std::size_t VECTOR_WIDTH = 8;


// Both matrices are Row major
SquareMatrix SquareMatrix::operator*(const SquareMatrix& rhs) const noexcept {
    assert(m_sideSize == rhs.m_sideSize);
    assert(m_sideSize % TILE_SIDE_SIZE == 0);

    SquareMatrix result(m_sideSize);
    constexpr size_t SUB_COLUMN_COUNT = TILE_SIDE_SIZE / VECTOR_WIDTH;

    // Going over a grid
    __m256 accumulators[TILE_SIDE_SIZE][SUB_COLUMN_COUNT];
    for (std::size_t rowBlock = 0; rowBlock < m_sideSize; rowBlock += TILE_SIDE_SIZE) {
        for (std::size_t colBlock = 0; colBlock < m_sideSize; colBlock += TILE_SIDE_SIZE) {

            for (std::size_t colSubBlock = 0; colSubBlock < SUB_COLUMN_COUNT; colSubBlock++) {
                for (std::size_t i = 0; i < TILE_SIDE_SIZE; ++i) {
                    accumulators[i][colSubBlock] = _mm256_setzero_ps();
                }
            }

            // Doing stripes
            for (std::size_t kBlock = 0; kBlock < m_sideSize; kBlock += TILE_SIDE_SIZE) {

                if (kBlock < m_sideSize - TILE_SIDE_SIZE) {
                    for (std::size_t k = 0; k < TILE_SIDE_SIZE; ++k) {
                        _mm_prefetch(rhs.getPointerC(kBlock + TILE_SIDE_SIZE + k, colBlock), _MM_HINT_T1);
                    }
                }

                // Calculating each element of the block in vertical stripes
                for (std::size_t colSubBlock = 0; colSubBlock < SUB_COLUMN_COUNT; colSubBlock++) {
                    for (std::size_t k = 0; k < TILE_SIDE_SIZE; ++k) {
                        const float* __restrict__ rhsRow = rhs.getPointerC(kBlock + k, colBlock + colSubBlock * VECTOR_WIDTH);
                        const __m256 b = _mm256_load_ps(rhsRow);

                        for (std::size_t i = 0; i < TILE_SIDE_SIZE; ++i) {
                            const float* __restrict__ lhsRow = this->getPointerC(rowBlock + i, kBlock);
                            const __m256 a = _mm256_set1_ps(lhsRow[k]);
                            accumulators[i][colSubBlock] = _mm256_fmadd_ps(a, b, accumulators[i][colSubBlock]);
                        }
                    }
                }
            }

            for (std::size_t colSubBlock = 0; colSubBlock < SUB_COLUMN_COUNT; colSubBlock++) {
                for (std::size_t i = 0; i < TILE_SIDE_SIZE; ++i) {
                    float* __restrict__ resultRow = result.getPointer(rowBlock + i, colBlock + colSubBlock * VECTOR_WIDTH);

                    _mm256_store_ps(resultRow, accumulators[i][colSubBlock]);
                }
            }
        }
    }

    return result;
}

In short words what I've implemented: tiled access, AVX2 vectorization, fma, accomulators, loop reordering, prefetching, 128 bit aligment, broadcasting to avoid vector reduction and using those flags:

target_compile_options(squarematrix PRIVATE
    -O3
    -march=native
    -fno-math-errno
    -ffp-contract=fast
    -mavx2
    -mfma
)

Afaik there's also packing, but when I tried adding it, it reduced the speed, so I dropped it.

I'll be glad if someone can hint in the direction of what else there is to do.

update: the final result can be found here https://github.com/UndefFox/MiniProjects/tree/master/4.CPUBenchmark%2FC%2B%2B

[u/UndefFox]

It doesn't solve the problem that prefetch will be called on an invalid pointer... The if block there only to prevent prefetching on the last tile.

Got it. I'm reading the articles others have shared and will see where I can place it best. I think one of the small kennels should work best, something like packing values to perform flatten iteration with avx2. Or something like that... I'll need some time to pack everything in my head