▶ 目录 ← 首页
CS61C

DGEMM: SIMD with Intel Intrinsics

Learning Outcomes

  • Compare SIMD-extended matrix multiplication with sequential programs.
  • Compare performance using different gcc optimization flags: -O0, -O2, -O3.
🎥 Lecture Video

Before continuing, we recommend reviewing:

Three matrix diagram showing DGEMM nested-loop structure and index mapping over i, j, and k dimensions. Multiplication of row 0 in matrix A with column 0 of matrix B results in element 0,0 in matrix C.

DGEMM for 8×88 \times 8 square matrices A and B. C00C_{00} is computed as the dot product of row i = 0 of A and column j = 0 of B.

DGEMM 7: Naive SIMD DGEMM

With the ability to perform SIMD multiplication, it is tempting to use our new 256-bit-wide registers to load in blocks of row A and blocks of column B for the element-wise multiplication necessary for dot products, as shown in #fig-dgemm-simd-naive.

Diagram of naive SIMD DGEMM using three matrix panels labeled A, B, and C. Row 0 of A and several columns of B are outlined or shaded to show which operands load into wide SIMD registers for one dot-product step; element C[0,0] is emphasized as the accumulator slot. The figure illustrates packing several independent multiplies from the same dot product into one SIMD instruction stream before summing into C.

“Naive” SIMD DGEMM that leverages SIMD architecture registers to parallelize multiplications within a single dot product. The outlined boxes indicate which values are loaded into the 256-bit-wide registers.1

This naive SIMD implementation does better than our naive DGEMM but worse than specifying narrow registers with the register keyword in C:

C               0.768672 seconds
registers:      0.277462 seconds
simd,naive:     0.416584 seconds

The non-compulsory cache miss persists in our naive SIMD implementation—we are still loading in multiple rows of B instead of columns of B.

DGEMM 8: Transpose SIMD DGEMM

We next apply cache blocking by first transposing B, then leveraging our 256-bit-wide (four-double) registers for element-wise multiplication:

Same three-matrix layout as the naive SIMD case but B is shown transposed so consecutive memory accesses align with SIMD loads. Row 0 of A and contiguous strips along the transposed B layout are highlighted to indicate 256-bit-wide register fills; C[0,0] remains the highlighted output cell. The diagram stresses that transposing B turns column walks into row-major-friendly blocks for the SIMD kernel.

“Transposed” SIMD DGEMM that uses a transposed B to load in columns of B to streamline memory accesses. The outlined boxes indicate which values are loaded into the 256-bit-wide registers.

C               0.768672 seconds
registers:      0.277462 seconds
simd,naive:     0.416584 seconds
simd,transpose: 0.275622 seconds

This version is certainly speedier, but it does not prove huge benefits beyond our register keyword approach.

DGEMM 9: Tiled SIMD DGEMM

Recall that cache blocking is any re-design of our algorithm to adjust memory accesses. Earlier, we discussed a submatrix tiling approach to matrix multiplication—where we compute multiple elements of C with the current set of rows of A and set of columns of B.

In #fig-dgemm-simd-block, we assume that the product of a scalar with a vector can be computed as a vector operation, provided that the scalar is copied to each element in a vector of the same length.

Tiled SIMD matrix-multiply illustration split into left and right regions. Left: a row of C with four adjacent elements accumulated in parallel, starting at C[i,j]. Right: cloud-shaped grouping where four elements from row k of B are scaled by the scalar A[i,k] and summed across k iterations; arrows or indices suggest the inner k loop. Together the panels show how a tile reuses A and B data while updating a short vector of C entries per step.

Compute four elements of C (starting with CijC_{ij}) by iteratively adding the result of scaling four elements of the kk-th row of B (starting with BkjB_{kj}) with AikA_{ik}.

We can extend this idea to a tiled SIMD approach shown in #fig-simd-dgemm-animate.

