| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #pragma once |
| 7 | |
| 8 | #if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__) |
| 9 | #define CPU_CAPABILITY_AVX512 |
| 10 | #endif |
| 11 | |
| 12 | #if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) |
| 13 | #define CPU_CAPABILITY_RVV |
| 14 | #endif |
| 15 | |
| 16 | #include <ATen/cpu/vec/functional.h> |
| 17 | #include <ATen/cpu/vec/vec.h> |
| 18 | #if defined(CPU_CAPABILITY_AVX512) |
| 19 | #include <immintrin.h> |
| 20 | #endif |
| 21 | |
| 22 | #if defined(CPU_CAPABILITY_RVV) |
| 23 | #include "../cpu_types_riscv_defs.hpp" |
| 24 | #endif |
| 25 | namespace { |
| 26 | |
| 27 | using namespace at::vec; |
| 28 | |
| 29 | template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0> |
| 30 | inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) { |
| 31 | return at::vec::convert_from_float<scalar_t>(a, b); |
| 32 | } |
| 33 | |
| 34 | template <typename scalar_t> |
| 35 | inline void store_from_float_ext(scalar_t* out, const Vectorized<float>& a) { |
| 36 | float out_buffer[Vectorized<float>::size()]; |
| 37 | a.store(out_buffer); |
| 38 | for (int i = 0; i < Vectorized<float>::size(); ++i) { |
| 39 | out[i] = static_cast<scalar_t>(out_buffer[i]); |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | // allow f16, bf16 |
| 44 | template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1> |
| 45 | inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) { |
| 46 | using bVec = at::vec::Vectorized<scalar_t>; |
| 47 | using fVec = at::vec::Vectorized<float>; |
| 48 | bVec x_vec = bVec::loadu(data); |
| 49 | fVec x0, x1; |
| 50 | std::tie(x0, x1) = at::vec::convert_to_float(x_vec); |
| 51 | return std::make_tuple(x0, x1); |
| 52 | } |
| 53 | |
| 54 | // allow f32 |
| 55 | inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const float* __restrict__ data) { |
| 56 | using fVec = at::vec::Vectorized<float>; |
| 57 | fVec x0 = fVec::loadu(data); |
| 58 | fVec x1 = fVec::loadu(data + fVec::size()); |
| 59 | return std::make_tuple(x0, x1); |
| 60 | } |
| 61 | |
| 62 | template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1> |
| 63 | inline at::vec::Vectorized<float> load_float_vec(const scalar_t* __restrict__ data) { |
| 64 | at::vec::Vectorized<float> out; |
| 65 | if constexpr (std::is_same_v<scalar_t, at::BFloat16>) { |
| 66 | at::vec::load_fp32_from_bf16(data, out); |
| 67 | } else { |
| 68 | at::vec::load_fp32_from_fp16(data, out); |
| 69 | } |
| 70 | return out; |
| 71 | } |
| 72 | |
| 73 | #if defined(CPU_CAPABILITY_AVX512) |
| 74 | |
| 75 | // `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics |
| 76 | // use native instruction for bfloat16->float32 conversion |
| 77 | template <> |
| 78 | inline Vectorized<at::BFloat16> |
| 79 | convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) { |
| 80 | return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a))); |
| 81 | } |
| 82 | |
| 83 | template <> |
| 84 | inline void store_from_float_ext<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) { |
| 85 | _mm256_storeu_si256(reinterpret_cast<__m256i*>(out), (__m256i)(_mm512_cvtneps_pbh(__m512(a)))); |
| 86 | } |
| 87 | |
| 88 | template <> |
| 89 | inline void store_from_float_ext<at::Half>(at::Half* out, const Vectorized<float>& a) { |
| 90 | _mm256_storeu_si256( |
| 91 | reinterpret_cast<__m256i*>(out), _mm512_cvtps_ph(__m512(a), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); |
| 92 | } |
| 93 | |
| 94 | #define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16)) |
| 95 | |
| 96 | #define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a) |
| 97 | |
| 98 | // this doesn't handle NaN. |
| 99 | inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) { |
| 100 | const __m512i x = _mm512_cvtepu8_epi16(fp8_vec); |
| 101 | __m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780)); |
| 102 | combined = _mm512_slli_epi16(combined, 4); |
| 103 | combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0)); |
| 104 | combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00)); |
| 105 | |
| 106 | const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512()); |
| 107 | return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined); |
| 108 | } |
| 109 | |
| 110 | inline __m512bh cvt_e4m3_bf16_intrinsic_without_denorm(__m256i fp8_vec) { |
| 111 | // The following conversion is without denorm behavior, that is to say, |
| 112 | // Max subnorm : S.0000.111 = 0.875 ∗ 2**(−6) |
| 113 | // Min subnorm : S.0000.001 = 2**(−9) |
| 114 | // 0.0019 ~ 0.0137 cannot be converted correctly. |
| 115 | __m512i x = _mm512_cvtepu8_epi16(fp8_vec); |
| 116 | auto mask = _mm512_cmpneq_epi16_mask( |
| 117 | _mm512_and_si512(x, _mm512_set1_epi16(127)), |
| 118 | _mm512_setzero_si512()); // mask = x & 0x7f |
| 119 | auto mask_nan = _mm512_cmpneq_epi16_mask( |
| 120 | _mm512_and_si512(x, _mm512_set1_epi16(127)), |
| 121 | _mm512_set1_epi16(127)); // mask_nan = x & 0x7f |
| 122 | auto mantissa = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4); // mantissa = (x & 7) << 4 |
| 123 | auto exponent = _mm512_add_epi16( |
| 124 | _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), |
| 125 | _mm512_set1_epi16(120)); // exponent = (((x >> 3) & 15) + 120) |
| 126 | auto nonsign = _mm512_maskz_mov_epi16(mask, _mm512_or_si512(mantissa, _mm512_slli_epi16(exponent, 7))); |
| 127 | nonsign = _mm512_mask_mov_epi16(_mm512_set1_epi16(0x7fff), mask_nan, nonsign); // deal with Nan |
| 128 | return (__m512bh)(_mm512_or_si512( |
| 129 | nonsign, |
| 130 | _mm512_slli_epi16( |
| 131 | _mm512_and_si512(x, _mm512_set1_epi16(128)), |
| 132 | 8))); // add sign (x & 128) << 8 |
| 133 | } |
| 134 | |
| 135 | inline __m512bh cvt_e4m3_bf16_intrinsic_with_denorm(__m256i fp8_vec) { |
| 136 | __m512i x = _mm512_cvtepu8_epi16(fp8_vec); |
| 137 | __m512i lg2mant = _mm512_mask_mov_epi16( |
| 138 | _mm512_mask_mov_epi16( |
| 139 | _mm512_setzero_si512(), _mm512_test_epi16_mask(x, _mm512_set1_epi16(2)), _mm512_set1_epi16(1)), |
| 140 | _mm512_test_epi16_mask(x, _mm512_set1_epi16(4)), |
| 141 | _mm512_set1_epi16(2)); |
| 142 | return (__m512bh)(_mm512_or_si512( |
| 143 | _mm512_maskz_mov_epi16( |
| 144 | _mm512_cmpneq_epi16_mask(_mm512_and_si512(x, _mm512_set1_epi16(127)), _mm512_setzero_si512()), |
| 145 | _mm512_mask_blend_epi16( |
| 146 | _mm512_test_epi16_mask(x, _mm512_set1_epi16(120)), |
| 147 | _mm512_or_si512( |
| 148 | _mm512_and_si512( |
| 149 | _mm512_sllv_epi16( |
| 150 | _mm512_and_si512(x, _mm512_set1_epi16(3)), _mm512_sub_epi16(_mm512_set1_epi16(7), lg2mant)), |
| 151 | _mm512_set1_epi16(0x007f)), |
| 152 | _mm512_slli_epi16(_mm512_add_epi16(lg2mant, _mm512_set1_epi16(118)), 7)), |
| 153 | _mm512_or_si512( |
| 154 | _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4), |
| 155 | _mm512_slli_epi16( |
| 156 | _mm512_add_epi16( |
| 157 | _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), _mm512_set1_epi16(120)), |
| 158 | 7)))), |
| 159 | _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(128)), 8))); |
| 160 | } |
| 161 | |
| 162 | inline __m512bh CVT_FP8_TO_BF16(__m256i a) { |
| 163 | #ifdef SGLANG_CPU_FP8_CVT_FTZ |
| 164 | return cvt_e4m3_bf16_intrinsic_no_nan(a); |
| 165 | #else |
| 166 | return cvt_e4m3_bf16_intrinsic_with_denorm(a); |
| 167 | #endif |
| 168 | } |
| 169 | // faster version of float8_e4m3fn conversion to bfloat16 |
| 170 | // |
| 171 | // we mapped cuda implementation from below link and vectorized with avx512: |
| 172 | // https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387 |
| 173 | // |
| 174 | inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) { |
| 175 | const __m512i mask0 = _mm512_set1_epi16(0x80); // sign bit |
| 176 | const __m512i mask1 = _mm512_set1_epi16(0x7F); // exponent and mantissa |
| 177 | const __m512i mask2 = _mm512_set1_epi16(0x4000); |
| 178 | |
| 179 | __m512i x = _mm512_cvtepu8_epi16(a); |
| 180 | __m512i vsign = _mm512_and_si512(x, mask0); |
| 181 | vsign = _mm512_slli_epi16(vsign, 8); |
| 182 | |
| 183 | __m512i vexp_and_mant = _mm512_and_si512(x, mask1); |
| 184 | vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4); |
| 185 | |
| 186 | // _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110 |
| 187 | return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110)); |
| 188 | } |
| 189 | |
| 190 | // bias for conversion of fp8 to bf16 1/256 in float32 |
| 191 | #define kFP8_BIAS 0x3b800000 |
| 192 | |
| 193 | // remove warning: ignoring attributes on template argument ‘__m512bh’ [-Wignored-attributes] |
| 194 | #pragma GCC diagnostic push |
| 195 | #pragma GCC diagnostic ignored "-Wignored-attributes" |
| 196 | |
| 197 | #define MXFP4_VALUES \ |
| 198 | -6.0f, -4.0f, -3.0f, -2.0f, -1.5f, -1.0f, -0.5f, -0.0f, 6.0f, 4.0f, 3.0f, 2.0f, 1.5f, 1.0f, 0.5f, 0.0f |
| 199 | |
| 200 | // convert 64 mxfp4 to 2x bf16 vectors, expect input 32-way packing |
| 201 | inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i a, __m512i s0, __m512i s1) { |
| 202 | // LUT |
| 203 | const __m512 values = _mm512_set_ps(MXFP4_VALUES); |
| 204 | const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values)); |
| 205 | |
| 206 | const __m512i abs_mask = _mm512_set1_epi16(0x7FFF); |
| 207 | |
| 208 | // expand values to 16-bit integers |
| 209 | __m512i x0 = _mm512_cvtepu8_epi16(a); |
| 210 | __m512i x1 = _mm512_srli_epi32(x0, 4); |
| 211 | |
| 212 | // LUT to convert mxfp4 values to bf16 |
| 213 | x0 = _mm512_permutexvar_epi16(x0, lut); |
| 214 | x1 = _mm512_permutexvar_epi16(x1, lut); |
| 215 | |
| 216 | // Emulate the bf16 multiply by the E8M0 scale as an integer add on the |
| 217 | // exponent field, forcing zeros to stay zero. |
| 218 | // |
| 219 | // vptestmw sets a lane's mask bit when (x & 0x7FFF) != 0, i.e. for the lanes |
| 220 | // that are not +0.0 or -0.0. That is the *complement* of the mask the old |
| 221 | // and/cmp/add/blend sequence built, which selected the zero lanes; the |
| 222 | // zero-masking form of vpaddw then writes 0 where the mask is clear, so the |
| 223 | // two agree lane by lane. That includes the LUT's -0.0 entry, which both |
| 224 | // versions turn into +0.0. |
| 225 | // |
| 226 | // 2 instructions per vector instead of 4, no extra ISA requirement: |
| 227 | // vptestmw is AVX512BW, which vpermw above already needs. |
| 228 | __mmask32 mask0 = _mm512_test_epi16_mask(x0, abs_mask); |
| 229 | __mmask32 mask1 = _mm512_test_epi16_mask(x1, abs_mask); |
| 230 | |
| 231 | x0 = _mm512_maskz_add_epi16(mask0, x0, s0); |
| 232 | x1 = _mm512_maskz_add_epi16(mask1, x1, s1); |
| 233 | |
| 234 | return std::make_tuple(__m512bh(x0), __m512bh(x1)); |
| 235 | } |
| 236 | |
| 237 | #define CVT_MXFP4_TO_BF16(a, s0, s1) cvt_mxfp4_e2m1_bf16_intrinsic_lut(a, s0, s1) |
| 238 | |
| 239 | #pragma GCC diagnostic pop |
| 240 | |
| 241 | #endif |
| 242 | |
| 243 | // vector to scalar reduction |
| 244 | #if defined(CPU_CAPABILITY_AVX512) |
| 245 | inline float vec_reduce_sum(const Vectorized<float>& a) { |
| 246 | return _mm512_reduce_add_ps(__m512(a)); |
| 247 | } |
| 248 | |
| 249 | inline float vec_reduce_max(const Vectorized<float>& a) { |
| 250 | return _mm512_reduce_max_ps(__m512(a)); |
| 251 | } |
| 252 | #else |
| 253 | inline float vec_reduce_sum(const Vectorized<float>& a) { |
| 254 | return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return x + y; }, a); |
| 255 | } |
| 256 | |
| 257 | inline float vec_reduce_max(const Vectorized<float>& a) { |
| 258 | return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return maximum(x, y); }, a); |
| 259 | } |
| 260 | #endif |
| 261 | |
| 262 | // https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 |
| 263 | template <typename scalar_t> |
| 264 | inline void |
| 265 | quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) { |
| 266 | float amax = 0.f; // absolute max |
| 267 | for (int64_t k = 0; k < K; ++k) { |
| 268 | const float val = static_cast<float>(A[k]); |
| 269 | amax = std::max(amax, std::abs(val)); |
| 270 | } |
| 271 | |
| 272 | amax = std::max(amax, eps); |
| 273 | const float scale = amax / 127; |
| 274 | const float inv_scale = 127 / amax; |
| 275 | |
| 276 | for (int64_t k = 0; k < K; ++k) { |
| 277 | const float val = static_cast<float>(A[k]) * inv_scale; |
| 278 | Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128); |
| 279 | } |
| 280 | As = scale; |
| 281 | } |
| 282 | |
| 283 | #if defined(CPU_CAPABILITY_AVX512) |
| 284 | template <> |
| 285 | inline void quantize_row_int8<at::BFloat16>( |
| 286 | uint8_t* __restrict__ Aq, float& As, const at::BFloat16* __restrict__ A, int64_t K, float eps) { |
| 287 | const __m512 signBit = _mm512_set1_ps(-0.0f); |
| 288 | const __m512i off = _mm512_set1_epi32(128); |
| 289 | |
| 290 | // K is 32x, no remainder |
| 291 | float amax = 0.f; |
| 292 | __m512 vamax0 = _mm512_set1_ps(0.f); |
| 293 | __m512 vamax1 = _mm512_set1_ps(0.f); |
| 294 | for (int64_t k = 0; k < K; k += 32) { |
| 295 | __m512i va = _mm512_loadu_si512((void*)(A + k)); |
| 296 | __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0)); |
| 297 | __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1)); |
| 298 | vamax0 = _mm512_max_ps(vamax0, _mm512_andnot_ps(signBit, va0)); |
| 299 | vamax1 = _mm512_max_ps(vamax1, _mm512_andnot_ps(signBit, va1)); |
| 300 | } |
| 301 | amax = _mm512_reduce_max_ps(_mm512_max_ps(vamax0, vamax1)); |
| 302 | amax = std::max(amax, eps); |
| 303 | const float scale = amax / 127; |
| 304 | const float inv_scale = 127 / amax; |
| 305 | const __m512 vd = _mm512_set1_ps(inv_scale); |
| 306 | |
| 307 | for (int64_t k = 0; k < K; k += 32) { |
| 308 | __m512i va = _mm512_loadu_si512((void*)(A + k)); |
| 309 | __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0)); |
| 310 | __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1)); |
| 311 | va0 = _mm512_mul_ps(va0, vd); |
| 312 | va1 = _mm512_mul_ps(va1, vd); |
| 313 | va0 = _mm512_roundscale_ps(va0, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); |
| 314 | va1 = _mm512_roundscale_ps(va1, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); |
| 315 | __m128i i0 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va0), off)); |
| 316 | __m128i i1 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va1), off)); |
| 317 | _mm256_storeu_si256(reinterpret_cast<__m256i*>(Aq + k), _mm256_set_m128i(i1, i0)); |
| 318 | } |
| 319 | As = scale; |
| 320 | } |
| 321 | #endif |
| 322 | |
| 323 | // transpose utils |
| 324 | // taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998 |
| 325 | #if defined(CPU_CAPABILITY_AVX512) |
| 326 | inline void transpose_16x16_16bit(__m256i* v) { |
| 327 | __m256i v1[16]; |
| 328 | v1[0] = _mm256_unpacklo_epi16(v[0], v[1]); |
| 329 | v1[1] = _mm256_unpackhi_epi16(v[0], v[1]); |
| 330 | v1[2] = _mm256_unpacklo_epi16(v[2], v[3]); |
| 331 | v1[3] = _mm256_unpackhi_epi16(v[2], v[3]); |
| 332 | v1[4] = _mm256_unpacklo_epi16(v[4], v[5]); |
| 333 | v1[5] = _mm256_unpackhi_epi16(v[4], v[5]); |
| 334 | v1[6] = _mm256_unpacklo_epi16(v[6], v[7]); |
| 335 | v1[7] = _mm256_unpackhi_epi16(v[6], v[7]); |
| 336 | v1[8] = _mm256_unpacklo_epi16(v[8], v[9]); |
| 337 | v1[9] = _mm256_unpackhi_epi16(v[8], v[9]); |
| 338 | v1[10] = _mm256_unpacklo_epi16(v[10], v[11]); |
| 339 | v1[11] = _mm256_unpackhi_epi16(v[10], v[11]); |
| 340 | v1[12] = _mm256_unpacklo_epi16(v[12], v[13]); |
| 341 | v1[13] = _mm256_unpackhi_epi16(v[12], v[13]); |
| 342 | v1[14] = _mm256_unpacklo_epi16(v[14], v[15]); |
| 343 | v1[15] = _mm256_unpackhi_epi16(v[14], v[15]); |
| 344 | |
| 345 | v[0] = _mm256_unpacklo_epi32(v1[0], v1[2]); |
| 346 | v[1] = _mm256_unpackhi_epi32(v1[0], v1[2]); |
| 347 | v[2] = _mm256_unpacklo_epi32(v1[1], v1[3]); |
| 348 | v[3] = _mm256_unpackhi_epi32(v1[1], v1[3]); |
| 349 | v[4] = _mm256_unpacklo_epi32(v1[4], v1[6]); |
| 350 | v[5] = _mm256_unpackhi_epi32(v1[4], v1[6]); |
| 351 | v[6] = _mm256_unpacklo_epi32(v1[5], v1[7]); |
| 352 | v[7] = _mm256_unpackhi_epi32(v1[5], v1[7]); |
| 353 | v[8] = _mm256_unpacklo_epi32(v1[8], v1[10]); |
| 354 | v[9] = _mm256_unpackhi_epi32(v1[8], v1[10]); |
| 355 | v[10] = _mm256_unpacklo_epi32(v1[9], v1[11]); |
| 356 | v[11] = _mm256_unpackhi_epi32(v1[9], v1[11]); |
| 357 | v[12] = _mm256_unpacklo_epi32(v1[12], v1[14]); |
| 358 | v[13] = _mm256_unpackhi_epi32(v1[12], v1[14]); |
| 359 | v[14] = _mm256_unpacklo_epi32(v1[13], v1[15]); |
| 360 | v[15] = _mm256_unpackhi_epi32(v1[13], v1[15]); |
| 361 | |
| 362 | v1[0] = _mm256_unpacklo_epi64(v[0], v[4]); |
| 363 | v1[1] = _mm256_unpackhi_epi64(v[0], v[4]); |
| 364 | v1[2] = _mm256_unpacklo_epi64(v[1], v[5]); |
| 365 | v1[3] = _mm256_unpackhi_epi64(v[1], v[5]); |
| 366 | v1[4] = _mm256_unpacklo_epi64(v[2], v[6]); |
| 367 | v1[5] = _mm256_unpackhi_epi64(v[2], v[6]); |
| 368 | v1[6] = _mm256_unpacklo_epi64(v[3], v[7]); |
| 369 | v1[7] = _mm256_unpackhi_epi64(v[3], v[7]); |
| 370 | v1[8] = _mm256_unpacklo_epi64(v[8], v[12]); |
| 371 | v1[9] = _mm256_unpackhi_epi64(v[8], v[12]); |
| 372 | v1[10] = _mm256_unpacklo_epi64(v[9], v[13]); |
| 373 | v1[11] = _mm256_unpackhi_epi64(v[9], v[13]); |
| 374 | v1[12] = _mm256_unpacklo_epi64(v[10], v[14]); |
| 375 | v1[13] = _mm256_unpackhi_epi64(v[10], v[14]); |
| 376 | v1[14] = _mm256_unpacklo_epi64(v[11], v[15]); |
| 377 | v1[15] = _mm256_unpackhi_epi64(v[11], v[15]); |
| 378 | |
| 379 | v[0] = _mm256_permute2x128_si256(v1[0], v1[8], 0x20); |
| 380 | v[1] = _mm256_permute2x128_si256(v1[1], v1[9], 0x20); |
| 381 | v[2] = _mm256_permute2x128_si256(v1[2], v1[10], 0x20); |
| 382 | v[3] = _mm256_permute2x128_si256(v1[3], v1[11], 0x20); |
| 383 | v[4] = _mm256_permute2x128_si256(v1[4], v1[12], 0x20); |
| 384 | v[5] = _mm256_permute2x128_si256(v1[5], v1[13], 0x20); |
| 385 | v[6] = _mm256_permute2x128_si256(v1[6], v1[14], 0x20); |
| 386 | v[7] = _mm256_permute2x128_si256(v1[7], v1[15], 0x20); |
| 387 | v[8] = _mm256_permute2x128_si256(v1[0], v1[8], 0x31); |
| 388 | v[9] = _mm256_permute2x128_si256(v1[1], v1[9], 0x31); |
| 389 | v[10] = _mm256_permute2x128_si256(v1[2], v1[10], 0x31); |
| 390 | v[11] = _mm256_permute2x128_si256(v1[3], v1[11], 0x31); |
| 391 | v[12] = _mm256_permute2x128_si256(v1[4], v1[12], 0x31); |
| 392 | v[13] = _mm256_permute2x128_si256(v1[5], v1[13], 0x31); |
| 393 | v[14] = _mm256_permute2x128_si256(v1[6], v1[14], 0x31); |
| 394 | v[15] = _mm256_permute2x128_si256(v1[7], v1[15], 0x31); |
| 395 | } |
| 396 | |
| 397 | inline void transpose_16x16_32bit(__m512i* v) { |
| 398 | __m512i v1[16]; |
| 399 | v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); |
| 400 | v1[1] = _mm512_unpackhi_epi32(v[0], v[1]); |
| 401 | v1[2] = _mm512_unpacklo_epi32(v[2], v[3]); |
| 402 | v1[3] = _mm512_unpackhi_epi32(v[2], v[3]); |
| 403 | v1[4] = _mm512_unpacklo_epi32(v[4], v[5]); |
| 404 | v1[5] = _mm512_unpackhi_epi32(v[4], v[5]); |
| 405 | v1[6] = _mm512_unpacklo_epi32(v[6], v[7]); |
| 406 | v1[7] = _mm512_unpackhi_epi32(v[6], v[7]); |
| 407 | v1[8] = _mm512_unpacklo_epi32(v[8], v[9]); |
| 408 | v1[9] = _mm512_unpackhi_epi32(v[8], v[9]); |
| 409 | v1[10] = _mm512_unpacklo_epi32(v[10], v[11]); |
| 410 | v1[11] = _mm512_unpackhi_epi32(v[10], v[11]); |
| 411 | v1[12] = _mm512_unpacklo_epi32(v[12], v[13]); |
| 412 | v1[13] = _mm512_unpackhi_epi32(v[12], v[13]); |
| 413 | v1[14] = _mm512_unpacklo_epi32(v[14], v[15]); |
| 414 | v1[15] = _mm512_unpackhi_epi32(v[14], v[15]); |
| 415 | |
| 416 | v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]); |
| 417 | v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]); |
| 418 | v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]); |
| 419 | v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]); |
| 420 | v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]); |
| 421 | v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]); |
| 422 | v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]); |
| 423 | v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]); |
| 424 | v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]); |
| 425 | v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]); |
| 426 | v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]); |
| 427 | v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]); |
| 428 | v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]); |
| 429 | v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]); |
| 430 | v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]); |
| 431 | v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]); |
| 432 | |
| 433 | v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88); |
| 434 | v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88); |
| 435 | v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88); |
| 436 | v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88); |
| 437 | v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd); |
| 438 | v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd); |
| 439 | v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd); |
| 440 | v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd); |
| 441 | v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88); |
| 442 | v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88); |
| 443 | v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88); |
| 444 | v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88); |
| 445 | v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd); |
| 446 | v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd); |
| 447 | v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd); |
| 448 | v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd); |
| 449 | |
| 450 | v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); |
| 451 | v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); |
| 452 | v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); |
| 453 | v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); |
| 454 | v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); |
| 455 | v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); |
| 456 | v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); |
| 457 | v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); |
| 458 | v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); |
| 459 | v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); |
| 460 | v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); |
| 461 | v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); |
| 462 | v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); |
| 463 | v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); |
| 464 | v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); |
| 465 | v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); |
| 466 | } |
| 467 | |
| 468 | // remove warning : ignoring attributes on template argument ‘__m512i’ [-Wignored-attributes] |
| 469 | #pragma GCC diagnostic push |
| 470 | #pragma GCC diagnostic ignored "-Wignored-attributes" |
| 471 | |
| 472 | // transpose from [2, 32] to [32, 2] |
| 473 | inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) { |
| 474 | // r0: {a0, a1, ..., a31} |
| 475 | // r1: {b0, b1, ..., b31} |
| 476 | // |
| 477 | // d0: {a0, b0, ..., a15, b15} |
| 478 | // d1: {a16, b16, ..., a31, b31} |
| 479 | // |
| 480 | __m512i d0 = _mm512_unpacklo_epi16(r0, r1); |
| 481 | __m512i d1 = _mm512_unpackhi_epi16(r0, r1); |
| 482 | r0 = _mm512_shuffle_i32x4(d0, d1, 0x88); |
| 483 | r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd); |
| 484 | d0 = _mm512_shuffle_i32x4(r0, r1, 0x88); |
| 485 | d1 = _mm512_shuffle_i32x4(r0, r1, 0xdd); |
| 486 | return std::make_tuple(d0, d1); |
| 487 | } |
| 488 | #pragma GCC diagnostic pop |
| 489 | |
| 490 | // Note: mapped from aten exp_u20 |
| 491 | inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) { |
| 492 | const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); |
| 493 | const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); |
| 494 | const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); |
| 495 | const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); |
| 496 | const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); |
| 497 | const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e) |
| 498 | const __m512 vec_half = _mm512_set1_ps(0.5f); |
| 499 | const __m512 vec_one = _mm512_set1_ps(1.f); |
| 500 | const __m512 vec_zero = _mm512_set1_ps(0.f); |
| 501 | const __m512 vec_two = _mm512_set1_ps(2.f); |
| 502 | const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); |
| 503 | const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); |
| 504 | const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); |
| 505 | const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); |
| 506 | const int n_mantissa_bits = 23; |
| 507 | |
| 508 | // exp(x) = |
| 509 | // = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem |
| 510 | // = 2^n * exp(r) // simplify the exp(n*ln(2)) expression |
| 511 | |
| 512 | auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); |
| 513 | auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); |
| 514 | vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); |
| 515 | |
| 516 | // fx = floorf(x * log2ef + 0.5) |
| 517 | auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); |
| 518 | auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); |
| 519 | vec_fx = _mm512_cvtepi32_ps(vec_fx_i); |
| 520 | |
| 521 | // x = x - fx * ln2 |
| 522 | auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); |
| 523 | |
| 524 | // compute polynomial |
| 525 | auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); |
| 526 | vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); |
| 527 | vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); |
| 528 | vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); |
| 529 | vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); |
| 530 | |
| 531 | // compute 2^(n-1) |
| 532 | auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); |
| 533 | auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); |
| 534 | auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); |
| 535 | vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); |
| 536 | auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); |
| 537 | vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero); |
| 538 | |
| 539 | // y = y * 2^n |
| 540 | vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); |
| 541 | vec_res = _mm512_mul_ps(vec_res, vec_two); |
| 542 | return vec_res; |
| 543 | } |
| 544 | |
| 545 | // Note: mapped from aten fexp_u20 |
| 546 | inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) { |
| 547 | const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f); |
| 548 | const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f); |
| 549 | const __m512 vec_c2 = _mm512_set1_ps(-0.22433836478672356); |
| 550 | const __m512 vec_c3 = _mm512_set1_ps(-0.079204240219773236); |
| 551 | |
| 552 | const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e) |
| 553 | |
| 554 | const __m512 vec_a = _mm512_set1_ps(std::pow(2, 23) / std::log2(2)); |
| 555 | const __m512 vec_b = _mm512_set1_ps(std::pow(2, 23) * 127.f); |
| 556 | |
| 557 | const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); |
| 558 | const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); |
| 559 | __m512i vec_infinity = _mm512_set1_epi32(0x7F800000); |
| 560 | __m512i vec_zero = _mm512_setzero_epi32(); |
| 561 | |
| 562 | // Fast Exponential Computation on SIMD Architectures |
| 563 | // A. Cristiano I. Malossi, Yves Ineichen, Costas Bekas, and Alessandro |
| 564 | // Curioni exp(x) = 2**(x * log2(e)) |
| 565 | // = 2**xi * 2**xf - TIPS we are using the EEEE floating point |
| 566 | // representation with identification to the exponent and the |
| 567 | // mentissa |
| 568 | // 2**xf will be approximated to a polynomial of degree 3 computed with |
| 569 | // Horner method |
| 570 | // mask for the boundary condition |
| 571 | auto min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, _CMP_LT_OS); |
| 572 | auto max_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_max, _CMP_GT_OS); |
| 573 | |
| 574 | // transformation with log2(e) |
| 575 | auto vec_src = _mm512_mul_ps(values, vec_exp_log2ef); |
| 576 | auto vec_fractional = _mm512_sub_ps(vec_src, _mm512_floor_ps(vec_src)); |
| 577 | |
| 578 | // compute polynomial using Horner Scheme, for superscalar processor |
| 579 | auto vec_res = _mm512_fmadd_ps(vec_fractional, vec_c3, vec_c2); |
| 580 | vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c1); |
| 581 | vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c0); |
| 582 | |
| 583 | vec_src = _mm512_sub_ps(vec_src, vec_res); |
| 584 | // the tips is here, headache in perspective |
| 585 | auto tmp = _mm512_fmadd_ps(vec_a, vec_src, vec_b); |
| 586 | // headache bis - we loose precision with the cast but it "fits", but ok |
| 587 | // after f32 -> f16 later |
| 588 | __m512i casted_integer = _mm512_cvttps_epi32(tmp); |
| 589 | // boundary condition, lower than the min -> 0 |
| 590 | casted_integer = _mm512_mask_mov_epi32(casted_integer, min_mask, vec_zero); |
| 591 | // boundary condition, larger than the max -> +oo |
| 592 | casted_integer = _mm512_mask_mov_epi32(casted_integer, max_mask, vec_infinity); |
| 593 | // final interpretation to float |
| 594 | return _mm512_castsi512_ps(casted_integer); |
| 595 | } |
| 596 | |
| 597 | // sigmoid(x) = 1 / (1 + exp(-x)); avoid vdivps via rcp14 |
| 598 | inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_ps(__m512 x) { |
| 599 | __m512 minus_x = _mm512_xor_ps(_mm512_set1_ps(-0.f), x); |
| 600 | __m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_x), _mm512_set1_ps(1.f)); |
| 601 | return _mm512_rcp14_ps(denom); |
| 602 | } |
| 603 | |
| 604 | // SiLU(x) = x * sigmoid(x) |
| 605 | inline __attribute__((always_inline)) __m512 _mm512_rcp14_silu_ps(__m512 x) { |
| 606 | return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(x)); |
| 607 | } |
| 608 | |
| 609 | // x * sigmoid(x * alpha) for clamped SwiGLU |
| 610 | inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512 x, __m512 alpha) { |
| 611 | __m512 xa = _mm512_mul_ps(x, alpha); |
| 612 | return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa)); |
| 613 | } |
| 614 | |
| 615 | #endif |
| 616 | |
| 617 | inline at::vec::Vectorized<float> fast_sigmoid(const at::vec::Vectorized<float>& x) { |
| 618 | #if defined(CPU_CAPABILITY_AVX512) |
| 619 | return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_ps(x)); |
| 620 | #else |
| 621 | const auto one = at::vec::Vectorized<float>(1.f); |
| 622 | return one / (one + x.neg().exp_u20()); |
| 623 | #endif |
| 624 | } |
| 625 | |
| 626 | inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x) { |
| 627 | #if defined(CPU_CAPABILITY_AVX512) |
| 628 | return at::vec::Vectorized<float>(_mm512_rcp14_silu_ps(x)); |
| 629 | #else |
| 630 | return x * fast_sigmoid(x); |
| 631 | #endif |
| 632 | } |
| 633 | |
| 634 | inline at::vec::Vectorized<float> |
| 635 | fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) { |
| 636 | #if defined(CPU_CAPABILITY_AVX512) |
| 637 | return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_glu_ps(x, alpha)); |
| 638 | #else |
| 639 | return x * fast_sigmoid(x * alpha); |
| 640 | #endif |
| 641 | } |
| 642 | |
| 643 | } // anonymous namespace |