代码

// mmmd1k: C = A*B, n = 1024, row-major doubles.
// Tuned config: MC=384, KC=1024 (single k pass), NC=1024; A packed 6-row k-major,
// B packed 8-wide k-major panels; AVX2+FMA 6x8 microkernel (12 accumulators, k unroll 4,
// software prefetch ap+384/bp+96); C written once directly (no memset) with streaming
// stores when C is 32B aligned and n%4==0 (falls back to normal stores otherwise).
// wsbench (judge HW): 46.70ms vs 51.56ms for the previous MC=120/KC=256 version.
#pragma GCC target("avx2,fma")
#include <immintrin.h>
#include <cstring>

#define MR 4
#define MR6 6
#define NR 8
#define MCVAL 384
#define KCVAL 1024
#define NCVAL 1024
#define PREF 1
#define APF 384
#define BPF 96
#define CSTREAMVAL 1
#define CSTREAM CSTREAMVAL

#define MC MCVAL
#define KC KCVAL
#define NC NCVAL

static double Apack[(size_t)MC * KC + 64] __attribute__((aligned(64)));
static double Bpack[(size_t)(KC + 4) * (NC + NR) + 64] __attribute__((aligned(64)));

// Pack 6 rows of an A block, k-major: Apack[(ig/6)*kc*6 + k*6 + i]
static void packA6(const double *A, int n, int ic, int pc, int mc, int kc) {
    double *dst = Apack;
    for (int ig = 0; ig < mc; ig += MR6) {
        const double *r0 = A + (size_t)(ic + ig) * n + pc;
        const double *r1 = r0 + n, *r2 = r1 + n, *r3 = r2 + n, *r4 = r3 + n, *r5 = r4 + n;
        for (int k = 0; k < kc; k++) {
            dst[0] = r0[k]; dst[1] = r1[k]; dst[2] = r2[k];
            dst[3] = r3[k]; dst[4] = r4[k]; dst[5] = r5[k];
            dst += MR6;
        }
    }
}

// Pack B into NR-wide panels: panel p holds kc rows of NR contiguous doubles.
static void packB(const double *B, int n, int pc, int jc, int kc, int nc) {
    double *dst = Bpack;
    for (int jr = 0; jr < nc; jr += NR) {
        const double *src = B + (size_t)pc * n + jc + jr;
        for (int k = 0; k < kc; k++) {
            _mm256_store_pd(dst, _mm256_loadu_pd(src + (size_t)k * n));
            _mm256_store_pd(dst + 4, _mm256_loadu_pd(src + (size_t)k * n + 4));
            dst += NR;
        }
    }
}

static void packA4(const double *A, int n, int ic, int pc, int mc, int kc) {
    double *dst = Apack;
    for (int ig = 0; ig < mc; ig += MR) {
        const double *r0 = A + (size_t)(ic + ig) * n + pc;
        const double *r1 = r0 + n, *r2 = r1 + n, *r3 = r2 + n;
        for (int k = 0; k < kc; k++) {
            dst[0] = r0[k]; dst[1] = r1[k]; dst[2] = r2[k]; dst[3] = r3[k];
            dst += MR;
        }
    }
}

