@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 372 lines 11.3 KB
Raw
1 #pragma once
2
3 #include <cstdint>
4 #include <hip/hip_runtime.h>
5 #include <hip/hip_fp16.h>
6 #include <hip/hip_bf16.h>
7
8 #define __quickreduce_device_inline__ __device__ __forceinline__
9 #define __quickreduce_launch_bounds_two_shot__ __launch_bounds__(256, 4)
10 #define __quickreduce_launch_bounds_one_shot__ __launch_bounds__(512, 4)
11
12 namespace quickreduce {
13
14 typedef __hip_bfloat16 nv_bfloat16;
15 typedef __hip_bfloat162 nv_bfloat162;
16
17 using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
18 using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
19
20 // Setup acquire-release semantics for vector memory reads (mubuf instruction)
21 // as per architecture.
22 #if defined(__gfx942__)
23 // CDNA3: Scope bits sc0, sc1
24 #define MUBUF_ACQUIRE 16
25 #define MUBUF_RELEASE 16
26 #elif (defined(__gfx908__) || defined(__gfx90a__))
27 // CDNA1 and CDNA2 - glc bit
28 #define MUBUF_ACQUIRE 1
29 #define MUBUF_RELEASE 0
30 #endif
31
32 static constexpr int kNegOne = 0xBC00BC00; // {-1, -1}, fp16x2_t
33
34 // Number of atoms (4xf16x2_t) processed by a single thread
35 static constexpr int kAtoms = 8;
36
37 // We use a workgroup of 256 threads
38 static constexpr int kBlockSize = 256;
39 static constexpr int kAtomStride = kBlockSize;
40
41 // Size and atom stride of source/destination data that the block will
42 // process.
43 // Workgroup scope = Tile = (256 threads x 8 atoms x 16B)
44 static constexpr int kTileSize = kBlockSize * kAtoms * sizeof(int32x4_t);
45
46 // Max number of blocks. 304 CUs on MI300
47 static constexpr int kMaxNumBlocks = 304 * 4;
48
49 // Standard CDNA wavefront size.
50 static constexpr int kWavefront = 64;
51
52 // 256 thread, 4 wavefronts.
53 static dim3 constexpr kBlockTwoShot = {kWavefront, kBlockSize / kWavefront, 1};
54
55 // Number of threads in a group for quantization
56 // It corresponds to 32 F16 elements in quantization block
57 static constexpr int kThreadGroupSize = 8;
58
59 // Methods
60 __quickreduce_device_inline__ __host__ unsigned long divceil(unsigned long x,
61 unsigned long y) {
62 return ((x + y - 1) / y);
63 }
64
65 union BufferResource {
66 __quickreduce_device_inline__ constexpr BufferResource()
67 : config(0x00020000U) {}
68
69 __quickreduce_device_inline__ constexpr BufferResource(void* buffer_address,
70 uint32_t buffer_size)
71 : address(buffer_address), range(buffer_size), config(0x00020000U) {}
72
73 int32x4_t descriptor;
74 struct {
75 void* address; // 8B, out of which first 48b is address, and 16b is stride
76 // (unused)
77 uint32_t range; // Byte range for the buffer resource
78 uint32_t config; // Constant, DFMT=32b
79 };
80 };
81
82 #if !defined(__gfx1250__)
83 __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
84 int32x4_t srsrc, int32_t voffset, int32_t soffset,
85 int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
86
87 __quickreduce_device_inline__ static void buffer_store_dwordx4(
88 int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
89 int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
90 #else
91 __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
92 int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {}
93
94 __quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data,
95 int32x4_t srsrc,
96 int32_t voffset,
97 int32_t soffset,
98 int32_t aux) {}
99 #endif
100
101 __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
102 #if defined(__gfx942__)
103 if (value) {
104 asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
105 } else {
106 asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::);
107 }
108 #endif
109 }
110 union bf162_int_union {
111 int i;
112 nv_bfloat162 bf2;
113 };
114
115 template <typename T>
116 __quickreduce_device_inline__ void packed_assign_add(int32x4_t* A,
117 int32x4_t* B);
118
119 template <>
120 __quickreduce_device_inline__ void packed_assign_add<half>(int32x4_t* A,
121 int32x4_t* B) {
122 int32x4_t& tR_fragment = A[0];
123 int32x4_t& tA_fragment = B[0];
124
125 asm volatile("v_pk_add_f16 %0, %1, %2"
126 : "=v"(tR_fragment[0])
127 : "v"(tR_fragment[0]), "v"(tA_fragment[0]));
128 asm volatile("v_pk_add_f16 %0, %1, %2"
129 : "=v"(tR_fragment[1])
130 : "v"(tR_fragment[1]), "v"(tA_fragment[1]));
131 asm volatile("v_pk_add_f16 %0, %1, %2"
132 : "=v"(tR_fragment[2])
133 : "v"(tR_fragment[2]), "v"(tA_fragment[2]));
134 asm volatile("v_pk_add_f16 %0, %1, %2"
135 : "=v"(tR_fragment[3])
136 : "v"(tR_fragment[3]), "v"(tA_fragment[3]));
137 }
138
139 template <>
140 __quickreduce_device_inline__ void packed_assign_add<nv_bfloat16>(
141 int32x4_t* A, int32x4_t* B) {
142 nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(A);
143 nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(B);
144 #pragma unroll
145 for (int i = 0; i < 4; i++) {
146 tA[i] = __hadd2(tA[i], tB[i]);
147 }
148 }
149
150 template <typename T>
151 __quickreduce_device_inline__ int packed_max(int a, int b);
152
153 template <>
154 __quickreduce_device_inline__ int packed_max<half>(int a, int b) {
155 int result;
156 asm volatile("v_pk_max_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
157 return result;
158 }
159
160 template <>
161 __quickreduce_device_inline__ int packed_max<nv_bfloat16>(int a, int b) {
162 bf162_int_union A, B, R;
163 A.i = a;
164 B.i = b;
165 R.bf2 = __hmax2(A.bf2, B.bf2);
166 return R.i;
167 }
168
169 template <typename T>
170 __quickreduce_device_inline__ int packed_min(int a, int b);
171
172 template <>
173 __quickreduce_device_inline__ int packed_min<half>(int a, int b) {
174 int result;
175 asm volatile("v_pk_min_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
176 return result;
177 }
178
179 template <>
180 __quickreduce_device_inline__ int packed_min<nv_bfloat16>(int a, int b) {
181 bf162_int_union A, B, R;
182 A.i = a;
183 B.i = b;
184 R.bf2 = __hmin2(A.bf2, B.bf2);
185 return R.i;
186 }
187
188 template <typename T>
189 __quickreduce_device_inline__ int packed_abs_max(int a, int b);
190
191 template <>
192 __quickreduce_device_inline__ int packed_abs_max<half>(int a, int b) {
193 half2 wmaxh2 = __builtin_bit_cast(half2, a);
194 half2 wminh2 = __builtin_bit_cast(half2, b);
195 half2 wblockmaxh2;
196
197 wblockmaxh2.x =
198 __hgt(__habs(wmaxh2.x), __habs(wminh2.x)) ? wmaxh2.x : wminh2.x;
199 wblockmaxh2.y =
200 __hgt(__habs(wmaxh2.y), __habs(wminh2.y)) ? wmaxh2.y : wminh2.y;
201 return __builtin_bit_cast(int, wblockmaxh2);
202 }
203
204 template <>
205 __quickreduce_device_inline__ int packed_abs_max<nv_bfloat16>(int a, int b) {
206 bf162_int_union A, B, R;
207 A.i = a;
208 B.i = b;
209 R.bf2.x = __hgt(__habs(A.bf2.x), __habs(B.bf2.x)) ? A.bf2.x : B.bf2.x;
210 R.bf2.y = __hgt(__habs(A.bf2.y), __habs(B.bf2.y)) ? A.bf2.y : B.bf2.y;
211 return R.i;
212 }
213
214 template <typename T>
215 __quickreduce_device_inline__ int packed_add(int a, int b);
216
217 template <>
218 __quickreduce_device_inline__ int packed_add<half>(int a, int b) {
219 int result;
220 asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
221 return result;
222 }
223
224 template <>
225 __quickreduce_device_inline__ int packed_add<nv_bfloat16>(int a, int b) {
226 bf162_int_union A, B, R;
227 A.i = a;
228 B.i = b;
229 R.bf2 = __hadd2(A.bf2, B.bf2);
230 return R.i;
231 }
232
233 template <>
234 __quickreduce_device_inline__ int packed_add<int16_t>(int a, int b) {
235 int result;
236 asm volatile("v_pk_add_i16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
237 return result;
238 }
239
240 template <typename T>
241 __quickreduce_device_inline__ int packed_sub(int a, int b);
242
243 template <>
244 __quickreduce_device_inline__ int packed_sub<half>(int a, int b) {
245 int result;
246
247 // MI300 lacks packed fp16 sub instruction. So we do -1 * min + max
248 asm volatile("v_pk_fma_f16 %0, %1, %2 %3"
249 : "=v"(result)
250 : "v"(kNegOne), "v"(b), "v"(a));
251 return result;
252 }
253
254 template <>
255 __quickreduce_device_inline__ int packed_sub<nv_bfloat16>(int a, int b) {
256 bf162_int_union A, B, R;
257 A.i = a;
258 B.i = b;
259 R.bf2 = __hsub2(A.bf2, B.bf2);
260 return R.i;
261 }
262
263 template <typename T>
264 __quickreduce_device_inline__ int packed_mul(int a, int b);
265
266 template <>
267 __quickreduce_device_inline__ int packed_mul<half>(int a, int b) {
268 int result;
269 asm volatile("v_pk_mul_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
270 return result;
271 }
272
273 template <>
274 __quickreduce_device_inline__ int packed_mul<nv_bfloat16>(int a, int b) {
275 nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(&a);
276 nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(&b);
277 nv_bfloat162 tR = __hmul2(*tA, *tB);
278 return *(reinterpret_cast<int*>(&tR));
279 }
280
281 template <typename T>
282 __quickreduce_device_inline__ int packed_rcp(int a);
283
284 template <>
285 __quickreduce_device_inline__ int packed_rcp<half>(int a) {
286 return __builtin_bit_cast(int, h2rcp(__builtin_bit_cast(half2, a)));
287 }
288
289 template <>
290 __quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
291 bf162_int_union A, R;
292 A.i = a;
293 R.bf2 = h2rcp(A.bf2);
294 return R.i;
295 }
296
297 template <typename T>
298 __quickreduce_device_inline__ int packed_from_int16_pair(int16_t low,
299 int16_t high);
300
301 template <>
302 __quickreduce_device_inline__ int packed_from_int16_pair<half>(int16_t low,
303 int16_t high) {
304 // Convert two signed integers to one fp16x2 packed 32-bit lane.
305 half2 h = __halves2half2(__int2half_rn(static_cast<int>(low)),
306 __int2half_rn(static_cast<int>(high)));
307 return __builtin_bit_cast(int, h);
308 }
309
310 template <>
311 __quickreduce_device_inline__ int packed_from_int16_pair<nv_bfloat16>(
312 int16_t low, int16_t high) {
313 // Convert two signed integers to one bf16x2 packed 32-bit lane.
314 nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
315 nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
316 nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
317 return *reinterpret_cast<int*>(&bf2);
318 }
319
320 // changes dtype
321 __quickreduce_device_inline__ float T2float_cast(half a) {
322 return __half2float(a);
323 }
324
325 __quickreduce_device_inline__ float T2float_cast(nv_bfloat16 a) {
326 return __bfloat162float(a);
327 }
328
329 template <typename T>
330 __quickreduce_device_inline__ int group_abs_max(int32x4_t atom) {
331 const int group_leader = (threadIdx.x / kThreadGroupSize) * kThreadGroupSize;
332
333 int wmax, wmin, wblockmax;
334 int a, b;
335 a = packed_max<T>(atom[0], atom[1]);
336 b = packed_max<T>(atom[2], atom[3]);
337
338 wmax = packed_max<T>(a, b);
339
340 a = packed_min<T>(atom[0], atom[1]);
341 b = packed_min<T>(atom[2], atom[3]);
342
343 wmin = packed_min<T>(a, b);
344
345 // Reduce the max among a group of threads
346 // Note: This is basically 2 blocks of values setup as the
347 // upper/lower halves of the f16x2_t
348 for (int i = 1; i < kThreadGroupSize; i <<= 1) {
349 int x = __shfl_down(wmax, i);
350 wmax = packed_max<T>(wmax, x);
351
352 int y = __shfl_down(wmin, i);
353 wmin = packed_min<T>(wmin, y);
354 }
355 wblockmax = packed_abs_max<T>(wmax, wmin);
356 // Share with the cohort
357 wblockmax = __shfl(wblockmax, group_leader);
358 return wblockmax;
359 }
360
361 __quickreduce_device_inline__ void set_sync_flag(uint32_t* flag_ptr,
362 uint32_t flag) {
363 __atomic_store_n(flag_ptr, flag, __ATOMIC_RELEASE);
364 }
365
366 __quickreduce_device_inline__ void wait_sync_flag(uint32_t* flag_ptr,
367 uint32_t flag) {
368 while (__atomic_load_n(flag_ptr, __ATOMIC_RELAXED) != flag) {
369 }
370 }
371
372 } // namespace quickreduce