4
\$\begingroup\$

I need help to make matrix multiplication in C run as fast as possible. On my AMD Phenom(tm) II X6 1090T, my program multiplies two square singe precision 4096x4096 matrices in about 6.9 seconds. NumPy, which is based on OpenBLAS, multiplies the same sized matrices in about 1.5 seconds. So I think it should be possible to double the speed of my program. Since my processor is old, only 128-bit SIMD instructions are available.

Original program

#include <math.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define SIZE        4096
#define N_ELS       (SIZE * SIZE)
#define N_BYTES     (sizeof(float) * N_ELS)
#define MIN(a, b)   ((a > b) ? (b) : (a))
#define N_THREADS   6
#define TILE_SIZE   32

static void
mul_tile(int i0, int i1, int j0, int j1, int k0, int k1,
         float * restrict A, float * restrict B, float * restrict C) {
    for (int i = i0; i < i1; i++) {
        for (int k = k0; k < k1; k++) {
            float a = A[SIZE * i + k];
            __builtin_prefetch(&C[SIZE * i + j0]);
            __builtin_prefetch(&B[SIZE * (k + 1) + j0]);
            __builtin_prefetch(&B[SIZE * (k + 2) + j0]);
            for (int j = j0; j < j1; j++) {
                float b = B[SIZE * k + j];
                C[SIZE * i + j] += a * b;
            }
        }
    }
}

typedef struct {
    float *A, *B, *C;
    int start_i, end_i;
} mul_job_t;

void *mul_thread(void *arg) {
    mul_job_t *job = (mul_job_t *)arg;
    int start_i = job->start_i;
    int end_i = job->end_i;
    for (int i = start_i; i < end_i; i += TILE_SIZE) {
        int imax = MIN(i + TILE_SIZE, SIZE);
        for (int j = 0; j < SIZE; j += TILE_SIZE) {
            int jmax = MIN(j + TILE_SIZE, SIZE);
            for (int k = 0; k < SIZE; k += TILE_SIZE) {
                int kmax = MIN(k + TILE_SIZE, SIZE);
                mul_tile(i, imax, j, jmax, k, kmax,
                         job->A, job->B, job->C);
            }
        }
    }
    return 0;
}

static void
mul(float * restrict A, float * restrict B, float * restrict C) {
    pthread_t threads[N_THREADS];
    mul_job_t jobs[N_THREADS];
    memset(C, 0, N_BYTES);
    int n_i_tiles = (int)ceil((float)SIZE / (float)TILE_SIZE);
    int tiles_per_thread = (int)ceil((float)n_i_tiles / (float)N_THREADS);
    for (int i = 0; i < N_THREADS; i++) {
        int start = TILE_SIZE * i * tiles_per_thread;
        int end = MIN(start + TILE_SIZE * tiles_per_thread, SIZE);
        jobs[i] = (mul_job_t){A, B, C, start, end};
        pthread_create(&threads[i], NULL, mul_thread, &jobs[i]);
    }
    for (int i = 0; i < N_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }
}

int
main(int argc, char *argv[]) {
    float *A = (float *)malloc(N_BYTES);
    float *B = (float *)malloc(N_BYTES);
    float *C = (float *)malloc(N_BYTES);
    for (int i = 0; i < SIZE * SIZE; i++) {
        A[i] = ((float)rand() / (float)RAND_MAX) * 10;
        B[i] = ((float)rand() / (float)RAND_MAX) * 10;
    }
    struct timespec begin, end;
    clock_gettime(CLOCK_MONOTONIC_RAW, &begin);
    mul(A, B, C);
    clock_gettime(CLOCK_MONOTONIC_RAW, &end);
    double delta = (end.tv_nsec - begin.tv_nsec) / 1000000000.0 +
        (end.tv_sec  - begin.tv_sec);
    printf("%.6lfs\n", delta);
    free(A);
    free(B);
    free(C);
}

For best results, compile with clang -o mul mul.c -O3 -fomit-frame-pointer -march=native -mtune=native -lpthread.

Here is my cache hierarchy:

enter image description here

Any answer that significantly improves the performance is acceptable to me. I'm only interested in performance and not other issues the code might have.

Notes and additions:

  • gcc 12.2 is about 0.4 seconds slower than clang 14.0 at 7.2 seconds.
  • The machine has a Radeon HD 4200 card, so if someone could optimize the code using GPU instructions that would be cool too.
  • The code is veeery cache sensitive; setting TILE_SIZE to 64 quadruples the execution time, for example.

