Added AVX512 Support

Bench: 3291373

Author: jhonnold

src/makefile

@@ -18,6 +18,7 @@ endif
 
 SSSE = -mssse3 -msse3 -msse2 -msse -mpopcnt
 AVX2 = -mavx2 -mavx -mfma -msse4.1 $(SSSE)
+AVX512 = -mavx512f -mavx512bw -mavx512dq $(AVX2)
 PEXT = -mbmi2 -DUSE_PEXT
 
 all:
@@ -44,4 +45,10 @@ avx2:
 avx2-pext:
 	$(CC) $(RFLAGS) $(DEFS) $(SRC) $(LIBS) $(AVX2) $(PEXT) -o $(EXE)-$(VERSION)-x64-avx2-pext$(EXT)
 
-release: ssse pext avx2 avx2-pext
+avx512:
+	$(CC) $(RFLAGS) $(DEFS) $(SRC) $(LIBS) $(AVX512) -o $(EXE)-$(VERSION)-x64-avx512$(EXT)
+
+avx512-pext:
+	$(CC) $(RFLAGS) $(DEFS) $(SRC) $(LIBS) $(AVX512) $(PEXT) -o $(EXE)-$(VERSION)-x64-avx512-pext$(EXT)
+
+release: ssse pext avx2 avx2-pext avx512 avx512-pext

src/nn.c

@@ -64,7 +64,35 @@ inline void RefreshAccumulator(Accumulator accumulator, Board* board, const int
   }
 }
 
-#if defined(__AVX2__)
+#if defined(__AVX512F__)
+const size_t WIDTH = sizeof(__m512i) / sizeof(int16_t);
+const size_t CHUNKS = N_HIDDEN / WIDTH;
+
+int OutputLayer(Accumulator stm, Accumulator xstm) {
+  int result = OUTPUT_BIAS * QUANTIZATION_PRECISION_IN;
+
+  const __m512i zero = _mm512_setzero_si512();
+  __m512i s0 = _mm512_setzero_si512();
+  __m512i s1 = _mm512_setzero_si512();
+
+  for (size_t j = 0; j < CHUNKS; j++) {
+    const __m512i ac0 = _mm512_max_epi16(*(__m512i*)&stm[j * WIDTH], zero);
+    const __m512i ac1 = _mm512_max_epi16(*(__m512i*)&xstm[j * WIDTH], zero);
+
+    s0 = _mm512_add_epi32(s0, _mm512_madd_epi16(ac0, *(__m512i*)&HIDDEN_WEIGHTS[j * WIDTH]));
+    s1 = _mm512_add_epi32(s1, _mm512_madd_epi16(ac1, *(__m512i*)&HIDDEN_WEIGHTS[j * WIDTH + N_HIDDEN]));
+  }
+
+  const __m512i r16 = _mm512_add_epi32(s0, s1);
+  const __m256i r8 = _mm256_add_epi32(_mm512_castsi512_si256(r16), _mm512_extracti32x8_epi32(r16, 1));
+  const __m128i r4 = _mm_add_epi32(_mm256_castsi256_si128(r8), _mm256_extractf128_si256(r8, 1));
+  const __m128i r2 = _mm_add_epi32(r4, _mm_srli_si128(r4, 8));
+  const __m128i r1 = _mm_add_epi32(r2, _mm_srli_si128(r2, 4));
+
+  result += _mm_cvtsi128_si32(r1);
+  return result / QUANTIZATION_PRECISION_IN / QUANTIZATION_PRECISION_OUT;
+}
+#elif defined(__AVX2__)
 const size_t WIDTH = sizeof(__m256i) / sizeof(int16_t);
 const size_t CHUNKS = N_HIDDEN / WIDTH;