ℹ️ Select 'Choose Exercise', or randomize 'Next Random Exercise' in selected language.

Choose Exercise:
Timer 00:00
WPM --
Score --
Acc --
Correct chars --

ASM x86 SIMD: Vectorized Matrix Multiplication (AVX2)

ASM x86

Goal -- WPM

Ready
Exercise Algorithm Area
1; Optimized Matrix Multiplication using AVX2
2; Computes C = A * B, where A, B, and C are matrices of size N x N.
3; Assumes N is a multiple of 8 for optimal AVX2 vectorization.
4
5section .data
6; Example matrices (replace with actual data or dynamic allocation)
7matrix_a: times 64*64*8 db 0 ; 64x64 matrix of doubles (8 bytes each)
8matrix_b: times 64*64*8 db 0
9matrix_c: times 64*64*8 db 0
10
11section .text
12global _start
13
14_start:
15; --- Initialization (example: fill matrices with dummy data) ---
16; In a real scenario, matrices would be loaded or computed.
17; For demonstration, we'll assume they are populated.
18; ... (code to populate matrix_a, matrix_b) ...
19
20; --- Matrix Multiplication Call ---
21; Assume N = 64 (matrix dimension)
22mov rdi, matrix_a ; Pointer to matrix A
23mov rsi, matrix_b ; Pointer to matrix B
24mov rdx, matrix_c ; Pointer to matrix C
25mov rcx, 64 ; Dimension N
26call multiply_matrices_avx2
27
28; --- Exit Program ---
29mov rax, 60
30mov rdi, 0
31syscall
32
33; Function: multiply_matrices_avx2
34; Computes C = A * B using AVX2 instructions.
35; Input:
36; RDI = Pointer to matrix A (double precision)
37; RSI = Pointer to matrix B (double precision)
38; RDX = Pointer to matrix C (double precision)
39; RCX = Matrix dimension N (must be multiple of 8)
40; Output:
41; Matrix C is updated with the result.
42multiply_matrices_avx2:
43; --- Prologue ---
44push rbp
45mov rbp, rsp
46push rbx
47push r12
48push r13
49push r14
50push r15
51; Save YMM registers if they might be used by called functions
52; (not strictly necessary here as we are the leaf function)
53
54; Register allocation:
55; RDI: ptrA
56; RSI: ptrB
57; RDX: ptrC
58; RCX: N
59; RBX: outer loop counter (row i)
60; R12: middle loop counter (col j)
61; R13: inner loop counter (k)
62; R14: ptrA_row_i
63; R15: ptrB_col_j
64
65mov rbx, 0 ; i = 0 (row index for C and A)
66mov r14, rdi ; ptrA_row_i = ptrA
67
68.outer_loop:
69cmp rbx, rcx ; if i >= N
70jge .end_outer_loop
71
72mov r12, 0 ; j = 0 (column index for C and B)
73mov r15, rsi ; ptrB_col_j = ptrB (start of column j)
74; Calculate the starting address for column j of B
75; This calculation is tricky: ptrB + j * N * sizeof(double)
76; Since N is in RCX, and we are iterating j, we need to adjust ptrB.
77; For simplicity, let's assume we can re-calculate ptrB_col_j each time.
78; A more optimized approach would pre-calculate column pointers or transpose B.
79
80.middle_loop:
81cmp r12, rcx ; if j >= N
82jge .end_middle_loop
83
84; Initialize sum vector for C[i][j] to zero
85; We will accumulate 8 doubles (one AVX vector) at a time.
86vpxor ymm0, ymm0, ymm0 ; sum = {0.0, 0.0, ..., 0.0}
87
88mov r13, 0 ; k = 0 (inner loop counter)
89mov r14, rdi ; Reset ptrA_row_i for each row i
90add r14, rbx ; Move to the correct row in A
91; Calculate offset for row i: i * N * sizeof(double)
92mov rax, rbx ; rax = i
93imul rax, rcx ; rax = i * N
94imul rax, 8 ; rax = i * N * sizeof(double)
95add r14, rax ; r14 now points to the start of row i in A
96
97; Calculate the starting address for column j of B
98; ptrB_col_j = ptrB + j * N * sizeof(double)
99mov rax, r12 ; rax = j
100imul rax, rcx ; rax = j * N
101imul rax, 8 ; rax = j * N * sizeof(double)
102mov r15, rsi ; r15 = ptrB
103add r15, rax ; r15 now points to the start of column j in B
104
105.inner_loop:
106cmp r13, rcx ; if k >= N
107jge .end_inner_loop
108
109; Load 8 doubles from A (row i, starting at k)
110; Address: r14 + k * sizeof(double)
111vmovupd ymm1, [r14 + r13 * 8]
112
113; Load 8 doubles from B (column j, starting at k)
114; Address: r15 + k * sizeof(double)
115vmovupd ymm2, [r15 + r13 * 8]
116
117; Multiply vectors: ymm3 = ymm1 * ymm2
118vmulpd ymm3, ymm1, ymm2
119
120; Accumulate into sum: ymm0 = ymm0 + ymm3
121vaddpd ymm0, ymm0, ymm3
122
123add r13, 8 ; Increment k by 8 (vector width)
124jmp .inner_loop
125
126.end_inner_loop:
127; Store the accumulated sum vector into C[i][j]
128; Address: ptrC + i * N * sizeof(double) + j * sizeof(double)
129; We need to calculate the address for C[i][j]
130mov rax, rbx ; rax = i
131imul rax, rcx ; rax = i * N
132imul rax, 8 ; rax = i * N * sizeof(double)
133add rax, r12 ; rax = i * N * sizeof(double) + j
134imul rax, 8 ; rax = (i * N + j) * sizeof(double)
135mov rdi, rdx ; rdi = ptrC
136add rdi, rax ; rdi points to C[i][j]
137vmovupd [rdi], ymm0 ; Store the 8 doubles
138
139add r12, 8 ; Increment j by 8 (vector width)
140jmp .middle_loop
141
142.end_middle_loop:
143inc rbx ; Increment i
144jmp .outer_loop
145
146.end_outer_loop:
147
148; --- Epilogue ---
149pop r15
150pop r14
151pop r13
152pop r12
153pop rbx
154pop rbp
155ret
Algorithm description viewbox