Optimized program

#include <math.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define SIZE        4096
#define N_ELS       (SIZE * SIZE)
#define N_BYTES     (sizeof(float) * N_ELS)
#define MIN(a, b)   ((a > b) ? (b) : (a))
#define N_THREADS   6
#define TILE_SIZE   32

typedef unsigned int uint_t;

static void
mul_tile(uint_t i0, uint_t i1, uint_t j0, uint_t j1, uint_t k0, uint_t k1,
         float * restrict A, float * restrict B, float * restrict C
) {
    for (uint_t i = i0; i < i1; i ++) {
        for (uint_t k = k0; k < k1; k++) {
            float a = A[SIZE * i + k];
            for (uint_t j = j0; j < j1; j++) {
                float b = B[SIZE * k + j];
                C[SIZE * i + j] += a * b;
            }
        }
    }
}

typedef struct {
    float *A, *B, *C;
    uint_t start_i, end_i;
} mul_job_t;

void *mul_thread(void *arg) {
    mul_job_t *job = (mul_job_t *)arg;
    uint_t start_i = job->start_i;
    uint_t end_i = job->end_i;
    for (uint_t i = start_i; i < end_i; i += TILE_SIZE) {
        uint_t imax = MIN(i + TILE_SIZE, SIZE);
        for (uint_t j = 0; j < SIZE; j += TILE_SIZE) {
            uint_t jmax = MIN(j + TILE_SIZE, SIZE);
            for (uint_t k = 0; k < SIZE; k += TILE_SIZE) {
                uint_t kmax = MIN(k + TILE_SIZE, SIZE);
                mul_tile(i, imax, j, jmax, k, kmax,
                         job->A, job->B, job->C);

            }
        }
    }
    return 0;
}

static void
mul(float * restrict A, float * restrict B, float * restrict C) {
    pthread_t threads[N_THREADS];
    mul_job_t jobs[N_THREADS];
    memset(C, 0, N_BYTES);
    int n_i_tiles = (int)ceil((float)SIZE / (float)TILE_SIZE);
    int tiles_per_thread = (int)ceil((float)n_i_tiles / (float)N_THREADS);
    for (int i = 0; i < N_THREADS; i++) {
        int start = TILE_SIZE * i * tiles_per_thread;
        int end = MIN(start + TILE_SIZE * tiles_per_thread, SIZE);
        jobs[i] = (mul_job_t){A, B, C, start, end};
        pthread_create(&threads[i], NULL, mul_thread, &jobs[i]);
    }
    for (int i = 0; i < N_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }
}

int
main(int argc, char *argv[]) {
    float *A = (float *)malloc(N_BYTES);
    float *B = (float *)malloc(N_BYTES);
    float *C = (float *)malloc(N_BYTES);
    for (int i = 0; i < SIZE * SIZE; i++) {
        A[i] = ((float)rand() / (float)RAND_MAX) * 10;
        B[i] = ((float)rand() / (float)RAND_MAX) * 10;
    }
    struct timespec begin, end;
    clock_gettime(CLOCK_MONOTONIC_RAW, &begin);
    mul(A, B, C);
    clock_gettime(CLOCK_MONOTONIC_RAW, &end);
    double delta = (end.tv_nsec - begin.tv_nsec) / 1000000000.0 +
        (end.tv_sec  - begin.tv_sec);
    printf("%.6lfs\n", delta);
    free(A);
    free(B);
    free(C);
}
\$\endgroup\$
6
  • 1
    \$\begingroup\$ Why are you doing this yourself? Just call into Atlas. \$\endgroup\$
    – Reinderien
    Aug 29, 2022 at 2:49
  • 1
    \$\begingroup\$ Because it's fun \$\endgroup\$ Aug 29, 2022 at 6:54
  • 2
    \$\begingroup\$ Will the code run as 32-bit code or 64-bit code? Usually that's not so important, but now it is because it strongly affects the number of vector registers \$\endgroup\$
    – user555045
    Aug 29, 2022 at 7:33
  • 3
    \$\begingroup\$ You may be "only interested in performance and not other issues the code might have", but you should know that posting your code for review here invites critique of any and all aspects of the code regardless of what you ask. \$\endgroup\$ Aug 29, 2022 at 8:24
  • \$\begingroup\$ I have rolled back Rev 5 → 4. Please see What to do when someone answers. \$\endgroup\$ Sep 11, 2022 at 0:12