SIMD `dgemm` tiled matrix multiplication. The outlined boxes indicate which values are loaded into the 256-bit-wide registers. Use the menu bar to trace through the animation or access the [original Google Slides](https://docs.google.com/presentation/d/1luqaX7cXBd158mvN9ZJDBcNa5O2MK4aWIZcrm1wsXeo/edit?usp=sharing).
Show Explanation

#fig-simd-dgemm-animate shows an 8-by-8 matrix multiplication.

  1. i = 0. Compute elements C0jC_{0j} for j=0,1,2,3j = 0, 1, 2, 3; let this be c⃗\vec{c}. Let b⃗\vec{b} be the corresponding part of the kk-th row of BB.

    1a through 1h: Loop through k=0,...,7k = 0, ..., 7. Compute the result of Aik⋅A_{ik} \cdot \vec{b}andaddtothefourelementsofand add to the four elements of\vec{c}$.

  2. i = 0, still. Next, compute elements C0jC_{0j} for j=4,5,6,7j = 4, 5, 6, 7; let this be c⃗\vec{c}. Let b⃗\vec{b} be the corresponding part of the kk-th row of BB. Again, loop through the index kk to add the resulting scalar-by-vector to the elements of c⃗\vec{c}.

  3. i = 1. Compute elements C1jC_{1j} for j=0,1,2,3j = 0, 1, 2, 3.

  4. i = 1. Compute elements C1jC_{1j} for j=4,5,6,7j = 4, 5, 6, 7.

The “tiled” cache blocking SIMD approach performs slightly worse than the transposed SIMD approach.

C               0.768672 seconds
registers:      0.277462 seconds
simd,naive:     0.416584 seconds
simd,transpose: 0.275622 seconds
simd,tiled:     0.356344 seconds

DGEMM 10: GCC Optimization

Finally, let’s compile using different gcc optimization flags. In #tab-dgemm-gcc-simd, we can see that even with mild gcc optimizations like -O1, SIMD vastly outperforms any SISD approach.

Generally speaking, most of the speedup comes not from doing multiple math operations at a time, but instead from doing a large memory load/store at a time.

DGEMM SIMD Code

The three algorithms discussed in this section leverage Intel SIMD extensions. You will see in the code below that the SIMD instructions are written in C as Intel Intrinsics. More next!

DGEMM, SIMD naive 


matrix_d_t *dgemm_simd(matrix_d_t *A, matrix_d_t *B) {
  if (A->ncols!=B->nrows) return NULL; 
  matrix_d_t *C = init_mat_d(A->nrows, B->ncols);
  for (int i = 0; i < A->nrows; i++) {
    for (int j = 0; j < B->ncols; j+=4) { // 4 doubles at a time
            avx256_t v_C = avx_load(C->data + i*C->ncols +j);
            for (int k = 0; k < A->ncols; k++) {
                avx256_t s_A = avx_set_num(A->data[(i*A->ncols)+k]);
                avx256_t v_B = avx_load(B->data+k*B->ncols+j);
                // C_ij += a_ik * B_jk (for j = 0...3)
                v_C = avx_mul_add(s_A, v_B, v_C);
            }
            avx_store(C->data+i*C->ncols+j, v_C);
    }
  }
  return C;
}
DGEMM, SIMD Transpose 


matrix_d_t *dgemm_simd_transpose(matrix_d_t *A, matrix_d_t *B) {
  if (A->ncols!=B->nrows) return NULL;
  matrix_d_t *C = init_mat_d(A->nrows, B->ncols);
  matrix_d_t *B_T = transpose_mat_d(B);
  for (int i = 0; i < A->nrows; i++) {
    for (int j = 0; j < B->ncols; j++) {
      double *ptr_A = A->data+(i*A->ncols);
      double *ptr_B_T = B_T->data+(j*B_T->ncols);
      int k = 0;
  
            // 4 doubles at a time
            avx256_t v_C = avx_set_num(0);
            for (; k < A->ncols/4*4; k+= 4) {
                avx256_t v_A = avx_load(ptr_A+k);
                avx256_t v_B = avx_load(ptr_B_T+k);
                v_C = avx_mul_add(v_A, v_B, v_C);
            }
  
      double mem[4] __attribute__ ((aligned (64)));
            avx_store(mem, v_C);
            double sum = mem[0] + mem[1] + mem[2] + mem[3];

            // tail case
            for(; k < A->ncols; k++) {
                sum += ptr_A[k]*ptr_B_T[k];
            }
            C->data[i*C->ncols+j] = sum;
    }
  }
  free_mat_d(B_T);
  return C;
}
DGEMM, SIMD Tiled 


static inline void matmul_simd_tile(int si, int sj, int sk,
                 matrix_d_t *A, matrix_d_t *B, matrix_d_t *C) {
    for (int i = si; i < si + BLOCKSIZE; i++) {
        for (int j = sj; j < sj + BLOCKSIZE; j+=4) { // 4 doubles at a time
            avx256_t v_C = avx_load(C->data+i*C->ncols+j);

            for (int k = sk; k < sk + BLOCKSIZE; k++) {
                avx256_t s_A = avx_set_num(A->data[(i*A->ncols)+k]);
                avx256_t v_B = avx_load(B->data+k*B->ncols+j);
                v_C = avx_mul_add(s_A, v_B, v_C);
            }
            avx_store(C->data+i*C->ncols+j, v_C);
        }
    }
}
matrix_d_t *dgemm_simd_block(matrix_d_t *A, matrix_d_t *B) {
  if (A->ncols!=B->nrows) return NULL;
  matrix_d_t *C = init_mat_d(A->nrows, B->ncols);
  for (int si = 0; si < A->nrows; si += BLOCKSIZE) {
    for (int sj = 0; sj < B->ncols; sj += BLOCKSIZE) {
            for (int sk = 0; sk < A->ncols; sk+= BLOCKSIZE) {
                matmul_simd_tile(si, sj, sk, A, B, C);
            }
    }
  }
  return C;
}
DGEMM (original)

见 #code-dgemm

AVX intrinsics 


typedef __m256d avx256_t;

// Loads 8 doubles at memory address A into a avx256_t
static inline avx256_t avx_load(double *A) {
  return _mm256_load_pd(A);
}

// Stores the avx256_t at SRC to DST. Each avx256_t element gets stored in a 
// different index of the array passed into DST.
static inline void avx_store(double *dst, avx256_t src) {
  _mm256_store_pd(dst, src);
}

// Creates a avx256_t where every element is equal to num
static inline avx256_t avx_set_num(double num) {
  return _mm256_set1_pd(num);
}

// A * B + C
static inline avx256_t avx_mul_add(avx256_t A, avx256_t B, avx256_t C) {
  return _mm256_fmadd_pd(A, B, C);
}

Footnotes

  1. In #fig-dgemm-simd-naive, we assume a cache that has 256-bit blocks. This carries over the cache assumption from our cache blocking discussion. ↩