The things that destroy performance tend to be things that are much easier to type, such as division operations, poor cache usage, bad instruction scheduling (which is a topic by itself), and mispredicted branches.

Về việc nhân ma trận, Fast Matrix Multiplication là 1 bài được sinh ra để "thử thách" người đọc. Bài này bảo bạn thực hiện ~2^{10}~ lần nhân hai ma trận ~256\times 256~ trong 800ms, điều này là có thể làm được, nhưng cần biết rõ bạn đang làm gì.

Nhận thấy việc CPU này có AVX2 là điều khá quan trọng. CPU ở đây là Ryzen 3600X (Zen 2), 3.8 GHz.

Phần 1. Nhân ma trận cơ bản.

Về việc nhân ma trận cơ bản, các bạn đã có thuật toán cổ điển mà đơn thuần chúng ta sẽ tối ưu:

for (int i = 0; i < 256; i++)
    for (int j = 0; j < 256; j++)
        for (int k = 0; k < 256; k++)
            c[i][k] += a[i][j] * b[j][k];

Trước tiên, chúng ta hãy xem thuật toán này thực hiện ~2^{10}~ lần gọi trong thời gian bao nhiêu.

Test case #1: WA [1,495s, 5,48 MB] (0/10)

Không hẳn là tệ nhưng ta có thể làm tốt hơn. Kết quả cuối cùng là ta sẽ có được chương trình nhanh gấp đôi (chỉ 700ms!)

Phần 2. Các khái niệm.

Ta sẽ sử dụng kiểu dữ liệu __m256d để lưu 4 số kiểu double trong 256-bit. Để sử dụng kiểu dữ liệu này và các hàm, cần thêm thư viện immintrin.h.

Về các thao tác trên __m256d có 1 vài thao tác chính sau:

_mm256_fmadd_pd(a, b, c): trả về a * b + c (vector 4 phần tử chứa giá trị của a[i] * b[i] + c[i] với ~0\leq i<4~).

_mm256_broadcast_sd(x): Gán giá trị *x cho cả 4 phần tử.

_mm256_load_pd(x): Lấy 256-bit (4 phần tử *(x+i) với ~0\leq i<4~) từ x thành 1 vector.

Ta sẽ sử dụng hàm sau để tính tổng trong vector (chi tiết này thực ra không quan trọng, bạn chỉ cần làm s[0] + s[1] + s[2] + s[3] là compiler tự sinh code đủ tốt rồi):

inline double hsum_double_avx(__m256d v)
{
    __m128d vlow = _mm256_castpd256_pd128(v);
    __m128d vhigh = _mm256_extractf128_pd(v, 1); // high 128
    vlow = _mm_add_pd(vlow, vhigh);              // reduce down to 128
    __m128d high64 = _mm_unpackhi_pd(vlow, vlow);
    return _mm_cvtsd_f64(_mm_add_sd(vlow, high64)); // reduce to scalar
}l

Phần 3: Chiến thuật tối ưu

Về phần này, ta phải nhận xét về độ trễ và thông lượng của các hàm trên. Hiểu cơ bản là độ trễ là thời gian (tính theo chu kì) từ khi lệnh đó thực hiện đến khi lệnh khác dùng được kết quả. Thông lượng là số lệnh tối đa có thể thực hiện được mỗi chu kì với các lệnh độc lập.

Do lệnh vbroadcastsd thực hiện trên các cổng độc lập so với vfmadd nên thông lượng ở đây là không quan trọng lắm. Để hạn chế việc độ trễ của vfmadd làm ảnh hưởng đến thông lượng thực tế, ta cần tối thiểu 10 vfmadd độc lập.

Để làm được điều đó, ta sẽ làm nhân ma trận ~256\times 256\times 260~.

Ta còn phải xét có thể tải được bao nhiêu vector mỗi chu kì từ cache. L1d trên Zen 2 là 32KB, L2 là 512KB.

Khi thực hiện nhân ma trận như trên, chiến thuật sẽ như sau:

Đặt biến __m256d s{} làm biến kết quả.

Với mỗi khối ~2\times 20~ của C (kết quả):

• Ta đặt ~Cr_i~ với ~0\leq i < 10~ làm 10 vector kết quả của C.

• Thực hiện nhân ma trận ~2\times 256\times 20~ (lấy ~2\times 256~ phần tử từ A, ~256\times 20~ phần tử từ B.)

• Cộng ~Cr_i \times Cr_i~ vào s.

Kết quả sẽ là hsum_double_avx(s).

Với việc thực hiện nhân ma trận ~2\times 256\times 20~ như sau:

Lặp 256 lần (~k~ chạy từ ~0~ đến ~255~):

Phân tán ~a[i][k]~ và ~a[i+1][k]~ ra thành các vector ~ai_1~ và ~ai_2~.

Với ~x~ chạy từ ~0~ đến ~4~:

• ~Cr_x = Cr_x + ai_1 \times \text{_mm256_load_ps}(B + j \times 256 + k \times 20 + x\times 4).~
• ~Cr_{x+5} = Cr_{x+5} + ai_2 \times \text{_mm256_load_ps}(B + j \times 256 + k \times 20 + x \times 4).~

