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, 16:00 UTC
[CPU][IBM Z] Fix BF16 support and vectorize math operations for s390x (#28926)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
R3hankhan committed
Nov 24, 2025 at 17:38 UTC
4de87866a8fdc9395ccd40e00c0f9075439b07ae
2 files changed
+531
-57
csrc/cpu/cpu_attn_impl.hpp
+1
-1
@@ -847,7 +847,7 @@ struct VecTypeTrait<c10::BFloat16> {
847
};
848
#endif
849
850
-#if !defined(__powerpc__)
850
+#if !defined(__powerpc__) && !defined(__s390x__)
851
template <>
852
struct VecTypeTrait<c10::Half> {
853
using vec_t = vec_op::FP16Vec16;
csrc/cpu/cpu_types_vxe.hpp
+530
-56
@@ -4,6 +4,7 @@
4
5
#include <vecintrin.h>
6
#include <cmath>
7
+#include <limits>
8
#include <torch/all.h>
9
namespace vec_op {
10
@@ -174,8 +175,9 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
175
}
176
177
explicit FP32Vec8(const BF16Vec8& v) {
177
- reg.val[0] = (__vector float)vec_mergeh(zero, v.reg);
178
- reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
178
+ // On big-endian s390x, place BF16 first to get correct byte order
179
+ reg.val[0] = (__vector float)vec_mergeh(v.reg, zero);
180
+ reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
181
}
182
183
float reduce_sum() const {
@@ -189,51 +191,257 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
191
}
192
193
FP32Vec8 exp() const {
192
- // TODO: Vectorize this
193
- AliasReg ar;
194
- ar.reg = reg;
195
- f32x4x4_t ret;
196
- ret.val[0][0] = std::exp(ar.values[0]);
197
- ret.val[0][1] = std::exp(ar.values[1]);
198
- ret.val[0][2] = std::exp(ar.values[2]);
199
- ret.val[0][3] = std::exp(ar.values[3]);
200
- ret.val[1][0] = std::exp(ar.values[4]);
201
- ret.val[1][1] = std::exp(ar.values[5]);
202
- ret.val[1][2] = std::exp(ar.values[6]);
203
- ret.val[1][3] = std::exp(ar.values[7]);
204
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
194
+ f32x4x2_t out;
195
+
196
+ const __vector float log2e = vec_splats(1.44269504088896341f);
197
+ const __vector float one = vec_splats(1.0f);
198
+ const __vector float min_x = vec_splats(-87.3f);
199
+ const __vector float max_x = vec_splats(88.7f);
200
+
201
+ // 5th-degree minimax polynomial for 2^r (r in [0,1))
202
+ const __vector float c1 = vec_splats(0.6931471805599453f);
203
+ const __vector float c2 = vec_splats(0.240226506959101f);
204
+ const __vector float c3 = vec_splats(0.05550410866482158f);
205
+ const __vector float c4 = vec_splats(0.009618129107628477f);
206
+ const __vector float c5 = vec_splats(0.0013333558146428443f);
207
+
208
+ for (int i = 0; i < 2; i++) {
209
+ __vector float x = reg.val[i];
210
+
211
+ x = vec_max(x, min_x);
212
+ x = vec_min(x, max_x);
213
+
214
+ __vector float y = vec_mul(x, log2e);
215
+
216
+ __vector float kf = vec_floor(y);
217
+ __vector float r = vec_sub(y, kf);
218
+
219
+ __vector signed int k = vec_signed(kf);
220
+ const __vector signed int min_k = vec_splats((signed int)-126);
221
+ const __vector signed int max_k = vec_splats((signed int)127);
222
+ k = vec_min(vec_max(k, min_k), max_k);
223
+
224
+ // Build 2^k from exponent bits
225
+ __vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
226
+ __vector unsigned int bits = (__vector unsigned int)exp_int;
227
+ bits = vec_sl(bits, vec_splats((unsigned int)23));
228
+ __vector float pow2k = (__vector float)bits;
229
+
230
+ // Improved minimax polynomial
231
+ __vector float poly = vec_madd(c5, r, c4);
232
+ poly = vec_madd(poly, r, c3);
233
+ poly = vec_madd(poly, r, c2);
234
+ poly = vec_madd(poly, r, c1);
235
+ poly = vec_madd(poly, r, one);
236
+
237
+ out.val[i] = vec_mul(pow2k, poly);
238
+ }
239
+
240
+ return FP32Vec8(out);
241
}
242
243
FP32Vec8 tanh() const {
208
- // TODO: Vectorize this
209
- AliasReg ar;
210
- ar.reg = reg;
211
- f32x4x4_t ret;
212
- ret.val[0][0] = std::tanh(ar.values[0]);
213
- ret.val[0][1] = std::tanh(ar.values[1]);
214
- ret.val[0][2] = std::tanh(ar.values[2]);
215
- ret.val[0][3] = std::tanh(ar.values[3]);
216
- ret.val[1][0] = std::tanh(ar.values[4]);
217
- ret.val[1][1] = std::tanh(ar.values[5]);
218
- ret.val[1][2] = std::tanh(ar.values[6]);
219
- ret.val[1][3] = std::tanh(ar.values[7]);
220
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
244
+ // tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
245
+ const __vector float one = vec_splats(1.0f);
246
+ const __vector float two = vec_splats(2.0f);
247
+ const __vector float zero = vec_splats(0.0f);
248
+ const __vector float sat =
249
+ vec_splats(9.0f); // beyond this, tanh(x) ~ sign(x)
250
+
251
+ f32x4x2_t out;
252
+
253
+ for (int i = 0; i < 2; i++) {
254
+ __vector float x = reg.val[i];
255
+ __vector float ax = vec_abs(x);
256
+
257
+ // sign(x): +1 or -1
258
+ __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
259
+
260
+ // saturation mask: |x| > sat
261
+ __vector __bool int saturated = vec_cmpgt(ax, sat);
262
+
263
+ // 2x
264
+ __vector float two_x = vec_mul(x, two);
265
+
266
+ // Build a temporary FP32Vec8 with both lanes = 2x, reuse exp()
267
+ f32x4x2_t tmp;
268
+ tmp.val[0] = two_x;
269
+ tmp.val[1] = two_x;
270
+ FP32Vec8 exp_2x_vec(tmp);
271
+
272
+ FP32Vec8 e2x = exp_2x_vec.exp();
273
+ __vector float e = e2x.reg.val[i];
274
+
275
+ // tanh(x) = (e - 1) / (e + 1)
276
+ __vector float num = vec_sub(e, one);
277
+ __vector float den = vec_add(e, one);
278
+
279
+ __vector float t = vec_div(num, den);
280
+
281
+ // For large |x|, clamp to sign(x)
282
+ out.val[i] = vec_sel(t, sign, saturated);
283
+ }
284
+
285
+ return FP32Vec8(out);
286
}
287
288
FP32Vec8 er() const {
224
- // TODO: Vectorize this
225
- AliasReg ar;
226
- ar.reg = reg;
227
- f32x4x4_t ret;
228
- ret.val[0][0] = std::erf(ar.values[0]);
229
- ret.val[0][1] = std::erf(ar.values[1]);
230
- ret.val[0][2] = std::erf(ar.values[2]);
231
- ret.val[0][3] = std::erf(ar.values[3]);
232
- ret.val[1][0] = std::erf(ar.values[4]);
233
- ret.val[1][1] = std::erf(ar.values[5]);
234
- ret.val[1][2] = std::erf(ar.values[6]);
235
- ret.val[1][3] = std::erf(ar.values[7]);
236
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
289
+ // A&S 7.1.26 approximation:
290
+ // erf(x) = sign(x) * (1 - ((((a5*t + a4)*t + a3)*t + a2)*t + a1) * t *
291
+ // exp(-x^2)) t = 1 / (1 + p*|x|), p = 0.3275911
292
+
293
+ const __vector float one = vec_splats(1.0f);
294
+ const __vector float zero = vec_splats(0.0f);
295
+ const __vector float p = vec_splats(0.3275911f);
296
+
297
+ // Polynomial coeffs
298
+ const __vector float a1 = vec_splats(0.254829592f);
299
+ const __vector float a2 = vec_splats(-0.284496736f);
300
+ const __vector float a3 = vec_splats(1.421413741f);
301
+ const __vector float a4 = vec_splats(-1.453152027f);
302
+ const __vector float a5 = vec_splats(1.061405429f);
303
+
304
+ // Threshold where erf(x) ~ sign(x)
305
+ const __vector float sat = vec_splats(6.0f);
306
+
307
+ f32x4x2_t out;
308
+
309
+ for (int lane = 0; lane < 2; lane++) {
310
+ __vector float x = reg.val[lane];
311
+ __vector float ax = vec_abs(x);
312
+
313
+ // sign(x)
314
+ __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
315
+
316
+ // |x| > 6 → erf(x) = ±1
317
+ __vector __bool int saturated = vec_cmpgt(ax, sat);
318
+
319
+ // t = 1 / (1 + p * |x|)
320
+ __vector float t = vec_madd(p, ax, one);
321
+ t = vec_div(one, t);
322
+
323
+ // poly = a5
324
+ __vector float poly = a5;
325
+ poly = vec_madd(poly, t, a4);
326
+ poly = vec_madd(poly, t, a3);
327
+ poly = vec_madd(poly, t, a2);
328
+ poly = vec_madd(poly, t, a1);
329
+
330
+ // full polynomial: poly = poly * t
331
+ poly = vec_mul(poly, t);
332
+
333
+ // Compute exp(-x^2)
334
+ __vector float x2 = vec_mul(x, x);
335
+ __vector float neg_x2 = vec_neg(x2);
336
+
337
+ f32x4x2_t tmp;
338
+ tmp.val[0] = neg_x2;
339
+ tmp.val[1] = neg_x2;
340
+ FP32Vec8 exp_neg_x2(tmp);
341
+
342
+ FP32Vec8 e = exp_neg_x2.exp();
343
+ __vector float ex = e.reg.val[lane];
344
+
345
+ // erf(x) = sign * (1 - poly * exp(-x^2))
346
+ __vector float term = vec_mul(poly, ex);
347
+ __vector float y = vec_sub(one, term);
348
+ y = vec_mul(y, sign);
349
+
350
+ // saturated → ±1
351
+ __vector float sat_val = vec_mul(sign, one);
352
+ out.val[lane] = vec_sel(y, sat_val, saturated);
353
+ }
354
+
355
+ return FP32Vec8(out);
356
+ }
357
+ // Elementwise sigmoid(x) = 1 / (1 + exp(-x))
358
+ FP32Vec8 sigmoid() const {
359
+ const __vector float one = vec_splats(1.0f);
360
+
361
+ f32x4x2_t neg;
362
+ for (int i = 0; i < 2; ++i) {
363
+ neg.val[i] = vec_neg(reg.val[i]);
364
+ }
365
+
366
+ FP32Vec8 neg_x(neg);
367
+ FP32Vec8 e = neg_x.exp(); // exp(-x)
368
+
369
+ f32x4x2_t denom;
370
+ for (int i = 0; i < 2; ++i) {
371
+ denom.val[i] = vec_add(one, e.reg.val[i]);
372
+ }
373
+
374
+ FP32Vec8 denom_vec(denom);
375
+ FP32Vec8 one_vec(1.0f);
376
+
377
+ return one_vec / denom_vec;
378
+ }
379
+
380
+ // Tanh-based GELU:
381
+ // gelu(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3)))
382
+ FP32Vec8 gelu_tanh() const {
383
+ const __vector float k_s2pi = vec_splats(0.7978845608028654f); // √(2/π)
384
+ const __vector float k_0_0447 = vec_splats(0.044715f);
385
+
386
+ f32x4x2_t x2, x3, inner;
387
+ for (int i = 0; i < 2; ++i) {
388
+ __vector float x = reg.val[i];
389
+ x2.val[i] = vec_mul(x, x); // x^2
390
+ x3.val[i] = vec_mul(x2.val[i], x); // x^3
391
+ __vector float t = vec_madd(k_0_0447, x3.val[i], x); // x + 0.044715*x^3
392
+ inner.val[i] = vec_mul(k_s2pi, t); // √(2/π)*(...)
393
+ }
394
+
395
+ FP32Vec8 inner_vec(inner);
396
+ FP32Vec8 t = inner_vec.tanh(); // tanh part
397
+
398
+ FP32Vec8 one_vec(1.0f);
399
+ FP32Vec8 half_vec(0.5f);
400
+
401
+ FP32Vec8 x_vec(*this);
402
+ return x_vec * half_vec * (one_vec + t);
403
+ }
404
+
405
+ // Erf-based GELU:
406
+ // gelu(x) = 0.5 * x * (1 + erf(x / √2))
407
+ FP32Vec8 gelu_erf() const {
408
+ const __vector float inv_sqrt2 = vec_splats(0.7071067811865476f); // 1/√2
409
+ FP32Vec8 x_vec(*this);
410
+
411
+ f32x4x2_t scaled;
412
+ for (int i = 0; i < 2; ++i) {
413
+ scaled.val[i] = vec_mul(reg.val[i], inv_sqrt2);
414
+ }
415
+ FP32Vec8 x_scaled(scaled);
416
+
417
+ FP32Vec8 erf_x = x_scaled.er();
418
+
419
+ FP32Vec8 one_vec(1.0f);
420
+ FP32Vec8 half_vec(0.5f);
421
+
422
+ return x_vec * half_vec * (one_vec + erf_x);
423
+ }
424
+
425
+ // Elementwise reciprocal: 1/x (scalar per lane, for correctness)
426
+ FP32Vec8 rcp() const {
427
+ AliasReg in, out;
428
+ in.reg = reg;
429
+
430
+ for (int i = 0; i < VEC_ELEM_NUM; ++i) {
431
+ out.values[i] = 1.0f / in.values[i];
432
+ }
433
+ return FP32Vec8(out.reg);
434
+ }
435
+
436
+ // Elementwise rsqrt(x) = 1 / sqrt(x) (scalar per lane, for correctness)
437
+ FP32Vec8 rsqrt() const {
438
+ AliasReg in, out;
439
+ in.reg = reg;
440
+
441
+ for (int i = 0; i < VEC_ELEM_NUM; ++i) {
442
+ out.values[i] = 1.0f / std::sqrt(in.values[i]);
443
+ }
444
+ return FP32Vec8(out.reg);
445
}
446
447
FP32Vec8 operator*(const FP32Vec8& b) const {
@@ -316,10 +524,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
524
}
525
526
explicit FP32Vec16(const BF16Vec16& v) {
319
- reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
320
- reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
321
- reg.val[2] = (__vector float)vec_mergeh(zero, v.reg.val[1]);
322
- reg.val[3] = (__vector float)vec_mergel(zero, v.reg.val[1]);
527
+ // On big-endian s390x, place BF16 first to get correct byte order
528
+ reg.val[0] = (__vector float)vec_mergeh(v.reg.val[0], zero);
529
+ reg.val[1] = (__vector float)vec_mergel(v.reg.val[0], zero);
530
+ reg.val[2] = (__vector float)vec_mergeh(v.reg.val[1], zero);
531
+ reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
532
}
533
534
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -376,6 +585,23 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
585
return result;
586
}
587
588
+ FP32Vec16 max(const FP32Vec16& b) const {
589
+ return FP32Vec16(f32x4x4_t({vec_max(reg.val[0], b.reg.val[0]),
590
+ vec_max(reg.val[1], b.reg.val[1]),
591
+ vec_max(reg.val[2], b.reg.val[2]),
592
+ vec_max(reg.val[3], b.reg.val[3])}));
593
+ }
594
+
595
+ float reduce_max() const {
596
+ AliasReg ar;
597
+ ar.reg = reg;
598
+ float result = ar.values[0];
599
+ unroll_loop<int, VEC_ELEM_NUM>([&result, &ar](int i) {
600
+ if (ar.values[i] > result) result = ar.values[i];
601
+ });
602
+ return result;
603
+ }
604
+
605
void save(float* ptr) const {
606
vec_xst(reg.val[0], 0, ptr);
607
vec_xst(reg.val[1], 16, ptr);
@@ -402,15 +628,14 @@ struct VecType<c10::BFloat16> {
628
using vec_type = BF16Vec8;
629
};
630
631
+// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
632
+using FP16Vec16 = FP32Vec16;
633
+
634
template <typename T>
635
void storeFP32(float v, T* ptr) {
636
*ptr = v;
637
}
638
410
-inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
411
- acc = acc + a * b;
412
-}
413
-
639
namespace c10 {
640
struct BFloat16 {
641
uint16_t value; // Assume BFloat16 is defined as a struct containing a 16-bit
@@ -429,6 +654,79 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
654
#define __VEC_CLASS_FP_NAN (1 << 6)
655
#endif
656
657
+// Optimized FMA (Fused Multiply-Add) implementations using IBM Z vector
658
+// intrinsics
659
+
660
+// FP32Vec4 FMA: acc = acc + (a * b) or equivalently acc = fma(a, b, acc)
661
+FORCE_INLINE void fma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
662
+ acc.reg = vec_madd(a.reg, b.reg, acc.reg);
663
+}
664
+
665
+// FP32Vec8 FMA: acc = acc + (a * b)
666
+FORCE_INLINE void fma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
667
+ acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
668
+ acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
669
+}
670
+
671
+// FP32Vec16 FMA: acc = acc + (a * b)
672
+FORCE_INLINE void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
673
+ acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
674
+ acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
675
+ acc.reg.val[2] = vec_madd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
676
+ acc.reg.val[3] = vec_madd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
677
+}
678
+
679
+// Multiply-Subtract: acc = acc - (a * b)
680
+FORCE_INLINE void fms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
681
+ acc.reg = vec_msub(a.reg, b.reg, acc.reg);
682
+}
683
+
684
+FORCE_INLINE void fms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
685
+ acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
686
+ acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
687
+}
688
+
689
+FORCE_INLINE void fms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
690
+ acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
691
+ acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
692
+ acc.reg.val[2] = vec_msub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
693
+ acc.reg.val[3] = vec_msub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
694
+}
695
+
696
+// Negative Multiply-Add: acc = -(a * b) + acc
697
+FORCE_INLINE void nfma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
698
+ acc.reg = vec_nmadd(a.reg, b.reg, acc.reg);
699
+}
700
+
701
+FORCE_INLINE void nfma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
702
+ acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
703
+ acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
704
+}
705
+
706
+FORCE_INLINE void nfma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
707
+ acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
708
+ acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
709
+ acc.reg.val[2] = vec_nmadd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
710
+ acc.reg.val[3] = vec_nmadd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
711
+}
712
+
713
+// Negative Multiply-Subtract: acc = -(a * b) - acc
714
+FORCE_INLINE void nfms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
715
+ acc.reg = vec_nmsub(a.reg, b.reg, acc.reg);
716
+}
717
+
718
+FORCE_INLINE void nfms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
719
+ acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
720
+ acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
721
+}
722
+
723
+FORCE_INLINE void nfms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
724
+ acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
725
+ acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
726
+ acc.reg.val[2] = vec_nmsub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
727
+ acc.reg.val[3] = vec_nmsub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
728
+}
729
+
730
const static __vector unsigned char omask = {2, 3, 6, 7, 10, 11, 14, 15,
731
18, 19, 22, 23, 26, 27, 30, 31};
732
const static __vector unsigned int bias = {0x00007fff, 0x00007fff, 0x00007fff,
@@ -441,13 +739,24 @@ const static __vector unsigned int one = {1, 1, 1, 1};
739
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
740
__vector unsigned int inp0 = (__vector unsigned int)(v.reg.val[0]);
741
__vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
742
+ __vector unsigned int lsb0 = inp0 >> sh16;
743
+ __vector unsigned int lsb1 = inp1 >> sh16;
744
+ lsb0 = lsb0 & one;
745
+ lsb1 = lsb1 & one;
746
+ __vector unsigned int rnd0 = lsb0 + bias;
747
+ __vector unsigned int rnd1 = lsb1 + bias;
748
+ inp0 = inp0 + rnd0;
749
+ inp1 = inp1 + rnd1;
750
int cc;
751
__vector __bool int sel0 =
752
vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
753
__vector __bool int sel1 =
754
vec_fp_test_data_class(v.reg.val[1], __VEC_CLASS_FP_NAN, &cc);
449
- inp0 = vec_sel(inp0, nan, sel0) >> sh16;
450
- inp1 = vec_sel(inp1, nan, sel1) >> sh16;
755
+ inp0 = vec_sel(inp0, nan, sel0);
756
+ inp1 = vec_sel(inp1, nan, sel1);
757
+ inp0 = inp0 >> sh16;
758
+ inp1 = inp1 >> sh16;
759
+
760
reg = (__vector signed short)vec_perm(inp0, inp1, omask);
761
}
762
@@ -456,6 +765,22 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
765
__vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
766
__vector unsigned int inp2 = (__vector unsigned int)(v.reg.val[2]);
767
__vector unsigned int inp3 = (__vector unsigned int)(v.reg.val[3]);
768
+ __vector unsigned int lsb0 = inp0 >> sh16;
769
+ __vector unsigned int lsb1 = inp1 >> sh16;
770
+ __vector unsigned int lsb2 = inp2 >> sh16;
771
+ __vector unsigned int lsb3 = inp3 >> sh16;
772
+ lsb0 = lsb0 & one;
773
+ lsb1 = lsb1 & one;
774
+ lsb2 = lsb2 & one;
775
+ lsb3 = lsb3 & one;
776
+ __vector unsigned int rnd0 = lsb0 + bias;
777
+ __vector unsigned int rnd1 = lsb1 + bias;
778
+ __vector unsigned int rnd2 = lsb2 + bias;
779
+ __vector unsigned int rnd3 = lsb3 + bias;
780
+ inp0 = inp0 + rnd0;
781
+ inp1 = inp1 + rnd1;
782
+ inp2 = inp2 + rnd2;
783
+ inp3 = inp3 + rnd3;
784
int cc;
785
__vector __bool int sel0 =
786
vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
@@ -465,15 +790,164 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
790
vec_fp_test_data_class(v.reg.val[2], __VEC_CLASS_FP_NAN, &cc);
791
__vector __bool int sel3 =
792
vec_fp_test_data_class(v.reg.val[3], __VEC_CLASS_FP_NAN, &cc);
468
- inp0 = vec_sel(inp0, nan, sel0) >> sh16;
469
- inp1 = vec_sel(inp1, nan, sel1) >> sh16;
470
- inp2 = vec_sel(inp2, nan, sel2) >> sh16;
471
- inp3 = vec_sel(inp3, nan, sel3) >> sh16;
793
+ inp0 = vec_sel(inp0, nan, sel0);
794
+ inp1 = vec_sel(inp1, nan, sel1);
795
+ inp2 = vec_sel(inp2, nan, sel2);
796
+ inp3 = vec_sel(inp3, nan, sel3);
797
+ inp0 = inp0 >> sh16;
798
+ inp1 = inp1 >> sh16;
799
+ inp2 = inp2 >> sh16;
800
+ inp3 = inp3 >> sh16;
801
+
802
reg.val[0] = (__vector signed short)vec_perm(inp0, inp1, omask);
803
reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
804
}
805
476
-inline void prefetch(const void* addr) { void __dcbt(const void* addr); }
806
+// 1D softmax over `n` elements in `input`, writes result to `output`.
807
+// Uses FP32Vec8 for main body, scalar tail handling.
808
+// Requirement: n > 0
809
+FORCE_INLINE void softmax_fp32vec8(float* output, const float* input, int n) {
810
+ if (n <= 0) return;
811
+
812
+ // ---------- Pass 1: find max ----------
813
+ float max_val = -std::numeric_limits<float>::infinity();
814
+ int i = 0;
815
+
816
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
817
+ FP32Vec8 v(input + i);
818
+ FP32Vec8::AliasReg ar;
819
+ ar.reg = v.reg;
820
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
821
+ if (ar.values[j] > max_val) max_val = ar.values[j];
822
+ }
823
+ }
824
+ for (; i < n; ++i) {
825
+ if (input[i] > max_val) max_val = input[i];
826
+ }
827
+
828
+ // ---------- Pass 2: compute exp(x - max) and sum ----------
829
+ float sum = 0.0f;
830
+ i = 0;
831
+
832
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
833
+ float tmp[FP32Vec8::VEC_ELEM_NUM];
834
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
835
+ tmp[j] = input[i + j] - max_val;
836
+ }
837
+
838
+ FP32Vec8 v(tmp);
839
+ FP32Vec8 e = v.exp();
840
+
841
+ FP32Vec8::AliasReg ar;
842
+ ar.reg = e.reg;
843
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
844
+ output[i + j] = ar.values[j];
845
+ sum += ar.values[j];
846
+ }
847
+ }
848
+
849
+ // Tail
850
+ for (; i < n; ++i) {
851
+ float x = input[i] - max_val;
852
+ float ex = std::exp(x); // scalar tail
853
+ output[i] = ex;
854
+ sum += ex;
855
+ }
856
+
857
+ // ---------- Pass 3: normalize ----------
858
+ float inv_sum = 1.0f / sum;
859
+ i = 0;
860
+
861
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
862
+ float tmp[FP32Vec8::VEC_ELEM_NUM];
863
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
864
+ tmp[j] = output[i + j] * inv_sum;
865
+ }
866
+ FP32Vec8 v(tmp);
867
+ v.save(output + i);
868
+ }
869
+
870
+ for (; i < n; ++i) {
871
+ output[i] *= inv_sum;
872
+ }
873
+}
874
+
875
+// 1D RMSNorm kernel:
876
+// input: x[0..n-1]
877
+// weight: w[0..n-1] (gamma), may be nullptr
878
+// output: y[i] = x[i] * inv_rms * (weight[i] if weight != nullptr else 1)
879
+// eps: small epsilon for numerical stability
880
+FORCE_INLINE void rmsnorm_fp32vec8(float* output, const float* input,
881
+ const float* weight, int n, float eps) {
882
+ if (n <= 0) return;
883
+
884
+ // ---------- Pass 1: compute sum of squares ----------
885
+ float sum_sq = 0.0f;
886
+ int i = 0;
887
+
888
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
889
+ FP32Vec8 x_vec(input + i);
890
+
891
+ FP32Vec8 sq = x_vec * x_vec;
892
+
893
+ FP32Vec8::AliasReg ar;
894
+ ar.reg = sq.reg;
895
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
896
+ sum_sq += ar.values[j];
897
+ }
898
+ }
899
+
900
+ // Tail
901
+ for (; i < n; ++i) {
902
+ float v = input[i];
903
+ sum_sq += v * v;
904
+ }
905
+
906
+ float mean_sq = sum_sq / static_cast<float>(n);
907
+ float inv_rms = 1.0f / std::sqrt(mean_sq + eps);
908
+
909
+ // ---------- Pass 2: scale (and apply weight if given) ----------
910
+ const float inv_rms_f = inv_rms;
911
+ i = 0;
912
+
913
+ if (weight) {
914
+ // with gamma
915
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
916
+ FP32Vec8 x_vec(input + i);
917
+
918
+ float wtmp[FP32Vec8::VEC_ELEM_NUM];
919
+ for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
920
+ wtmp[j] = weight[i + j];
921
+ }
922
+ FP32Vec8 w_vec(wtmp);
923
+
924
+ FP32Vec8 scale_vec(inv_rms_f);
925
+ FP32Vec8 y = x_vec * scale_vec * w_vec;
926
+ y.save(output + i);
927
+ }
928
+
929
+ for (; i < n; ++i) {
930
+ output[i] = input[i] * inv_rms_f * weight[i];
931
+ }
932
+ } else {
933
+ // without gamma
934
+ for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
935
+ FP32Vec8 x_vec(input + i);
936
+ FP32Vec8 scale_vec(inv_rms_f);
937
+ FP32Vec8 y = x_vec * scale_vec;
938
+ y.save(output + i);
939
+ }
940
+
941
+ for (; i < n; ++i) {
942
+ output[i] = input[i] * inv_rms_f;
943
+ }
944
+ }
945
+}
946
+
947
+// Prefetch data to cache for better memory access performance
948
+FORCE_INLINE void prefetch(const void* addr) {
949
+ __builtin_prefetch(addr, 0, 3); // 0=read, 3=high temporal locality
950
+}
951
952
}; // namespace vec_op
953