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