char-fork/llama.cpp

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

Max Krasnyanskyhexagon: support for multi-device model split (aka row-split) (#28589)eafe15a5e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
10.8 KiB222 linesraw
1// HMX tile-level inline helpers (FP16 32x32 tile operations).
2// Ported from htp-ops-lib/include/dsp/hmx_utils.h. (https://github.com/haozixu/htp-ops-lib)
3
4#ifndef HMX_UTILS_H
5#define HMX_UTILS_H
6
7#include "hvx-base.h"
8
9#include <assert.h>
10#include <hexagon_types.h>
11#include <stddef.h>
12
13#define HMX_FP16_TILE_N_ROWS 32
14#define HMX_FP16_TILE_N_COLS 32
15#define HMX_FP16_TILE_N_ELMS 1024
16#define HMX_FP16_TILE_SIZE   2048
17
18// Initialise aligned 256-byte area with scale vector + zero padding.
19static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale) {
20    volatile HVX_Vector *pv = (HVX_Vector *) out_scales;
21    pv[0] = v_scale;
22    pv[1] = Q6_V_vzero();
23}
24
25// --- Shared scatter offsets and interleave helper ---
26
27// vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile.
28// word[i] = i*128 maps K-row-pair i to byte offset i*128.
29// Column offset (n*4) is added at runtime.  Entries 0..15 cover one tile (region 2047);
30// entries 16..31 cover the next adjacent tile (region 4095) - pick region size at the
31// call site to scatter into one tile (masked) or two contiguous tiles (unmasked).
32static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = {
33    0 * 128,  1 * 128,  2 * 128,  3 * 128,  4 * 128,  5 * 128,  6 * 128,  7 * 128,  8 * 128,  9 * 128,  10 * 128,
34    11 * 128, 12 * 128, 13 * 128, 14 * 128, 15 * 128, 16 * 128, 17 * 128, 18 * 128, 19 * 128, 20 * 128, 21 * 128,
35    22 * 128, 23 * 128, 24 * 128, 25 * 128, 26 * 128, 27 * 128, 28 * 128, 29 * 128, 30 * 128, 31 * 128,
36};
37
38// Scatter row-major FP16 data (in VTCM scratch) into transposed [K][N] tiles.
39// vtcm_src: [n_cols][src_stride] row-major fp16 (only first k elements per row are used)
40// vtcm_dst: [n_col_tiles][n_k_tiles][HMX_FP16_TILE_N_ELMS] tile-major interleaved fp16
41// Processes rows [start_row, end_row) for multi-thread slicing.
42// Full range: start_row=0, end_row=n_cols.
43static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
44                                            const __fp16 * restrict vtcm_src,
45                                            uint32_t n_cols,
46                                            uint32_t k,
47                                            size_t src_stride,
48                                            uint32_t start_row,
49                                            uint32_t end_row) {
50    assert(k % HMX_FP16_TILE_N_COLS == 0);
51
52    const uint32_t       n_k_tiles     = k / HMX_FP16_TILE_N_COLS;
53    const HVX_Vector     v_scat_base   = hvx_vmem(hmx_transpose_scatter_offsets);
54    const HVX_Vector     v_scat_step   = Q6_V_vsplat_R(4);
55    const HVX_VectorPred q_mask64      = Q6_Q_vsetq_R(64);
56    // Each hvx_vmemu load brings 64 fp16 = 128 bytes covering 2 adjacent K-tiles.
57    // When n_k_tiles is even, scatter into 2 K-tiles per call (region 4095, no mask)
58    // using the upper half of hmx_transpose_scatter_offsets.  Tail one K-tile (when
59    // n_k_tiles is odd) falls back to single-tile masked scatter.
60    const bool           pair_scatter  = (n_k_tiles & 1) == 0;
61    const size_t         pair_region   = (size_t) (2 * HMX_FP16_TILE_SIZE - 1);
62    const size_t         single_region = (size_t) (HMX_FP16_TILE_SIZE - 1);
63    __builtin_assume(k > 0);
64    __builtin_assume(end_row > start_row);
65
66    if (pair_scatter) {
67        // Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter.
68        const uint32_t c_step      = 2 * HMX_FP16_TILE_N_COLS;
69        const size_t   c_byte_step = (size_t) c_step * sizeof(__fp16);
70        const size_t   dst_step    = 2 * (size_t) HMX_FP16_TILE_N_ELMS;
71        const uint32_t n_c_iters   = k / c_step;
72
73        for (uint32_t r = start_row; r < end_row; r += 2) {
74            const uint32_t   ct             = r / HMX_FP16_TILE_N_ROWS;
75            const uint32_t   local_r        = r % HMX_FP16_TILE_N_ROWS;
76            const bool       next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
77            const HVX_Vector v_off0         = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
78            const HVX_Vector v_off1         = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
79
80            __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
81            const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
82            const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
83
84            assert(hex_is_aligned(p0, 128));
85            assert(hex_is_aligned(p1, 128));
86            assert(c_byte_step % 128 == 0);
87
88            if (p1) {
89                for (uint32_t i = 0; i < n_c_iters; ++i) {
90                    HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
91                    HVX_Vector v1 = hvx_vmem(p1); p1 += c_byte_step;
92                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
93                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1);
94                    tile_base += dst_step;
95                }
96            } else {
97                const HVX_Vector vzero = Q6_V_vzero();
98                for (uint32_t i = 0; i < n_c_iters; ++i) {
99                    HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
100                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
101                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero);
102                    tile_base += dst_step;
103                }
104            }
105        }
106    } else {
107        // Fallback: scatter one K-tile per call (region 2047, masked).
108        const uint32_t c_step      = HMX_FP16_TILE_N_COLS;
109        const size_t   c_byte_step = (size_t) c_step * sizeof(__fp16);
110        const size_t   dst_step    = (size_t) HMX_FP16_TILE_N_ELMS;
111        const uint32_t n_c_iters   = k / c_step;
112
113        for (uint32_t r = start_row; r < end_row; r += 2) {
114            const uint32_t   ct             = r / HMX_FP16_TILE_N_ROWS;
115            const uint32_t   local_r        = r % HMX_FP16_TILE_N_ROWS;
116            const bool       next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
117            const HVX_Vector v_off0         = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
118            const HVX_Vector v_off1         = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
119
120            __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
121            const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
122            const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
123
124            if (p1) {
125                for (uint32_t i = 0; i < n_c_iters; ++i) {
126                    HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
127                    HVX_Vector v1 = hvx_vmemu(p1); p1 += c_byte_step;
128                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
129                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1);
130                    tile_base += dst_step;
131                }
132            } else {
133                const HVX_Vector vzero = Q6_V_vzero();
134                for (uint32_t i = 0; i < n_c_iters; ++i) {
135                    HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
136                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
137                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero);
138                    tile_base += dst_step;
139                }
140            }
141        }
142    }
143}
144
145// Interleave row-major FP16 data into column-major tile format.
146// Input: [n_rows, head_dim] row-major.  Output: tile[dim_tile][row_tile].
147// Processes rows [start_row, end_row) for multi-thread slicing.
148// Full range: start_row=0, end_row=n_rows.
149static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
150                                            const __fp16 * restrict src,
151                                            uint32_t n_rows,
152                                            uint32_t head_dim,
153                                            size_t src_stride,
154                                            uint32_t n_row_tiles,
155                                            uint32_t start_row,
156                                            uint32_t end_row) {
157    __builtin_assume(head_dim > 0);
158    const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS;
159
160    for (uint32_t r = start_row; r < end_row; r += 2) {
161        const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows;
162
163        const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride);
164        const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL;
165
166        // Row-pair invariants hoisted out of the c loop.
167        const uint32_t r0      = r / HMX_FP16_TILE_N_ROWS;
168        const uint32_t r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2;
169
170        // tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0).
171        // Each c step (+= 64) advances both by 2 dim-tiles worth of fp16.
172        __fp16 *     tb0     = tiles_out + (size_t) r0 * HMX_FP16_TILE_N_ELMS;
173        __fp16 *     tb1     = tb0 + tile_stride_elms;
174        const size_t tb_step = 2 * tile_stride_elms;
175
176        if (pv_in1) {
177            for (uint32_t c = 0; c < head_dim; c += 64) {
178                HVX_Vector     v0             = *pv_in0++;
179                HVX_Vector     v1             = *pv_in1++;
180                HVX_VectorPair vp             = Q6_W_vshuff_VVR(v1, v0, -2);
181                ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp);
182                ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp);
183                tb0 += tb_step;
184                tb1 += tb_step;
185            }
186        } else {
187            const HVX_Vector vzero = Q6_V_vzero();
188            for (uint32_t c = 0; c < head_dim; c += 64) {
189                HVX_Vector     v0             = *pv_in0++;
190                HVX_VectorPair vp             = Q6_W_vshuff_VVR(vzero, v0, -2);
191                ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp);
192                ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp);
193                tb0 += tb_step;
194                tb1 += tb_step;
195            }
196        }
197    }
198}
199
200// --- HMX inline asm macros for load-store packetization ---
201#define HMX_LOAD_MPY_F16(act, wt, range)              \
202    "{\n"                                             \
203    "    activation.hf = mxmem(" act ", " range ")\n" \
204    "    weight.hf = mxmem(" wt ", " range ")\n"      \
205    "}\n"
206
207#define HMX_LOAD_MPY_DEEP_F16(act, wt, range)              \
208    "{\n"                                                  \
209    "    activation.hf = mxmem(" act ", " range "):deep\n" \
210    "    weight.hf = mxmem(" wt ", " range ")\n"           \
211    "}\n"
212
213#define HMX_STORE_AFTER_F16(out, scale_reg) \
214    "mxmem(" out ", " scale_reg "):after.hf = acc\n"
215
216#define HMX_SET_BIAS(scales) \
217    "bias = mxmem2(" scales ")\n"
218
219#define HMX_CLRACC_F16() \
220    "mxclracc.hf\n"
221
222#endif // HMX_UTILS_H