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 4, 22:30 UTC
[CPU][RISC-V] Support multiple RVV VLEN targets via compile-time dispatch (#39478)
Signed-off-by: velonica0 <like@mail.nankai.edu.cn>
velonica0 committed
Apr 20, 2026 at 14:36 UTC
ec7aafc02ae86c41edf2f245df793a37bd097fc4
4 files changed
+1046
-831
cmake/cpu_extension.cmake
+41
-8
@@ -161,16 +161,49 @@ elseif (S390_FOUND)
161
"-mtune=native")
162
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
163
message(STATUS "RISC-V detected")
164
- if(RVV_BF16_FOUND)
165
- message(STATUS "BF16 extension detected")
166
- set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
167
- add_compile_definitions(RISCV_BF16_SUPPORT)
168
- elseif (RVV_FP16_FOUND)
169
- message(WARNING "BF16 functionality is not available")
170
- set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
164
+ # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
165
+ # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
166
+ if(NOT DEFINED VLLM_RVV_VLEN)
167
+ # Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
168
+ if(EXISTS /proc/cpuinfo)
169
+ file(READ /proc/cpuinfo _cpuinfo)
170
+ set(_best 0)
171
+ foreach(_n IN ITEMS 128 256 512 1024)
172
+ if(_cpuinfo MATCHES "zvl${_n}b")
173
+ set(_best ${_n})
174
+ endif()
175
+ endforeach()
176
+ if(_best GREATER 0)
177
+ set(VLLM_RVV_VLEN ${_best})
178
+ endif()
179
+ endif()
180
+ # If auto-detect failed (no /proc/cpuinfo or no zvl<N>b reported)
181
+ # but the compiler supports RVV, require explicit specification.
182
+ if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
183
+ message(FATAL_ERROR
184
+ "RISC-V RVV is available but VLEN could not be auto-detected. "
185
+ "Please specify VLEN explicitly:\n"
186
+ " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
187
+ " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n"
188
+ " -DVLLM_RVV_VLEN=0 (force scalar, no RVV)")
189
+ endif()
190
+ endif()
191
+ if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
192
+ message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
193
+ if(RVV_BF16_FOUND)
194
+ message(STATUS "BF16 extension detected")
195
+ set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
196
+ add_compile_definitions(RISCV_BF16_SUPPORT)
197
+ elseif(RVV_FP16_FOUND)
198
+ message(WARNING "BF16 functionality is not available")
199
+ set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
200
+ else()
201
+ message(STATUS "compile riscv with scalar (no FP16/BF16)")
202
+ set(MARCH_FLAGS -march=rv64gc)
203
+ endif()
204
else()
205
message(STATUS "compile riscv with scalar")
173
- list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
206
+ set(MARCH_FLAGS -march=rv64gc)
207
endif()
208
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
209
else()
csrc/cpu/cpu_types_riscv.hpp
+16
-823
@@ -1,832 +1,25 @@
1
#ifndef CPU_TYPES_RISCV_HPP
2
#define CPU_TYPES_RISCV_HPP
3
4
-#include <algorithm>
5
-#include <cmath>
6
-#include <cstring>
7
-#include <iostream>
8
-#include <limits>
9
-#include <riscv_vector.h>
10
-#include <torch/all.h>
4
+// RISC-V Vector (RVV) CPU type definitions for vLLM.
5
+//
6
+// Supports multiple VLENs via compile-time dispatch. The compiler defines
7
+// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
8
+// maps VLEN to the correct LMUL suffixes, and the impl header provides
9
+// VLEN-independent class implementations.
10
+//
11
+// To add support for a new VLEN, add the LMUL mapping in
12
+// cpu_types_riscv_defs.hpp (the impl header needs no changes).
13
12
-// ============================================================================
13
-// Vector Register Type Definitions (VLEN=128 bits)
14
-// ============================================================================
15
-
16
-typedef vfloat16m1_t fixed_vfloat16m1_t
17
- __attribute__((riscv_rvv_vector_bits(128)));
18
-typedef vfloat16m2_t fixed_vfloat16m2_t
19
- __attribute__((riscv_rvv_vector_bits(256)));
20
-
21
-typedef vfloat32m1_t fixed_vfloat32m1_t
22
- __attribute__((riscv_rvv_vector_bits(128)));
23
-typedef vfloat32m2_t fixed_vfloat32m2_t
24
- __attribute__((riscv_rvv_vector_bits(256)));
25
-typedef vfloat32m4_t fixed_vfloat32m4_t
26
- __attribute__((riscv_rvv_vector_bits(512)));
27
-typedef vfloat32m8_t fixed_vfloat32m8_t
28
- __attribute__((riscv_rvv_vector_bits(1024)));
29
-
30
-typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
31
-typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
32
-
33
-typedef vuint16m1_t fixed_vuint16m1_t
34
- __attribute__((riscv_rvv_vector_bits(128)));
35
-typedef vuint16m2_t fixed_vuint16m2_t
36
- __attribute__((riscv_rvv_vector_bits(256)));
37
-typedef vuint16m4_t fixed_vuint16m4_t
38
- __attribute__((riscv_rvv_vector_bits(512)));
39
-
40
-#ifdef RISCV_BF16_SUPPORT
41
-typedef vbfloat16m1_t fixed_vbfloat16m1_t
42
- __attribute__((riscv_rvv_vector_bits(128)));
43
-typedef vbfloat16m2_t fixed_vbfloat16m2_t
44
- __attribute__((riscv_rvv_vector_bits(256)));
45
-typedef vbfloat16m4_t fixed_vbfloat16m4_t
46
- __attribute__((riscv_rvv_vector_bits(512)));
47
-#endif
48
-
49
-namespace vec_op {
50
-
51
-#ifdef RISCV_BF16_SUPPORT
52
- #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
53
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
54
- AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
55
- AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
56
-#else
57
- #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
58
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
59
- AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
14
+#ifndef __riscv_vector
15
+ #error "cpu_types_riscv.hpp included in a non-RVV translation unit"
16
#endif
17
62
-#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
63
- AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
64
-
65
-#define FORCE_INLINE __attribute__((always_inline)) inline
66
-
67
-namespace {
68
-template <typename T, T... indexes, typename F>
69
-constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
70
- (f(std::integral_constant<T, indexes>{}), ...);
71
-};
72
-} // namespace
73
-
74
-template <typename T, T count, typename F,
75
- typename = std::enable_if_t<std::is_invocable_v<F, T>>>
76
-constexpr void unroll_loop(F&& f) {
77
- unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
78
-}
79
-
80
-template <typename T>
81
-struct Vec {
82
- constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
83
-};
84
-
85
-struct FP32Vec8;
86
-struct FP32Vec16;
87
-
88
-// ============================================================================
89
-// FP16 Implementation
90
-// ============================================================================
91
-
92
-struct FP16Vec8 : public Vec<FP16Vec8> {
93
- constexpr static int VEC_ELEM_NUM = 8;
94
- fixed_vfloat16m1_t reg;
95
-
96
- explicit FP16Vec8(const void* ptr)
97
- : reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
98
- VEC_ELEM_NUM)) {};
99
-
100
- explicit FP16Vec8(const FP32Vec8&);
101
-
102
- void save(void* ptr) const {
103
- __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
104
- }
105
- void save(void* ptr, int elem_num) const {
106
- __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
107
- }
108
- void save_strided(void* ptr, ptrdiff_t stride) const {
109
- ptrdiff_t byte_stride = stride * sizeof(_Float16);
110
- __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
111
- VEC_ELEM_NUM);
112
- }
113
-};
114
-
115
-struct FP16Vec16 : public Vec<FP16Vec16> {
116
- constexpr static int VEC_ELEM_NUM = 16;
117
- fixed_vfloat16m2_t reg;
118
-
119
- explicit FP16Vec16(const void* ptr)
120
- : reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
121
- VEC_ELEM_NUM)) {};
122
-
123
- explicit FP16Vec16(const FP32Vec16& vec);
124
-
125
- void save(void* ptr) const {
126
- __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
127
- }
128
- void save(void* ptr, int elem_num) const {
129
- __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
130
- }
131
- void save_strided(void* ptr, ptrdiff_t stride) const {
132
- ptrdiff_t byte_stride = stride * sizeof(_Float16);
133
- __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
134
- VEC_ELEM_NUM);
135
- }
136
-};
137
-
138
-// ============================================================================
139
-// BF16 Implementation
140
-// ============================================================================
141
-
142
-#ifdef RISCV_BF16_SUPPORT
143
-
144
-FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
145
- return __riscv_vreinterpret_v_bf16m1_u16m1(v);
146
-}
147
-FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
148
- return __riscv_vreinterpret_v_bf16m2_u16m2(v);
149
-}
150
-FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
151
- return __riscv_vreinterpret_v_bf16m4_u16m4(v);
152
-}
153
-
154
-struct BF16Vec8 : public Vec<BF16Vec8> {
155
- constexpr static int VEC_ELEM_NUM = 8;
156
- fixed_vbfloat16m1_t reg;
157
-
158
- explicit BF16Vec8(const void* ptr)
159
- : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
160
- reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
161
-
162
- explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
163
- explicit BF16Vec8(const FP32Vec8&);
164
-
165
- void save(void* ptr) const {
166
- __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
167
- VEC_ELEM_NUM);
168
- }
169
- void save(void* ptr, int elem_num) const {
170
- __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
171
- elem_num);
172
- }
173
- void save_strided(void* ptr, ptrdiff_t stride) const {
174
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
175
- __riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
176
- bf16_to_u16(reg), VEC_ELEM_NUM);
177
- }
178
-};
179
-
180
-struct BF16Vec16 : public Vec<BF16Vec16> {
181
- constexpr static int VEC_ELEM_NUM = 16;
182
- fixed_vbfloat16m2_t reg;
183
-
184
- explicit BF16Vec16(const void* ptr)
185
- : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
186
- reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
187
-
188
- explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
189
- explicit BF16Vec16(const FP32Vec16&);
190
-
191
- void save(void* ptr) const {
192
- __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
193
- VEC_ELEM_NUM);
194
- }
195
- void save(void* ptr, int elem_num) const {
196
- __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
197
- elem_num);
198
- }
199
- void save_strided(void* ptr, ptrdiff_t stride) const {
200
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
201
- __riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
202
- bf16_to_u16(reg), VEC_ELEM_NUM);
203
- }
204
-};
205
-
206
-struct BF16Vec32 : public Vec<BF16Vec32> {
207
- constexpr static int VEC_ELEM_NUM = 32;
208
- fixed_vbfloat16m4_t reg;
209
-
210
- explicit BF16Vec32(const void* ptr)
211
- : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
212
- reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
213
-
214
- explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
215
-
216
- explicit BF16Vec32(const BF16Vec8& v) {
217
- fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
218
- fixed_vuint16m4_t u16_combined =
219
- __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
220
- reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
221
- };
222
-
223
- void save(void* ptr) const {
224
- __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
225
- VEC_ELEM_NUM);
226
- }
227
- void save(void* ptr, int elem_num) const {
228
- __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
229
- elem_num);
230
- }
231
- void save_strided(void* ptr, ptrdiff_t stride) const {
232
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
233
- __riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
234
- bf16_to_u16(reg), VEC_ELEM_NUM);
235
- }
236
-};
237
-
238
-#else
239
-// ============================================================================
240
-// BF16 Fallback Implementation (FP32 Simulation)
241
-// ============================================================================
242
-
243
-struct BF16Vec8 : public Vec<BF16Vec8> {
244
- constexpr static int VEC_ELEM_NUM = 8;
245
- fixed_vfloat32m2_t reg_fp32;
246
- explicit BF16Vec8(const void* ptr) {
247
- const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
248
- float tmp[8];
249
- for (int i = 0; i < 8; ++i) {
250
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
251
- std::memcpy(&tmp[i], &v, 4);
252
- }
253
- reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
254
- }
255
- explicit BF16Vec8(const FP32Vec8&);
256
- void save(void* ptr) const {
257
- float tmp[8];
258
- __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
259
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
260
- for (int i = 0; i < 8; ++i) {
261
- uint32_t v;
262
- std::memcpy(&v, &tmp[i], 4);
263
- u16[i] = static_cast<uint16_t>(v >> 16);
264
- }
265
- }
266
- void save(void* ptr, int elem_num) const {
267
- float tmp[8];
268
- __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
269
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
270
- for (int i = 0; i < elem_num; ++i) {
271
- uint32_t v;
272
- std::memcpy(&v, &tmp[i], 4);
273
- u16[i] = static_cast<uint16_t>(v >> 16);
274
- }
275
- }
276
- void save_strided(void* ptr, ptrdiff_t stride) const {
277
- float tmp[8];
278
- __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
279
- uint8_t* u8 = static_cast<uint8_t*>(ptr);
280
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
281
- for (int i = 0; i < 8; ++i) {
282
- uint32_t v;
283
- std::memcpy(&v, &tmp[i], 4);
284
- uint16_t val = static_cast<uint16_t>(v >> 16);
285
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
286
- }
287
- }
288
-};
289
-
290
-struct BF16Vec16 : public Vec<BF16Vec16> {
291
- constexpr static int VEC_ELEM_NUM = 16;
292
- fixed_vfloat32m4_t reg_fp32;
293
- explicit BF16Vec16(const void* ptr) {
294
- const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
295
- float tmp[16];
296
- for (int i = 0; i < 16; ++i) {
297
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
298
- std::memcpy(&tmp[i], &v, 4);
299
- }
300
- reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
301
- }
302
- explicit BF16Vec16(const FP32Vec16&);
303
- void save(void* ptr) const {
304
- float tmp[16];
305
- __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
306
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
307
- for (int i = 0; i < 16; ++i) {
308
- uint32_t v;
309
- std::memcpy(&v, &tmp[i], 4);
310
- u16[i] = static_cast<uint16_t>(v >> 16);
311
- }
312
- }
313
- void save(void* ptr, int elem_num) const {
314
- float tmp[16];
315
- __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
316
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
317
- for (int i = 0; i < elem_num; ++i) {
318
- uint32_t v;
319
- std::memcpy(&v, &tmp[i], 4);
320
- u16[i] = static_cast<uint16_t>(v >> 16);
321
- }
322
- }
323
- void save_strided(void* ptr, ptrdiff_t stride) const {
324
- float tmp[16];
325
- __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
326
- uint8_t* u8 = static_cast<uint8_t*>(ptr);
327
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
328
- for (int i = 0; i < 16; ++i) {
329
- uint32_t v;
330
- std::memcpy(&v, &tmp[i], 4);
331
- uint16_t val = static_cast<uint16_t>(v >> 16);
332
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
333
- }
334
- }
335
-};
336
-
337
-struct BF16Vec32 : public Vec<BF16Vec32> {
338
- constexpr static int VEC_ELEM_NUM = 32;
339
- fixed_vfloat32m8_t reg_fp32;
340
-
341
- explicit BF16Vec32(const void* ptr) {
342
- const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
343
- float tmp[32];
344
- for (int i = 0; i < 32; ++i) {
345
- uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
346
- std::memcpy(&tmp[i], &v, 4);
347
- }
348
- reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
349
- }
350
-
351
- explicit BF16Vec32(const BF16Vec8& v) {
352
- float tmp_small[8];
353
- __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
354
- float tmp_large[32];
355
- for (int i = 0; i < 4; ++i) {
356
- std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
357
- }
358
- reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
359
- }
360
-
361
- void save(void* ptr) const {
362
- float tmp[32];
363
- __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
364
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
365
- for (int i = 0; i < 32; ++i) {
366
- uint32_t v;
367
- std::memcpy(&v, &tmp[i], 4);
368
- u16[i] = static_cast<uint16_t>(v >> 16);
369
- }
370
- }
371
-
372
- void save(void* ptr, int elem_num) const {
373
- float tmp[32];
374
- __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
375
- uint16_t* u16 = static_cast<uint16_t*>(ptr);
376
- for (int i = 0; i < elem_num; ++i) {
377
- uint32_t v;
378
- std::memcpy(&v, &tmp[i], 4);
379
- u16[i] = static_cast<uint16_t>(v >> 16);
380
- }
381
- }
382
-
383
- void save_strided(void* ptr, ptrdiff_t stride) const {
384
- float tmp[32];
385
- __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
386
- uint8_t* u8 = static_cast<uint8_t*>(ptr);
387
- ptrdiff_t byte_stride = stride * sizeof(uint16_t);
388
- for (int i = 0; i < 32; ++i) {
389
- uint32_t v;
390
- std::memcpy(&v, &tmp[i], 4);
391
- uint16_t val = static_cast<uint16_t>(v >> 16);
392
- *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
393
- }
394
- }
395
-};
396
-#endif
397
-
398
-// ============================================================================
399
-// FP32 Implementation
400
-// ============================================================================
401
-
402
-struct FP32Vec4 : public Vec<FP32Vec4> {
403
- constexpr static int VEC_ELEM_NUM = 4;
404
- fixed_vfloat32m1_t reg;
405
- explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
406
- explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
407
- explicit FP32Vec4(const float* ptr)
408
- : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
409
- explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
410
- explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
411
- void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
412
- void save(float* ptr, int elem_num) const {
413
- __riscv_vse32_v_f32m1(ptr, reg, elem_num);
414
- }
415
-};
416
-
417
-struct FP32Vec8 : public Vec<FP32Vec8> {
418
- constexpr static int VEC_ELEM_NUM = 8;
419
- fixed_vfloat32m2_t reg;
420
-
421
- explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
422
- explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
423
- explicit FP32Vec8(const float* ptr)
424
- : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
425
- explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
426
- explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
427
- explicit FP32Vec8(const FP16Vec8& v)
428
- : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
429
- explicit FP32Vec8(fixed_vfloat16m1_t v)
430
- : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
431
-
432
-#ifdef RISCV_BF16_SUPPORT
433
- explicit FP32Vec8(fixed_vbfloat16m1_t v)
434
- : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
435
- explicit FP32Vec8(const BF16Vec8& v)
436
- : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
437
-#else
438
- explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
18
+#ifndef __riscv_v_min_vlen
19
+ #error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
20
#endif
21
441
- float reduce_sum() const {
442
- fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
443
- scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
444
- return __riscv_vfmv_f_s_f32m1_f32(scalar);
445
- }
446
-
447
- FP32Vec8 operator*(const FP32Vec8& b) const {
448
- return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
449
- }
450
- FP32Vec8 operator+(const FP32Vec8& b) const {
451
- return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
452
- }
453
- FP32Vec8 operator-(const FP32Vec8& b) const {
454
- return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
455
- }
456
- FP32Vec8 operator/(const FP32Vec8& b) const {
457
- return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
458
- }
459
-
460
- FP32Vec8 min(const FP32Vec8& b) const {
461
- return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
462
- }
463
- FP32Vec8 max(const FP32Vec8& b) const {
464
- return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
465
- }
466
- FP32Vec8 abs() const {
467
- return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
468
- }
469
-
470
- FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
471
- return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
472
- }
473
- FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
474
- return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
475
- }
476
-
477
- FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
478
- fixed_vfloat32m2_t temp =
479
- __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
480
- return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
481
- }
482
-
483
- void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
484
- void save(float* ptr, int elem_num) const {
485
- __riscv_vse32_v_f32m2(ptr, reg, elem_num);
486
- }
487
- void save_strided(float* ptr, ptrdiff_t stride) const {
488
- ptrdiff_t byte_stride = stride * sizeof(float);
489
- __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
490
- }
491
-
492
- FP32Vec8 exp() const {
493
- const float inv_ln2 = 1.44269504088896341f;
494
- fixed_vfloat32m2_t x_scaled =
495
- __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
496
- fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
497
- fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
498
-
499
- fixed_vfloat32m2_t r =
500
- __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
501
-
502
- fixed_vfloat32m2_t poly =
503
- __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
504
- poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
505
- poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
506
- poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
507
- poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
508
- poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
509
- poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
510
- poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
511
- poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
512
- poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
513
- poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
514
-
515
- fixed_vint32m2_t biased_exp =
516
- __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
517
- biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
518
- fixed_vint32m2_t exponent_bits =
519
- __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
520
- fixed_vfloat32m2_t scale =
521
- __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
522
-
523
- return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
524
- }
525
-
526
- FP32Vec8 tanh() const {
527
- fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
528
- __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
529
- fixed_vfloat32m2_t x2 =
530
- __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
531
- FP32Vec8 exp_val = FP32Vec8(x2).exp();
532
- fixed_vfloat32m2_t num =
533
- __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
534
- fixed_vfloat32m2_t den =
535
- __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
536
- return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
537
- }
538
-
539
- FP32Vec8 er() const {
540
- const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
541
- a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
542
- fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
543
-
544
- fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
545
- __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
546
- t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
547
-
548
- fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
549
- poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
550
- a4, VEC_ELEM_NUM);
551
- poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
552
- a3, VEC_ELEM_NUM);
553
- poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
554
- a2, VEC_ELEM_NUM);
555
- poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
556
- a1, VEC_ELEM_NUM);
557
- poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
558
-
559
- fixed_vfloat32m2_t exp_val =
560
- FP32Vec8(__riscv_vfneg_v_f32m2(
561
- __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
562
- VEC_ELEM_NUM))
563
- .exp()
564
- .reg;
565
- fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
566
- __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
567
- VEC_ELEM_NUM);
568
-
569
- vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
570
- return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
571
- }
572
-};
573
-
574
-struct FP32Vec16 : public Vec<FP32Vec16> {
575
- constexpr static int VEC_ELEM_NUM = 16;
576
- fixed_vfloat32m4_t reg;
577
-
578
- explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
579
- explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
580
- explicit FP32Vec16(const float* ptr)
581
- : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
582
- explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
583
- explicit FP32Vec16(const FP32Vec8& data)
584
- : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
585
- explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
586
- explicit FP32Vec16(const FP16Vec16& v);
587
-
588
-#ifdef RISCV_BF16_SUPPORT
589
- explicit FP32Vec16(fixed_vbfloat16m2_t v)
590
- : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
591
- explicit FP32Vec16(const BF16Vec16& v)
592
- : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
593
-#else
594
- explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
595
-#endif
596
-
597
- FP32Vec16 operator+(const FP32Vec16& b) const {
598
- return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
599
- }
600
- FP32Vec16 operator-(const FP32Vec16& b) const {
601
- return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
602
- }
603
- FP32Vec16 operator*(const FP32Vec16& b) const {
604
- return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
605
- }
606
- FP32Vec16 operator/(const FP32Vec16& b) const {
607
- return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
608
- }
609
-
610
- FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
611
- return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
612
- }
613
-
614
- float reduce_sum() const {
615
- fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
616
- scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
617
- return __riscv_vfmv_f_s_f32m1_f32(scalar);
618
- }
619
-
620
- float reduce_max() const {
621
- fixed_vfloat32m1_t scalar =
622
- __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
623
- scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
624
- return __riscv_vfmv_f_s_f32m1_f32(scalar);
625
- }
626
-
627
- float reduce_min() const {
628
- fixed_vfloat32m1_t scalar =
629
- __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
630
- scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
631
- return __riscv_vfmv_f_s_f32m1_f32(scalar);
632
- }
633
-
634
- template <int group_size>
635
- float reduce_sub_sum(int idx) {
636
- static_assert(VEC_ELEM_NUM % group_size == 0);
637
- const int start = idx * group_size;
638
- vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
639
- vbool8_t mask = __riscv_vmand_mm_b8(
640
- __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
641
- __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
642
- VEC_ELEM_NUM);
643
- fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
644
- scalar =
645
- __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
646
- return __riscv_vfmv_f_s_f32m1_f32(scalar);
647
- };
648
-
649
- FP32Vec16 max(const FP32Vec16& b) const {
650
- return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
651
- }
652
- FP32Vec16 min(const FP32Vec16& b) const {
653
- return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
654
- }
655
- FP32Vec16 abs() const {
656
- return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
657
- }
658
-
659
- FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
660
- return FP32Vec16(__riscv_vfmin_vv_f32m4(
661
- max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
662
- VEC_ELEM_NUM));
663
- }
664
-
665
- void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
666
- void save(float* ptr, int elem_num) const {
667
- __riscv_vse32_v_f32m4(ptr, reg, elem_num);
668
- }
669
- void save_strided(float* ptr, ptrdiff_t stride) const {
670
- ptrdiff_t byte_stride = stride * sizeof(float);
671
- __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
672
- }
673
-
674
- FP32Vec16 exp() const {
675
- const float inv_ln2 = 1.44269504088896341f;
676
- fixed_vfloat32m4_t x_scaled =
677
- __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
678
- fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
679
- fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
680
- fixed_vfloat32m4_t r =
681
- __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
682
-
683
- fixed_vfloat32m4_t poly =
684
- __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
685
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
686
- 0.009618129107628f, VEC_ELEM_NUM);
687
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
688
- 0.055504108664821f, VEC_ELEM_NUM);
689
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
690
- 0.240226506959101f, VEC_ELEM_NUM);
691
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
692
- 0.693147180559945f, VEC_ELEM_NUM);
693
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
694
- 1.0f, VEC_ELEM_NUM);
695
-
696
- fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
697
- __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
698
- fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
699
- __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
700
-
701
- return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
702
- }
703
-
704
- FP32Vec16 tanh() const {
705
- fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
706
- __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
707
- FP32Vec16 exp_val =
708
- FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
709
- return FP32Vec16(__riscv_vfdiv_vv_f32m4(
710
- __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
711
- __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
712
- }
713
-
714
- FP32Vec16 er() const {
715
- const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
716
- a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
717
- fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
718
- fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
719
- __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
720
- 1.0f, VEC_ELEM_NUM),
721
- 1.0f, VEC_ELEM_NUM);
722
-
723
- fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
724
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
725
- a4, VEC_ELEM_NUM);
726
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
727
- a3, VEC_ELEM_NUM);
728
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
729
- a2, VEC_ELEM_NUM);
730
- poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
731
- a1, VEC_ELEM_NUM);
732
- poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
733
-
734
- fixed_vfloat32m4_t exp_val =
735
- FP32Vec16(__riscv_vfneg_v_f32m4(
736
- __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
737
- VEC_ELEM_NUM))
738
- .exp()
739
- .reg;
740
- fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
741
- __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
742
- VEC_ELEM_NUM);
743
-
744
- vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
745
- return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
746
- }
747
-};
748
-
749
-// ============================================================================
750
-// Type Traits & Global Helpers
751
-// ============================================================================
752
-
753
-template <typename T>
754
-struct VecType {
755
- using vec_type = void;
756
- using vec_t = void;
757
-};
758
-
759
-template <typename T>
760
-using vec_t = typename VecType<T>::vec_type;
761
-
762
-template <>
763
-struct VecType<float> {
764
- using vec_type = FP32Vec8;
765
- using vec_t = FP32Vec8;
766
-};
767
-template <>
768
-struct VecType<c10::Half> {
769
- using vec_type = FP16Vec8;
770
- using vec_t = FP16Vec8;
771
-};
772
-template <>
773
-struct VecType<c10::BFloat16> {
774
- using vec_type = BF16Vec8;
775
- using vec_t = BF16Vec8;
776
-};
777
-
778
-template <typename T>
779
-void storeFP32(float v, T* ptr) {
780
- *ptr = v;
781
-}
782
-template <>
783
-inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
784
- *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
785
-}
786
-
787
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
788
- reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
789
-}
790
-inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
791
- reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
792
-}
793
-inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
794
- reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
795
-}
796
-inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
797
- acc = acc.fma(a, b);
798
-}
799
-
800
-#ifdef RISCV_BF16_SUPPORT
801
-template <>
802
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
803
- *ptr = static_cast<__bf16>(v);
804
-};
805
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
806
- : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
807
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
808
- : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
809
-#else
810
-template <>
811
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
812
- uint32_t val;
813
- std::memcpy(&val, &v, 4);
814
- *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
815
-}
816
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
817
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
818
-#endif
819
-
820
-inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
821
-
822
-} // namespace vec_op
823
-
824
-#ifndef CPU_KERNEL_GUARD_IN
825
- #define CPU_KERNEL_GUARD_IN(NAME)
826
-#endif
827
-
828
-#ifndef CPU_KERNEL_GUARD_OUT
829
- #define CPU_KERNEL_GUARD_OUT(NAME)
830
-#endif
22
+#include "cpu_types_riscv_defs.hpp"
23
+#include "cpu_types_riscv_impl.hpp"
24
832
-#endif // CPU_TYPES_RISCV_HPP
\ No newline at end of file
25
+#endif // CPU_TYPES_RISCV_HPP
csrc/cpu/cpu_types_riscv_defs.hpp
new
+98
@@ -0,0 +1,98 @@
1
+#ifndef CPU_TYPES_RISCV_DEFS_HPP
2
+#define CPU_TYPES_RISCV_DEFS_HPP
3
+
4
+// VLEN-to-LMUL mapping for RISC-V Vector extension.
5
+//
6
+// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
7
+// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
8
+// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
9
+
10
+#include <riscv_vector.h>
11
+
12
+#if __riscv_v_min_vlen == 128
13
+ #define LMUL_128 m1
14
+ #define LMUL_256 m2
15
+ #define LMUL_512 m4
16
+ #define LMUL_1024 m8
17
+ #define BOOL_256 b16
18
+ #define BOOL_512 b8
19
+#elif __riscv_v_min_vlen == 256
20
+ #define LMUL_128 mf2
21
+ #define LMUL_256 m1
22
+ #define LMUL_512 m2
23
+ #define LMUL_1024 m4
24
+ #define BOOL_256 b32
25
+ #define BOOL_512 b16
26
+#else
27
+ #error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
28
+#endif
29
+
30
+// Token-paste helpers.
31
+#define _RVV_P2(a, b) a##b
32
+#define _RVV_P3(a, b, c) a##b##c
33
+#define _RVV_P4(a, b, c, d) a##b##c##d
34
+#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
35
+#define RVVI(base, lmul) _RVV_P2(base, lmul)
36
+#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
37
+#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
38
+// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
39
+#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
40
+#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
41
+
42
+// ---- Semantic fixed-vector typedefs (named by element count) ----
43
+
44
+// float16
45
+typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
46
+ __attribute__((riscv_rvv_vector_bits(128)));
47
+typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
48
+ __attribute__((riscv_rvv_vector_bits(256)));
49
+
50
+// float32
51
+typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
52
+ __attribute__((riscv_rvv_vector_bits(128)));
53
+typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
54
+ __attribute__((riscv_rvv_vector_bits(256)));
55
+typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
56
+ __attribute__((riscv_rvv_vector_bits(512)));
57
+typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
58
+ __attribute__((riscv_rvv_vector_bits(1024)));
59
+
60
+// int32
61
+typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
62
+ __attribute__((riscv_rvv_vector_bits(256)));
63
+typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
64
+ __attribute__((riscv_rvv_vector_bits(512)));
65
+
66
+// uint16
67
+typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
68
+ __attribute__((riscv_rvv_vector_bits(128)));
69
+typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
70
+ __attribute__((riscv_rvv_vector_bits(256)));
71
+typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
72
+ __attribute__((riscv_rvv_vector_bits(512)));
73
+
74
+// bfloat16
75
+#ifdef RISCV_BF16_SUPPORT
76
+typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
77
+ __attribute__((riscv_rvv_vector_bits(128)));
78
+typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
79
+ __attribute__((riscv_rvv_vector_bits(256)));
80
+typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
81
+ __attribute__((riscv_rvv_vector_bits(512)));
82
+#endif
83
+
84
+// ---- Reduction accumulator type (always m1 = one register of f32) ----
85
+// Used for scalar reductions; only element [0] is meaningful.
86
+typedef vfloat32m1_t rvv_f32_accum_t
87
+ __attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
88
+
89
+// ---- Mask types for f32 elements ----
90
+#if __riscv_v_min_vlen == 128
91
+typedef vbool16_t rvv_mask_f32x8_t;
92
+typedef vbool8_t rvv_mask_f32x16_t;
93
+#elif __riscv_v_min_vlen == 256
94
+typedef vbool32_t rvv_mask_f32x8_t;
95
+typedef vbool16_t rvv_mask_f32x16_t;
96
+#endif
97
+
98
+#endif // CPU_TYPES_RISCV_DEFS_HPP
csrc/cpu/cpu_types_riscv_impl.hpp
new
+891
@@ -0,0 +1,891 @@
1
+#ifndef CPU_TYPES_RISCV_IMPL_HPP
2
+#define CPU_TYPES_RISCV_IMPL_HPP
3
+
4
+// Shared implementation of RVV vector-type wrapper classes.
5
+// This file is VLEN-independent: it uses the semantic type names and
6
+// RVVI() intrinsic macros from cpu_types_riscv_defs.hpp.
7
+//
8
+// DO NOT include this file directly; include cpu_types_riscv.hpp instead.
9
+
10
+#include <algorithm>
11
+#include <cmath>
12
+#include <cstring>
13
+#include <iostream>
14
+#include <limits>
15
+#include <torch/all.h>
16
+namespace vec_op {
17
+
18
+#ifdef RISCV_BF16_SUPPORT
19
+ #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
20
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
21
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
22
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
23
+#else
24
+ #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
25
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
26
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
27
+#endif
28
+
29
+#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
30
+ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
31
+
32
+#define FORCE_INLINE __attribute__((always_inline)) inline
33
+
34
+namespace {
35
+template <typename T, T... indexes, typename F>
36
+constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
37
+ (f(std::integral_constant<T, indexes>{}), ...);
38
+};
39
+} // namespace
40
+
41
+template <typename T, T count, typename F,
42
+ typename = std::enable_if_t<std::is_invocable_v<F, T>>>
43
+constexpr void unroll_loop(F&& f) {
44
+ unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
45
+}
46
+
47
+template <typename T>
48
+struct Vec {
49
+ constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
50
+};
51
+
52
+struct FP32Vec8;
53
+struct FP32Vec16;
54
+
55
+// ============================================================================
56
+// FP16 Implementation
57
+// ============================================================================
58
+
59
+struct FP16Vec8 : public Vec<FP16Vec8> {
60
+ constexpr static int VEC_ELEM_NUM = 8;
61
+ fixed_fp16x8_t reg;
62
+
63
+ explicit FP16Vec8(const void* ptr)
64
+ : reg(RVVI(__riscv_vle16_v_f16, LMUL_128)(
65
+ static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
66
+
67
+ explicit FP16Vec8(const FP32Vec8&);
68
+
69
+ void save(void* ptr) const {
70
+ RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
71
+ VEC_ELEM_NUM);
72
+ }
73
+ void save(void* ptr, int elem_num) const {
74
+ RVVI(__riscv_vse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr), reg,
75
+ elem_num);
76
+ }
77
+ void save_strided(void* ptr, ptrdiff_t stride) const {
78
+ ptrdiff_t byte_stride = stride * sizeof(_Float16);
79
+ RVVI(__riscv_vsse16_v_f16, LMUL_128)(static_cast<_Float16*>(ptr),
80
+ byte_stride, reg, VEC_ELEM_NUM);
81
+ }
82
+};
83
+
84
+struct FP16Vec16 : public Vec<FP16Vec16> {
85
+ constexpr static int VEC_ELEM_NUM = 16;
86
+ fixed_fp16x16_t reg;
87
+
88
+ explicit FP16Vec16(const void* ptr)
89
+ : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
90
+ static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
91
+
92
+ explicit FP16Vec16(const FP32Vec16& vec);
93
+
94
+ void save(void* ptr) const {
95
+ RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
96
+ VEC_ELEM_NUM);
97
+ }
98
+ void save(void* ptr, int elem_num) const {
99
+ RVVI(__riscv_vse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr), reg,
100
+ elem_num);
101
+ }
102
+ void save_strided(void* ptr, ptrdiff_t stride) const {
103
+ ptrdiff_t byte_stride = stride * sizeof(_Float16);
104
+ RVVI(__riscv_vsse16_v_f16, LMUL_256)(static_cast<_Float16*>(ptr),
105
+ byte_stride, reg, VEC_ELEM_NUM);
106
+ }
107
+};
108
+
109
+// ============================================================================
110
+// BF16 Implementation
111
+// ============================================================================
112
+
113
+#ifdef RISCV_BF16_SUPPORT
114
+
115
+FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
116
+ return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
117
+}
118
+FORCE_INLINE fixed_u16x16_t bf16_to_u16(fixed_bf16x16_t v) {
119
+ return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_256, _u16, LMUL_256)(v);
120
+}
121
+FORCE_INLINE fixed_u16x32_t bf16_to_u16(fixed_bf16x32_t v) {
122
+ return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_512, _u16, LMUL_512)(v);
123
+}
124
+
125
+struct BF16Vec8 : public Vec<BF16Vec8> {
126
+ constexpr static int VEC_ELEM_NUM = 8;
127
+ fixed_bf16x8_t reg;
128
+
129
+ explicit BF16Vec8(const void* ptr)
130
+ : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16,
131
+ LMUL_128)(RVVI(__riscv_vle16_v_u16, LMUL_128)(
132
+ reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
133
+
134
+ explicit BF16Vec8(fixed_bf16x8_t data) : reg(data) {};
135
+ explicit BF16Vec8(const FP32Vec8&);
136
+
137
+ void save(void* ptr) const {
138
+ RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
139
+ bf16_to_u16(reg), VEC_ELEM_NUM);
140
+ }
141
+ void save(void* ptr, int elem_num) const {
142
+ RVVI(__riscv_vse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
143
+ bf16_to_u16(reg), elem_num);
144
+ }
145
+ void save_strided(void* ptr, ptrdiff_t stride) const {
146
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
147
+ RVVI(__riscv_vsse16_v_u16, LMUL_128)(reinterpret_cast<uint16_t*>(ptr),
148
+ byte_stride, bf16_to_u16(reg),
149
+ VEC_ELEM_NUM);
150
+ }
151
+};
152
+
153
+struct BF16Vec16 : public Vec<BF16Vec16> {
154
+ constexpr static int VEC_ELEM_NUM = 16;
155
+ fixed_bf16x16_t reg;
156
+
157
+ explicit BF16Vec16(const void* ptr)
158
+ : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16,
159
+ LMUL_256)(RVVI(__riscv_vle16_v_u16, LMUL_256)(
160
+ reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
161
+
162
+ explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
163
+ explicit BF16Vec16(const FP32Vec16&);
164
+
165
+ void save(void* ptr) const {
166
+ RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
167
+ bf16_to_u16(reg), VEC_ELEM_NUM);
168
+ }
169
+ void save(void* ptr, int elem_num) const {
170
+ RVVI(__riscv_vse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
171
+ bf16_to_u16(reg), elem_num);
172
+ }
173
+ void save_strided(void* ptr, ptrdiff_t stride) const {
174
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
175
+ RVVI(__riscv_vsse16_v_u16, LMUL_256)(reinterpret_cast<uint16_t*>(ptr),
176
+ byte_stride, bf16_to_u16(reg),
177
+ VEC_ELEM_NUM);
178
+ }
179
+};
180
+
181
+struct BF16Vec32 : public Vec<BF16Vec32> {
182
+ constexpr static int VEC_ELEM_NUM = 32;
183
+ fixed_bf16x32_t reg;
184
+
185
+ explicit BF16Vec32(const void* ptr)
186
+ : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
187
+ LMUL_512)(RVVI(__riscv_vle16_v_u16, LMUL_512)(
188
+ reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
189
+
190
+ explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
191
+
192
+ explicit BF16Vec32(const BF16Vec8& v) {
193
+ fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
194
+ fixed_u16x32_t u16_combined =
195
+ RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
196
+ u16_val, u16_val, u16_val, u16_val);
197
+ reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
198
+ LMUL_512)(u16_combined);
199
+ };
200
+
201
+ void save(void* ptr) const {
202
+ RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
203
+ bf16_to_u16(reg), VEC_ELEM_NUM);
204
+ }
205
+ void save(void* ptr, int elem_num) const {
206
+ RVVI(__riscv_vse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
207
+ bf16_to_u16(reg), elem_num);
208
+ }
209
+ void save_strided(void* ptr, ptrdiff_t stride) const {
210
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
211
+ RVVI(__riscv_vsse16_v_u16, LMUL_512)(reinterpret_cast<uint16_t*>(ptr),
212
+ byte_stride, bf16_to_u16(reg),
213
+ VEC_ELEM_NUM);
214
+ }
215
+};
216
+
217
+#else
218
+// ============================================================================
219
+// BF16 Fallback Implementation (FP32 Simulation)
220
+// ============================================================================
221
+
222
+struct BF16Vec8 : public Vec<BF16Vec8> {
223
+ constexpr static int VEC_ELEM_NUM = 8;
224
+ fixed_fp32x8_t reg_fp32;
225
+ explicit BF16Vec8(const void* ptr) {
226
+ const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
227
+ float tmp[8];
228
+ for (int i = 0; i < 8; ++i) {
229
+ uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
230
+ std::memcpy(&tmp[i], &v, 4);
231
+ }
232
+ reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
233
+ }
234
+ explicit BF16Vec8(const FP32Vec8&);
235
+ void save(void* ptr) const {
236
+ float tmp[8];
237
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
238
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
239
+ for (int i = 0; i < 8; ++i) {
240
+ uint32_t v;
241
+ std::memcpy(&v, &tmp[i], 4);
242
+ u16[i] = static_cast<uint16_t>(v >> 16);
243
+ }
244
+ }
245
+ void save(void* ptr, int elem_num) const {
246
+ float tmp[8];
247
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
248
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
249
+ for (int i = 0; i < elem_num; ++i) {
250
+ uint32_t v;
251
+ std::memcpy(&v, &tmp[i], 4);
252
+ u16[i] = static_cast<uint16_t>(v >> 16);
253
+ }
254
+ }
255
+ void save_strided(void* ptr, ptrdiff_t stride) const {
256
+ float tmp[8];
257
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
258
+ uint8_t* u8 = static_cast<uint8_t*>(ptr);
259
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
260
+ for (int i = 0; i < 8; ++i) {
261
+ uint32_t v;
262
+ std::memcpy(&v, &tmp[i], 4);
263
+ uint16_t val = static_cast<uint16_t>(v >> 16);
264
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
265
+ }
266
+ }
267
+};
268
+
269
+struct BF16Vec16 : public Vec<BF16Vec16> {
270
+ constexpr static int VEC_ELEM_NUM = 16;
271
+ fixed_fp32x16_t reg_fp32;
272
+ explicit BF16Vec16(const void* ptr) {
273
+ const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
274
+ float tmp[16];
275
+ for (int i = 0; i < 16; ++i) {
276
+ uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
277
+ std::memcpy(&tmp[i], &v, 4);
278
+ }
279
+ reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
280
+ }
281
+ explicit BF16Vec16(const FP32Vec16&);
282
+ void save(void* ptr) const {
283
+ float tmp[16];
284
+ RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
285
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
286
+ for (int i = 0; i < 16; ++i) {
287
+ uint32_t v;
288
+ std::memcpy(&v, &tmp[i], 4);
289
+ u16[i] = static_cast<uint16_t>(v >> 16);
290
+ }
291
+ }
292
+ void save(void* ptr, int elem_num) const {
293
+ float tmp[16];
294
+ RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
295
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
296
+ for (int i = 0; i < elem_num; ++i) {
297
+ uint32_t v;
298
+ std::memcpy(&v, &tmp[i], 4);
299
+ u16[i] = static_cast<uint16_t>(v >> 16);
300
+ }
301
+ }
302
+ void save_strided(void* ptr, ptrdiff_t stride) const {
303
+ float tmp[16];
304
+ RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
305
+ uint8_t* u8 = static_cast<uint8_t*>(ptr);
306
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
307
+ for (int i = 0; i < 16; ++i) {
308
+ uint32_t v;
309
+ std::memcpy(&v, &tmp[i], 4);
310
+ uint16_t val = static_cast<uint16_t>(v >> 16);
311
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
312
+ }
313
+ }
314
+};
315
+
316
+struct BF16Vec32 : public Vec<BF16Vec32> {
317
+ constexpr static int VEC_ELEM_NUM = 32;
318
+ fixed_fp32x32_t reg_fp32;
319
+
320
+ explicit BF16Vec32(const void* ptr) {
321
+ const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
322
+ float tmp[32];
323
+ for (int i = 0; i < 32; ++i) {
324
+ uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
325
+ std::memcpy(&tmp[i], &v, 4);
326
+ }
327
+ reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
328
+ }
329
+
330
+ explicit BF16Vec32(const BF16Vec8& v) {
331
+ float tmp_small[8];
332
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
333
+ float tmp_large[32];
334
+ for (int i = 0; i < 4; ++i) {
335
+ std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
336
+ }
337
+ reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp_large, 32);
338
+ }
339
+
340
+ void save(void* ptr) const {
341
+ float tmp[32];
342
+ RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
343
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
344
+ for (int i = 0; i < 32; ++i) {
345
+ uint32_t v;
346
+ std::memcpy(&v, &tmp[i], 4);
347
+ u16[i] = static_cast<uint16_t>(v >> 16);
348
+ }
349
+ }
350
+
351
+ void save(void* ptr, int elem_num) const {
352
+ float tmp[32];
353
+ RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
354
+ uint16_t* u16 = static_cast<uint16_t*>(ptr);
355
+ for (int i = 0; i < elem_num; ++i) {
356
+ uint32_t v;
357
+ std::memcpy(&v, &tmp[i], 4);
358
+ u16[i] = static_cast<uint16_t>(v >> 16);
359
+ }
360
+ }
361
+
362
+ void save_strided(void* ptr, ptrdiff_t stride) const {
363
+ float tmp[32];
364
+ RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
365
+ uint8_t* u8 = static_cast<uint8_t*>(ptr);
366
+ ptrdiff_t byte_stride = stride * sizeof(uint16_t);
367
+ for (int i = 0; i < 32; ++i) {
368
+ uint32_t v;
369
+ std::memcpy(&v, &tmp[i], 4);
370
+ uint16_t val = static_cast<uint16_t>(v >> 16);
371
+ *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
372
+ }
373
+ }
374
+};
375
+#endif
376
+
377
+// ============================================================================
378
+// FP32 Implementation
379
+// ============================================================================
380
+
381
+struct FP32Vec4 : public Vec<FP32Vec4> {
382
+ constexpr static int VEC_ELEM_NUM = 4;
383
+ fixed_fp32x4_t reg;
384
+ explicit FP32Vec4(float v)
385
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(v, VEC_ELEM_NUM)) {};
386
+ explicit FP32Vec4()
387
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_128)(0.0f, VEC_ELEM_NUM)) {};
388
+ explicit FP32Vec4(const float* ptr)
389
+ : reg(RVVI(__riscv_vle32_v_f32, LMUL_128)(ptr, VEC_ELEM_NUM)) {};
390
+ explicit FP32Vec4(fixed_fp32x4_t data) : reg(data) {};
391
+ explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
392
+ void save(float* ptr) const {
393
+ RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
394
+ }
395
+ void save(float* ptr, int elem_num) const {
396
+ RVVI(__riscv_vse32_v_f32, LMUL_128)(ptr, reg, elem_num);
397
+ }
398
+};
399
+
400
+struct FP32Vec8 : public Vec<FP32Vec8> {
401
+ constexpr static int VEC_ELEM_NUM = 8;
402
+ fixed_fp32x8_t reg;
403
+
404
+ explicit FP32Vec8(float v)
405
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
406
+ explicit FP32Vec8()
407
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, VEC_ELEM_NUM)) {};
408
+ explicit FP32Vec8(const float* ptr)
409
+ : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {};
410
+ explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {};
411
+ explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
412
+ explicit FP32Vec8(const FP16Vec8& v)
413
+ : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {};
414
+ explicit FP32Vec8(fixed_fp16x8_t v)
415
+ : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
416
+
417
+#ifdef RISCV_BF16_SUPPORT
418
+ explicit FP32Vec8(fixed_bf16x8_t v)
419
+ : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
420
+ explicit FP32Vec8(const BF16Vec8& v)
421
+ : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
422
+ };
423
+#else
424
+ explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
425
+#endif
426
+
427
+ float reduce_sum() const {
428
+ rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
429
+ scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_256, _f32m1)(reg, scalar,
430
+ VEC_ELEM_NUM);
431
+ return __riscv_vfmv_f_s_f32m1_f32(scalar);
432
+ }
433
+
434
+ FP32Vec8 operator*(const FP32Vec8& b) const {
435
+ return FP32Vec8(
436
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
437
+ }
438
+ FP32Vec8 operator+(const FP32Vec8& b) const {
439
+ return FP32Vec8(
440
+ RVVI(__riscv_vfadd_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
441
+ }
442
+ FP32Vec8 operator-(const FP32Vec8& b) const {
443
+ return FP32Vec8(
444
+ RVVI(__riscv_vfsub_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
445
+ }
446
+ FP32Vec8 operator/(const FP32Vec8& b) const {
447
+ return FP32Vec8(
448
+ RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
449
+ }
450
+
451
+ FP32Vec8 min(const FP32Vec8& b) const {
452
+ return FP32Vec8(
453
+ RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
454
+ }
455
+ FP32Vec8 max(const FP32Vec8& b) const {
456
+ return FP32Vec8(
457
+ RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, VEC_ELEM_NUM));
458
+ }
459
+ FP32Vec8 abs() const {
460
+ return FP32Vec8(RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM));
461
+ }
462
+
463
+ FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
464
+ return FP32Vec8(RVVI(__riscv_vfmin_vv_f32, LMUL_256)(reg, b.reg, elem_num));
465
+ }
466
+ FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
467
+ return FP32Vec8(RVVI(__riscv_vfmax_vv_f32, LMUL_256)(reg, b.reg, elem_num));
468
+ }
469
+
470
+ FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
471
+ fixed_fp32x8_t temp =
472
+ RVVI(__riscv_vfmax_vv_f32, LMUL_256)(min_v.reg, reg, VEC_ELEM_NUM);
473
+ return FP32Vec8(
474
+ RVVI(__riscv_vfmin_vv_f32, LMUL_256)(max_v.reg, temp, VEC_ELEM_NUM));
475
+ }
476
+
477
+ void save(float* ptr) const {
478
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, VEC_ELEM_NUM);
479
+ }
480
+ void save(float* ptr, int elem_num) const {
481
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(ptr, reg, elem_num);
482
+ }
483
+ void save_strided(float* ptr, ptrdiff_t stride) const {
484
+ ptrdiff_t byte_stride = stride * sizeof(float);
485
+ RVVI(__riscv_vsse32_v_f32, LMUL_256)(ptr, byte_stride, reg, VEC_ELEM_NUM);
486
+ }
487
+
488
+ FP32Vec8 exp() const {
489
+ const float inv_ln2 = 1.44269504088896341f;
490
+ fixed_fp32x8_t x_scaled =
491
+ RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM);
492
+ fixed_i32x8_t n_int =
493
+ RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
494
+ fixed_fp32x8_t n_float =
495
+ RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(n_int, VEC_ELEM_NUM);
496
+
497
+ fixed_fp32x8_t r =
498
+ RVVI(__riscv_vfsub_vv_f32, LMUL_256)(x_scaled, n_float, VEC_ELEM_NUM);
499
+
500
+ fixed_fp32x8_t poly =
501
+ RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.001333355810164f, VEC_ELEM_NUM);
502
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
503
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.009618129107628f,
504
+ VEC_ELEM_NUM);
505
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
506
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.055504108664821f,
507
+ VEC_ELEM_NUM);
508
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
509
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.240226506959101f,
510
+ VEC_ELEM_NUM);
511
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
512
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 0.693147180559945f,
513
+ VEC_ELEM_NUM);
514
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, r, VEC_ELEM_NUM);
515
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(poly, 1.0f, VEC_ELEM_NUM);
516
+
517
+ fixed_i32x8_t biased_exp =
518
+ RVVI(__riscv_vadd_vx_i32, LMUL_256)(n_int, 127, VEC_ELEM_NUM);
519
+ biased_exp =
520
+ RVVI(__riscv_vmax_vx_i32, LMUL_256)(biased_exp, 0, VEC_ELEM_NUM);
521
+ fixed_i32x8_t exponent_bits =
522
+ RVVI(__riscv_vsll_vx_i32, LMUL_256)(biased_exp, 23, VEC_ELEM_NUM);
523
+ fixed_fp32x8_t scale = RVVI4(__riscv_vreinterpret_v_i32, LMUL_256, _f32,
524
+ LMUL_256)(exponent_bits);
525
+
526
+ return FP32Vec8(
527
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, scale, VEC_ELEM_NUM));
528
+ }
529
+
530
+ FP32Vec8 tanh() const {
531
+ fixed_fp32x8_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
532
+ RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
533
+ VEC_ELEM_NUM);
534
+ fixed_fp32x8_t x2 =
535
+ RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x_clamped, 2.0f, VEC_ELEM_NUM);
536
+ FP32Vec8 exp_val = FP32Vec8(x2).exp();
537
+ fixed_fp32x8_t num =
538
+ RVVI(__riscv_vfsub_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
539
+ fixed_fp32x8_t den =
540
+ RVVI(__riscv_vfadd_vf_f32, LMUL_256)(exp_val.reg, 1.0f, VEC_ELEM_NUM);
541
+ return FP32Vec8(
542
+ RVVI(__riscv_vfdiv_vv_f32, LMUL_256)(num, den, VEC_ELEM_NUM));
543
+ }
544
+
545
+ FP32Vec8 er() const {
546
+ const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
547
+ a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
548
+ fixed_fp32x8_t abs_x =
549
+ RVVI(__riscv_vfabs_v_f32, LMUL_256)(reg, VEC_ELEM_NUM);
550
+
551
+ fixed_fp32x8_t t = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
552
+ RVVI(__riscv_vfmul_vf_f32, LMUL_256)(abs_x, p, VEC_ELEM_NUM), 1.0f,
553
+ VEC_ELEM_NUM);
554
+ t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_256)(t, 1.0f, VEC_ELEM_NUM);
555
+
556
+ fixed_fp32x8_t poly =
557
+ RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(a5, VEC_ELEM_NUM);
558
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
559
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a4,
560
+ VEC_ELEM_NUM);
561
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
562
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a3,
563
+ VEC_ELEM_NUM);
564
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
565
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a2,
566
+ VEC_ELEM_NUM);
567
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_256)(
568
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM), a1,
569
+ VEC_ELEM_NUM);
570
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, t, VEC_ELEM_NUM);
571
+
572
+ fixed_fp32x8_t exp_val = FP32Vec8(RVVI(__riscv_vfneg_v_f32, LMUL_256)(
573
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(
574
+ abs_x, abs_x, VEC_ELEM_NUM),
575
+ VEC_ELEM_NUM))
576
+ .exp()
577
+ .reg;
578
+ fixed_fp32x8_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_256)(
579
+ RVVI(__riscv_vfmul_vv_f32, LMUL_256)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
580
+ VEC_ELEM_NUM);
581
+
582
+ rvv_mask_f32x8_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_256, BOOL_256)(
583
+ reg, 0.0f, VEC_ELEM_NUM);
584
+ return FP32Vec8(
585
+ RVVI3(__riscv_vfneg_v_f32, LMUL_256, _m)(mask, res, VEC_ELEM_NUM));
586
+ }
587
+};
588
+
589
+struct FP32Vec16 : public Vec<FP32Vec16> {
590
+ constexpr static int VEC_ELEM_NUM = 16;
591
+ fixed_fp32x16_t reg;
592
+
593
+ explicit FP32Vec16(float v)
594
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
595
+ explicit FP32Vec16()
596
+ : reg(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.0f, VEC_ELEM_NUM)) {};
597
+ explicit FP32Vec16(const float* ptr)
598
+ : reg(RVVI(__riscv_vle32_v_f32, LMUL_512)(ptr, VEC_ELEM_NUM)) {};
599
+ explicit FP32Vec16(fixed_fp32x16_t data) : reg(data) {};
600
+ explicit FP32Vec16(const FP32Vec8& data)
601
+ : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
602
+ data.reg, data.reg)) {};
603
+ explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
604
+ explicit FP32Vec16(const FP16Vec16& v);
605
+
606
+#ifdef RISCV_BF16_SUPPORT
607
+ explicit FP32Vec16(fixed_bf16x16_t v)
608
+ : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
609
+ explicit FP32Vec16(const BF16Vec16& v)
610
+ : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM)) {
611
+ };
612
+#else
613
+ explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
614
+#endif
615
+
616
+ FP32Vec16 operator+(const FP32Vec16& b) const {
617
+ return FP32Vec16(
618
+ RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
619
+ }
620
+ FP32Vec16 operator-(const FP32Vec16& b) const {
621
+ return FP32Vec16(
622
+ RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
623
+ }
624
+ FP32Vec16 operator*(const FP32Vec16& b) const {
625
+ return FP32Vec16(
626
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
627
+ }
628
+ FP32Vec16 operator/(const FP32Vec16& b) const {
629
+ return FP32Vec16(
630
+ RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
631
+ }
632
+
633
+ FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
634
+ return FP32Vec16(
635
+ RVVI(__riscv_vfmacc_vv_f32, LMUL_512)(reg, a.reg, b.reg, VEC_ELEM_NUM));
636
+ }
637
+
638
+ float reduce_sum() const {
639
+ rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
640
+ scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1)(reg, scalar,
641
+ VEC_ELEM_NUM);
642
+ return __riscv_vfmv_f_s_f32m1_f32(scalar);
643
+ }
644
+
645
+ float reduce_max() const {
646
+ rvv_f32_accum_t scalar =
647
+ __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
648
+ scalar = RVVI3(__riscv_vfredmax_vs_f32, LMUL_512, _f32m1)(reg, scalar,
649
+ VEC_ELEM_NUM);
650
+ return __riscv_vfmv_f_s_f32m1_f32(scalar);
651
+ }
652
+
653
+ float reduce_min() const {
654
+ rvv_f32_accum_t scalar =
655
+ __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
656
+ scalar = RVVI3(__riscv_vfredmin_vs_f32, LMUL_512, _f32m1)(reg, scalar,
657
+ VEC_ELEM_NUM);
658
+ return __riscv_vfmv_f_s_f32m1_f32(scalar);
659
+ }
660
+
661
+ template <int group_size>
662
+ float reduce_sub_sum(int idx) {
663
+ static_assert(VEC_ELEM_NUM % group_size == 0);
664
+ const int start = idx * group_size;
665
+ auto indices = RVVI(__riscv_vid_v_u32, LMUL_512)(VEC_ELEM_NUM);
666
+ rvv_mask_f32x16_t mask = RVVI(__riscv_vmand_mm_, BOOL_512)(
667
+ RVVIB(__riscv_vmsgeu_vx_u32, LMUL_512, BOOL_512)(indices, start,
668
+ VEC_ELEM_NUM),
669
+ RVVIB(__riscv_vmsltu_vx_u32, LMUL_512, BOOL_512)(
670
+ indices, start + group_size, VEC_ELEM_NUM),
671
+ VEC_ELEM_NUM);
672
+ rvv_f32_accum_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
673
+ scalar = RVVI3(__riscv_vfredusum_vs_f32, LMUL_512, _f32m1_m)(
674
+ mask, reg, scalar, VEC_ELEM_NUM);
675
+ return __riscv_vfmv_f_s_f32m1_f32(scalar);
676
+ };
677
+
678
+ FP32Vec16 max(const FP32Vec16& b) const {
679
+ return FP32Vec16(
680
+ RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
681
+ }
682
+ FP32Vec16 min(const FP32Vec16& b) const {
683
+ return FP32Vec16(
684
+ RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
685
+ }
686
+ FP32Vec16 abs() const {
687
+ return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
688
+ }
689
+
690
+ FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
691
+ return FP32Vec16(RVVI(__riscv_vfmin_vv_f32, LMUL_512)(
692
+ max_v.reg,
693
+ RVVI(__riscv_vfmax_vv_f32, LMUL_512)(min_v.reg, reg, VEC_ELEM_NUM),
694
+ VEC_ELEM_NUM));
695
+ }
696
+
697
+ void save(float* ptr) const {
698
+ RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, VEC_ELEM_NUM);
699
+ }
700
+ void save(float* ptr, int elem_num) const {
701
+ RVVI(__riscv_vse32_v_f32, LMUL_512)(ptr, reg, elem_num);
702
+ }
703
+ void save_strided(float* ptr, ptrdiff_t stride) const {
704
+ ptrdiff_t byte_stride = stride * sizeof(float);
705
+ RVVI(__riscv_vsse32_v_f32, LMUL_512)(ptr, byte_stride, reg, VEC_ELEM_NUM);
706
+ }
707
+
708
+ FP32Vec16 exp() const {
709
+ const float inv_ln2 = 1.44269504088896341f;
710
+ fixed_fp32x16_t x_scaled =
711
+ RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM);
712
+ fixed_i32x16_t n_int =
713
+ RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
714
+ fixed_fp32x16_t n_float =
715
+ RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_512)(n_int, VEC_ELEM_NUM);
716
+ fixed_fp32x16_t r =
717
+ RVVI(__riscv_vfsub_vv_f32, LMUL_512)(x_scaled, n_float, VEC_ELEM_NUM);
718
+
719
+ fixed_fp32x16_t poly =
720
+ RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(0.001333355810164f, VEC_ELEM_NUM);
721
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
722
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
723
+ 0.009618129107628f, VEC_ELEM_NUM);
724
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
725
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
726
+ 0.055504108664821f, VEC_ELEM_NUM);
727
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
728
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
729
+ 0.240226506959101f, VEC_ELEM_NUM);
730
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
731
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM),
732
+ 0.693147180559945f, VEC_ELEM_NUM);
733
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
734
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, r, VEC_ELEM_NUM), 1.0f,
735
+ VEC_ELEM_NUM);
736
+
737
+ fixed_i32x16_t biased_exp = RVVI(__riscv_vmax_vx_i32, LMUL_512)(
738
+ RVVI(__riscv_vadd_vx_i32, LMUL_512)(n_int, 127, VEC_ELEM_NUM), 0,
739
+ VEC_ELEM_NUM);
740
+ fixed_fp32x16_t scale =
741
+ RVVI4(__riscv_vreinterpret_v_i32, LMUL_512, _f32, LMUL_512)(
742
+ RVVI(__riscv_vsll_vx_i32, LMUL_512)(biased_exp, 23, VEC_ELEM_NUM));
743
+
744
+ return FP32Vec16(
745
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, scale, VEC_ELEM_NUM));
746
+ }
747
+
748
+ FP32Vec16 tanh() const {
749
+ fixed_fp32x16_t x_clamped = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
750
+ RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, -9.0f, VEC_ELEM_NUM), 9.0f,
751
+ VEC_ELEM_NUM);
752
+ FP32Vec16 exp_val = FP32Vec16(RVVI(__riscv_vfmul_vf_f32, LMUL_512)(
753
+ x_clamped, 2.0f, VEC_ELEM_NUM))
754
+ .exp();
755
+ return FP32Vec16(RVVI(__riscv_vfdiv_vv_f32, LMUL_512)(
756
+ RVVI(__riscv_vfsub_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
757
+ RVVI(__riscv_vfadd_vf_f32, LMUL_512)(exp_val.reg, 1.0f, VEC_ELEM_NUM),
758
+ VEC_ELEM_NUM));
759
+ }
760
+
761
+ FP32Vec16 er() const {
762
+ const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
763
+ a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
764
+ fixed_fp32x16_t abs_x =
765
+ RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM);
766
+ fixed_fp32x16_t t = RVVI(__riscv_vfrdiv_vf_f32, LMUL_512)(
767
+ RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
768
+ RVVI(__riscv_vfmul_vf_f32, LMUL_512)(abs_x, p, VEC_ELEM_NUM), 1.0f,
769
+ VEC_ELEM_NUM),
770
+ 1.0f, VEC_ELEM_NUM);
771
+
772
+ fixed_fp32x16_t poly =
773
+ RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(a5, VEC_ELEM_NUM);
774
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
775
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a4,
776
+ VEC_ELEM_NUM);
777
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
778
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a3,
779
+ VEC_ELEM_NUM);
780
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
781
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a2,
782
+ VEC_ELEM_NUM);
783
+ poly = RVVI(__riscv_vfadd_vf_f32, LMUL_512)(
784
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM), a1,
785
+ VEC_ELEM_NUM);
786
+ poly = RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, t, VEC_ELEM_NUM);
787
+
788
+ fixed_fp32x16_t exp_val =
789
+ FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(
790
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(abs_x, abs_x,
791
+ VEC_ELEM_NUM),
792
+ VEC_ELEM_NUM))
793
+ .exp()
794
+ .reg;
795
+ fixed_fp32x16_t res = RVVI(__riscv_vfrsub_vf_f32, LMUL_512)(
796
+ RVVI(__riscv_vfmul_vv_f32, LMUL_512)(poly, exp_val, VEC_ELEM_NUM), 1.0f,
797
+ VEC_ELEM_NUM);
798
+
799
+ rvv_mask_f32x16_t mask = RVVIB(__riscv_vmflt_vf_f32, LMUL_512, BOOL_512)(
800
+ reg, 0.0f, VEC_ELEM_NUM);
801
+ return FP32Vec16(
802
+ RVVI3(__riscv_vfneg_v_f32, LMUL_512, _m)(mask, res, VEC_ELEM_NUM));
803
+ }
804
+};
805
+
806
+// ============================================================================
807
+// Type Traits & Global Helpers
808
+// ============================================================================
809
+
810
+template <typename T>
811
+struct VecType {
812
+ using vec_type = void;
813
+ using vec_t = void;
814
+};
815
+
816
+template <typename T>
817
+using vec_t = typename VecType<T>::vec_type;
818
+
819
+template <>
820
+struct VecType<float> {
821
+ using vec_type = FP32Vec8;
822
+ using vec_t = FP32Vec8;
823
+};
824
+template <>
825
+struct VecType<c10::Half> {
826
+ using vec_type = FP16Vec8;
827
+ using vec_t = FP16Vec8;
828
+};
829
+template <>
830
+struct VecType<c10::BFloat16> {
831
+ using vec_type = BF16Vec8;
832
+ using vec_t = BF16Vec8;
833
+};
834
+
835
+template <typename T>
836
+void storeFP32(float v, T* ptr) {
837
+ *ptr = v;
838
+}
839
+template <>
840
+inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
841
+ *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
842
+}
843
+
844
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
845
+ reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_256)(v.reg, VEC_ELEM_NUM);
846
+}
847
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
848
+ reg = RVVI(__riscv_vfncvt_f_f_w_f16, LMUL_128)(v.reg, VEC_ELEM_NUM);
849
+}
850
+inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
851
+ reg = RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_512)(v.reg, VEC_ELEM_NUM);
852
+}
853
+inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
854
+ acc = acc.fma(a, b);
855
+}
856
+
857
+#ifdef RISCV_BF16_SUPPORT
858
+template <>
859
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
860
+ *ptr = static_cast<__bf16>(v);
861
+};
862
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
863
+ : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_128)(v.reg, VEC_ELEM_NUM)) {
864
+ };
865
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
866
+ : reg(RVVI(__riscv_vfncvtbf16_f_f_w_bf16, LMUL_256)(v.reg, VEC_ELEM_NUM)) {
867
+ };
868
+#else
869
+template <>
870
+inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
871
+ uint32_t val;
872
+ std::memcpy(&val, &v, 4);
873
+ *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
874
+}
875
+inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
876
+inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
877
+#endif
878
+
879
+inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
880
+
881
+} // namespace vec_op
882
+
883
+#ifndef CPU_KERNEL_GUARD_IN
884
+ #define CPU_KERNEL_GUARD_IN(NAME)
885
+#endif
886
+
887
+#ifndef CPU_KERNEL_GUARD_OUT
888
+ #define CPU_KERNEL_GUARD_OUT(NAME)
889
+#endif
890
+
891
+#endif // CPU_TYPES_RISCV_IMPL_HPP