A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 23:30 UTC
refactor(cpu_types_scalar.hpp): Unify scalar loop implementations using unroll_loop (#28847)
Signed-off-by: ihb2032 <1355790728@qq.com> Co-authored-by: lyd1992 <liuyudong@iscas.ac.cn>
ihb2032 committed
Nov 19, 2025 at 19:05 UTC
815160958327d601933139b9e76a01eb6d2bc5cf
1 file changed
+87
-135
csrc/cpu/cpu_types_scalar.hpp
+87
-135
@@ -26,10 +26,6 @@ namespace vec_op {
26
27
#define FORCE_INLINE __attribute__((always_inline)) inline
28
29
-#define __max(a, b) ((a) > (b) ? (a) : (b))
30
-#define __min(a, b) ((a) < (b) ? (a) : (b))
31
-#define __abs(a) ((a) < (0) ? (0 - a) : (a))
32
-
29
typedef struct f16x8_t {
30
uint16_t val[8];
31
} f16x8_t;
@@ -99,7 +95,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
95
void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
96
97
void save(void* ptr, const int elem_num) const {
102
- int num = __min(elem_num, VEC_ELEM_NUM);
98
+ int num = std::min(elem_num, VEC_ELEM_NUM);
99
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
100
}
101
};
@@ -128,7 +124,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
124
void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
125
126
void save(void* ptr, const int elem_num) const {
131
- int num = __min(elem_num, VEC_ELEM_NUM);
127
+ int num = std::min(elem_num, VEC_ELEM_NUM);
128
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
129
}
130
};
@@ -143,9 +139,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
139
explicit BF16Vec32(f16x32_t data) : reg(data) {};
140
141
explicit BF16Vec32(BF16Vec8& vec8_data) {
146
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
142
+ unroll_loop<int, VEC_ELEM_NUM>([&vec8_data, this](int i) {
143
reg.val[i] = vec8_data.reg.val[i % BF16Vec8::VEC_ELEM_NUM];
148
- }
144
+ });
145
}
146
147
void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
@@ -157,15 +153,11 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
153
f32x4_t reg;
154
155
explicit FP32Vec4(float v) {
160
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
161
- reg.val[i] = v;
162
- }
156
+ unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
157
}
158
159
explicit FP32Vec4() {
166
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
167
- reg.val[i] = 0.0f;
168
- }
160
+ unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
161
}
162
163
explicit FP32Vec4(const float* ptr)
@@ -182,15 +174,11 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
174
f32x8_t reg;
175
176
explicit FP32Vec8(float v) {
185
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
186
- reg.val[i] = v;
187
- }
177
+ unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
178
}
179
180
explicit FP32Vec8() {
191
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
192
- reg.val[i] = 0.0f;
193
- }
181
+ unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
182
}
183
184
explicit FP32Vec8(const float* ptr)
@@ -201,78 +189,68 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
189
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
190
191
explicit FP32Vec8(const FP16Vec8& v) {
204
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
205
- reg.val[i] = fp16_to_float(v.reg.val[i]);
206
- }
192
+ unroll_loop<int, VEC_ELEM_NUM>(
193
+ [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
194
}
195
196
FP32Vec8(const BF16Vec8& v) {
210
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
211
- reg.val[i] = bf16_to_float(v.reg.val[i]);
212
- }
197
+ unroll_loop<int, VEC_ELEM_NUM>(
198
+ [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
199
}
200
201
float reduce_sum() const {
202
float result = 0;
217
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
218
- result += reg.val[i];
219
- }
203
+ unroll_loop<int, VEC_ELEM_NUM>(
204
+ [&result, this](int i) { result += reg.val[i]; });
205
return result;
206
}
207
208
FP32Vec8 exp() const {
209
f32x8_t ret;
225
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
226
- ret.val[i] = expf(reg.val[i]);
227
- }
210
+ unroll_loop<int, VEC_ELEM_NUM>(
211
+ [&ret, this](int i) { ret.val[i] = expf(reg.val[i]); });
212
return FP32Vec8(ret);
213
}
214
215
FP32Vec8 tanh() const {
216
f32x8_t ret;
233
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
234
- ret.val[i] = tanhf(reg.val[i]);
235
- }
217
+ unroll_loop<int, VEC_ELEM_NUM>(
218
+ [&ret, this](int i) { ret.val[i] = tanhf(reg.val[i]); });
219
return FP32Vec8(ret);
220
}
221
222
FP32Vec8 er() const {
223
f32x8_t ret;
241
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
242
- ret.val[i] = erf(reg.val[i]);
243
- }
224
+ unroll_loop<int, VEC_ELEM_NUM>(
225
+ [&ret, this](int i) { ret.val[i] = erf(reg.val[i]); });
226
return FP32Vec8(ret);
227
}
228
229
FP32Vec8 operator*(const FP32Vec8& b) const {
230
f32x8_t ret;
249
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
250
- ret.val[i] = reg.val[i] * b.reg.val[i];
251
- }
231
+ unroll_loop<int, VEC_ELEM_NUM>(
232
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
233
return FP32Vec8(ret);
234
}
235
236
FP32Vec8 operator+(const FP32Vec8& b) const {
237
f32x8_t ret;
257
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
258
- ret.val[i] = reg.val[i] + b.reg.val[i];
259
- }
238
+ unroll_loop<int, VEC_ELEM_NUM>(
239
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
240
return FP32Vec8(ret);
241
}
242
243
FP32Vec8 operator-(const FP32Vec8& b) const {
244
f32x8_t ret;
265
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
266
- ret.val[i] = reg.val[i] - b.reg.val[i];
267
- }
245
+ unroll_loop<int, VEC_ELEM_NUM>(
246
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
247
return FP32Vec8(ret);
248
}
249
250
FP32Vec8 operator/(const FP32Vec8& b) const {
251
f32x8_t ret;
273
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
274
- ret.val[i] = reg.val[i] / b.reg.val[i];
275
- }
252
+ unroll_loop<int, VEC_ELEM_NUM>(
253
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
254
return FP32Vec8(ret);
255
}
256
@@ -284,15 +262,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
262
f32x16_t reg;
263
264
explicit FP32Vec16(float v) {
287
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
288
- reg.val[i] = v;
289
- }
265
+ unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
266
}
267
268
explicit FP32Vec16() {
293
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
294
- reg.val[i] = 0.0f;
295
- }
269
+ unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
270
}
271
272
explicit FP32Vec16(const float* ptr)
@@ -301,29 +275,27 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
275
explicit FP32Vec16(f32x16_t data) : reg(data) {};
276
277
FP32Vec16(const FP32Vec4& data) {
304
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
278
+ unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
279
reg.val[i] = data.reg.val[i % FP32Vec4::VEC_ELEM_NUM];
306
- }
280
+ });
281
}
282
283
FP32Vec16(const FP32Vec8& data) {
310
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
284
+ unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
285
reg.val[i] = data.reg.val[i % FP32Vec8::VEC_ELEM_NUM];
312
- }
286
+ });
287
}
288
289
FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
290
291
explicit FP32Vec16(const FP16Vec16& v) {
318
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
319
- reg.val[i] = fp16_to_float(v.reg.val[i]);
320
- }
292
+ unroll_loop<int, VEC_ELEM_NUM>(
293
+ [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
294
}
295
296
explicit FP32Vec16(const BF16Vec16& v) {
324
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
325
- reg.val[i] = bf16_to_float(v.reg.val[i]);
326
- }
297
+ unroll_loop<int, VEC_ELEM_NUM>(
298
+ [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
299
}
300
301
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
@@ -331,82 +303,74 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
303
FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
304
305
FP32Vec16 operator*(const FP32Vec16& b) const {
334
- FP32Vec16 result(0.0f);
335
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
336
- result.reg.val[i] = reg.val[i] * b.reg.val[i];
337
- }
338
- return result;
306
+ f32x16_t ret;
307
+ unroll_loop<int, VEC_ELEM_NUM>(
308
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
309
+ return FP32Vec16(ret);
310
}
311
312
FP32Vec16 operator+(const FP32Vec16& b) const {
342
- FP32Vec16 result(0.0f);
343
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
344
- result.reg.val[i] = reg.val[i] + b.reg.val[i];
345
- }
346
- return result;
313
+ f32x16_t ret;
314
+ unroll_loop<int, VEC_ELEM_NUM>(
315
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
316
+ return FP32Vec16(ret);
317
}
318
319
FP32Vec16 operator-(const FP32Vec16& b) const {
350
- FP32Vec16 result(0.0f);
351
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
352
- result.reg.val[i] = reg.val[i] - b.reg.val[i];
353
- }
354
- return result;
320
+ f32x16_t ret;
321
+ unroll_loop<int, VEC_ELEM_NUM>(
322
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
323
+ return FP32Vec16(ret);
324
}
325
326
FP32Vec16 operator/(const FP32Vec16& b) const {
358
- FP32Vec16 result(0.0f);
359
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
360
- result.reg.val[i] = reg.val[i] / b.reg.val[i];
361
- }
362
- return result;
327
+ f32x16_t ret;
328
+ unroll_loop<int, VEC_ELEM_NUM>(
329
+ [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
330
+ return FP32Vec16(ret);
331
}
332
333
FP32Vec16 max(const FP32Vec16& b) const {
366
- FP32Vec16 result(0.0f);
367
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
368
- result.reg.val[i] = __max(reg.val[i], b.reg.val[i]);
369
- }
370
- return result;
334
+ f32x16_t ret;
335
+ unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
336
+ ret.val[i] = std::max(reg.val[i], b.reg.val[i]);
337
+ });
338
+ return FP32Vec16(ret);
339
}
340
341
FP32Vec16 min(const FP32Vec16& b) const {
374
- FP32Vec16 result(0.0f);
375
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
376
- result.reg.val[i] = __min(reg.val[i], b.reg.val[i]);
377
- }
378
- return result;
342
+ f32x16_t ret;
343
+ unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
344
+ ret.val[i] = std::min(reg.val[i], b.reg.val[i]);
345
+ });
346
+ return FP32Vec16(ret);
347
}
348
349
FP32Vec16 abs() const {
382
- FP32Vec16 result(0.0f);
383
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
384
- result.reg.val[i] = __abs(reg.val[i]);
385
- }
386
- return result;
350
+ f32x16_t ret;
351
+ unroll_loop<int, VEC_ELEM_NUM>(
352
+ [&ret, this](int i) { ret.val[i] = std::abs(reg.val[i]); });
353
+ return FP32Vec16(ret);
354
}
355
356
float reduce_sum() const {
357
float result = 0.0f;
391
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
392
- result += reg.val[i];
393
- }
358
+ unroll_loop<int, VEC_ELEM_NUM>(
359
+ [&result, this](int i) { result += reg.val[i]; });
360
return result;
361
}
362
363
float reduce_max() const {
398
- float result = reg.val[0];
399
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
400
- result = __max(reg.val[i], result);
401
- }
364
+ float result = std::numeric_limits<float>::lowest();
365
+ unroll_loop<int, VEC_ELEM_NUM>(
366
+ [&result, this](int i) { result = std::max(reg.val[i], result); });
367
return result;
368
}
369
370
float reduce_min() const {
406
- float result = reg.val[0];
407
- for (int i = 0; i < VEC_ELEM_NUM; ++i) {
408
- result = __min(reg.val[i], result);
409
- }
371
+ float result = std::numeric_limits<float>::max();
372
+ unroll_loop<int, VEC_ELEM_NUM>(
373
+ [&result, this](int i) { result = std::min(reg.val[i], result); });
374
return result;
375
}
376
@@ -414,13 +378,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
378
float reduce_sub_sum(int idx) {
379
static_assert(VEC_ELEM_NUM % group_size == 0);
380
float sum = 0.0;
417
- int start = idx * group_size;
418
- int end = (idx + 1) * group_size;
419
-
420
- for (; (start < VEC_ELEM_NUM) && (start < end); ++start) {
421
- sum += reg.val[start];
422
- }
423
-
381
+ const int start = idx * group_size;
382
+ unroll_loop<int, group_size>(
383
+ [&sum, &start, this](int i) { sum += reg.val[start + i]; });
384
return sum;
385
}
386
@@ -477,17 +437,13 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
437
}
438
439
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
480
- int i = 0;
481
- for (i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
482
- reg.val[i] = float_to_fp16(v.reg.val[i]);
483
- }
440
+ unroll_loop<int, FP16Vec16::VEC_ELEM_NUM>(
441
+ [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
442
}
443
444
inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
487
- int i = 0;
488
- for (i = 0; i < FP16Vec8::VEC_ELEM_NUM; ++i) {
489
- reg.val[i] = float_to_fp16(v.reg.val[i]);
490
- }
445
+ unroll_loop<int, FP16Vec8::VEC_ELEM_NUM>(
446
+ [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
447
}
448
449
inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
@@ -495,17 +451,13 @@ inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
451
}
452
453
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
498
- int i = 0;
499
- for (i = 0; i < BF16Vec8::VEC_ELEM_NUM; ++i) {
500
- reg.val[i] = float_to_bf16(v.reg.val[i]);
501
- }
454
+ unroll_loop<int, BF16Vec8::VEC_ELEM_NUM>(
455
+ [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
456
}
457
458
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
505
- int i = 0;
506
- for (i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
507
- reg.val[i] = float_to_bf16(v.reg.val[i]);
508
- }
459
+ unroll_loop<int, BF16Vec16::VEC_ELEM_NUM>(
460
+ [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
461
}
462
463
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 3); }