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 | |
15 | namespace 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 |
22 | void 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 |
44 | partition_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 | |
172 | partition_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 | |