char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
f027c4f1b
char/aliceai-t5-moe
1#ifndef HVX_SQRT_H 2#define HVX_SQRT_H 3 4#include <stdbool.h> 5#include <stdint.h> 6 7#include "hex-utils.h" 8 9#include "hvx-base.h" 10 11#define RSQRT_CONST 0x5f3759df // Constant for fast inverse square root calculation 12#define RSQRT_ONE_HALF 0x3f000000 // 0.5 13#define RSQRT_THREE_HALVES 0x3fc00000 // 1.5 14 15#if __HVX_ARCH__ < 79 16#define HVX_OP_MUL (a ,b ) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) 17#else 18#define HVX_OP_MUL (a ,b ) Q6_Vsf_vmpy_VsfVsf(a, b) 19#endif 20 21static inline HVX_Vector hvx_vec_rsqrt_f32 (HVX_Vector in_vec ) { 22//Algorithm : 23// x2 = input*0.5 24// y = * (long *) &input 25// y = 0x5f3759df - (y>>1) 26// y = y*(threehalfs - x2*y*y) 27 28HVX_Vector rsqrtconst = Q6_V_vsplat_R (RSQRT_CONST ); 29HVX_Vector onehalf = Q6_V_vsplat_R (RSQRT_ONE_HALF ); 30HVX_Vector threehalfs = Q6_V_vsplat_R (RSQRT_THREE_HALVES ); 31 32HVX_Vector x2 ,y ,ypower2 ,temp ; 33 34x2 = Q6_Vqf32_vmpy_VsfVsf (in_vec ,onehalf ); 35x2 = Q6_Vqf32_vadd_Vqf32Vsf (x2 ,Q6_V_vzero ()); 36 37y = Q6_Vw_vasr_VwR (in_vec ,1 ); 38y = Q6_Vw_vsub_VwVw (rsqrtconst ,y ); 39 40// 1st iteration 41ypower2 = Q6_Vqf32_vmpy_VsfVsf (y ,y ); 42ypower2 = Q6_Vqf32_vadd_Vqf32Vsf (ypower2 ,Q6_V_vzero ()); 43temp = Q6_Vqf32_vmpy_Vqf32Vqf32 (x2 ,ypower2 ); 44temp = Q6_Vqf32_vsub_VsfVsf (threehalfs ,Q6_Vsf_equals_Vqf32 (temp )); 45temp = Q6_Vqf32_vmpy_VsfVsf (y ,Q6_Vsf_equals_Vqf32 (temp )); 46 47// 2nd iteration 48y = Q6_Vqf32_vadd_Vqf32Vsf (temp ,Q6_V_vzero ()); 49ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32 (y ,y ); 50ypower2 = Q6_Vqf32_vadd_Vqf32Vsf (ypower2 ,Q6_V_vzero ()); 51temp = Q6_Vqf32_vmpy_Vqf32Vqf32 (x2 ,ypower2 ); 52temp = Q6_Vqf32_vsub_VsfVsf (threehalfs ,Q6_Vsf_equals_Vqf32 (temp )); 53temp = Q6_Vqf32_vmpy_Vqf32Vqf32 (y ,temp ); 54 55// 3rd iteration 56y = Q6_Vqf32_vadd_Vqf32Vsf (temp ,Q6_V_vzero ()); 57ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32 (y ,y ); 58ypower2 = Q6_Vqf32_vadd_Vqf32Vsf (ypower2 ,Q6_V_vzero ()); 59temp = Q6_Vqf32_vmpy_Vqf32Vqf32 (x2 ,ypower2 ); 60temp = Q6_Vqf32_vsub_VsfVsf (threehalfs ,Q6_Vsf_equals_Vqf32 (temp )); 61temp = Q6_Vqf32_vmpy_Vqf32Vqf32 (y ,temp ); 62 63return Q6_Vsf_equals_Vqf32 (temp ); 64} 65 66// Compute sqrt(x) as x*inv_sqrt(x) 67#define hvx_sqrt_f32_loop_body (dst_type ,src_type ,vec_store ) \ 68 do { \ 69 dst_type * restrict vdst = (dst_type *) dst; \ 70 src_type * restrict vsrc = (src_type *) src; \ 71 \ 72 const uint32_t nvec = n / VLEN_FP32; \ 73 const uint32_t nloe = n % VLEN_FP32; \ 74 \ 75 uint32_t i = 0; \ 76 \ 77 _Pragma("unroll(4)") \ 78 for (; i < nvec; i++) { \ 79 HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ 80 HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ 81 vdst[i] = sqrt_res; \ 82 } \ 83 if (nloe) { \ 84 HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ 85 HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ 86 vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, sqrt_res); \ 87 } \ 88 } while(0) 89 90static inline void hvx_sqrt_f32_aa (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 91assert ((unsigned long )dst %128 == 0 ); 92assert ((unsigned long )src %128 == 0 ); 93hvx_sqrt_f32_loop_body (HVX_Vector ,HVX_Vector ,hvx_vec_store_a ); 94} 95 96static inline void hvx_sqrt_f32_au (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 97assert ((unsigned long )dst %128 == 0 ); 98hvx_sqrt_f32_loop_body (HVX_Vector ,HVX_UVector ,hvx_vec_store_a ); 99} 100 101static inline void hvx_sqrt_f32_ua (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 102assert ((unsigned long )src %128 == 0 ); 103hvx_sqrt_f32_loop_body (HVX_UVector ,HVX_Vector ,hvx_vec_store_u ); 104} 105 106static inline void hvx_sqrt_f32_uu (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 107hvx_sqrt_f32_loop_body (HVX_UVector ,HVX_UVector ,hvx_vec_store_u ); 108} 109 110static inline void hvx_sqrt_f32 (uint8_t * restrictdst ,const uint8_t * restrictsrc ,const int num_elems ) { 111if ((unsigned long )dst %128 == 0 ) { 112if ((unsigned long )src %128 == 0 ) { 113hvx_sqrt_f32_aa (dst ,src ,num_elems ); 114 }else { 115hvx_sqrt_f32_au (dst ,src ,num_elems ); 116 } 117 }else { 118if ((unsigned long )src %128 == 0 ) { 119hvx_sqrt_f32_ua (dst ,src ,num_elems ); 120 }else { 121hvx_sqrt_f32_uu (dst ,src ,num_elems ); 122 } 123 } 124} 125 126// Compute sqrt(x) for f16 by promoting to f32, applying hvx_vec_rsqrt_f32, and narrowing back. 127#define hvx_sqrt_f16_loop_body (dst_type ,src_type ,vec_store ) \ 128 do { \ 129 dst_type * restrict vdst = (dst_type *) dst; \ 130 src_type * restrict vsrc = (src_type *) src; \ 131 \ 132 const uint32_t nvec = n / VLEN_FP16; \ 133 const uint32_t nloe = n % VLEN_FP16; \ 134 \ 135 uint32_t i = 0; \ 136 \ 137 _Pragma("unroll(4)") \ 138 for (; i < nvec; i++) { \ 139 HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]); \ 140 HVX_Vector r0 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_lo_W(p)), Q6_V_lo_W(p)); \ 141 HVX_Vector r1 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_hi_W(p)), Q6_V_hi_W(p)); \ 142 vdst[i] = hvx_vec_f32_to_f16(r0, r1); \ 143 } \ 144 if (nloe) { \ 145 HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]); \ 146 HVX_Vector r0 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_lo_W(p)), Q6_V_lo_W(p)); \ 147 HVX_Vector r1 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_hi_W(p)), Q6_V_hi_W(p)); \ 148 HVX_Vector v = hvx_vec_f32_to_f16(r0, r1); \ 149 vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v); \ 150 } \ 151 } while(0) 152 153static inline void hvx_sqrt_f16_aa (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 154assert ((unsigned long )dst %128 == 0 ); 155assert ((unsigned long )src %128 == 0 ); 156hvx_sqrt_f16_loop_body (HVX_Vector ,HVX_Vector ,hvx_vec_store_a ); 157} 158 159static inline void hvx_sqrt_f16_au (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 160assert ((unsigned long )dst %128 == 0 ); 161hvx_sqrt_f16_loop_body (HVX_Vector ,HVX_UVector ,hvx_vec_store_a ); 162} 163 164static inline void hvx_sqrt_f16_ua (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 165assert ((unsigned long )src %128 == 0 ); 166hvx_sqrt_f16_loop_body (HVX_UVector ,HVX_Vector ,hvx_vec_store_u ); 167} 168 169static inline void hvx_sqrt_f16_uu (uint8_t * restrictdst ,const uint8_t * restrictsrc ,uint32_t n ) { 170hvx_sqrt_f16_loop_body (HVX_UVector ,HVX_UVector ,hvx_vec_store_u ); 171} 172 173static inline void hvx_sqrt_f16 (uint8_t * restrictdst ,const uint8_t * restrictsrc ,const int num_elems ) { 174if ((unsigned long )dst %128 == 0 ) { 175if ((unsigned long )src %128 == 0 ) { 176hvx_sqrt_f16_aa (dst ,src ,num_elems ); 177 }else { 178hvx_sqrt_f16_au (dst ,src ,num_elems ); 179 } 180 }else { 181if ((unsigned long )src %128 == 0 ) { 182hvx_sqrt_f16_ua (dst ,src ,num_elems ); 183 }else { 184hvx_sqrt_f16_uu (dst ,src ,num_elems ); 185 } 186 } 187} 188 189#endif /* HVX_SQRT_H */