123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153 |
- // See LICENSE for license details.
- #include "common.h"
- #include <assert.h>
- #include <math.h>
- #include <stdint.h>
- #include <alloca.h>
- #define MIN(a, b) ((a) < (b) ? (a) : (b))
- static void mm_naive(size_t m, size_t n, size_t p,
- t* a, size_t lda, t* b, size_t ldb, t* c, size_t ldc)
- {
- for (size_t i = 0; i < m; i++)
- {
- for (size_t j = 0; j < n; j++)
- {
- t s0 = c[i*ldc+j], s1 = 0, s2 = 0, s3 = 0;
- for (size_t k = 0; k < p/4*4; k+=4)
- {
- s0 = fma(a[i*lda+k+0], b[(k+0)*ldb+j], s0);
- s1 = fma(a[i*lda+k+1], b[(k+1)*ldb+j], s1);
- s2 = fma(a[i*lda+k+2], b[(k+2)*ldb+j], s2);
- s3 = fma(a[i*lda+k+3], b[(k+3)*ldb+j], s3);
- }
- for (size_t k = p/4*4; k < p; k++)
- s0 = fma(a[i*lda+k], b[k*ldb+j], s0);
- c[i*ldc+j] = (s0 + s1) + (s2 + s3);
- }
- }
- }
- static inline void mm_rb(size_t m, size_t n, size_t p,
- t* a, size_t lda, t* b, size_t ldb, t* c, size_t ldc)
- {
- size_t mb = m/RBM*RBM, nb = n/RBN*RBN;
- for (size_t i = 0; i < mb; i += RBM)
- {
- for (size_t j = 0; j < nb; j += RBN)
- kloop(p, a+i*lda, lda, b+j, ldb, c+i*ldc+j, ldc);
- mm_naive(RBM, n - nb, p, a+i*lda, lda, b+nb, ldb, c+i*ldc+nb, ldc);
- }
- mm_naive(m - mb, n, p, a+mb*lda, lda, b, ldb, c+mb*ldc, ldc);
- }
- static inline void repack(t* a, size_t lda, const t* a0, size_t lda0, size_t m, size_t p)
- {
- for (size_t i = 0; i < m; i++)
- {
- for (size_t j = 0; j < p/8*8; j+=8)
- {
- t t0 = a0[i*lda0+j+0];
- t t1 = a0[i*lda0+j+1];
- t t2 = a0[i*lda0+j+2];
- t t3 = a0[i*lda0+j+3];
- t t4 = a0[i*lda0+j+4];
- t t5 = a0[i*lda0+j+5];
- t t6 = a0[i*lda0+j+6];
- t t7 = a0[i*lda0+j+7];
- a[i*lda+j+0] = t0;
- a[i*lda+j+1] = t1;
- a[i*lda+j+2] = t2;
- a[i*lda+j+3] = t3;
- a[i*lda+j+4] = t4;
- a[i*lda+j+5] = t5;
- a[i*lda+j+6] = t6;
- a[i*lda+j+7] = t7;
- }
- for (size_t j = p/8*8; j < p; j++)
- a[i*lda+j] = a0[i*lda0+j];
- }
- }
- static void mm_cb(size_t m, size_t n, size_t p,
- t* a, size_t lda, t* b, size_t ldb, t* c, size_t ldc)
- {
- size_t nmb = m/CBM, nnb = n/CBN, npb = p/CBK;
- size_t mb = nmb*CBM, nb = nnb*CBN, pb = npb*CBK;
- //t a1[mb*pb], b1[pb*nb], c1[mb*nb];
- t* a1 = (t*)alloca_aligned(sizeof(t)*mb*pb, 8192);
- t* b1 = (t*)alloca_aligned(sizeof(t)*pb*nb, 8192);
- t* c1 = (t*)alloca_aligned(sizeof(t)*mb*nb, 8192);
- for (size_t i = 0; i < mb; i += CBM)
- for (size_t j = 0; j < pb; j += CBK)
- repack(a1 + (npb*(i/CBM) + j/CBK)*(CBM*CBK), CBK, a + i*lda + j, lda, CBM, CBK);
- for (size_t i = 0; i < pb; i += CBK)
- for (size_t j = 0; j < nb; j += CBN)
- repack(b1 + (nnb*(i/CBK) + j/CBN)*(CBK*CBN), CBN, b + i*ldb + j, ldb, CBK, CBN);
- for (size_t i = 0; i < mb; i += CBM)
- for (size_t j = 0; j < nb; j += CBN)
- repack(c1 + (nnb*(i/CBM) + j/CBN)*(CBM*CBN), CBN, c + i*ldc + j, ldc, CBM, CBN);
- for (size_t i = 0; i < mb; i += CBM)
- {
- for (size_t j = 0; j < nb; j += CBN)
- {
- for (size_t k = 0; k < pb; k += CBK)
- {
- mm_rb(CBM, CBN, CBK,
- a1 + (npb*(i/CBM) + k/CBK)*(CBM*CBK), CBK,
- b1 + (nnb*(k/CBK) + j/CBN)*(CBK*CBN), CBN,
- c1 + (nnb*(i/CBM) + j/CBN)*(CBM*CBN), CBN);
- }
- if (pb < p)
- {
- mm_rb(CBM, CBN, p - pb,
- a + i*lda + pb, lda,
- b + pb*ldb + j, ldb,
- c1 + (nnb*(i/CBM) + j/CBN)*(CBM*CBN), CBN);
- }
- }
- if (nb < n)
- {
- for (size_t k = 0; k < p; k += CBK)
- {
- mm_rb(CBM, n - nb, MIN(p - k, CBK),
- a + i*lda + k, lda,
- b + k*ldb + nb, ldb,
- c + i*ldc + nb, ldc);
- }
- }
- }
- if (mb < m)
- {
- for (size_t j = 0; j < n; j += CBN)
- {
- for (size_t k = 0; k < p; k += CBK)
- {
- mm_rb(m - mb, MIN(n - j, CBN), MIN(p - k, CBK),
- a + mb*lda + k, lda,
- b + k*ldb + j, ldb,
- c + mb*ldc + j, ldc);
- }
- }
- }
- for (size_t i = 0; i < mb; i += CBM)
- for (size_t j = 0; j < nb; j += CBN)
- repack(c + i*ldc + j, ldc, c1 + (nnb*(i/CBM) + j/CBN)*(CBM*CBN), CBN, CBM, CBN);
- }
- void mm(size_t m, size_t n, size_t p,
- t* a, size_t lda, t* b, size_t ldb, t* c, size_t ldc)
- {
- if (__builtin_expect(m <= 2*CBM && n <= 2*CBN && p <= 2*CBK, 1))
- mm_rb(m, n, p, a, lda, b, ldb, c, ldc);
- else
- mm_cb(m, n, p, a, lda, b, ldb, c, ldc);
- }
|