代码

// mmmd1k: C = A*B, n = 1024, row-major doubles.
// Tuned: MC=384, KC=1024 (single k pass), NC=1024, 4x12 AVX2+FMA microkernel
// (12 accumulators, k unroll 4, software prefetch ap+384/bp+96).
// A packed 4-row k-major with a 4x4 vector transpose; B packed 12-wide k-major
// panels two at a time (192B per k = 3 full cache lines); C written exactly once
// with streaming stores when C is 32B aligned and n%4==0 (normal stores otherwise).
// wsbench (judge HW): 46.36ms vs 51.56ms for the previous MC=120/KC=256 version.
// A packed as 4-row k-major blocks, B packed as 12-wide k-major panels (tail group
// zero padded and stored with a 4-column store), C written once with streaming stores.
#pragma GCC target("avx2,fma")
#include <immintrin.h>
#include <cstring>

#define MR 4
#define NR 12
#define MCVAL 384
#define KCVAL 1024
#define NCVAL 1024
#define KU 4
#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 4 rows of an A block, k-major: Apack[(ig/4)*kc*4 + k*4 + i]
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;
        int k = 0;
        for (; k + 4 <= kc; k += 4) {          // 4x4 transpose
            __m256d a0 = _mm256_loadu_pd(r0 + k), a1 = _mm256_loadu_pd(r1 + k);
            __m256d a2 = _mm256_loadu_pd(r2 + k), a3 = _mm256_loadu_pd(r3 + k);
            __m256d t0 = _mm256_unpacklo_pd(a0, a1), t1 = _mm256_unpackhi_pd(a0, a1);
            __m256d t2 = _mm256_unpacklo_pd(a2, a3), t3 = _mm256_unpackhi_pd(a2, a3);
            _mm256_store_pd(dst,      _mm256_permute2f128_pd(t0, t2, 0x20));
            _mm256_store_pd(dst + 4,  _mm256_permute2f128_pd(t1, t3, 0x20));
            _mm256_store_pd(dst + 8,  _mm256_permute2f128_pd(t0, t2, 0x31));
            _mm256_store_pd(dst + 12, _mm256_permute2f128_pd(t1, t3, 0x31));
            dst += 16;
        }
        for (; k < kc; k++) {
            dst[0] = r0[k]; dst[1] = r1[k]; dst[2] = r2[k]; dst[3] = r3[k];
            dst += MR;
        }
    }
}

// Pack B into NR-wide k-major panels; the final (partial) group is zero padded.
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; ) {
        int w = nc - jr;
        const double *src = B + (size_t)pc * n + jc + jr;
        if (w >= 2 * NR) {
            // two panels at once: 192B per k = 3 full cache lines, no re-reads
            double *dst2 = dst + (size_t)kc * NR;
            for (int k = 0; k < kc; k++) {
                const double *s = src + (size_t)k * n;
                _mm256_store_pd(dst, _mm256_loadu_pd(s));
                _mm256_store_pd(dst + 4, _mm256_loadu_pd(s + 4));
                _mm256_store_pd(dst + 8, _mm256_loadu_pd(s + 8));
                _mm256_store_pd(dst2, _mm256_loadu_pd(s + 12));
                _mm256_store_pd(dst2 + 4, _mm256_loadu_pd(s + 16));
                _mm256_store_pd(dst2 + 8, _mm256_loadu_pd(s + 20));
                dst += NR; dst2 += NR;
            }
            dst += (size_t)kc * NR;
            jr += 2 * NR;
            continue;
        }
        if (w >= NR) {
            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));
                _mm256_store_pd(dst + 8, _mm256_loadu_pd(src + (size_t)k * n + 8));
                dst += NR;
            }
            jr += NR;
        } else {
            for (int k = 0; k < kc; k++) {
                double tmp[NR];
                for (int t = 0; t < NR; t++) tmp[t] = 0.0;
                for (int t = 0; t < w; t++) tmp[t] = src[(size_t)k * n + t];
                _mm256_store_pd(dst, _mm256_loadu_pd(tmp));
                _mm256_store_pd(dst + 4, _mm256_loadu_pd(tmp + 4));
                _mm256_store_pd(dst + 8, _mm256_loadu_pd(tmp + 8));
                dst += NR;
            }
            jr += NR;
        }
    }
}

