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_MMX)
  35 # include <mmintrin.h>
  36
  37 #elif defined(USE_NEON)
  38 # include <arm_neon.h>
  39 #endif
  40
  41 namespace Stockfish::Simd {
  42
  43 #if defined (USE_AVX512)
  44
  45     [[maybe_unused]] static int m512_hadd(__m512i sum, int bias) {
  46       return _mm512_reduce_add_epi32(sum) + bias;
  47     }
  48
  49     /*
  50       Parameters:
  51         sum0 = [zmm0.i128[0], zmm0.i128[1], zmm0.i128[2], zmm0.i128[3]]
  52         sum1 = [zmm1.i128[0], zmm1.i128[1], zmm1.i128[2], zmm1.i128[3]]
  53         sum2 = [zmm2.i128[0], zmm2.i128[1], zmm2.i128[2], zmm2.i128[3]]
  54         sum3 = [zmm3.i128[0], zmm3.i128[1], zmm3.i128[2], zmm3.i128[3]]
  55
  56       Returns:
  57         ret = [
  58           reduce_add_epi32(zmm0.i128[0]), reduce_add_epi32(zmm1.i128[0]), reduce_add_epi32(zmm2.i128[0]), reduce_add_epi32(zmm3.i128[0]),
  59           reduce_add_epi32(zmm0.i128[1]), reduce_add_epi32(zmm1.i128[1]), reduce_add_epi32(zmm2.i128[1]), reduce_add_epi32(zmm3.i128[1]),
  60           reduce_add_epi32(zmm0.i128[2]), reduce_add_epi32(zmm1.i128[2]), reduce_add_epi32(zmm2.i128[2]), reduce_add_epi32(zmm3.i128[2]),
  61           reduce_add_epi32(zmm0.i128[3]), reduce_add_epi32(zmm1.i128[3]), reduce_add_epi32(zmm2.i128[3]), reduce_add_epi32(zmm3.i128[3])
  62         ]
  63     */
  64     [[maybe_unused]] static __m512i m512_hadd128x16_interleave(
  65         __m512i sum0, __m512i sum1, __m512i sum2, __m512i sum3) {
  66
  67       __m512i sum01a = _mm512_unpacklo_epi32(sum0, sum1);
  68       __m512i sum01b = _mm512_unpackhi_epi32(sum0, sum1);
  69
  70       __m512i sum23a = _mm512_unpacklo_epi32(sum2, sum3);
  71       __m512i sum23b = _mm512_unpackhi_epi32(sum2, sum3);
  72
  73       __m512i sum01 = _mm512_add_epi32(sum01a, sum01b);
  74       __m512i sum23 = _mm512_add_epi32(sum23a, sum23b);
  75
  76       __m512i sum0123a = _mm512_unpacklo_epi64(sum01, sum23);
  77       __m512i sum0123b = _mm512_unpackhi_epi64(sum01, sum23);
  78
  79       return _mm512_add_epi32(sum0123a, sum0123b);
  80     }
  81
  82     [[maybe_unused]] static void m512_add_dpbusd_epi32(
  83         __m512i& acc,
  84         __m512i a,
  85         __m512i b) {
  86
  87 # if defined (USE_VNNI)
  88       acc = _mm512_dpbusd_epi32(acc, a, b);
  89 # else
  90       __m512i product0 = _mm512_maddubs_epi16(a, b);
  91       product0 = _mm512_madd_epi16(product0, _mm512_set1_epi16(1));
  92       acc = _mm512_add_epi32(acc, product0);
  93 # endif
  94     }
  95
  96     [[maybe_unused]] static void m512_add_dpbusd_epi32x2(
  97         __m512i& acc,
  98         __m512i a0, __m512i b0,
  99         __m512i a1, __m512i b1) {
 100
 101 # if defined (USE_VNNI)
 102       acc = _mm512_dpbusd_epi32(acc, a0, b0);
 103       acc = _mm512_dpbusd_epi32(acc, a1, b1);
 104 # else
 105       __m512i product0 = _mm512_maddubs_epi16(a0, b0);
 106       __m512i product1 = _mm512_maddubs_epi16(a1, b1);
 107       product0 = _mm512_madd_epi16(product0, _mm512_set1_epi16(1));
 108       product1 = _mm512_madd_epi16(product1, _mm512_set1_epi16(1));
 109       acc = _mm512_add_epi32(acc, _mm512_add_epi32(product0, product1));
 110 # endif
 111     }
 112
 113 #endif
 114
 115 #if defined (USE_AVX2)
 116
 117     [[maybe_unused]] static int m256_hadd(__m256i sum, int bias) {
 118       __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1));
 119       sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_BADC));
 120       sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_CDAB));
 121       return _mm_cvtsi128_si32(sum128) + bias;
 122     }
 123
 124     [[maybe_unused]] static void m256_add_dpbusd_epi32(
 125         __m256i& acc,
 126         __m256i a,
 127         __m256i b) {
 128
 129 # if defined (USE_VNNI)
 130       acc = _mm256_dpbusd_epi32(acc, a, b);
 131 # else
 132       __m256i product0 = _mm256_maddubs_epi16(a, b);
 133       product0 = _mm256_madd_epi16(product0, _mm256_set1_epi16(1));
 134       acc = _mm256_add_epi32(acc, product0);
 135 # endif
 136     }
 137
 138     [[maybe_unused]] static void m256_add_dpbusd_epi32x2(
 139         __m256i& acc,
 140         __m256i a0, __m256i b0,
 141         __m256i a1, __m256i b1) {
 142
 143 # if defined (USE_VNNI)
 144       acc = _mm256_dpbusd_epi32(acc, a0, b0);
 145       acc = _mm256_dpbusd_epi32(acc, a1, b1);
 146 # else
 147       __m256i product0 = _mm256_maddubs_epi16(a0, b0);
 148       __m256i product1 = _mm256_maddubs_epi16(a1, b1);
 149       product0 = _mm256_madd_epi16(product0, _mm256_set1_epi16(1));
 150       product1 = _mm256_madd_epi16(product1, _mm256_set1_epi16(1));
 151       acc = _mm256_add_epi32(acc, _mm256_add_epi32(product0, product1));
 152 # endif
 153     }
 154
 155 #endif
 156
 157 #if defined (USE_SSSE3)
 158
 159     [[maybe_unused]] static int m128_hadd(__m128i sum, int bias) {
 160       sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0x4E)); //_MM_PERM_BADC
 161       sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0xB1)); //_MM_PERM_CDAB
 162       return _mm_cvtsi128_si32(sum) + bias;
 163     }
 164
 165     [[maybe_unused]] static void m128_add_dpbusd_epi32(
 166         __m128i& acc,
 167         __m128i a,
 168         __m128i b) {
 169
 170       __m128i product0 = _mm_maddubs_epi16(a, b);
 171       product0 = _mm_madd_epi16(product0, _mm_set1_epi16(1));
 172       acc = _mm_add_epi32(acc, product0);
 173     }
 174
 175     [[maybe_unused]] static void m128_add_dpbusd_epi32x2(
 176         __m128i& acc,
 177         __m128i a0, __m128i b0,
 178         __m128i a1, __m128i b1) {
 179
 180       __m128i product0 = _mm_maddubs_epi16(a0, b0);
 181       __m128i product1 = _mm_maddubs_epi16(a1, b1);
 182       product0 = _mm_madd_epi16(product0, _mm_set1_epi16(1));
 183       product1 = _mm_madd_epi16(product1, _mm_set1_epi16(1));
 184       acc = _mm_add_epi32(acc, _mm_add_epi32(product0, product1));
 185     }
 186
 187 #endif
 188
 189 #if defined (USE_NEON_DOTPROD)
 190
 191     [[maybe_unused]] static void dotprod_m128_add_dpbusd_epi32x2(
 192         int32x4_t& acc,
 193         int8x16_t a0, int8x16_t b0,
 194         int8x16_t a1, int8x16_t b1) {
 195
 196         acc = vdotq_s32(acc, a0, b0);
 197         acc = vdotq_s32(acc, a1, b1);
 198     }
 199
 200     [[maybe_unused]] static void dotprod_m128_add_dpbusd_epi32(
 201         int32x4_t& acc,
 202         int8x16_t a, int8x16_t b) {
 203
 204         acc = vdotq_s32(acc, a, b);
 205     }
 206 #endif
 207
 208 #if defined (USE_NEON)
 209
 210     [[maybe_unused]] static int neon_m128_reduce_add_epi32(int32x4_t s) {
 211 #   if USE_NEON >= 8
 212       return vaddvq_s32(s);
 213 #   else
 214       return s[0] + s[1] + s[2] + s[3];
 215 #   endif
 216     }
 217
 218     [[maybe_unused]] static int neon_m128_hadd(int32x4_t sum, int bias) {
 219       return neon_m128_reduce_add_epi32(sum) + bias;
 220     }
 221
 222     [[maybe_unused]] static void neon_m128_add_dpbusd_epi32x2(
 223         int32x4_t& acc,
 224         int8x8_t a0, int8x8_t b0,
 225         int8x8_t a1, int8x8_t b1) {
 226
 227       int16x8_t product = vmull_s8(a0, b0);
 228       product = vmlal_s8(product, a1, b1);
 229       acc = vpadalq_s16(acc, product);
 230     }
 231 #endif
 232
 233 #if USE_NEON >= 8
 234     [[maybe_unused]] static void neon_m128_add_dpbusd_epi32(
 235         int32x4_t& acc,
 236         int8x16_t a, int8x16_t b) {
 237
 238       int16x8_t product0 = vmull_s8(vget_low_s8(a), vget_low_s8(b));
 239       int16x8_t product1 = vmull_high_s8(a, b);
 240       int16x8_t sum = vpaddq_s16(product0, product1);
 241       acc = vpadalq_s16(acc, sum);
 242     }
 243 #endif
 244 }
 245
 246 #endif // STOCKFISH_SIMD_H_INCLUDED