1 Answer 1

6
\$\begingroup\$

High level

The code implements tiling, which is good, but there is only one TILE_SIZE. In my experience, there is benefit from allowing tiles to be non-square. The A tiles and B tiles do not have to be the same size either, they only need to be compatible. So there are really 3 different tile dimensions to choose/tune.

The code doesn't implement tile repacking (copying a tile into contiguous memory), which can be useful to reduce TLB thrashing, depending on whether that is happening and how big the impact is. In the current code it's probably not relevant, but it may become so after other changes are done.

Low level

AMD Phenom(tm) II X6 1090T

So the K10 family, here are some relevant performance parameters:

  • 2 movaps (load) per cycle
  • 1 movaps (store) and it costs 2 mops
  • 1 mulps per cycle
  • 1 addps per cycle (latency: 4)
  • 2 shufps per cycle

Not all of them at the same time, the shuffle can take the place of an addition or multiplication or both. IDK about the load, there's a question mark in the execution unit column. Additionally, K10 can only decode up to 3 instructions per cycle.

I suggest the design goals should be:

  • reduce the number of stores and loads, making the code consist as much as possible/reasonable of mulps and addps
  • juggle at least 4 independent dependency chains of additions, to avoid bottlenecking on their latency (not a problem in the original code, but it will be when rearranging the code so that stores are avoided in the inner loop)
  • no horizontal sums (which are not SIMD-friendly)
  • do not exceed the number of vector registers (spill/reload would conflict with the first design goal)

Reducing the number of stores can be done by making the k-loop the inner loop, it will then sum up a bunch of products that are then summed into one entry of C. So, no stores in the inner loop. That degrades the access pattern to B, but with tile repacking that doesn't matter, actually repacking makes the access pattern even better than it is now: a purely sequential scan (rather than several sequential pieces). Unfortunately this also ruins the "load a outside of the loop"-strategy, but that will be fixed in a different way. This also creates a loop-carried dependency through the addition, but that will be fixed by using multiple accumulators.

To avoid horizontal sums, the factor-of-4 unroll for SIMD won't be across the k-loop, but across the j-loop. So SIMD will be used to compute 4 entries to sum into C, not to compute one entry faster (which would require a horizontal sum). Though the unroll factor will get bigger.

To avoid loads (obviously we cannot avoid all of them), the result of a load must be reused several times. This can be done by loading several entries from B (horizontally), and several entries from A (vertically), and then computing every relevant product between those entries. For example by loading 4 entries from B (4 SSE vectors actually, so 16 floats) and 4 entries from A (broadcasted into full vectors with shufps) then there are 16 products to compute. 4x4 doesn't fit in the register budget, but 3x4 and 2x4 are possible. 3 is an annoying number so let's go for 2x4 (by contrast, on Haswell the choice of 3x4 is more or less forced: it requires at least 10 independent dependency chains, due to having an FMA with a latency of 5 and a throughput of 2 per cycle). Important: for 32-bit code, even 2x4 is too much.

So 2 loads from A, vertically (note that this is an unroll of the i-loop by a factor of 2), and then broadcast using shufps (in the code this will be represented by _mm_set1_ps). And 4 loads from B, horizontally, which is already 16 floats at once. It looks likely that the B tiles will "want" to be wider than 32, otherwise the j-loop makes only 2 iterations. 32x32 is really tiny anyway, only 4KB per tile.

If I haven't made any mistakes, the important part of the codes comes out roughly like this (treat this more as a sketch of a solution rather than copy-and-paste code though, I just wrote this roughly without testing)

#include <xmmintrin.h>