// FIRST=true: C is written (first k block); FIRST=false: accumulate into C.
template <bool FIRST, bool STREAM>
static void kernel6x8(const double *Ap, const double *Bp, int kc, double *C, int ldc) {
    __m256d c0=_mm256_setzero_pd(),c1=_mm256_setzero_pd(),c2=_mm256_setzero_pd();
    __m256d c3=_mm256_setzero_pd(),c4=_mm256_setzero_pd(),c5=_mm256_setzero_pd();
    __m256d c6=_mm256_setzero_pd(),c7=_mm256_setzero_pd(),c8=_mm256_setzero_pd();
    __m256d c9=_mm256_setzero_pd(),c10=_mm256_setzero_pd(),c11=_mm256_setzero_pd();
    const double *ap = Ap, *bp = Bp;
#define K6STEP() {                                                              \
        __m256d b0 = _mm256_load_pd(bp);                                        \
        __m256d b1 = _mm256_load_pd(bp + 4);                                    \
        __m256d a;                                                              \
        a = _mm256_broadcast_sd(ap + 0); c0  = _mm256_fmadd_pd(a, b0, c0);  c6  = _mm256_fmadd_pd(a, b1, c6); \
        a = _mm256_broadcast_sd(ap + 1); c1  = _mm256_fmadd_pd(a, b0, c1);  c7  = _mm256_fmadd_pd(a, b1, c7); \
        a = _mm256_broadcast_sd(ap + 2); c2  = _mm256_fmadd_pd(a, b0, c2);  c8  = _mm256_fmadd_pd(a, b1, c8); \
        a = _mm256_broadcast_sd(ap + 3); c3  = _mm256_fmadd_pd(a, b0, c3);  c9  = _mm256_fmadd_pd(a, b1, c9); \
        a = _mm256_broadcast_sd(ap + 4); c4  = _mm256_fmadd_pd(a, b0, c4);  c10 = _mm256_fmadd_pd(a, b1, c10); \
        a = _mm256_broadcast_sd(ap + 5); c5  = _mm256_fmadd_pd(a, b0, c5);  c11 = _mm256_fmadd_pd(a, b1, c11); \
        ap += MR6; bp += NR;                                                    \
    }
    int k = 0;
#if PREF
    for (; k + 4 <= kc; k += 4) {
        _mm_prefetch((const char *)(ap + APF), _MM_HINT_T0);
        _mm_prefetch((const char *)(bp + BPF), _MM_HINT_T0);
        K6STEP() K6STEP() K6STEP() K6STEP()
    }
#else
    for (; k + 4 <= kc; k += 4) { K6STEP() K6STEP() K6STEP() K6STEP() }
#endif
    for (; k < kc; k++) K6STEP()
    double *cp = C;
#define STROW(R, A, B) if (FIRST) { \
        if (STREAM) { _mm256_stream_pd(cp, A); _mm256_stream_pd(cp + 4, B); } \
        else { _mm256_storeu_pd(cp, A); _mm256_storeu_pd(cp + 4, B); } } \
    else { __m256d t = _mm256_loadu_pd(cp); _mm256_storeu_pd(cp, _mm256_add_pd(t, A)); \
           t = _mm256_loadu_pd(cp + 4); _mm256_storeu_pd(cp + 4, _mm256_add_pd(t, B)); } cp += ldc;
    STROW(0, c0, c6) STROW(1, c1, c7) STROW(2, c2, c8)
    STROW(3, c3, c9) STROW(4, c4, c10) STROW(5, c5, c11)
}

template <bool FIRST, bool STREAM>
static void kernel4x8(const double *Ap, const double *Bp, int kc, int nc, double *C, int ldc) {
    __m256d c0=_mm256_setzero_pd(),c1=_mm256_setzero_pd(),c2=_mm256_setzero_pd(),c3=_mm256_setzero_pd();
    __m256d c4=_mm256_setzero_pd(),c5=_mm256_setzero_pd(),c6=_mm256_setzero_pd(),c7=_mm256_setzero_pd();
    const double *ap = Ap, *bp = Bp;
    for (int k = 0; k < kc; k++) {
        __m256d b0 = _mm256_load_pd(bp), b1 = _mm256_load_pd(bp + 4), a;
        a = _mm256_broadcast_sd(ap + 0); c0 = _mm256_fmadd_pd(a, b0, c0); c4 = _mm256_fmadd_pd(a, b1, c4);
        a = _mm256_broadcast_sd(ap + 1); c1 = _mm256_fmadd_pd(a, b0, c1); c5 = _mm256_fmadd_pd(a, b1, c5);
        a = _mm256_broadcast_sd(ap + 2); c2 = _mm256_fmadd_pd(a, b0, c2); c6 = _mm256_fmadd_pd(a, b1, c6);
        a = _mm256_broadcast_sd(ap + 3); c3 = _mm256_fmadd_pd(a, b0, c3); c7 = _mm256_fmadd_pd(a, b1, c7);
        ap += MR; bp += NR;
    }
    double *cp = C;
#define STROW4(R, A, B) if (FIRST) { \
        if (STREAM) { _mm256_stream_pd(cp, A); _mm256_stream_pd(cp + 4, B); } \
        else { _mm256_storeu_pd(cp, A); _mm256_storeu_pd(cp + 4, B); } } \
    else { __m256d t = _mm256_loadu_pd(cp); _mm256_storeu_pd(cp, _mm256_add_pd(t, A)); \
           t = _mm256_loadu_pd(cp + 4); _mm256_storeu_pd(cp + 4, _mm256_add_pd(t, B)); } cp += ldc;
    STROW4(0, c0, c4) STROW4(1, c1, c5) STROW4(2, c2, c6) STROW4(3, c3, c7)
}

