char-fork/llama.cpp

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

Todor Boinovskihexagon: add RELU and LEAKY_RELU ops (#28585)050dde50c

char/aliceai-t5-moe
5.9 KiB178 linesraw
1#ifndef HTP_UNARY_OPS_H
2#define HTP_UNARY_OPS_H
3
4#include "hex-common.h"
5#include "htp-ops.h"
6
7// Op-specific struct for precomputed unary params
8struct htp_unary_kernel_params {
9    uint32_t  n_threads;
10    uint32_t  col_tile;
11    uint32_t  vtcm_row_per_thread;
12    uint32_t  block;
13    uint32_t  broadcast_weight;
14
15    uint32_t  vtcm_src0_size_per_thread;
16    uint32_t  vtcm_src1_size_per_thread;
17    uint32_t  vtcm_dst_size_per_thread;
18
19    uint32_t  vtcm_src0_size;
20    uint32_t  vtcm_src1_size;
21    uint32_t  vtcm_dst_size;
22
23    uint32_t  src0_row_size_aligned;
24    uint32_t  src1_row_size_aligned;
25    uint32_t  dst_row_size_aligned;
26
27    uint32_t  vtcm_size;
28
29    // Fastdiv helpers
30    struct fastdiv_values div_ne01;
31    struct fastdiv_values div_ne02;
32    struct fastdiv_values div_ne012;
33    struct fastdiv_values div_tpr;
34};
35
36#if defined(__cplusplus)
37static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
38#else
39_Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
40#endif
41
42static inline bool htp_op_is_unary(uint32_t opcode) {
43    switch (opcode) {
44        case HTP_OP_CLAMP:
45        case HTP_OP_LEAKY_RELU:
46        case HTP_OP_NORM:
47        case HTP_OP_RMS_NORM:
48        case HTP_OP_RMS_NORM_MUL:
49        case HTP_OP_SCALE:
50        case HTP_OP_SQR:
51        case HTP_OP_SQRT:
52        case HTP_OP_UNARY_NEG:
53        case HTP_OP_UNARY_EXP:
54        case HTP_OP_UNARY_SIGMOID:
55        case HTP_OP_UNARY_SILU:
56        case HTP_OP_UNARY_GELU:
57        case HTP_OP_UNARY_SOFTPLUS:
58        case HTP_OP_UNARY_TANH:
59        case HTP_OP_UNARY_ABS:
60        case HTP_OP_UNARY_LOG:
61        case HTP_OP_UNARY_RELU:
62        case HTP_OP_L2_NORM:
63        case HTP_OP_TRI:
64            return true;
65        default:
66            return false;
67    }
68}
69
70struct htp_unary_vtcm_layout {
71    size_t total_bytes;
72    size_t off_src0;
73    size_t off_src1;
74    size_t off_dst;
75
76    size_t src0_bytes;
77    size_t src1_bytes;
78    size_t dst_bytes;
79};
80
81static inline void htp_unary_vtcm_layout_build(
82    struct htp_unary_vtcm_layout * L,
83    uint32_t op,
84    uint32_t ne00,
85    uint32_t ne10,
86    uint32_t ne11,
87    bool broadcast_weight,
88    uint32_t n_threads,
89    size_t vtcm_size,
90    size_t elem_size,
91    uint32_t * out_col_tile,
92    uint32_t * out_vtcm_row_per_thread
93) {
94    const size_t src0_data_row_size = ne00 * elem_size;
95    const size_t dst_data_row_size  = ne10 * elem_size;
96
97    const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
98    const size_t dst_row_size_aligned  = hex_round_up(dst_data_row_size,  128);
99
100    size_t src1_row_size_aligned = 0;
101    if (op == HTP_OP_RMS_NORM_MUL) {
102        // RMS_NORM_MUL fusion is F32-only; its weight tensor is always F32.
103        const size_t src1_data_row_size = ne11 * sizeof(float);
104        src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
105    }
106
107    size_t vtcm_size_per_row = 0;
108    size_t vtcm_row_per_thread = 0;
109
110    if (op == HTP_OP_RMS_NORM_MUL) {
111        if (broadcast_weight) {
112            size_t available_vtcm = vtcm_size;
113            size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
114            if (available_vtcm > src1_vtcm_total) {
115                available_vtcm -= src1_vtcm_total;
116            } else {
117                available_vtcm = 0;
118            }
119            vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
120            vtcm_row_per_thread = available_vtcm / (n_threads * vtcm_size_per_row);
121        } else {
122            vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned + src1_row_size_aligned);
123            vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
124        }
125    } else {
126        vtcm_size_per_row   = 2 * (src0_row_size_aligned + dst_row_size_aligned);
127        vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
128    }
129
130    const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
131                               op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM);
132    // The tiled fallback path below only has F32 task functions (unary_task_f32_tiled_*);
133    // F16 has no tiled kernels, so it must stay on the row-block path like reduction ops.
134    // NOTE: if F16 ends up with vtcm_row_per_thread == 0 here (row too large for the VTCM
135    // budget), execute_op_unary() will see BLOCK == 0 and skip computation for that op
136    // (logged via FARF(ERROR, ...)) since there is no F16 tiled fallback. This is a known
137    // limitation; supporting it would require adding F16 tiled kernels.
138    const bool is_f16 = (elem_size == sizeof(_Float16));
139    uint32_t col_tile = 0;
140
141    if (vtcm_row_per_thread == 0 && !is_reduction && !is_f16) {
142        const size_t per_thread_budget = vtcm_size / n_threads;
143        const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
144        col_tile = (uint32_t) (col_tile_bytes / elem_size);
145
146        L->src0_bytes = col_tile_bytes * 2;
147        L->dst_bytes  = col_tile_bytes * 2;
148        L->src1_bytes = 0;
149    } else {
150        L->src0_bytes = src0_row_size_aligned * vtcm_row_per_thread * 2;
151        L->dst_bytes  = dst_row_size_aligned * vtcm_row_per_thread * 2;
152        if (op == HTP_OP_RMS_NORM_MUL) {
153            if (broadcast_weight) {
154                L->src1_bytes = src1_row_size_aligned;
155            } else {
156                L->src1_bytes = src1_row_size_aligned * vtcm_row_per_thread * 2;
157            }
158        } else {
159            L->src1_bytes = 0;
160        }
161    }
162
163    L->off_src0 = 0;
164    if (op == HTP_OP_RMS_NORM_MUL) {
165        L->off_src1 = L->off_src0 + L->src0_bytes * n_threads;
166        L->off_dst  = L->off_src1 + L->src1_bytes * n_threads;
167    } else {
168        L->off_src1 = 0;
169        L->off_dst  = L->off_src0 + L->src0_bytes * n_threads;
170    }
171
172    L->total_bytes = L->off_dst + L->dst_bytes * n_threads;
173
174    *out_col_tile = col_tile;
175    *out_vtcm_row_per_thread = vtcm_row_per_thread;
176}
177
178#endif /* HTP_UNARY_OPS_H */