@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 280 lines 9.23 KB
Raw
1 #ifndef CPU_MICRO_GEMM_AMX_HPP
2 #define CPU_MICRO_GEMM_AMX_HPP
3 #include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
4
5 namespace cpu_micro_gemm {
6 namespace {
7 // AMX specific
8 constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
9 constexpr static int64_t AMX_TILE_ROW_NUM = 16;
10 constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
11
12 typedef struct __tile_config {
13 uint8_t palette_id = 1;
14 uint8_t start_row = 0;
15 uint8_t reserved_0[14] = {0};
16 uint16_t colsb[16] = {0};
17 uint8_t rows[16] = {0};
18 } __tilecfg;
19
20 // 2-2-4 pattern, for 16 < m <= 32
21 // TILE 0, 1: load A matrix, row num should be 16, m - 16
22 // TILE 2, 3: load B matrix, row num should be 16
23 // TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
24 // - 16
25 template <typename scalar_t>
26 class TileGemm224 {
27 public:
28 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
29 TORCH_CHECK(false, "Unsupported data type for TileGemm224");
30 }
31
32 FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
33 TORCH_CHECK(false, "Unsupported data type for TileGemm224");
34 }
35 };
36
37 template <>
38 class TileGemm224<c10::BFloat16> {
39 public:
40 using scalar_t = c10::BFloat16;
41 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
42 const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
43 c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
44 c10::BFloat16* __restrict__ a_tile_1 = a_ptr + lda * AMX_TILE_ROW_NUM;
45 const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
46
47 // B is always packed as 16 output channels block
48 c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
49 c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
50 const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
51
52 float* __restrict__ c_tile_4 = c_ptr;
53 float* __restrict__ c_tile_5 =
54 c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
55 float* __restrict__ c_tile_6 = c_ptr + AMX_TILE_ROW_NUM * ldc;
56 float* __restrict__ c_tile_7 =
57 c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
58 const int32_t c_tile_stride = ldc * sizeof(float);
59
60 if (accum_c) {
61 _tile_loadd(4, c_tile_4, c_tile_stride);
62 _tile_loadd(5, c_tile_5, c_tile_stride);
63 _tile_loadd(6, c_tile_6, c_tile_stride);
64 _tile_loadd(7, c_tile_7, c_tile_stride);
65 } else {
66 _tile_zero(4);
67 _tile_zero(5);
68 _tile_zero(6);
69 _tile_zero(7);
70 }
71
72 for (int32_t k = 0; k < k_times; ++k) {
73 _tile_loadd(0, a_tile_0, a_tile_stride);
74 _tile_stream_loadd(2, b_tile_2, b_tile_stride);
75 _tile_dpbf16ps(4, 0, 2);
76 _tile_stream_loadd(3, b_tile_3, b_tile_stride);
77 _tile_dpbf16ps(5, 0, 3);
78 _tile_loadd(1, a_tile_1, a_tile_stride);
79 _tile_dpbf16ps(6, 1, 2);
80 _tile_dpbf16ps(7, 1, 3);
81
82 // update ptrs
83 a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
84 a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
85 b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
86 b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
87 }
88
89 _tile_stored(4, c_tile_4, c_tile_stride);
90 _tile_stored(5, c_tile_5, c_tile_stride);
91 _tile_stored(6, c_tile_6, c_tile_stride);
92 _tile_stored(7, c_tile_7, c_tile_stride);
93 }
94
95 FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
96 const int32_t m_0 = AMX_TILE_ROW_NUM;
97 const int32_t m_1 = m - AMX_TILE_ROW_NUM;
98 config.rows[0] = m_0;
99 config.rows[1] = m_1;
100 config.rows[2] = AMX_TILE_ROW_NUM;
101 config.rows[3] = AMX_TILE_ROW_NUM;
102 config.rows[4] = m_0;
103 config.rows[5] = m_0;
104 config.rows[6] = m_1;
105 config.rows[7] = m_1;
106 _tile_loadconfig(&config);
107 }
108 };
109
110 // 1-2-2 pattern, for 0 < m <= 16
111 // TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
112 // m, m
113 // TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
114 // num should be 16
115 // TILE 6, 7, (6, 7): store results C matrix, row num should be
116 // m
117 template <typename scalar_t>
118 class TileGemm122 {
119 public:
120 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
121 TORCH_CHECK(false, "Unsupported data type for TileGemm122");
122 }
123
124 FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
125 TORCH_CHECK(false, "Unsupported data type for TileGemm122");
126 }
127 };
128
129 template <>
130 class TileGemm122<c10::BFloat16> {
131 public:
132 using scalar_t = c10::BFloat16;
133 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
134 c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
135 c10::BFloat16* __restrict__ a_tile_1 =
136 a_ptr + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
137 const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
138
139 c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
140 c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
141 c10::BFloat16* __restrict__ b_tile_4 =
142 b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
143 c10::BFloat16* __restrict__ b_tile_5 =
144 b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
145 int64_t b_stride = AMX_TILE_ROW_BYTES;
146
147 float* __restrict__ c_tile_6 = c_ptr;
148 float* __restrict__ c_tile_7 = c_ptr + AMX_TILE_ROW_BYTES / sizeof(float);
149 int64_t c_stride = ldc * sizeof(float);
150
151 const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
152 const int32_t k_group_times = k_times / 2;
153 const bool has_tail = (k_times % 2 == 1);
154
155 if (accum_c) {
156 _tile_loadd(6, c_tile_6, c_stride);
157 _tile_loadd(7, c_tile_7, c_stride);
158 } else {
159 _tile_zero(6);
160 _tile_zero(7);
161 }
162
163 for (int32_t k = 0; k < k_group_times; ++k) {
164 _tile_loadd(0, a_tile_0, a_tile_stride);
165 _tile_stream_loadd(2, b_tile_2, b_stride);
166 _tile_dpbf16ps(6, 0, 2);
167 _tile_stream_loadd(3, b_tile_3, b_stride);
168 _tile_dpbf16ps(7, 0, 3);
169 _tile_loadd(1, a_tile_1, a_tile_stride);
170 _tile_stream_loadd(4, b_tile_4, b_stride);
171 _tile_dpbf16ps(6, 1, 4);
172 _tile_stream_loadd(5, b_tile_5, b_stride);
173 _tile_dpbf16ps(7, 1, 5);
174
175 // update ptrs
176 a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
177 a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
178 b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
179 b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
180 b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
181 b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
182 }
183
184 if (has_tail) {
185 _tile_loadd(0, a_tile_0, a_tile_stride);
186 _tile_stream_loadd(2, b_tile_2, b_stride);
187 _tile_dpbf16ps(6, 0, 2);
188 _tile_stream_loadd(3, b_tile_3, b_stride);
189 _tile_dpbf16ps(7, 0, 3);
190 }
191
192 _tile_stored(6, c_tile_6, c_stride);
193 _tile_stored(7, c_tile_7, c_stride);
194 }
195
196 FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
197 config.rows[0] = m;
198 config.rows[1] = m;
199 config.rows[2] = AMX_TILE_ROW_NUM;
200 config.rows[3] = AMX_TILE_ROW_NUM;
201 config.rows[4] = AMX_TILE_ROW_NUM;
202 config.rows[5] = AMX_TILE_ROW_NUM;
203 config.rows[6] = m;
204 config.rows[7] = m;
205 _tile_loadconfig(&config);
206 }
207 };
208 } // namespace
209
210 // Gemm kernel uses AMX, requires B matrix to be packed
211 template <typename scalar_t>
212 class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
213 public:
214 static constexpr int32_t MaxMSize = 32;
215 static constexpr int32_t NSize = 32;
216 static constexpr int32_t WeightOCGroupSize = 16;
217 static constexpr bool PackA = false;
218
219 public:
220 MicroGemm() : curr_m_(-1) {
221 vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
222 }
223
224 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
225 if (m > AMX_TILE_ROW_NUM) {
226 if (m != curr_m_) {
227 curr_m_ = m;
228 TileGemm224<scalar_t>::init_tile_config(m, amx_tile_config_);
229 }
230 TileGemm224<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
231 } else {
232 if (m != curr_m_) {
233 curr_m_ = m;
234 TileGemm122<scalar_t>::init_tile_config(m, amx_tile_config_);
235 }
236 TileGemm122<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
237 }
238 }
239
240 static void pack_weight(const scalar_t* __restrict__ weight,
241 scalar_t* __restrict__ packed_weight,
242 const int32_t output_size, const int32_t input_size) {
243 constexpr int32_t elem_num_per_group = 4 / sizeof(scalar_t);
244 TORCH_CHECK_EQ(output_size % 16, 0);
245 TORCH_CHECK_EQ(input_size % (16 * elem_num_per_group), 0);
246
247 const int32_t output_group_num = output_size / 16;
248 const int32_t input_32b_num = input_size / elem_num_per_group;
249 for (int32_t output_group_idx = 0; output_group_idx < output_group_num;
250 ++output_group_idx) {
251 const int32_t* __restrict__ weight_32b =
252 reinterpret_cast<const int32_t*>(weight);
253 int32_t* __restrict__ packed_weight_32b =
254 reinterpret_cast<int32_t*>(packed_weight);
255 for (int32_t output_idx = 0; output_idx < 16; ++output_idx) {
256 for (int32_t weight_offset = 0, packed_offset = 0;
257 weight_offset < input_32b_num;
258 ++weight_offset, packed_offset += 16) {
259 packed_weight_32b[packed_offset] = weight_32b[weight_offset];
260 }
261
262 // update
263 weight_32b += input_32b_num;
264 packed_weight_32b += 1;
265 }
266
267 // update
268 weight += 16 * input_size;
269 packed_weight += 16 * input_size;
270 }
271 }
272
273 private:
274 alignas(64) __tilecfg amx_tile_config_;
275 int32_t curr_m_;
276 };
277
278 } // namespace cpu_micro_gemm
279
280 #endif