git.sesse.net Git - stockfish/blob - src/nnue/layers/simd.h

   1 /*
   2   Stockfish, a UCI chess playing engine derived from Glaurung 2.1
   3   Copyright (C) 2004-2023 The Stockfish developers (see AUTHORS file)
   4
   5   Stockfish is free software: you can redistribute it and/or modify
   6   it under the terms of the GNU General Public License as published by
   7   the Free Software Foundation, either version 3 of the License, or
   8   (at your option) any later version.
   9
  10   Stockfish is distributed in the hope that it will be useful,
  11   but WITHOUT ANY WARRANTY; without even the implied warranty of
  12   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13   GNU General Public License for more details.
  14
  15   You should have received a copy of the GNU General Public License
  16   along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 */
  18
  19 #ifndef STOCKFISH_SIMD_H_INCLUDED
  20 #define STOCKFISH_SIMD_H_INCLUDED
  21
  22 #if defined(USE_AVX2)
  23     #include <immintrin.h>
  24
  25 #elif defined(USE_SSE41)
  26     #include <smmintrin.h>
  27
  28 #elif defined(USE_SSSE3)
  29     #include <tmmintrin.h>
  30
  31 #elif defined(USE_SSE2)
  32     #include <emmintrin.h>
  33
  34 #elif defined(USE_NEON)
  35     #include <arm_neon.h>
  36 #endif
  37
  38 namespace Stockfish::Simd {
  39
  40 #if defined(USE_AVX512)
  41
  42 [[maybe_unused]] static int m512_hadd(__m512i sum, int bias) {
  43     return _mm512_reduce_add_epi32(sum) + bias;
  44 }
  45
  46 /*
  47       Parameters:
  48         sum0 = [zmm0.i128[0], zmm0.i128[1], zmm0.i128[2], zmm0.i128[3]]
  49         sum1 = [zmm1.i128[0], zmm1.i128[1], zmm1.i128[2], zmm1.i128[3]]
  50         sum2 = [zmm2.i128[0], zmm2.i128[1], zmm2.i128[2], zmm2.i128[3]]
  51         sum3 = [zmm3.i128[0], zmm3.i128[1], zmm3.i128[2], zmm3.i128[3]]
  52
  53       Returns:
  54         ret = [
  55           reduce_add_epi32(zmm0.i128[0]), reduce_add_epi32(zmm1.i128[0]), reduce_add_epi32(zmm2.i128[0]), reduce_add_epi32(zmm3.i128[0]),
  56           reduce_add_epi32(zmm0.i128[1]), reduce_add_epi32(zmm1.i128[1]), reduce_add_epi32(zmm2.i128[1]), reduce_add_epi32(zmm3.i128[1]),
  57           reduce_add_epi32(zmm0.i128[2]), reduce_add_epi32(zmm1.i128[2]), reduce_add_epi32(zmm2.i128[2]), reduce_add_epi32(zmm3.i128[2]),
  58           reduce_add_epi32(zmm0.i128[3]), reduce_add_epi32(zmm1.i128[3]), reduce_add_epi32(zmm2.i128[3]), reduce_add_epi32(zmm3.i128[3])
  59         ]
  60     */
  61 [[maybe_unused]] static __m512i
  62 m512_hadd128x16_interleave(__m512i sum0, __m512i sum1, __m512i sum2, __m512i sum3) {
  63
  64     __m512i sum01a = _mm512_unpacklo_epi32(sum0, sum1);
  65     __m512i sum01b = _mm512_unpackhi_epi32(sum0, sum1);
  66
  67     __m512i sum23a = _mm512_unpacklo_epi32(sum2, sum3);
  68     __m512i sum23b = _mm512_unpackhi_epi32(sum2, sum3);
  69
  70     __m512i sum01 = _mm512_add_epi32(sum01a, sum01b);
  71     __m512i sum23 = _mm512_add_epi32(sum23a, sum23b);
  72
  73     __m512i sum0123a = _mm512_unpacklo_epi64(sum01, sum23);
  74     __m512i sum0123b = _mm512_unpackhi_epi64(sum01, sum23);
  75
  76     return _mm512_add_epi32(sum0123a, sum0123b);
  77 }
  78
  79 [[maybe_unused]] static void m512_add_dpbusd_epi32(__m512i& acc, __m512i a, __m512i b) {
  80
  81     #if defined(USE_VNNI)
  82     acc = _mm512_dpbusd_epi32(acc, a, b);
  83     #else
  84     __m512i product0 = _mm512_maddubs_epi16(a, b);
  85     product0         = _mm512_madd_epi16(product0, _mm512_set1_epi16(1));
  86     acc              = _mm512_add_epi32(acc, product0);
  87     #endif
  88 }
  89
  90 [[maybe_unused]] static void
  91 m512_add_dpbusd_epi32x2(__m512i& acc, __m512i a0, __m512i b0, __m512i a1, __m512i b1) {
  92
  93     #if defined(USE_VNNI)
  94     acc = _mm512_dpbusd_epi32(acc, a0, b0);
  95     acc = _mm512_dpbusd_epi32(acc, a1, b1);
  96     #else
  97     __m512i product0 = _mm512_maddubs_epi16(a0, b0);
  98     __m512i product1 = _mm512_maddubs_epi16(a1, b1);
  99     product0         = _mm512_madd_epi16(product0, _mm512_set1_epi16(1));
 100     product1         = _mm512_madd_epi16(product1, _mm512_set1_epi16(1));
 101     acc              = _mm512_add_epi32(acc, _mm512_add_epi32(product0, product1));
 102     #endif
 103 }
 104
 105 #endif
 106
 107 #if defined(USE_AVX2)
 108
 109 [[maybe_unused]] static int m256_hadd(__m256i sum, int bias) {
 110     __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1));
 111     sum128         = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_BADC));
 112     sum128         = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_CDAB));
 113     return _mm_cvtsi128_si32(sum128) + bias;
 114 }
 115
 116 [[maybe_unused]] static void m256_add_dpbusd_epi32(__m256i& acc, __m256i a, __m256i b) {
 117
 118     #if defined(USE_VNNI)
 119     acc = _mm256_dpbusd_epi32(acc, a, b);
 120     #else
 121     __m256i product0 = _mm256_maddubs_epi16(a, b);
 122     product0         = _mm256_madd_epi16(product0, _mm256_set1_epi16(1));
 123     acc              = _mm256_add_epi32(acc, product0);
 124     #endif
 125 }
 126
 127 [[maybe_unused]] static void
 128 m256_add_dpbusd_epi32x2(__m256i& acc, __m256i a0, __m256i b0, __m256i a1, __m256i b1) {
 129
 130     #if defined(USE_VNNI)
 131     acc = _mm256_dpbusd_epi32(acc, a0, b0);
 132     acc = _mm256_dpbusd_epi32(acc, a1, b1);
 133     #else
 134     __m256i product0 = _mm256_maddubs_epi16(a0, b0);
 135     __m256i product1 = _mm256_maddubs_epi16(a1, b1);
 136     product0         = _mm256_madd_epi16(product0, _mm256_set1_epi16(1));
 137     product1         = _mm256_madd_epi16(product1, _mm256_set1_epi16(1));
 138     acc              = _mm256_add_epi32(acc, _mm256_add_epi32(product0, product1));
 139     #endif
 140 }
 141
 142 #endif
 143
 144 #if defined(USE_SSSE3)
 145
 146 [[maybe_unused]] static int m128_hadd(__m128i sum, int bias) {
 147     sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0x4E));  //_MM_PERM_BADC
 148     sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0xB1));  //_MM_PERM_CDAB
 149     return _mm_cvtsi128_si32(sum) + bias;
 150 }
 151
 152 [[maybe_unused]] static void m128_add_dpbusd_epi32(__m128i& acc, __m128i a, __m128i b) {
 153
 154     __m128i product0 = _mm_maddubs_epi16(a, b);
 155     product0         = _mm_madd_epi16(product0, _mm_set1_epi16(1));
 156     acc              = _mm_add_epi32(acc, product0);
 157 }
 158
 159 [[maybe_unused]] static void
 160 m128_add_dpbusd_epi32x2(__m128i& acc, __m128i a0, __m128i b0, __m128i a1, __m128i b1) {
 161
 162     __m128i product0 = _mm_maddubs_epi16(a0, b0);
 163     __m128i product1 = _mm_maddubs_epi16(a1, b1);
 164     product0         = _mm_madd_epi16(product0, _mm_set1_epi16(1));
 165     product1         = _mm_madd_epi16(product1, _mm_set1_epi16(1));
 166     acc              = _mm_add_epi32(acc, _mm_add_epi32(product0, product1));
 167 }
 168
 169 #endif
 170
 171 #if defined(USE_NEON_DOTPROD)
 172
 173 [[maybe_unused]] static void dotprod_m128_add_dpbusd_epi32x2(
 174   int32x4_t& acc, int8x16_t a0, int8x16_t b0, int8x16_t a1, int8x16_t b1) {
 175
 176     acc = vdotq_s32(acc, a0, b0);
 177     acc = vdotq_s32(acc, a1, b1);
 178 }
 179
 180 [[maybe_unused]] static void
 181 dotprod_m128_add_dpbusd_epi32(int32x4_t& acc, int8x16_t a, int8x16_t b) {
 182
 183     acc = vdotq_s32(acc, a, b);
 184 }
 185 #endif
 186
 187 #if defined(USE_NEON)
 188
 189 [[maybe_unused]] static int neon_m128_reduce_add_epi32(int32x4_t s) {
 190     #if USE_NEON >= 8
 191     return vaddvq_s32(s);
 192     #else
 193     return s[0] + s[1] + s[2] + s[3];
 194     #endif
 195 }
 196
 197 [[maybe_unused]] static int neon_m128_hadd(int32x4_t sum, int bias) {
 198     return neon_m128_reduce_add_epi32(sum) + bias;
 199 }
 200
 201 [[maybe_unused]] static void
 202 neon_m128_add_dpbusd_epi32x2(int32x4_t& acc, int8x8_t a0, int8x8_t b0, int8x8_t a1, int8x8_t b1) {
 203
 204     int16x8_t product = vmull_s8(a0, b0);
 205     product           = vmlal_s8(product, a1, b1);
 206     acc               = vpadalq_s16(acc, product);
 207 }
 208 #endif
 209
 210 #if USE_NEON >= 8
 211 [[maybe_unused]] static void neon_m128_add_dpbusd_epi32(int32x4_t& acc, int8x16_t a, int8x16_t b) {
 212
 213     int16x8_t product0 = vmull_s8(vget_low_s8(a), vget_low_s8(b));
 214     int16x8_t product1 = vmull_high_s8(a, b);
 215     int16x8_t sum      = vpaddq_s16(product0, product1);
 216     acc                = vpadalq_s16(acc, sum);
 217 }
 218 #endif
 219 }
 220
 221 #endif  // STOCKFISH_SIMD_H_INCLUDED