1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 * This source code is licensed under the BSD-style license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8#include "fbgemm/FbgemmFPCommon.h"
9#include "fbgemm/Fbgemm.h"
10
11#include <array>
12#include <cmath>
13#include <utility>
14
15namespace fbgemm {
16
17/// class that performs packing of matrix in
18/// row-major or col-major format into
19/// internal packed blocked-row major format
20
21/// Todo: make it fast with AVX2 transpose
22void PackA(int nrow, int ncol, const float* from, int ldim, float* to) {
23 // for (int r = 0; r < nrow; ++r) {
24 // for (int c = 0; c < ncol; ++c) {
25 // to[r + c * nrow] = from[r * ldim + c];
26 // }
27 // }
28 transpose_simd(nrow, ncol, from, ldim, to, nrow);
29}
30
31// Each kernel does the following computation that multiplies
32// mb x k A sub-matrix with k x b_block_cols*64 B sub-matrix
33// for (int j = 0; j < b_block_cols * 64; j += 64) {
34// for (int kk = 0; kk < k; ++k) {
35// for (int i = 0; i < mb; ++i) {
36// c[i][j:j+64] += a[i][kk] * b[kk][j:j+64]
37// }
38// }
39// }
40
41// autotuned kernel splits for various cases m = 1:mb_max
42// may need re-autotuning for new uarch
43// clang-format off
44partition_array_t partition_avx2 = {
45 // NOTE: clang-format wants to use a different formatting but the current
46 // formatting should be easier to read.
47 {
48 {{ { 0, 0 }, { 0, 0 } } }, // 0
49 {{ { 1, 1 }, { 0, 0 } } }, // 1
50 {{ { 2, 1 }, { 0, 0 } } }, // 2
51 {{ { 3, 1 }, { 0, 0 } } }, // 3
52 {{ { 4, 1 }, { 0, 0 } } }, // 4
53 {{ { 5, 1 }, { 0, 0 } } }, // 5
54 {{ { 6, 1 }, { 0, 0 } } }, // 6
55 {{ { 5, 1 }, { 2, 1 } } }, // 7
56 {{ { 4, 2 }, { 0, 0 } } }, // 8
57 {{ { 5, 1 }, { 4, 1 } } }, // 9
58 {{ { 5, 2 }, { 0, 0 } } }, // 10
59 {{ { 6, 1 }, { 5, 1 } } }, // 11
60 {{ { 6, 2 }, { 0, 0 } } }, // 12
61 {{ { 5, 2 }, { 3, 1 } } }, // 13
62 {{ { 6, 2 }, { 2, 1 } } }, // 14
63 {{ { 5, 3 }, { 0, 0 } } }, // 15
64 {{ { 6, 2 }, { 4, 1 } } }, // 16
65 {{ { 6, 2 }, { 5, 1 } } }, // 17
66 {{ { 6, 3 }, { 0, 0 } } }, // 18
67 {{ { 5, 3 }, { 4, 1 } } }, // 19
68 {{ { 5, 4 }, { 0, 0 } } }, // 20
69 {{ { 5, 3 }, { 6, 1 } } }, // 21
70 {{ { 6, 3 }, { 4, 1 } } }, // 22
71 {{ { 6, 3 }, { 5, 1 } } }, // 23
72 {{ { 6, 4 }, { 0, 0 } } }, // 24
73 {{ { 5, 5 }, { 0, 0 } } }, // 25
74 {{ { 5, 4 }, { 6, 1 } } }, // 26
75 {{ { 6, 4 }, { 3, 1 } } }, // 27
76 {{ { 6, 4 }, { 4, 1 } } }, // 28
77 {{ { 6, 4 }, { 5, 1 } } }, // 29
78 {{ { 6, 5 }, { 0, 0 } } }, // 30
79 {{ { 6, 5 }, { 1, 1 } } }, // 31
80 {{ { 6, 5 }, { 2, 1 } } }, // 32
81 {{ { 6, 5 }, { 3, 1 } } }, // 33
82 {{ { 6, 5 }, { 4, 1 } } }, // 34
83 {{ { 6, 5 }, { 5, 1 } } }, // 35
84 {{ { 6, 6 }, { 0, 0 } } }, // 36
85 {{ { 6, 6 }, { 1, 1 } } }, // 37
86 {{ { 6, 6 }, { 2, 1 } } }, // 38
87 {{ { 6, 6 }, { 3, 1 } } }, // 39
88 {{ { 6, 6 }, { 4, 1 } } }, // 40
89 {{ { 6, 6 }, { 5, 1 } } }, // 41
90 {{ { 6, 7 }, { 0, 0 } } }, // 42
91 {{ { 6, 7 }, { 1, 1 } } }, // 43
92 {{ { 6, 7 }, { 2, 1 } } }, // 44
93 {{ { 6, 7 }, { 3, 1 } } }, // 45
94 {{ { 6, 7 }, { 4, 1 } } }, // 46
95 {{ { 6, 7 }, { 5, 1 } } }, // 47
96 {{ { 6, 8 }, { 0, 0 } } }, // 48
97 {{ { 6, 8 }, { 1, 1 } } }, // 49
98 {{ { 6, 8 }, { 2, 1 } } }, // 50
99 {{ { 6, 8 }, { 3, 1 } } }, // 51
100 {{ { 6, 8 }, { 4, 1 } } }, // 52
101 {{ { 6, 8 }, { 5, 1 } } }, // 53
102 {{ { 6, 9 }, { 0, 0 } } }, // 54
103 {{ { 6, 9 }, { 1, 1 } } }, // 55
104 {{ { 6, 9 }, { 2, 1 } } }, // 56
105 {{ { 6, 9 }, { 3, 1 } } }, // 57
106 {{ { 6, 9 }, { 4, 1 } } }, // 58
107 {{ { 6, 9 }, { 5, 1 } } }, // 59
108 {{ { 6, 10 }, { 0, 0 } } }, // 60
109 {{ { 6, 10 }, { 1, 1 } } }, // 61
110 {{ { 6, 10 }, { 2, 1 } } }, // 62
111 {{ { 6, 10 }, { 3, 1 } } }, // 63
112 {{ { 6, 10 }, { 4, 1 } } }, // 64
113 {{ { 6, 10 }, { 5, 1 } } }, // 65
114 {{ { 6, 11 }, { 0, 0 } } }, // 66
115 {{ { 6, 11 }, { 1, 1 } } }, // 67
116 {{ { 6, 11 }, { 2, 1 } } }, // 68
117 {{ { 6, 11 }, { 3, 1 } } }, // 69
118 {{ { 6, 11 }, { 4, 1 } } }, // 70
119 {{ { 6, 11 }, { 5, 1 } } }, // 71
120 {{ { 6, 12 }, { 0, 0 } } }, // 72
121 {{ { 6, 12 }, { 1, 1 } } }, // 73
122 {{ { 6, 12 }, { 2, 1 } } }, // 74
123 {{ { 6, 12 }, { 3, 1 } } }, // 75
124 {{ { 6, 12 }, { 4, 1 } } }, // 76
125 {{ { 6, 12 }, { 5, 1 } } }, // 77
126 {{ { 6, 13 }, { 0, 0 } } }, // 78
127 {{ { 6, 13 }, { 1, 1 } } }, // 79
128 {{ { 6, 13 }, { 2, 1 } } }, // 80
129 {{ { 6, 13 }, { 3, 1 } } }, // 81
130 {{ { 6, 13 }, { 4, 1 } } }, // 82
131 {{ { 6, 13 }, { 5, 1 } } }, // 83
132 {{ { 6, 14 }, { 0, 0 } } }, // 84
133 {{ { 6, 14 }, { 1, 1 } } }, // 85
134 {{ { 6, 14 }, { 2, 1 } } }, // 86
135 {{ { 6, 14 }, { 3, 1 } } }, // 87
136 {{ { 6, 14 }, { 4, 1 } } }, // 88
137 {{ { 6, 14 }, { 5, 1 } } }, // 89
138 {{ { 6, 15 }, { 0, 0 } } }, // 90
139 {{ { 6, 15 }, { 1, 1 } } }, // 91
140 {{ { 6, 15 }, { 2, 1 } } }, // 92
141 {{ { 6, 15 }, { 3, 1 } } }, // 93
142 {{ { 6, 15 }, { 4, 1 } } }, // 94
143 {{ { 6, 15 }, { 5, 1 } } }, // 95
144 {{ { 6, 16 }, { 0, 0 } } }, // 96
145 {{ { 6, 16 }, { 1, 1 } } }, // 97
146 {{ { 6, 16 }, { 2, 1 } } }, // 98
147 {{ { 6, 16 }, { 3, 1 } } }, // 99
148 {{ { 6, 16 }, { 4, 1 } } }, // 100
149 {{ { 6, 16 }, { 5, 1 } } }, // 101
150 {{ { 6, 17 }, { 0, 0 } } }, // 102
151 {{ { 6, 17 }, { 1, 1 } } }, // 103
152 {{ { 6, 17 }, { 2, 1 } } }, // 104
153 {{ { 6, 17 }, { 3, 1 } } }, // 105
154 {{ { 6, 17 }, { 4, 1 } } }, // 106
155 {{ { 6, 17 }, { 5, 1 } } }, // 107
156 {{ { 6, 18 }, { 0, 0 } } }, // 108
157 {{ { 6, 18 }, { 1, 1 } } }, // 109
158 {{ { 6, 18 }, { 2, 1 } } }, // 110
159 {{ { 6, 18 }, { 3, 1 } } }, // 111
160 {{ { 6, 18 }, { 4, 1 } } }, // 112
161 {{ { 6, 18 }, { 5, 1 } } }, // 113
162 {{ { 6, 19 }, { 0, 0 } } }, // 114
163 {{ { 6, 19 }, { 1, 1 } } }, // 115
164 {{ { 6, 19 }, { 2, 1 } } }, // 116
165 {{ { 6, 19 }, { 3, 1 } } }, // 117
166 {{ { 6, 19 }, { 4, 1 } } }, // 118
167 {{ { 6, 19 }, { 5, 1 } } }, // 119
168 {{ { 6, 20 }, { 0, 0 } } }, // 120
169 }
170};
171
172partition_array_t partition_avx512 = {
173 // NOTE: clang-format wants to use a different formatting but the current
174 // formatting should be easier to read.
175 {
176 {{ { 0, 0 }, { 0, 0 } } }, // 0
177 {{ { 1, 1 }, { 0, 0 } } }, // 1
178 {{ { 2, 1 }, { 0, 0 } } }, // 2
179 {{ { 3, 1 }, { 0, 0 } } }, // 3
180 {{ { 4, 1 }, { 0, 0 } } }, // 4
181 {{ { 5, 1 }, { 0, 0 } } }, // 5
182 {{ { 6, 1 }, { 0, 0 } } }, // 6
183 {{ { 7, 1 }, { 0, 0 } } }, // 7
184 {{ { 8, 1 }, { 0, 0 } } }, // 8
185 {{ { 9, 1 }, { 0, 0 } } }, // 9
186 {{ { 10, 1 }, { 0, 0 } } }, // 10
187 {{ { 11, 1 }, { 0, 0 } } }, // 11
188 {{ { 12, 1 }, { 0, 0 } } }, // 12
189 {{ { 13, 1 }, { 0, 0 } } }, // 13
190 {{ { 14, 1 }, { 0, 0 } } }, // 14
191 {{ { 8, 1 }, { 7, 1 } } }, // 15
192 {{ { 8, 2 }, { 0, 0 } } }, // 16
193 {{ { 9, 1 }, { 8, 1 } } }, // 17
194 {{ { 9, 2 }, { 0, 0 } } }, // 18
195 {{ { 10, 1 }, { 9, 1 } } }, // 19
196 {{ { 10, 2 }, { 0, 0 } } }, // 20
197 {{ { 11, 1 }, { 10, 1 } } }, // 21
198 {{ { 11, 2 }, { 0, 0 } } }, // 22
199 {{ { 12, 1 }, { 11, 1 } } }, // 23
200 {{ { 12, 2 }, { 0, 0 } } }, // 24
201 {{ { 13, 1 }, { 12, 1 } } }, // 25
202 {{ { 13, 2 }, { 0, 0 } } }, // 26
203 {{ { 14, 1 }, { 13, 1 } } }, // 27
204 {{ { 14, 2 }, { 0, 0 } } }, // 28
205 {{ { 10, 2 }, { 9, 1 } } }, // 29
206 {{ { 10, 3 }, { 0, 0 } } }, // 30
207 {{ { 11, 2 }, { 9, 1 } } }, // 31
208 {{ { 11, 2 }, { 10, 1 } } }, // 32
209 {{ { 11, 3 }, { 0, 0 } } }, // 33
210 {{ { 12, 2 }, { 10, 1 } } }, // 34
211 {{ { 12, 2 }, { 11, 1 } } }, // 35
212 {{ { 12, 3 }, { 0, 0 } } }, // 36
213 {{ { 13, 2 }, { 11, 1 } } }, // 37
214 {{ { 13, 2 }, { 12, 1 } } }, // 38
215 {{ { 13, 3 }, { 0, 0 } } }, // 39
216 {{ { 14, 2 }, { 12, 1 } } }, // 40
217 {{ { 14, 2 }, { 13, 1 } } }, // 41
218 {{ { 14, 3 }, { 0, 0 } } }, // 42
219 {{ { 11, 3 }, { 10, 1 } } }, // 43
220 {{ { 11, 4 }, { 0, 0 } } }, // 44
221 {{ { 12, 3 }, { 9, 1 } } }, // 45
222 {{ { 12, 3 }, { 10, 1 } } }, // 46
223 {{ { 12, 3 }, { 11, 1 } } }, // 47
224 {{ { 12, 4 }, { 0, 0 } } }, // 48
225 {{ { 13, 3 }, { 10, 1 } } }, // 49
226 {{ { 13, 3 }, { 11, 1 } } }, // 50
227 {{ { 13, 3 }, { 12, 1 } } }, // 51
228 {{ { 13, 4 }, { 0, 0 } } }, // 52
229 {{ { 14, 3 }, { 11, 1 } } }, // 53
230 {{ { 14, 3 }, { 12, 1 } } }, // 54
231 {{ { 14, 3 }, { 13, 1 } } }, // 55
232 {{ { 14, 4 }, { 0, 0 } } }, // 56
233 {{ { 12, 4 }, { 9, 1 } } }, // 57
234 {{ { 12, 4 }, { 10, 1 } } }, // 58
235 {{ { 12, 4 }, { 11, 1 } } }, // 59
236 {{ { 12, 5 }, { 0, 0 } } }, // 60
237 {{ { 13, 4 }, { 9, 1 } } }, // 61
238 {{ { 13, 4 }, { 10, 1 } } }, // 62
239 {{ { 13, 4 }, { 11, 1 } } }, // 63
240 {{ { 13, 4 }, { 12, 1 } } }, // 64
241 {{ { 13, 5 }, { 0, 0 } } }, // 65
242 {{ { 14, 4 }, { 10, 1 } } }, // 66
243 {{ { 14, 4 }, { 11, 1 } } }, // 67
244 {{ { 14, 4 }, { 12, 1 } } }, // 68
245 {{ { 14, 4 }, { 13, 1 } } }, // 69
246 {{ { 14, 5 }, { 0, 0 } } }, // 70
247 {{ { 12, 5 }, { 11, 1 } } }, // 71
248 {{ { 12, 6 }, { 0, 0 } } }, // 72
249 {{ { 13, 5 }, { 8, 1 } } }, // 73
250 {{ { 13, 5 }, { 9, 1 } } }, // 74
251 {{ { 13, 5 }, { 10, 1 } } }, // 75
252 {{ { 13, 5 }, { 11, 1 } } }, // 76
253 {{ { 13, 5 }, { 12, 1 } } }, // 77
254 {{ { 13, 6 }, { 0, 0 } } }, // 78
255 {{ { 14, 5 }, { 9, 1 } } }, // 79
256 {{ { 14, 5 }, { 10, 1 } } }, // 80
257 {{ { 14, 5 }, { 11, 1 } } }, // 81
258 {{ { 14, 5 }, { 12, 1 } } }, // 82
259 {{ { 14, 5 }, { 13, 1 } } }, // 83
260 {{ { 14, 6 }, { 0, 0 } } }, // 84
261 {{ { 13, 6 }, { 7, 1 } } }, // 85
262 {{ { 13, 6 }, { 8, 1 } } }, // 86
263 {{ { 13, 6 }, { 9, 1 } } }, // 87
264 {{ { 13, 6 }, { 10, 1 } } }, // 88
265 {{ { 13, 6 }, { 11, 1 } } }, // 89
266 {{ { 13, 6 }, { 12, 1 } } }, // 90
267 {{ { 13, 7 }, { 0, 0 } } }, // 91
268 {{ { 14, 6 }, { 8, 1 } } }, // 92
269 {{ { 14, 6 }, { 9, 1 } } }, // 93
270 {{ { 14, 6 }, { 10, 1 } } }, // 94
271 {{ { 14, 6 }, { 11, 1 } } }, // 95
272 {{ { 14, 6 }, { 12, 1 } } }, // 96
273 {{ { 14, 6 }, { 13, 1 } } }, // 97
274 {{ { 14, 7 }, { 0, 0 } } }, // 98
275 {{ { 13, 7 }, { 8, 1 } } }, // 99
276 {{ { 13, 7 }, { 9, 1 } } }, // 100
277 {{ { 13, 7 }, { 10, 1 } } }, // 101
278 {{ { 13, 7 }, { 11, 1 } } }, // 102
279 {{ { 13, 7 }, { 12, 1 } } }, // 103
280 {{ { 13, 8 }, { 0, 0 } } }, // 104
281 {{ { 14, 7 }, { 7, 1 } } }, // 105
282 {{ { 14, 7 }, { 8, 1 } } }, // 106
283 {{ { 14, 7 }, { 9, 1 } } }, // 107
284 {{ { 14, 7 }, { 10, 1 } } }, // 108
285 {{ { 14, 7 }, { 11, 1 } } }, // 109
286 {{ { 14, 7 }, { 12, 1 } } }, // 110
287 {{ { 14, 7 }, { 13, 1 } } }, // 111
288 {{ { 14, 8 }, { 0, 0 } } }, // 112
289 {{ { 13, 8 }, { 9, 1 } } }, // 113
290 {{ { 13, 8 }, { 10, 1 } } }, // 114
291 {{ { 13, 8 }, { 11, 1 } } }, // 115
292 {{ { 13, 8 }, { 12, 1 } } }, // 116
293 {{ { 13, 9 }, { 0, 0 } } }, // 117
294 {{ { 14, 8 }, { 6, 1 } } }, // 118
295 {{ { 14, 8 }, { 7, 1 } } }, // 119
296 {{ { 14, 8 }, { 8, 1 } } }, // 120
297 }
298};
299
300}
301