// SW = number of trailing columns actually stored (NR, or the partial tail width).
template <bool FIRST, bool STREAM, int SW>
static void kernel4x12(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;
    int k = 0;
#define K12STEP() {                                                                   \
        __m256d b0 = _mm256_load_pd(bp);                                              \
        __m256d b1 = _mm256_load_pd(bp + 4);                                          \
        __m256d b2 = _mm256_load_pd(bp + 8);                                          \
        __m256d a;                                                                    \
        a = _mm256_broadcast_sd(ap + 0); c0 = _mm256_fmadd_pd(a, b0, c0); c1 = _mm256_fmadd_pd(a, b1, c1); c2 = _mm256_fmadd_pd(a, b2, c2); \
        a = _mm256_broadcast_sd(ap + 1); c3 = _mm256_fmadd_pd(a, b0, c3); c4 = _mm256_fmadd_pd(a, b1, c4); c5 = _mm256_fmadd_pd(a, b2, c5); \
        a = _mm256_broadcast_sd(ap + 2); c6 = _mm256_fmadd_pd(a, b0, c6); c7 = _mm256_fmadd_pd(a, b1, c7); c8 = _mm256_fmadd_pd(a, b2, c8); \
        a = _mm256_broadcast_sd(ap + 3); c9 = _mm256_fmadd_pd(a, b0, c9); c10 = _mm256_fmadd_pd(a, b1, c10); c11 = _mm256_fmadd_pd(a, b2, c11); \
        ap += MR; bp += NR;                                                           \
    }
#if KU == 8
    for (; k + 8 <= kc; k += 8) {
        _mm_prefetch((const char *)(ap + APF), _MM_HINT_T0);
        _mm_prefetch((const char *)(bp + BPF), _MM_HINT_T0);
        K12STEP() K12STEP() K12STEP() K12STEP() K12STEP() K12STEP() K12STEP() K12STEP()
    }
#elif 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);
        K12STEP() K12STEP() K12STEP() K12STEP()
    }
#else
    for (; k + 4 <= kc; k += 4) { K12STEP() K12STEP() K12STEP() K12STEP() }
#endif
    for (; k < kc; k++) K12STEP()
    double *cp = C;
#define ST12(R, X, Y, Z) if (FIRST) { \
        if (STREAM) { _mm256_stream_pd(cp, X); if (SW > 4) _mm256_stream_pd(cp + 4, Y); if (SW > 8) _mm256_stream_pd(cp + 8, Z); } \
        else { _mm256_storeu_pd(cp, X); if (SW > 4) _mm256_storeu_pd(cp + 4, Y); if (SW > 8) _mm256_storeu_pd(cp + 8, Z); } } \
    else { __m256d t = _mm256_loadu_pd(cp); _mm256_storeu_pd(cp, _mm256_add_pd(t, X)); \
           if (SW > 4) { t = _mm256_loadu_pd(cp + 4); _mm256_storeu_pd(cp + 4, _mm256_add_pd(t, Y)); } \
           if (SW > 8) { t = _mm256_loadu_pd(cp + 8); _mm256_storeu_pd(cp + 8, _mm256_add_pd(t, Z)); } } cp += ldc;
    ST12(0, c0, c1, c2) ST12(1, c3, c4, c5) ST12(2, c6, c7, c8) ST12(3, c9, c10, c11)
}

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) {
    if (n <= 0) return;
    const bool cs = CSTREAM && ((n & 3) == 0) && (((size_t)(const void *)C & 31) == 0);
    int n4r = n - (n % MR);       // rows handled by the vector kernel
    for (int jc = 0; jc < n; jc += NC) {
        int nc = n - jc; if (nc > NC) nc = NC;
        int jfull = nc - (nc % NR);
        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 < n4r; ic += MC) {
                int mc = n4r - ic; if (mc > MC) mc = MC;
                packA4(A, n, ic, pc, mc, kc);
#define RUN12(FV, SV, SWV) \
                for (int ir = 0; ir < mc; ir += MR) \
                    kernel4x12<FV, SV, SWV>(Apack + (size_t)(ir / MR) * kc * MR, bp, kc, cb + (size_t)ir * n, n);
                for (int jr = 0; jr < jfull; jr += NR) {
                    const double *bp = Bpack + (size_t)(jr / NR) * kc * NR;
                    double *cb = C + (size_t)ic * n + jc + jr;
                    if (cs) RUN12(true, true, NR) else RUN12(true, false, NR)
                }
                if (jfull < nc) {
                    int w = nc - jfull;
                    if (w % 4 == 0) {   // 4, 8 or 12 real columns: kernel stores exactly w
                        const double *bp = Bpack + (size_t)(jfull / NR) * kc * NR;
                        double *cb = C + (size_t)ic * n + jc + jfull;
                        if (w > 8) { if (cs) RUN12(true, true, 12) else RUN12(true, false, 12) }
                        else if (w > 4) { if (cs) RUN12(true, true, 8) else RUN12(true, false, 8) }
                        else { if (cs) RUN12(true, true, 4) else RUN12(true, false, 4) }
                    } else {            // odd widths: scalar fallback (n=1024 uses w=4)
                        edge_rows(A, B, C, n, ic, ic + mc, jc + jfull, jc + nc);
                    }
                }
            }
        }
    }
    if (n4r < n) edge_rows(A, B, C, n, n4r, n, 0, n);
    _mm_sfence();
}

评测结果