ASM x86 SIMD: Vectorized Matrix Multiplication (AVX2)

Algorithm description:

This ASM x86 code implements a matrix multiplication routine (C = A * B) optimized for speed using AVX2 SIMD instructions. It processes data in 256-bit (32-byte) chunks, performing 4 double-precision floating-point operations in parallel. This significantly accelerates computations for large matrices, making it suitable for scientific computing, machine learning, and graphics processing.

Algorithm explanation:

The `multiply_matrices_avx2` function computes C = A * B for N x N matrices, where N is a multiple of 8. It uses three nested loops: the outer loop iterates through rows of C (i), the middle loop iterates through columns of C (j), and the inner loop iterates through the dot product elements (k). The key optimization is the inner loop: instead of processing one element at a time, it loads 8 doubles from row i of A and 8 doubles from column j of B into YMM registers (`ymm1`, `ymm2`). These vectors are then multiplied (`vmulpd`) and accumulated into a sum vector (`ymm0`) using `vaddpd`. Finally, the accumulated sum vector is stored into C[i][j]. This vectorized approach leverages the parallelism of AVX2, performing 4 operations per clock cycle per YMM register. The time complexity is O(N^3), but the constant factor is significantly reduced compared to scalar multiplication. Space complexity is O(1) beyond input storage.

Pseudocode:

FUNCTION multiply_matrices_avx2(ptrA, ptrB, ptrC, N):
  SAVE registers (rbp, rbx, r12-r15)

  FOR i FROM 0 TO N-1:
    FOR j FROM 0 TO N-1 STEP 8:
      INITIALIZE sum_vector YMM0 to {0.0, ..., 0.0}
      ptrA_row_i = ptrA + i * N * sizeof(double)
      ptrB_col_j = ptrB + j * N * sizeof(double)

      FOR k FROM 0 TO N-1 STEP 8:
        LOAD 8 doubles from A[i][k..k+7] into YMM1
        LOAD 8 doubles from B[k..k+7][j] into YMM2
        MULTIPLY YMM1 BY YMM2, store in YMM3
        ADD YMM0 BY YMM3
      END FOR

      STORE YMM0 into C[i][j..j+7]
    END FOR
  END FOR

  RESTORE registers
  RETURN