char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Max Krasnyanskyhexagon: flash attention rework (optimizations, accuracy improvements, etc) (#25085)13e673863

char/aliceai-t5-moe
70.0 KiB1511 linesraw
1// Dynamic quantizers that produce flat (non-tiled) activations
2
3static inline void quantize_block_f32_q8_0_flat(
4    float * restrict x,
5    uint8_t * restrict y_quants,
6    __fp16 * restrict y_scales,
7    uint32_t block_idx
8) {
9    HVX_Vector * vx = (HVX_Vector *) x;
10    HVX_Vector zero = Q6_V_vzero();
11
12    HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0]));
13    HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1]));
14    HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2]));
15    HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3]));
16
17    HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero);
18    HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero);
19    HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero);
20    HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero);
21
22    HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero);
23    HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero);
24    HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero);
25    HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero);
26
27    HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf)));
28    HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf)));
29
30    HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf)));
31    HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf)));
32
33    HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
34    HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
35    HVX_Vector vd01_hf   = Q6_Vhf_equals_Vqf16(vd01_qf16);
36    HVX_Vector vd23_hf   = Q6_Vhf_equals_Vqf16(vd23_qf16);
37
38    HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf);
39    HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf);
40    vx01_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf));
41    vx23_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf));
42
43    HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf);
44    HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf);
45    HVX_Vector vx_i8    = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16);
46
47    * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8;
48
49    HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2);
50    HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2);
51    HVX_Vector v_scales = Q6_V_lo_W(vp2);
52    hvx_vec_store_u(y_scales + block_idx * 4, 8, v_scales);
53}
54
55static inline void quantize_block_f32_q8_1_flat(
56    float * restrict x,
57    uint8_t * restrict y_quants,
58    __fp16 * restrict y_scales,
59    uint32_t block_idx
60) {
61    HVX_Vector * vx = (HVX_Vector *) x;
62    HVX_Vector zero = Q6_V_vzero();
63
64    HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0]));
65    HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1]));
66    HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2]));
67    HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3]));
68
69    HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero);
70    HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero);
71    HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero);
72    HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero);
73
74    HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero);
75    HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero);
76    HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero);
77    HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero);
78
79    HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf)));
80    HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf)));
81
82    HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf)));
83    HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf)));
84
85    HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
86    HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
87    HVX_Vector vd01_hf   = Q6_Vhf_equals_Vqf16(vd01_qf16);
88    HVX_Vector vd23_hf   = Q6_Vhf_equals_Vqf16(vd23_qf16);
89
90    HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf);
91    HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf);
92    vx01_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf));
93    vx23_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf));
94
95    HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf);
96    HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf);
97    HVX_Vector vx_i8    = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16);
98
99    const HVX_Vector ones = Q6_Vb_vsplat_R(1);
100    HVX_Vector v_sums = Q6_Vw_vrmpy_VbVb(vx_i8, ones);
101    v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 4));
102    v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8));
103    v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16));
104
105    * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8;
106
107    HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2);
108    HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2);
109    HVX_Vector v_scales = Q6_V_lo_W(vp2);
110
111    HVX_VectorPair v_deal1 = Q6_W_vdeal_VVR(v_sums, v_sums, -4);
112    HVX_Vector v_even1 = Q6_V_lo_W(v_deal1);
113    HVX_VectorPair v_deal2 = Q6_W_vdeal_VVR(v_even1, v_even1, -4);
114    HVX_Vector v_even2 = Q6_V_lo_W(v_deal2);
115    HVX_VectorPair v_deal3 = Q6_W_vdeal_VVR(v_even2, v_even2, -4);
116    HVX_Vector v_sums_shuffled = Q6_V_lo_W(v_deal3);
117
118    HVX_Vector v_sums_sf = Q6_Vsf_equals_Vw(v_sums_shuffled);
119    HVX_Vector v_sums_hf = hvx_vec_f32_to_f16(v_sums_sf, Q6_V_vzero());
120
121    HVX_Vector v_prod = hvx_vec_mul_f16_f16(v_scales, v_sums_hf);
122
123    HVX_VectorPair vp_scales = Q6_W_vshuff_VVR(v_prod, v_scales, -2);
124    HVX_Vector v_final = Q6_V_lo_W(vp_scales);
125
126    hvx_vec_store_u(y_scales + block_idx * 8, 16, v_final);
127}
128
129static inline void quantize_row_f32_q8_0_flat(float * restrict x, uint8_t * restrict y, uint32_t k) {
130    assert(k % 32 == 0);
131    const uint32_t quants_size = hex_round_up(k, 128);
132    uint8_t * restrict y_quants = y;
133    __fp16 * restrict y_scales = (__fp16 *) (y + quants_size);
134
135    const uint32_t nb = (k + 127) / 128;
136    for (uint32_t i = 0; i < nb; i++) {
137        quantize_block_f32_q8_0_flat(x + i * 128, y_quants, y_scales, i);
138    }
139}
140
141static inline void quantize_row_f32_q8_1_flat(float * restrict x, uint8_t * restrict y, uint32_t k) {
142    assert(k % 32 == 0);
143    const uint32_t quants_size = hex_round_up(k, 128);
144    uint8_t * restrict y_quants = y;
145    __fp16 * restrict y_scales = (__fp16 *) (y + quants_size);
146
147    const uint32_t nb = (k + 127) / 128;
148    for (uint32_t i = 0; i < nb; i++) {
149        quantize_block_f32_q8_1_flat(x + i * 128, y_quants, y_scales, i);
150    }
151}
152
153static inline void quantize_f32_q8_0_flat_kernel(
154    const uint8_t * restrict src_data,
155    uint8_t * restrict dst_data,
156    uint8_t * restrict tmp_data,
157    uint32_t ne0,
158    uint32_t nrows,
159    size_t src_row_size,
160    size_t dst_row_size
161) {
162    const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float));
163    hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float));
164
165    for (uint32_t i = 0; i < nrows; ++i) {
166        hex_l2fetch(src_data, src_row_size, src_row_size, 2);
167        hvx_copy_f32_aa(tmp_data, src_data, ne0);
168
169        quantize_row_f32_q8_0_flat((float *) tmp_data, dst_data, ne0);
170        dst_data += dst_row_size;
171        src_data += src_row_size;
172    }
173}
174
175static inline void quantize_f32_q8_1_flat_kernel(
176    const uint8_t * restrict src_data,
177    uint8_t * restrict dst_data,
178    uint8_t * restrict tmp_data,
179    uint32_t ne0,
180    uint32_t nrows,
181    size_t src_row_size,
182    size_t dst_row_size
183) {
184    const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float));
185    hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float));
186
187    for (uint32_t i = 0; i < nrows; ++i) {
188        hex_l2fetch(src_data, src_row_size, src_row_size, 2);
189        hvx_copy_f32_aa(tmp_data, src_data, ne0);
190
191        quantize_row_f32_q8_1_flat((float *) tmp_data, dst_data, ne0);
192        dst_data += dst_row_size;
193        src_data += src_row_size;
194    }
195}
196
197static inline void quantize_f32_f32_flat_kernel(
198    const uint8_t * restrict src_data,
199    uint8_t * restrict dst_data,
200    uint8_t * restrict tmp_data,
201    uint32_t ne0,
202    uint32_t nrows,
203    size_t src_stride,
204    size_t dst_stride
205) {
206    (void) tmp_data;
207    const size_t src_row_size = ne0 * sizeof(float);
208    for (uint32_t i = 0; i < nrows; ++i) {
209        hex_l2fetch(src_data, src_row_size, src_stride, 2);
210        hvx_copy_f32_au(dst_data, src_data, ne0);
211
212        dst_data += dst_stride;
213        src_data += src_stride;
214    }
215}
216
217static inline void quantize_f32_f16_flat_kernel(
218    const uint8_t * restrict src_data,
219    uint8_t * restrict dst_data,
220    uint8_t * restrict tmp_data,
221    uint32_t ne0,
222    uint32_t nrows,
223    size_t src_stride,
224    size_t dst_stride
225) {
226    (void) tmp_data;
227    const size_t src_row_size = ne0 * sizeof(float);
228    for (uint32_t i = 0; i < nrows; ++i) {
229        hex_l2fetch(src_data, src_row_size, src_stride, 2);
230        hvx_copy_f16_f32_au(dst_data, src_data, ne0);
231
232        dst_data += dst_stride;
233        src_data += src_stride;
234    }
235}
236
237static inline void quantize_f16_f16_flat_kernel(
238    const uint8_t * restrict src_data,
239    uint8_t * restrict dst_data,
240    uint8_t * restrict tmp_data,
241    uint32_t ne0,
242    uint32_t nrows,
243    size_t src_stride,
244    size_t dst_stride
245) {
246    (void) tmp_data;
247    const size_t src_row_size = ne0 * sizeof(float);
248    for (uint32_t i = 0; i < nrows; ++i) {
249        hex_l2fetch(src_data, src_row_size, src_stride, 2);
250        hvx_copy_f16_au(dst_data, src_data, ne0);
251
252        dst_data += dst_stride;
253        src_data += src_stride;
254    }
255}
256
257// Dot kernels that consume flat (non-tiled) activations
258
259static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
260    const uint8_t * restrict tile_ptr = vx;
261    const uint8_t * restrict y_q = vy;
262
263    HVX_Vector v_sum_float = Q6_V_vzero();
264    HVX_Vector i8 = Q6_Vb_vsplat_R(8);
265
266    static const uint8_t __attribute__((aligned(128))) repl[128] = {
267        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
268        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
269        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
270        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
271        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
272        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
273        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
274        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
275    };
276    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
277
278    const uint32_t quants_size = hex_round_up(n, 128);
279    const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size);
280
281    uint32_t n_k_tiles = n / 32;
282    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
283        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
284
285        uint32_t block_idx = kt / 4;
286        uint32_t sub_idx = kt % 4;
287
288        HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128);
289        HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32);
290
291        HVX_Vector v_act_rep[8];
292        v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl);
293        v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl);
294        v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl);
295        v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl);
296        v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl);
297        v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl);
298        v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl);
299        v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl);
300
301        HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, i8);
302        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
303
304        HVX_Vector v_scale_w = vptr[4];
305
306        __fp16 scale_a_val = y_scales[kt];
307        HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val));
308
309        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a);
310        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
311
312        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
313    }
314
315    if (sz) {
316        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
317    } else {
318        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
319    }
320}
321
322static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
323    const uint8_t * restrict tile_ptr = vx;
324    const uint8_t * restrict y0_q = vy0;
325    const uint8_t * restrict y1_q = vy1;
326
327    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
328    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
329    HVX_Vector i8 = Q6_Vb_vsplat_R(8);
330
331    static const uint8_t __attribute__((aligned(128))) repl[128] = {
332        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
333        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
334        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
335        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
336        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
337        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
338        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
339        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
340    };
341    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
342
343    const uint32_t quants_size = hex_round_up(n, 128);
344    const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size);
345    const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size);
346
347    uint32_t n_k_tiles = n / 32;
348    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
349        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
350
351        uint32_t block_idx = kt / 4;
352        uint32_t sub_idx = kt % 4;
353
354        HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128);
355        HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128);
356
357        HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32);
358        HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32);
359
360        HVX_Vector v_act0_rep[8];
361        v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl);
362        v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl);
363        v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl);
364        v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl);
365        v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl);
366        v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl);
367        v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl);
368        v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl);
369
370        HVX_Vector v_act1_rep[8];
371        v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl);
372        v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl);
373        v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl);
374        v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl);
375        v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl);
376        v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl);
377        v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl);
378        v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl);
379
380        HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, i8);
381        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
382        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
383
384        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
385        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
386
387        HVX_Vector v_scale_w = vptr[4];
388
389        __fp16 scale_a0_val = y0_scales[kt];
390        __fp16 scale_a1_val = y1_scales[kt];
391        HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val));
392        HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val));
393
394        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0);
395        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1);
396
397        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
398        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
399
400        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
401        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
402    }
403
404    if (sz0) {
405        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
406    } else {
407        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
408    }
409    if (sz1) {
410        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
411    } else {
412        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
413    }
414}
415
416static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
417    const uint8_t * restrict tile_ptr = vx;
418    const uint8_t * restrict y_q = vy;
419
420    HVX_Vector v_sum_float = Q6_V_vzero();
421
422    static const uint8_t __attribute__((aligned(128))) repl[128] = {
423        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
424        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
425        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
426        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
427        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
428        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
429        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
430        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
431    };
432    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
433
434    const uint32_t quants_size = hex_round_up(n, 128);
435    const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size);
436
437    uint32_t n_k_tiles = n / 32;
438    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
439        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
440
441        uint32_t block_idx = kt / 4;
442        uint32_t sub_idx = kt % 4;
443
444        HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128);
445        HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32);
446
447        HVX_Vector v_act_rep[8];
448        v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl);
449        v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl);
450        v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl);
451        v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl);
452        v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl);
453        v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl);
454        v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl);
455        v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl);
456
457        HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, Q6_V_vzero());
458        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
459
460        HVX_Vector v_scale_offset = vptr[4];
461        HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2);
462        HVX_Vector v_scale = Q6_V_lo_W(p_deal);
463        HVX_Vector v_offset = Q6_V_hi_W(p_deal);
464
465        __fp16 scale_a_val = y_scales[kt * 2 + 0];
466        __fp16 sum_a_val   = y_scales[kt * 2 + 1];
467        HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val));
468        HVX_Vector v_sum_a   = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a_val));
469
470        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a);
471        HVX_Vector v_offset_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a);
472
473        HVX_Vector v_scaled_dot = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
474        HVX_Vector v_sum_scaled = hvx_vec_add_f32_f32(v_scaled_dot, v_offset_comb);
475
476        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
477    }
478
479    if (sz) {
480        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
481    } else {
482        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
483    }
484}
485
486static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
487    const uint8_t * restrict tile_ptr = vx;
488    const uint8_t * restrict y0_q = vy0;
489    const uint8_t * restrict y1_q = vy1;
490
491    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
492    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
493
494    static const uint8_t __attribute__((aligned(128))) repl[128] = {
495        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
496        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
497        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
498        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
499        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
500        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
501        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
502        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
503    };
504    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
505
506    const uint32_t quants_size = hex_round_up(n, 128);
507    const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size);
508    const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size);
509
510    uint32_t n_k_tiles = n / 32;
511    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
512        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
513
514        uint32_t block_idx = kt / 4;
515        uint32_t sub_idx = kt % 4;
516
517        HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128);
518        HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128);
519
520        HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32);
521        HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32);
522
523        HVX_Vector v_act0_rep[8];
524        v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl);
525        v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl);
526        v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl);
527        v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl);
528        v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl);
529        v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl);
530        v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl);
531        v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl);
532
533        HVX_Vector v_act1_rep[8];
534        v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl);
535        v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl);
536        v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl);
537        v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl);
538        v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl);
539        v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl);
540        v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl);
541        v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl);
542
543        HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, Q6_V_vzero());
544        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
545        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
546
547        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
548        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
549
550        HVX_Vector v_scale_offset = vptr[4];
551        HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2);
552        HVX_Vector v_scale = Q6_V_lo_W(p_deal);
553        HVX_Vector v_offset = Q6_V_hi_W(p_deal);
554
555        __fp16 scale_a0_val = y0_scales[kt * 2 + 0];
556        __fp16 sum_a0_val   = y0_scales[kt * 2 + 1];
557        __fp16 scale_a1_val = y1_scales[kt * 2 + 0];
558        __fp16 sum_a1_val   = y1_scales[kt * 2 + 1];
559
560        HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val));
561        HVX_Vector v_sum_a0   = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a0_val));
562        HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val));
563        HVX_Vector v_sum_a1   = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a1_val));
564
565        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a0);
566        HVX_Vector v_offset_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a0);
567        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a1);
568        HVX_Vector v_offset_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a1);
569
570        HVX_Vector v_scaled_dot_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
571        HVX_Vector v_sum_scaled_c0 = hvx_vec_add_f32_f32(v_scaled_dot_c0, v_offset_comb_c0);
572
573        HVX_Vector v_scaled_dot_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
574        HVX_Vector v_sum_scaled_c1 = hvx_vec_add_f32_f32(v_scaled_dot_c1, v_offset_comb_c1);
575
576        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
577        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
578    }
579
580    if (sz0) {
581        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
582    } else {
583        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
584    }
585    if (sz1) {
586        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
587    } else {
588        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
589    }
590}
591
592static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
593    const uint8_t * restrict tile_ptr = vx;
594    const uint8_t * restrict y_q = vy;
595
596    HVX_Vector v_sum_float = Q6_V_vzero();
597
598    static const uint8_t __attribute__((aligned(128))) repl[128] = {
599        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
600        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
601        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
602        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
603        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
604        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
605        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
606        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
607    };
608    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
609
610    const uint32_t quants_size = hex_round_up(n, 128);
611    const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size);
612
613    uint32_t n_k_tiles = n / 32;
614    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
615        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152);
616
617        uint32_t block_idx = kt / 4;
618        uint32_t sub_idx = kt % 4;
619
620        HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128);
621        HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32);
622
623        HVX_Vector v_act_rep[8];
624        v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl);
625        v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl);
626        v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl);
627        v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl);
628        v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl);
629        v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl);
630        v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl);
631        v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl);
632
633        HVX_Vector v_sum = accum_q8_0_32x1(vptr, v_act_rep);
634        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
635
636        HVX_Vector v_scale_w = vptr[8];
637
638        __fp16 scale_a_val = y_scales[kt];
639        HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val));
640
641        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a);
642        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
643
644        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
645    }
646
647    if (sz) {
648        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
649    } else {
650        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
651    }
652}
653
654static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
655    const uint8_t * restrict tile_ptr = vx;
656    const uint8_t * restrict y0_q = vy0;
657    const uint8_t * restrict y1_q = vy1;
658
659    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
660    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
661
662    static const uint8_t __attribute__((aligned(128))) repl[128] = {
663        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
664        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
665        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
666        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
667        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
668        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
669        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
670        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
671    };
672    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
673
674    const uint32_t quants_size = hex_round_up(n, 128);
675    const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size);
676    const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size);
677
678    uint32_t n_k_tiles = n / 32;
679    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
680        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152);
681
682        uint32_t block_idx = kt / 4;
683        uint32_t sub_idx = kt % 4;
684
685        HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128);
686        HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128);
687
688        HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32);
689        HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32);
690
691        HVX_Vector v_act0_rep[8];
692        v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl);
693        v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl);
694        v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl);
695        v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl);
696        v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl);
697        v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl);
698        v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl);
699        v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl);
700
701        HVX_Vector v_act1_rep[8];
702        v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl);
703        v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl);
704        v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl);
705        v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl);
706        v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl);
707        v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl);
708        v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl);
709        v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl);
710
711        HVX_VectorPair v_sums = accum_q8_0_32x2(vptr, v_act0_rep, v_act1_rep);
712        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
713        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
714
715        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
716        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
717
718        HVX_Vector v_scale_w = vptr[8];
719
720        __fp16 scale_a0_val = y0_scales[kt];
721        __fp16 scale_a1_val = y1_scales[kt];
722        HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val));
723        HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val));
724
725        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0);
726        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1);
727
728        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
729        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
730
731        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
732        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
733    }
734
735    if (sz0) {
736        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
737    } else {
738        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
739    }
740    if (sz1) {
741        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
742    } else {
743        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
744    }
745}
746
747static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
748    const uint8_t * restrict tile_ptr = vx;
749    const uint8_t * restrict y_q = vy;
750
751    HVX_Vector v_sum_float = Q6_V_vzero();
752    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
753    HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut;
754
755    static const uint8_t __attribute__((aligned(128))) repl[128] = {
756        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
757        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
758        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
759        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
760        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
761        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
762        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
763        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
764    };
765    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
766
767    const uint32_t quants_size = hex_round_up(n, 128);
768    const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size);
769
770    uint32_t n_k_tiles = n / 32;
771    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
772        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
773
774        uint32_t block_idx = kt / 4;
775        uint32_t sub_idx = kt % 4;
776
777        HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128);
778        HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32);
779
780        HVX_Vector v_act_rep[8];
781        v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl);
782        v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl);
783        v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl);
784        v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl);
785        v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl);
786        v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl);
787        v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl);
788        v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl);
789
790        HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut);
791        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
792
793        HVX_Vector v_scale_w = vptr[4];
794
795        __fp16 scale_a_val = y_scales[kt];
796        HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val));
797
798        HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a);
799        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
800
801        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
802    }
803
804    if (sz) {
805        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
806    } else {
807        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
808    }
809}
810
811static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
812    const uint8_t * restrict tile_ptr = vx;
813    const uint8_t * restrict y0_q = vy0;
814    const uint8_t * restrict y1_q = vy1;
815
816    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
817    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
818    HVX_Vector mask_h4        = Q6_Vb_vsplat_R(0x0F);
819    HVX_Vector lut            = *(const HVX_Vector *) kvalues_iq4nl_lut;
820
821    static const uint8_t __attribute__((aligned(128))) repl[128] = {
822        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
823        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
824        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
825        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
826        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
827        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
828        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
829        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
830    };
831    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
832
833    const uint32_t quants_size = hex_round_up(n, 128);
834    const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size);
835    const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size);
836
837    uint32_t n_k_tiles = n / 32;
838    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
839        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
840
841        uint32_t block_idx = kt / 4;
842        uint32_t sub_idx = kt % 4;
843
844        HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128);
845        HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128);
846
847        HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32);
848        HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32);
849
850        HVX_Vector v_act0_rep[8];
851        v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl);
852        v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl);
853        v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl);
854        v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl);
855        v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl);
856        v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl);
857        v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl);
858        v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl);
859
860        HVX_Vector v_act1_rep[8];
861        v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl);
862        v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl);
863        v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl);
864        v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl);
865        v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl);
866        v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl);
867        v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl);
868        v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl);
869
870        HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut);
871        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
872        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
873
874        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
875        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
876
877        HVX_Vector v_scale_w = vptr[4];
878
879        __fp16 scale_a0_val = y0_scales[kt];
880        __fp16 scale_a1_val = y1_scales[kt];
881        HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val));
882        HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val));
883
884        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0);
885        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1);
886
887        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
888        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
889
890        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
891        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
892    }
893
894    if (sz0) {
895        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
896    } else {
897        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
898    }
899    if (sz1) {
900        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
901    } else {
902        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
903    }
904}
905
906static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
907    const uint8_t * restrict tile_ptr = vx;
908    const uint8_t * restrict y_q = vy;
909
910    HVX_Vector v_sum_float = Q6_V_vzero();
911    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
912    HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut;
913    HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0;
914    HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff);
915
916    static const uint8_t __attribute__((aligned(128))) repl[128] = {
917        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
918        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
919        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
920        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
921        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
922        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
923        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
924        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
925    };
926    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
927
928    const uint32_t quants_size = hex_round_up(n, 128);
929    const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size);
930
931    uint32_t n_k_tiles = n / 32;
932    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
933        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
934
935        uint32_t block_idx = kt / 4;
936        uint32_t sub_idx = kt % 4;
937
938        HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128);
939        HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32);
940
941        HVX_Vector v_act_rep[8];
942        v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl);
943        v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl);
944        v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl);
945        v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl);
946        v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl);
947        v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl);
948        v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl);
949        v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl);
950
951        HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut);
952        HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum);
953
954        HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512);
955        HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand);
956        r0_d = Q6_V_vand_VV(r0_d, e8m0_mask);
957        HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23);
958
959        __fp16 scale_a_val = y_scales[kt];
960        HVX_Vector v_scale_a_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val));
961        HVX_VectorPair p_scale_a_f32 = hvx_vec_f16_to_f32(v_scale_a_f16);
962        HVX_Vector v_scale_a = Q6_V_lo_W(p_scale_a_f32);
963
964        HVX_Vector v_scale_comb = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a);
965        HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb);
966
967        v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
968    }
969
970    v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
971
972    if (sz) {
973        hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
974    } else {
975        hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
976    }
977}
978
979static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
980    const uint8_t * restrict tile_ptr = vx;
981    const uint8_t * restrict y0_q = vy0;
982    const uint8_t * restrict y1_q = vy1;
983
984    HVX_Vector v_sum_float_c0 = Q6_V_vzero();
985    HVX_Vector v_sum_float_c1 = Q6_V_vzero();
986    HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F);
987    HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut;
988    HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0;
989    HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff);
990
991    static const uint8_t __attribute__((aligned(128))) repl[128] = {
992        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
993        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
994        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
995        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
996        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
997        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
998        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
999        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
1000    };
1001    HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
1002
1003    const uint32_t quants_size = hex_round_up(n, 128);
1004    const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size);
1005    const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size);
1006
1007    uint32_t n_k_tiles = n / 32;
1008    for (uint32_t kt = 0; kt < n_k_tiles; kt++) {
1009        const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640);
1010
1011        uint32_t block_idx = kt / 4;
1012        uint32_t sub_idx = kt % 4;
1013
1014        HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128);
1015        HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128);
1016
1017        HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32);
1018        HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32);
1019
1020        HVX_Vector v_act0_rep[8];
1021        v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl);
1022        v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl);
1023        v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl);
1024        v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl);
1025        v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl);
1026        v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl);
1027        v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl);
1028        v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl);
1029
1030        HVX_Vector v_act1_rep[8];
1031        v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl);
1032        v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl);
1033        v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl);
1034        v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl);
1035        v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl);
1036        v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl);
1037        v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl);
1038        v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl);
1039
1040        HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut);
1041        HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums);
1042        HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums);
1043
1044        HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0);
1045        HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1);
1046
1047        HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512);
1048        HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand);
1049        r0_d = Q6_V_vand_VV(r0_d, e8m0_mask);
1050        HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23);
1051
1052        __fp16 scale_a0_val = y0_scales[kt];
1053        __fp16 scale_a1_val = y1_scales[kt];
1054        HVX_Vector v_scale_a0_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val));
1055        HVX_Vector v_scale_a1_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val));
1056        HVX_VectorPair p_scale_a0_f32 = hvx_vec_f16_to_f32(v_scale_a0_f16);
1057        HVX_VectorPair p_scale_a1_f32 = hvx_vec_f16_to_f32(v_scale_a1_f16);
1058        HVX_Vector v_scale_a0 = Q6_V_lo_W(p_scale_a0_f32);
1059        HVX_Vector v_scale_a1 = Q6_V_lo_W(p_scale_a1_f32);
1060
1061        HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a0);
1062        HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a1);
1063
1064        HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0);
1065        HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1);
1066
1067        v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0);
1068        v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
1069    }
1070
1071    v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
1072    v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
1073
1074    if (sz0) {
1075        hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
1076    } else {
1077        hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
1078    }
1079    if (sz1) {
1080        hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
1081    } else {
1082        hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
1083    }
1084}
1085
1086#if __HVX_ARCH__ < 79
1087#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
1088#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
1089#else
1090#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
1091#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
1092#endif
1093
1094static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
1095    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
1096    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
1097
1098    uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
1099    uint32_t nloe = n % VLEN_FP32; // leftover elements
1100
1101    HVX_Vector rsum = Q6_V_vzero();
1102
1103    uint32_t i = 0;
1104
1105    #pragma unroll(4)
1106    for (i = 0; i < nvec; i++) {
1107        HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
1108        rsum = HVX_OP_ADD_F32(rsum, prod);
1109    }
1110
1111    if (nloe) {
1112        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
1113        HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
1114        HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
1115        HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
1116        rsum = HVX_OP_ADD_F32(rsum, prod);
1117    }
1118
1119    *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
1120}
1121
1122static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
1123                                const void * restrict vx0, const void * restrict vx1,
1124                                const void * restrict vy0) {
1125    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
1126    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
1127    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
1128
1129    uint32_t nvec = n / VLEN_FP32;
1130    uint32_t nloe = n % VLEN_FP32;
1131
1132    HVX_Vector rsum0 = Q6_V_vzero();
1133    HVX_Vector rsum1 = Q6_V_vzero();
1134
1135    uint32_t i = 0;
1136
1137    #pragma unroll(2)
1138    for (i = 0; i < nvec; i++) {
1139        HVX_Vector y_sf = y[i];
1140        HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
1141        HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
1142        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
1143        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
1144    }
1145
1146    if (nloe) {
1147        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
1148        HVX_Vector y_sf  = Q6_V_vand_QV(bmask, y[i]);
1149        HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
1150        HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
1151        HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
1152        HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
1153        rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
1154        rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
1155    }
1156
1157    HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
1158    hvx_vec_store_u(s0, 8, rsum);
1159}
1160
1161static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
1162                                const void * restrict vx0, const void * restrict vx1,
1163                                const void * restrict vy0, const void * restrict vy1) {
1164    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
1165    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
1166    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
1167    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
1168
1169    uint32_t nvec = n / VLEN_FP32;
1170    uint32_t nloe = n % VLEN_FP32;
1171
1172    HVX_Vector r0_c0_sum = Q6_V_vzero();
1173    HVX_Vector r0_c1_sum = Q6_V_vzero();
1174    HVX_Vector r1_c0_sum = Q6_V_vzero();
1175    HVX_Vector r1_c1_sum = Q6_V_vzero();
1176
1177    uint32_t i = 0;
1178
1179    #pragma unroll(2)
1180    for (i = 0; i < nvec; i++) {
1181        HVX_Vector r0_sf = x0[i];
1182        HVX_Vector r1_sf = x1[i];
1183        HVX_Vector c0_sf = y0[i];
1184        HVX_Vector c1_sf = y1[i];
1185
1186        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
1187        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
1188        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
1189        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
1190    }
1191
1192    if (nloe) {
1193        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
1194
1195        HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
1196        HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
1197        HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
1198        HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
1199
1200        r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
1201        r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
1202        r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
1203        r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
1204    }
1205
1206    // Reduce and store results
1207    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
1208    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
1209
1210    hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
1211    hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
1212}
1213
1214static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
1215    const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
1216    const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
1217
1218    uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
1219    uint32_t nloe = n % VLEN_FP32; // leftover elements
1220
1221    HVX_Vector       rsum = Q6_V_vzero();
1222
1223    uint32_t i = 0;
1224
1225    #pragma unroll(2)
1226    for (i = 0; i < nvec; i++) {
1227        HVX_Vector x_sf = vx[i];
1228        HVX_Vector y_sf = vy[i];
1229
1230        rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
1231    }
1232
1233    if (nloe) {
1234        HVX_Vector x_sf = vx[i];
1235        HVX_Vector y_sf = vy[i];
1236
1237        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
1238        x_sf = Q6_V_vand_QV(bmask, x_sf);
1239        y_sf = Q6_V_vand_QV(bmask, y_sf);
1240
1241        rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
1242    }
1243
1244    rsum = hvx_vec_reduce_sum_f32(rsum);
1245    hvx_vec_store_u(&s[0], 4, rsum);
1246}
1247
1248#undef HVX_OP_ADD_F32
1249#undef HVX_OP_MUL_F32
1250
1251static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
1252    const HVX_Vector * restrict x = (const HVX_Vector *) vx;
1253    const HVX_Vector * restrict y = (const HVX_Vector *) vy;
1254
1255    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
1256    uint32_t nloe = n % VLEN_FP16; // leftover elements
1257
1258    HVX_VectorPair rsum_p = Q6_W_vzero();
1259
1260    uint32_t i = 0;
1261
1262    #pragma unroll(4)
1263    for (i = 0; i < nvec; i++) {
1264        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
1265    }
1266
1267    if (nloe) {
1268        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
1269        HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
1270        HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
1271        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
1272    }
1273
1274    HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
1275    hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
1276}
1277
1278static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
1279                                const void * restrict vx0, const void * restrict vx1,
1280                                const void * restrict vy0) {
1281    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
1282    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
1283    const HVX_Vector * restrict y  = (const HVX_Vector *) vy0;
1284
1285    uint32_t nvec = n / VLEN_FP16;
1286    uint32_t nloe = n % VLEN_FP16;
1287
1288    HVX_VectorPair rsum0_p = Q6_W_vzero();
1289    HVX_VectorPair rsum1_p = Q6_W_vzero();
1290
1291    uint32_t i = 0;
1292
1293    #pragma unroll(2)
1294    for (i = 0; i < nvec; i++) {
1295        HVX_Vector y_hf = y[i];
1296        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
1297        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
1298    }
1299
1300    if (nloe) {
1301        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
1302        HVX_Vector y_hf  = Q6_V_vand_QV(bmask, y[i]);
1303        HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
1304        HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
1305        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
1306        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
1307    }
1308
1309    HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
1310    HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
1311    HVX_Vector rsum  = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
1312    hvx_vec_store_u(s0, 8, rsum);
1313}
1314
1315static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
1316                                const void * restrict vx0, const void * restrict vx1,
1317                                const void * restrict vy0, const void * restrict vy1) {
1318    const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
1319    const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
1320    const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
1321    const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
1322
1323    uint32_t nvec = n / VLEN_FP16;
1324    uint32_t nloe = n % VLEN_FP16;
1325
1326    // Row sums (sf) - 4 accumulators for 2x2 tile
1327    HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
1328    HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
1329    HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
1330    HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
1331
1332    uint32_t i = 0;
1333
1334    #pragma unroll(2)
1335    for (i = 0; i < nvec; i++) {
1336        HVX_Vector r0_hf = x0[i];
1337        HVX_Vector r1_hf = x1[i];
1338        HVX_Vector c0_hf = y0[i];
1339        HVX_Vector c1_hf = y1[i];
1340
1341        // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1
1342        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
1343        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
1344        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
1345        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
1346    }
1347
1348    if (nloe) {
1349        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
1350
1351        HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
1352        HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
1353        HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
1354        HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
1355
1356        r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
1357        r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
1358        r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
1359        r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
1360    }
1361
1362    HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
1363    HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
1364    HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
1365    HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
1366
1367    // Reduce and store results
1368    HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
1369    HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
1370
1371    hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum);  // row0,col0 row1,col0
1372    hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum);  // row0,col1 row1,col1
1373}
1374
1375static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
1376    const HVX_UVector * restrict x = (const HVX_UVector *) vx;
1377    const HVX_UVector * restrict y = (const HVX_UVector *) vy;
1378
1379    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
1380    uint32_t nloe = n % VLEN_FP16; // leftover elements
1381
1382    HVX_Vector rsum = Q6_V_vzero();
1383
1384    uint32_t i = 0;
1385
1386    #pragma unroll(4)
1387    for (i = 0; i < nvec; i++) {
1388        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]);
1389        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
1390    }
1391
1392    if (nloe) {
1393        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
1394        HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
1395        HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
1396
1397        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
1398        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
1399    }
1400
1401    rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
1402    hvx_vec_store_u(&s[0], 4, rsum);
1403}
1404
1405static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
1406    const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
1407    const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
1408
1409    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
1410    uint32_t nloe = n % VLEN_FP16; // leftover elements
1411
1412    const HVX_Vector zero = Q6_V_vzero();
1413
1414    HVX_Vector       rsum = Q6_V_vzero();
1415
1416    uint32_t i = 0;
1417
1418    #pragma unroll(2)
1419    for (i = 0; i < nvec; i++) {
1420        // Load y (fp32) and convert into fp16
1421        HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero);  // 32 elements
1422        HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero);  // 32 elements
1423        HVX_Vector y_hf  = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
1424
1425        // Load x (fp16)
1426        HVX_Vector x_hf  = vx[i];
1427
1428        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
1429
1430        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
1431    }
1432
1433    if (nloe) {
1434        // Load y (fp32) and convert into fp16
1435        HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero);  // 32 elements
1436        HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero);  // 32 elements
1437        HVX_Vector y_hf  = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
1438
1439        // Load x (fp16)
1440        HVX_Vector x_hf  = vx[i];
1441
1442        // Zero-out unused elements
1443        // Note that we need to clear both x and y because they may contain NANs
1444        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
1445        x_hf = Q6_V_vand_QV(bmask, x_hf);
1446        y_hf = Q6_V_vand_QV(bmask, y_hf);
1447
1448        HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
1449
1450        rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf),  Q6_V_hi_W(xy_qf)));
1451    }
1452
1453    // Convert into fp32 and reduce
1454    rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
1455    hvx_vec_store_u(&s[0], 4, rsum);
1456}
1457
1458static inline void hvx_tensor_add_f32_grid(
1459    const struct htp_tensor * restrict dst,
1460    const struct htp_tensor * restrict src2,
1461    uint32_t start_row,
1462    uint32_t end_row,
1463    uint32_t start_col,
1464    uint32_t end_col,
1465    const struct fastdiv_values * div_ne11_12,
1466    const struct fastdiv_values * div_ne11
1467) {
1468    if (start_row >= end_row || start_col >= end_col) return;
1469    const uint32_t nb1 = dst->nb[1]; // row stride in bytes
1470
1471    const uint32_t ne11 = dst->ne[1];
1472    const uint32_t ne12 = dst->ne[2];
1473    const uint32_t ne11_12 = ne11 * ne12;
1474
1475    const bool is_broadcast1 = (src2->ne[1] == 1);
1476    const bool is_broadcast2 = (src2->ne[2] == 1);
1477    const bool is_broadcast3 = (src2->ne[3] == 1);
1478
1479    for (uint32_t r = start_row; r < end_row; r++) {
1480        float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1);
1481
1482        uint32_t i13 = fastdiv(r, div_ne11_12);
1483        uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11);
1484        uint32_t i11 = r - i13 * ne11_12 - i12 * ne11;
1485
1486        uint32_t i23 = is_broadcast3 ? 0 : i13;
1487        uint32_t i22 = is_broadcast2 ? 0 : i12;
1488        uint32_t i21 = is_broadcast1 ? 0 : i11;
1489
1490        const float * src2_row = (const float *) ((const uint8_t *) src2->data +
1491                                  i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]);
1492
1493        float * dst_ptr = &dst_row[start_col];
1494        const float * src2_ptr = &src2_row[start_col];
1495        int remaining = end_col - start_col;
1496        while (remaining >= 32) {
1497            HVX_Vector v_out = hvx_vmemu(dst_ptr);
1498            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
1499            hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z);
1500            dst_ptr += 32;
1501            src2_ptr += 32;
1502            remaining -= 32;
1503        }
1504        if (remaining > 0) {
1505            HVX_Vector v_out = hvx_vmemu(dst_ptr);
1506            HVX_Vector v_z   = hvx_vmemu(src2_ptr);
1507            hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
1508        }
1509    }
1510}
1511