A high-throughput and memory-efficient inference and serving engine for LLMs
Use std::bit_cast for type punning in CPU kernels (#45089)
Signed-off-by: Yuanyuan Chen <cyyever@outlook.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Yuanyuan Chen committed
Jun 11, 2026 at 13:07 UTC
0b995f860952a50fbeeda88d8a229c27a1fac2bb
3 files changed
+36
-56
csrc/cpu/cpu_types_riscv_impl.hpp
+20
-39
@@ -9,10 +9,14 @@
9
10
#include <algorithm>
11
#include <cmath>
12
+#include <cstdint>
13
#include <cstring>
14
#include <iostream>
15
#include <limits>
16
#include <torch/all.h>
17
+
18
+#include "float_convert.hpp"
19
+
20
namespace vec_op {
21
22
// FP8 KV cache is not supported on RISC-V. These tag types and the
@@ -245,8 +249,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
249
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
250
float tmp[8];
251
for (int i = 0; i < 8; ++i) {
248
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
249
- std::memcpy(&tmp[i], &v, 4);
252
+ tmp[i] = bf16_to_float(u16[i]);
253
}
254
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
255
}
@@ -256,9 +259,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
259
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
260
uint16_t* u16 = static_cast<uint16_t*>(ptr);
261
for (int i = 0; i < 8; ++i) {
259
- uint32_t v;
260
- std::memcpy(&v, &tmp[i], 4);
261
- u16[i] = static_cast<uint16_t>(v >> 16);
262
+ u16[i] = float_to_bf16(tmp[i]);
263
}
264
}
265
void save(void* ptr, int elem_num) const {
@@ -266,9 +267,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
267
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
268
uint16_t* u16 = static_cast<uint16_t*>(ptr);
269
for (int i = 0; i < elem_num; ++i) {
269
- uint32_t v;
270
- std::memcpy(&v, &tmp[i], 4);
271
- u16[i] = static_cast<uint16_t>(v >> 16);
270
+ u16[i] = float_to_bf16(tmp[i]);
271
}
272
}
273
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -277,10 +276,8 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
276
uint8_t* u8 = static_cast<uint8_t*>(ptr);
277
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
278
for (int i = 0; i < 8; ++i) {
280
- uint32_t v;
281
- std::memcpy(&v, &tmp[i], 4);
282
- uint16_t val = static_cast<uint16_t>(v >> 16);
283
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
279
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
280
+ float_to_bf16(tmp[i]);
281
}
282
}
283
};
@@ -292,8 +289,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
289
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
290
float tmp[16];
291
for (int i = 0; i < 16; ++i) {
295
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
296
- std::memcpy(&tmp[i], &v, 4);
292
+ tmp[i] = bf16_to_float(u16[i]);
293
}
294
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
295
}
@@ -306,9 +302,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
302
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
303
uint16_t* u16 = static_cast<uint16_t*>(ptr);
304
for (int i = 0; i < 16; ++i) {
309
- uint32_t v;
310
- std::memcpy(&v, &tmp[i], 4);
311
- u16[i] = static_cast<uint16_t>(v >> 16);
305
+ u16[i] = float_to_bf16(tmp[i]);
306
}
307
}
308
void save(void* ptr, int elem_num) const {
@@ -316,9 +310,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
310
RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
311
uint16_t* u16 = static_cast<uint16_t*>(ptr);
312
for (int i = 0; i < elem_num; ++i) {
319
- uint32_t v;
320
- std::memcpy(&v, &tmp[i], 4);
321
- u16[i] = static_cast<uint16_t>(v >> 16);
313
+ u16[i] = float_to_bf16(tmp[i]);
314
}
315
}
316
void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -327,10 +319,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
319
uint8_t* u8 = static_cast<uint8_t*>(ptr);
320
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
321
for (int i = 0; i < 16; ++i) {
330
- uint32_t v;
331
- std::memcpy(&v, &tmp[i], 4);
332
- uint16_t val = static_cast<uint16_t>(v >> 16);
333
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
322
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
323
+ float_to_bf16(tmp[i]);
324
}
325
}
326
};
@@ -343,8 +333,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
333
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
334
float tmp[32];
335
for (int i = 0; i < 32; ++i) {
346
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
347
- std::memcpy(&tmp[i], &v, 4);
336
+ tmp[i] = bf16_to_float(u16[i]);
337
}
338
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
339
}
@@ -371,9 +360,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
360
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
361
uint16_t* u16 = static_cast<uint16_t*>(ptr);
362
for (int i = 0; i < 32; ++i) {
374
- uint32_t v;
375
- std::memcpy(&v, &tmp[i], 4);
376
- u16[i] = static_cast<uint16_t>(v >> 16);
363
+ u16[i] = float_to_bf16(tmp[i]);
364
}
365
}
366
@@ -382,9 +369,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
369
RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
370
uint16_t* u16 = static_cast<uint16_t*>(ptr);
371
for (int i = 0; i < elem_num; ++i) {
385
- uint32_t v;
386
- std::memcpy(&v, &tmp[i], 4);
387
- u16[i] = static_cast<uint16_t>(v >> 16);
372
+ u16[i] = float_to_bf16(tmp[i]);
373
}
374
}
375
@@ -394,10 +379,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
379
uint8_t* u8 = static_cast<uint8_t*>(ptr);
380
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
381
for (int i = 0; i < 32; ++i) {
397
- uint32_t v;
398
- std::memcpy(&v, &tmp[i], 4);
399
- uint16_t val = static_cast<uint16_t>(v >> 16);
400
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
382
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
383
+ float_to_bf16(tmp[i]);
384
}
385
}
386
};
@@ -985,9 +968,7 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
968
#else
969
template <>
970
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
988
- uint32_t val;
989
- std::memcpy(&val, &v, 4);
990
- *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
971
+ *reinterpret_cast<uint16_t*>(ptr) = float_to_bf16(v);
972
}
973
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
974
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
csrc/cpu/cpu_types_vxe.hpp
+3
-2
@@ -3,7 +3,9 @@
3
#define CPU_TYPES_VXE_HPP
4
5
#include <vecintrin.h>
6
+#include <bit>
7
#include <cmath>
8
+#include <cstdint>
9
#include <limits>
10
#include <torch/all.h>
11
namespace vec_op {
@@ -817,8 +819,7 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
819
// intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
820
// produce incorrect results for some inputs. Process each of the 4 vectors
821
// separately.
820
- uint32_t in;
821
- std::memcpy(&in, &v, sizeof(in));
822
+ uint32_t in = std::bit_cast<uint32_t>(v);
823
824
uint32_t s = (in & 0x80000000) >> 16; // Sign
825
uint32_t e = (in & 0x7F800000) >> 23; // Exponent
csrc/cpu/float_convert.hpp
+13
-15
@@ -1,14 +1,15 @@
1
+#pragma once
2
2
-static float bf16_to_float(uint16_t bf16) {
3
+#include <bit>
4
+#include <cstdint>
5
+
6
+inline float bf16_to_float(uint16_t bf16) {
7
uint32_t bits = static_cast<uint32_t>(bf16) << 16;
4
- float fp32;
5
- std::memcpy(&fp32, &bits, sizeof(fp32));
6
- return fp32;
8
+ return std::bit_cast<float>(bits);
9
}
10
9
-static uint16_t float_to_bf16(float fp32) {
10
- uint32_t bits;
11
- std::memcpy(&bits, &fp32, sizeof(fp32));
11
+inline uint16_t float_to_bf16(float fp32) {
12
+ uint32_t bits = std::bit_cast<uint32_t>(fp32);
13
return static_cast<uint16_t>(bits >> 16);
14
}
15
@@ -18,14 +19,13 @@ static uint16_t float_to_bf16(float fp32) {
19
* Codes below copied from
20
* https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
21
*************************************************/
21
-static uint16_t float_to_fp16(float fp32) {
22
+inline uint16_t float_to_fp16(float fp32) {
23
uint16_t fp16;
24
24
- unsigned x;
25
unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
26
unsigned sign, exponent, mantissa;
27
28
- std::memcpy(&x, &fp32, sizeof(fp32));
28
+ uint32_t x = std::bit_cast<uint32_t>(fp32);
29
u = (x & 0x7fffffff);
30
31
// Get rid of +NaN/-NaN case first.
@@ -77,12 +77,11 @@ static uint16_t float_to_fp16(float fp32) {
77
return fp16;
78
}
79
80
-static float fp16_to_float(uint16_t fp16) {
80
+inline float fp16_to_float(uint16_t fp16) {
81
unsigned sign = ((fp16 >> 15) & 1);
82
unsigned exponent = ((fp16 >> 10) & 0x1f);
83
unsigned mantissa = ((fp16 & 0x3ff) << 13);
84
- int temp;
85
- float fp32;
84
+ uint32_t temp;
85
if (exponent == 0x1f) { /* NaN or Inf */
86
mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
87
exponent = 0xff;
@@ -101,6 +100,5 @@ static float fp16_to_float(uint16_t fp16) {
100
exponent += 0x70;
101
}
102
temp = ((sign << 31) | (exponent << 23) | mantissa);
104
- std::memcpy(&fp32, &temp, sizeof(temp));
105
- return fp32;
103
+ return std::bit_cast<float>(temp);
104
}