void mul_tile(int i0, int i1, int j0, int j1, int k0, int k1,
         float * restrict A, float * restrict B, float * restrict C) {
    for (int i = i0; i < i1; i += 2) {
        float *Bptr = B;
        for (int j = j0; j < j1; j += 16) {
            __m128 acc0 = _mm_load_ps(&C[SIZE * i + j]);
            __m128 acc1 = _mm_load_ps(&C[SIZE * i + (j + 4)]);
            __m128 acc2 = _mm_load_ps(&C[SIZE * i + (j + 8)]);
            __m128 acc3 = _mm_load_ps(&C[SIZE * i + (j + 12)]);
            __m128 acc4 = _mm_load_ps(&C[SIZE * (i + 1) + j]);
            __m128 acc5 = _mm_load_ps(&C[SIZE * (i + 1) + (j + 4)]);
            __m128 acc6 = _mm_load_ps(&C[SIZE * (i + 1) + (j + 8)]);
            __m128 acc7 = _mm_load_ps(&C[SIZE * (i + 1) + (j + 12)]);
            
            for (int k = k0; k < k1; k++) {
                // broadcast 2 entries from A
                __m128 A0 = _mm_set1_ps(A[SIZE * i + k]);
                __m128 A1 = _mm_set1_ps(A[SIZE * (i + 1) + k]);
                // assuming that B is a repacked tile, memory is accessed perfectly sequentially
                __m128 B0 = _mm_load_ps(Bptr);
                __m128 B1 = _mm_load_ps(Bptr + 4);
                __m128 B2 = _mm_load_ps(Bptr + 8);
                __m128 B3 = _mm_load_ps(Bptr + 12);
                Bptr += 16;
                // perform all products and sum them into the accumulators
                acc0 = _mm_add_ps(acc0, _mm_mul_ps(A0, B0));
                acc1 = _mm_add_ps(acc1, _mm_mul_ps(A0, B1));
                acc2 = _mm_add_ps(acc2, _mm_mul_ps(A0, B2));
                acc3 = _mm_add_ps(acc3, _mm_mul_ps(A0, B3));
                acc4 = _mm_add_ps(acc4, _mm_mul_ps(A1, B0));
                acc5 = _mm_add_ps(acc5, _mm_mul_ps(A1, B1));
                acc6 = _mm_add_ps(acc6, _mm_mul_ps(A1, B2));
                acc7 = _mm_add_ps(acc7, _mm_mul_ps(A1, B3));
            }

            // store results to C
            _mm_store_ps(&C[SIZE * i + j], acc0);
            _mm_store_ps(&C[SIZE * i + (j + 4)], acc1);
            _mm_store_ps(&C[SIZE * i + (j + 8)], acc2);
            _mm_store_ps(&C[SIZE * i + (j + 12)], acc3);
            _mm_store_ps(&C[SIZE * (i + 1) + j], acc4);
            _mm_store_ps(&C[SIZE * (i + 1) + (j + 4)], acc5);
            _mm_store_ps(&C[SIZE * (i + 1) + (j + 8)], acc6);
            _mm_store_ps(&C[SIZE * (i + 1) + (j + 12)], acc7);
        }
    }
}

I neglected to handle cases where the tile size is not divisible by the unroll factor. All wide loads are aligned loads, K10 isn't really that picky about using the unaligned load instruction as long as the actual address is aligned, but you should align the addresses (tile repacking fixes alignment as well, assuming you allocate the memory properly). I cannot predict the best way to prefetch, you would have to try it out. I did not assume that A is repacked, the access pattern isn't that bad and wouldn't cost that many TLB entries at once, you can repack it though.

The resulting asm looks reasonable to me, but I'm certainly no K10 expert.

Loop variable type

int doesn't work out too badly if it is the same size as a pointer, which in various cases it is, but notably not on 64-bit Windows. When int is smaller than a pointer, using an int to offset a pointer (ie accessing an array) likely results in a sign-extension. That may sound insignificant, but these sign-extensions are purely overhead, conferring absolutely no benefit. The cost can be significant especially in small loops. Using an unsigned integer or pointer-sized integer avoids this penalty.

\$\endgroup\$
2
  • \$\begingroup\$ Amazing! Only changing the ints to unsigned ints halved the execution time from 6.9 to 3.2!! I haven't yet had time to test your other suggestions but I bet they will reduce the execution time further. First I want to ask what is tile repacking? I'm not familiar with that technique. Transposing the B matrix? \$\endgroup\$ Aug 29, 2022 at 23:02
  • 1
    \$\begingroup\$ @BjörnLindqvist it's sort of a transpose but in general it's whatever rearrangement is necessary such that the final access pattern is linear. If you view B as having 16-float-wide elements (assuming the j-loop has a total unroll factor of 16, which in my example it does), then it is a transpose. BTW you can do this on-demand on a tile-by-tile basis, in a constant amount of memory. \$\endgroup\$
    – user555045
    Aug 30, 2022 at 14:22

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.