static void edge_rows(const double *A, const double *B, double *C, int n, int i0, int i1, int j0, int j1) {
    for (int i = i0; i < i1; i++) {
        for (int j = j0; j < j1; j++) {
            const double *ar = A + (size_t)i * n;
            double s = 0;
            for (int k = 0; k < n; k++) s += ar[k] * B[(size_t)k * n + j];
            C[(size_t)i * n + j] = s;
        }
    }
}

void matrix_multiply(int n, const double *A, const double *B, double *C) {
    // streaming stores need 32B-aligned targets: row stride n doubles must keep
    // (row*n + jr) 32B aligned, which needs n % 4 == 0 and C itself aligned.
    const bool cs = CSTREAM && ((n & 3) == 0) && (((size_t)(const void *)C & 31) == 0);
    int n6 = n - (n % MR6);
    int n4 = n - (n % NR);
    int ntail = n - n6;
    int n4r = n6 + (ntail - (ntail % MR));
    for (int jc = 0; jc < n4; jc += NC) {
        int nc = n4 - jc; if (nc > NC) nc = NC;
        for (int pc = 0; pc < n; pc += KC) {
            int kc = n - pc; if (kc > KC) kc = KC;
            packB(B, n, pc, jc, kc, nc);
            for (int ic = 0; ic < n6; ic += MC) {
                int mc = n6 - ic; if (mc > MC) mc = MC;
                packA6(A, n, ic, pc, mc, kc);
#define RUN6(FV, SV) for (int ir = 0; ir < mc; ir += MR6) \
        kernel6x8<FV, SV>(Apack + (size_t)(ir / MR6) * kc * MR6, bp, kc, cb + (size_t)ir * n, n);
                for (int jr = 0; jr < nc; jr += NR) {
                    const double *bp = Bpack + (size_t)(jr / NR) * kc * NR;
                    double *cb = C + (size_t)ic * n + jc + jr;
                    if (pc == 0) { if (cs) RUN6(true, true) else RUN6(true, false) }
                    else { if (cs) RUN6(false, true) else RUN6(false, false) }
                }
            }
            if (n6 < n4r) {
                int mc = n4r - n6;
                packA4(A, n, n6, pc, mc, kc);
#define RUN4(FV, SV) for (int ir = 0; ir < mc; ir += MR) \
        kernel4x8<FV, SV>(Apack + (size_t)(ir / MR) * kc * MR, bp, kc, nc, cb + (size_t)ir * n, n);
                for (int jr = 0; jr < nc; jr += NR) {
                    const double *bp = Bpack + (size_t)(jr / NR) * kc * NR;
                    double *cb = C + (size_t)n6 * n + jc + jr;
                    if (pc == 0) { if (cs) RUN4(true, true) else RUN4(true, false) }
                    else { if (cs) RUN4(false, true) else RUN4(false, false) }
                }
            }
        }
    }
    if (n4 < n) edge_rows(A, B, C, n, 0, n, n4, n);
    if (n4r < n) edge_rows(A, B, C, n, n4r, n, 0, n4);
    _mm_sfence();
}

评测结果