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 →
9.0 KiB181 linesraw
1#ifndef HVX_SIGMOID_H
2#define HVX_SIGMOID_H
3
4#include "hvx-base.h"
5#include "hvx-inverse.h"
6#include "hvx-exp.h"
7
8#define FAST_SIGMOID_LOG2F (0x3fb8aa3b)  // 1.442695022
9#define FAST_SIGMOID_C1    (0x3d009076)  // 0.03138777
10#define FAST_SIGMOID_C2    (0x3e8d74bd)  // 0.276281267
11#define FAST_SIGMOID_C3    (0x3f000000)  // 0.5
12
13static inline HVX_Vector hvx_vec_fast_sigmoid_f32(HVX_Vector v) {
14    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
15    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
16
17    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
18    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
19    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
20
21    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
22    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
23
24    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
25    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
26    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
27
28    HVX_Vector v3          = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
29    HVX_Vector v3_exponent = Q6_Vw_vasl_VwR(v3, 1);
30    v3_exponent            = Q6_Vuw_vlsr_VuwR(v3_exponent, 24);
31    v3_exponent            = Q6_Vw_vadd_VwVw(in_int, v3_exponent);
32    v3                     = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
33
34    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
35    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
36
37    HVX_Vector res = hvx_vec_inverse_f32(v5);
38    res            = Q6_Vqf32_vmpy_VsfVsf(v3, res);
39
40    return Q6_Vsf_equals_Vqf32(res);
41}
42
43static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v,
44                                                         HVX_Vector one,
45                                                         HVX_Vector max_exp,
46                                                         HVX_Vector min_exp) {
47    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
48    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
49
50    HVX_Vector out = hvx_vec_fast_sigmoid_f32(v);
51    out            = Q6_V_vmux_QVV(pred_max, out, one);
52    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
53}
54
55static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
56    // tanh(x) = 2 * sigmoid(2x) - 1
57    HVX_Vector two = hvx_vec_splat_f32(2.0f);
58    HVX_Vector one = hvx_vec_splat_f32(1.0f);
59    HVX_Vector x2  = Q6_Vqf32_vmpy_VsfVsf(x, two);
60
61    HVX_Vector max_exp = hvx_vec_splat_f32(87.f);
62    HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);
63
64    HVX_Vector sig2x = hvx_vec_fast_sigmoid_f32_guard(Q6_Vsf_equals_Vqf32(x2), one, max_exp, min_exp);
65
66    HVX_Vector res = Q6_Vqf32_vmpy_VsfVsf(sig2x, two);
67    res = Q6_Vqf32_vsub_Vqf32Vsf(res, one);
68    return Q6_Vsf_equals_Vqf32(res);
69}
70
71#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)                                  \
72    do {                                                                                      \
73        dst_type * restrict vdst = (dst_type *) dst;                                          \
74        src_type * restrict vsrc = (src_type *) src;                                          \
75                                                                                              \
76        const HVX_Vector one     = hvx_vec_splat_f32(1.f);                                    \
77        const HVX_Vector max_exp = hvx_vec_splat_f32(87.f);                                   \
78        const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);                                  \
79                                                                                              \
80        const uint32_t epv  = 128 / sizeof(float);                                            \
81        const uint32_t nvec = n / epv;                                                        \
82        const uint32_t nloe = n % epv;                                                        \
83                                                                                              \
84        uint32_t i = 0;                                                                       \
85                                                                                              \
86        _Pragma("unroll(4)")                                                                  \
87        for (; i < nvec; i++) {                                                               \
88             vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp);        \
89        }                                                                                     \
90        if (nloe) {                                                                           \
91             HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \
92             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp);                         \
93        }                                                                                     \
94    } while(0)
95
96#define hvx_tanh_loop_body(dst_type, src_type, vec_store)             \
97    do {                                                              \
98        dst_type * restrict vdst = (dst_type *) dst;                  \
99        src_type * restrict vsrc = (src_type *) src;                  \
100                                                                      \
101        const uint32_t epv  = 128 / sizeof(float);                    \
102        const uint32_t nvec = n / epv;                                \
103        const uint32_t nloe = n % epv;                                \
104                                                                      \
105        uint32_t i = 0;                                               \
106                                                                      \
107        _Pragma("unroll(4)")                                          \
108        for (; i < nvec; i++) {                                       \
109             vdst[i] = hvx_vec_tanh_f32(vsrc[i]);                     \
110        }                                                             \
111        if (nloe) {                                                   \
112             HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]);              \
113             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \
114        }                                                             \
115    } while(0)
116
117static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
118    assert((unsigned long) dst % 128 == 0);
119    assert((unsigned long) src % 128 == 0);
120    hvx_sigmoid_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
121}
122
123static inline void hvx_sigmoid_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
124    assert((unsigned long) dst % 128 == 0);
125    hvx_sigmoid_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
126}
127
128static inline void hvx_sigmoid_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
129    assert((unsigned long) src % 128 == 0);
130    hvx_sigmoid_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
131}
132
133static inline void hvx_sigmoid_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
134    hvx_sigmoid_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
135}
136
137static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
138    assert((unsigned long) dst % 128 == 0);
139    assert((unsigned long) src % 128 == 0);
140    hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
141}
142
143static inline HVX_Vector hvx_vec_fast_sigmoid_f16(HVX_Vector x_v) {
144    const HVX_Vector v_one       = hvx_vec_splat_f16(1.0f);
145    const HVX_Vector v_neg_log2e = hvx_vec_splat_f16(-EXP_LOG2E_F);
146    const HVX_Vector em_mask     = Q6_Vh_vsplat_R(0x7FFF);
147
148    // Compute absolute value of x_v
149    HVX_Vector abs_x = Q6_V_vand_VV(x_v, em_mask);
150
151    // Compute u = -abs_x * log2(e) <= 0.
152    HVX_Vector u = hvx_vec_mul_f16_f16(abs_x, v_neg_log2e);
153
154    // Clamp input to prevent underflow in exp2
155    const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
156    u = Q6_Vhf_vmax_VhfVhf(v_clamp_min, u);
157
158    HVX_Vector exp_val = hvx_vec_exp2_f16(u);
159    HVX_Vector denom   = hvx_vec_add_f16_f16(v_one, exp_val);
160    HVX_Vector sig_abs = hvx_vec_inverse_f16(denom);
161
162    // check if x_v < 0 (using integer comparison on absolute value)
163    HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(abs_x, x_v);
164
165    // If x_v < 0, return 1.0f - sig_abs
166    HVX_Vector sig_neg = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(v_one, sig_abs));
167    return Q6_V_vmux_QVV(is_neg, sig_neg, sig_abs);
168}
169
170static inline HVX_Vector hvx_vec_tanh_f16(HVX_Vector x) {
171    // tanh(x) = 2 * sigmoid(2x) - 1
172    const HVX_Vector v_two = hvx_vec_splat_f16(2.0f);
173
174    HVX_Vector x2 = hvx_vec_mul_f16_f16(x, v_two);
175    HVX_Vector sig2x = hvx_vec_fast_sigmoid_f16(x2);
176
177    const HVX_Vector v_neg_one = hvx_vec_splat_f16(-1.0f);
178    return hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(sig2x, v_two), v_neg_one);
179}
180
181#endif /* HVX_SIGMOID_H */