@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 228 lines 8.31 KB
Raw
1 #ifndef CPU_MICRO_GEMM_RVV_HPP
2 #define CPU_MICRO_GEMM_RVV_HPP
3
4 #include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
5
6 #if defined(__riscv_v)
7
8 namespace cpu_micro_gemm {
9 namespace {
10
11 constexpr int32_t RVV_MGEMM_N8 = 8;
12 constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
13
14 template <typename scalar_t>
15 FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
16
17 template <>
18 FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
19 return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
20 }
21
22 template <>
23 FORCE_INLINE fixed_fp32x8_t
24 load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
25 #if defined(__riscv_zvfh)
26 fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
27 reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
28 return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
29 #else
30 alignas(32) float values[RVV_MGEMM_N8];
31 for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
32 values[i] = static_cast<float>(ptr[i]);
33 }
34 return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
35 #endif
36 }
37
38 template <>
39 FORCE_INLINE fixed_fp32x8_t
40 load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
41 #if defined(__riscv_zvfbfmin)
42 fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
43 reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
44 fixed_bf16x8_t vec =
45 RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
46 return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
47 #else
48 fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
49 reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
50 auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
51 auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
52 return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
53 #endif
54 }
55
56 // Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
57 // with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
58 template <int32_t M, typename scalar_t>
59 FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
60 const scalar_t* __restrict__ b_ptr,
61 float* __restrict__ c_ptr,
62 const int64_t lda,
63 const int64_t ldc, const int32_t k,
64 const bool accum_c) {
65 static_assert(0 < M && M <= 8);
66
67 #define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
68 #define RVV_IF_M(i) if constexpr (M > (i))
69
70 #define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
71 RVV_ROWS_APPLY(RVV_DECL_A)
72 #undef RVV_DECL_A
73
74 #define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
75 RVV_ROWS_APPLY(RVV_DECL_ACC)
76 #undef RVV_DECL_ACC
77
78 #define RVV_INIT_ACC(i) \
79 RVV_IF_M(i) { \
80 if (accum_c) { \
81 acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
82 RVV_MGEMM_N8); \
83 } else { \
84 acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
85 } \
86 }
87 RVV_ROWS_APPLY(RVV_INIT_ACC)
88 #undef RVV_INIT_ACC
89
90 int32_t k_idx = 0;
91 for (; k_idx + 3 < k; k_idx += 4) {
92 #define RVV_FMA_ROW(i, K_OFFSET) \
93 RVV_IF_M(i) { \
94 acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
95 acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
96 RVV_MGEMM_N8); \
97 }
98
99 #define RVV_STEP_K(K_OFFSET) \
100 { \
101 fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
102 b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
103 RVV_FMA_ROW(0, K_OFFSET) \
104 RVV_FMA_ROW(1, K_OFFSET) \
105 RVV_FMA_ROW(2, K_OFFSET) \
106 RVV_FMA_ROW(3, K_OFFSET) \
107 RVV_FMA_ROW(4, K_OFFSET) \
108 RVV_FMA_ROW(5, K_OFFSET) \
109 RVV_FMA_ROW(6, K_OFFSET) \
110 RVV_FMA_ROW(7, K_OFFSET) \
111 }
112
113 RVV_STEP_K(0)
114 RVV_STEP_K(1)
115 RVV_STEP_K(2)
116 RVV_STEP_K(3)
117 #undef RVV_STEP_K
118 #undef RVV_FMA_ROW
119 }
120
121 for (; k_idx < k; ++k_idx) {
122 fixed_fp32x8_t b =
123 load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
124 #define RVV_TAIL_ROW(i) \
125 RVV_IF_M(i) { \
126 acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
127 acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
128 }
129 RVV_ROWS_APPLY(RVV_TAIL_ROW)
130 #undef RVV_TAIL_ROW
131 }
132
133 #define RVV_STORE_ROW(i) \
134 RVV_IF_M(i) { \
135 RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
136 RVV_MGEMM_N8); \
137 }
138 RVV_ROWS_APPLY(RVV_STORE_ROW)
139 #undef RVV_STORE_ROW
140
141 #undef RVV_ROWS_APPLY
142 #undef RVV_IF_M
143 }
144
145 template <int32_t M, typename scalar_t>
146 FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
147 static_assert(0 < M && M <= 8);
148 scalar_t* __restrict__ curr_b_0 = b_ptr;
149 scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
150
151 gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
152 gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
153 c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
154 gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
155 accum_c);
156 gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
157 ldc, k, accum_c);
158 }
159
160 class TileGemmRVV {
161 public:
162 template <typename scalar_t>
163 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
164 switch (m) {
165 case 1:
166 gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
167 break;
168 case 2:
169 gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
170 break;
171 case 3:
172 gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
173 break;
174 case 4:
175 gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
176 break;
177 case 5:
178 gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
179 break;
180 case 6:
181 gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
182 break;
183 case 7:
184 gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
185 break;
186 case 8:
187 gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
188 break;
189 }
190 }
191 };
192
193 } // namespace
194
195 template <typename scalar_t>
196 class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
197 public:
198 static constexpr int32_t MaxMSize = 8;
199 static constexpr int32_t NSize = 32;
200
201 public:
202 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
203 TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
204 }
205
206 static void pack_weight(const scalar_t* __restrict__ weight,
207 scalar_t* __restrict__ packed_weight,
208 const int32_t output_size, const int32_t input_size) {
209 TORCH_CHECK_EQ(output_size % 16, 0);
210 for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
211 const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
212 scalar_t* __restrict__ curr_packed_weight =
213 packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
214 for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
215 *curr_packed_weight = *curr_weight;
216
217 curr_packed_weight += 16;
218 ++curr_weight;
219 }
220 }
221 }
222 };
223
224 } // namespace cpu_micro_gemm
225
226 #endif // defined(__riscv_v)
227
228 #endif // CPU_MICRO_GEMM_RVV_HPP