Do chúng ta chỉ có 16 thanh ghi nên sẽ phải đọc B 6 lần. Tuy nhiên, do khối ~256\times 20~ của B nằm hoàn toàn trong L2 và 80% là trong L1, điều này không gây ra vấn đề lớn khi mỗi lần lặp k tốn ~\geq 5~ chu kì.

Nhận thấy 10 lệnh trên là độc lập nên ta không bị giới hạn về độ trễ của vfmadd nữa.

Về việc biến đổi ma trận B: Ta chuyển vị ma trận theo các khối ~1\times 20~, như sau:

for (int i = 0; i < 256; i++)
    for (int j = 0; j < 256; j++)
        B[j / 20 * 5120 + i * 20 + j % 20] = b[i][j];

Phần 4: Code

#pragma GCC optimize("O3,unroll-loops")
#include <bits/stdc++.h>
#include <immintrin.h>
using namespace std;
// doesn't matter if C's not in order; we can actually just square then sum
alignas(32) double B[260 * 256];
inline double hsum_double_avx(__m256d v)
{
    __m128d vlow = _mm256_castpd256_pd128(v);
    __m128d vhigh = _mm256_extractf128_pd(v, 1); // high 128
    vlow = _mm_add_pd(vlow, vhigh);              // reduce down to 128
    __m128d high64 = _mm_unpackhi_pd(vlow, vlow);
    return _mm_cvtsd_f64(_mm_add_sd(vlow, high64)); // reduce to scalar
}
double multiply_matrices(double (*__restrict__ a)[256], double (*__restrict__ b)[256])
{
    double *A = (double *)(a);
    // make the compiler's life easier
    for (int i = 0; i < 256; i++)
        for (int j = 0; j < 240; j += 20)
            for (int k = 0; k < 20; k++)
                B[j / 20 * 5120 + i * 20 + k] = b[i][j + k];
    for (int i = 0; i < 256; i++)
        for (int k = 0; k < 16; k++)
            B[61440 + i * 20 + k] = b[i][240 + k];
    __m256d s{};
    // 256 x 256 x 260 matrix multiplication
    // Zen 2 has 512K L2, B should fit in L2 for bandwidth demands (5-6x loads per iteration) - B is 520KB. 52-53 GFLOPS
    // w/o B copying.
    for (int j = 0; j < 260; j += 20)
        for (int i = 0; i < 256; i += 2)
        {
            __m256d Cr0{}, Cr1{}, Cr2{}, Cr3{}, Cr4{}, Cr5{}, Cr6{}, Cr7{}, Cr8{}, Cr9{};
            for (int k = 0; k < 256; k++)
            {
                auto ai = _mm256_broadcast_sd(A + i * 256 + k);
                auto ai2 = _mm256_broadcast_sd(A + (i + 1) * 256 + k);
                Cr0 = _mm256_fmadd_pd(ai, _mm256_load_pd(B + j * 256 + k * 20), Cr0);
                Cr1 = _mm256_fmadd_pd(ai, _mm256_load_pd(B + j * 256 + k * 20 + 4), Cr1);
                Cr2 = _mm256_fmadd_pd(ai, _mm256_load_pd(B + j * 256 + k * 20 + 8), Cr2);
                Cr3 = _mm256_fmadd_pd(ai, _mm256_load_pd(B + j * 256 + k * 20 + 12), Cr3);
                Cr4 = _mm256_fmadd_pd(ai, _mm256_load_pd(B + j * 256 + k * 20 + 16), Cr4);
                Cr5 = _mm256_fmadd_pd(ai2, _mm256_load_pd(B + j * 256 + k * 20), Cr5);
                Cr6 = _mm256_fmadd_pd(ai2, _mm256_load_pd(B + j * 256 + k * 20 + 4), Cr6);
                Cr7 = _mm256_fmadd_pd(ai2, _mm256_load_pd(B + j * 256 + k * 20 + 8), Cr7);
                Cr8 = _mm256_fmadd_pd(ai2, _mm256_load_pd(B + j * 256 + k * 20 + 12), Cr8);
                Cr9 = _mm256_fmadd_pd(ai2, _mm256_load_pd(B + j * 256 + k * 20 + 16), Cr9);
            }
            s = _mm256_fmadd_pd(Cr0, Cr0, s);
            s = _mm256_fmadd_pd(Cr1, Cr1, s);
            s = _mm256_fmadd_pd(Cr2, Cr2, s);
            s = _mm256_fmadd_pd(Cr3, Cr3, s);
            s = _mm256_fmadd_pd(Cr4, Cr4, s);
            s = _mm256_fmadd_pd(Cr5, Cr5, s);
            s = _mm256_fmadd_pd(Cr6, Cr6, s);
            s = _mm256_fmadd_pd(Cr7, Cr7, s);
            s = _mm256_fmadd_pd(Cr8, Cr8, s);
            s = _mm256_fmadd_pd(Cr9, Cr9, s);
        }
    return hsum_double_avx(s);
}

Và đây là thành quả:

Test case #2: WA [0,700s, 5,34 MB] (0/10)

Chúc bạn thành công!

https://www.youtube.com/watch?v=5qUYFApo0s4