diff --git a/src/nnue/features/half_ka_v2_hm.cpp b/src/nnue/features/half_ka_v2_hm.cpp index a82e89de4..09a5dc451 100644 --- a/src/nnue/features/half_ka_v2_hm.cpp +++ b/src/nnue/features/half_ka_v2_hm.cpp @@ -27,6 +27,75 @@ namespace Stockfish::Eval::NNUE::Features { +#if defined(USE_AVX512ICL) +void HalfKAv2_hm::write_indices(const std::array& oldPieces, + const std::array& newPieces, + Bitboard removedBB, + Bitboard addedBB, + Color perspective, + Square ksq, + IndexList& removed, + IndexList& added) { + + auto* write_removed = removed.make_space(popcount(removedBB)); + auto* write_added = added.make_space(popcount(addedBB)); + + const __m512i vecOldPieces = _mm512_loadu_si512(oldPieces.data()); + const __m512i vecNewPieces = _mm512_loadu_si512(newPieces.data()); + + static constexpr uint16_t psiTable[COLOR_NB][32] = { + {PS_NONE, PS_W_PAWN, PS_W_KNIGHT, PS_W_BISHOP, PS_W_ROOK, PS_W_QUEEN, PS_KING, PS_NONE, + PS_NONE, PS_B_PAWN, PS_B_KNIGHT, PS_B_BISHOP, PS_B_ROOK, PS_B_QUEEN, PS_KING, PS_NONE, + PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, + PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE}, + + {PS_NONE, PS_B_PAWN, PS_B_KNIGHT, PS_B_BISHOP, PS_B_ROOK, PS_B_QUEEN, PS_KING, PS_NONE, + PS_NONE, PS_W_PAWN, PS_W_KNIGHT, PS_W_BISHOP, PS_W_ROOK, PS_W_QUEEN, PS_KING, PS_NONE, + PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, + PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE, PS_NONE}}; + const __m512i psi = _mm512_loadu_si512(psiTable[perspective]); + + const __m512i allSquares = _mm512_set_epi8( + 63, 62, 61, 60, 59, 58, 57, 56, 55, 54, 53, 52, 51, 50, 49, 48, 47, 46, 45, 44, 43, 42, 41, + 40, 39, 38, 37, 36, 35, 34, 33, 32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, + 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); + + const uint16_t flip = 56 * perspective; + const __m512i orient = _mm512_set1_epi16((uint16_t) OrientTBL[ksq] ^ flip); + const __m512i bucket = _mm512_set1_epi16((uint16_t) KingBuckets[int(ksq) ^ flip]); + + __m512i removed_squares = _mm512_maskz_compress_epi8(removedBB, allSquares); + __m512i removed_pieces = _mm512_permutexvar_epi8(removed_squares, vecOldPieces); + removed_squares = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(removed_squares)); + removed_pieces = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(removed_pieces)); + const __m512i removed_psi = _mm512_permutexvar_epi16(removed_pieces, psi); + __m512i removed_indices = _mm512_xor_si512(removed_squares, orient); + removed_indices = _mm512_add_epi16(removed_indices, removed_psi); + removed_indices = _mm512_add_epi16(removed_indices, bucket); + + __m512i added_squares = _mm512_maskz_compress_epi8(addedBB, allSquares); + __m512i added_pieces = _mm512_permutexvar_epi8(added_squares, vecNewPieces); + added_squares = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(added_squares)); + added_pieces = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(added_pieces)); + const __m512i added_psi = _mm512_permutexvar_epi16(added_pieces, psi); + __m512i added_indices = _mm512_xor_si512(added_squares, orient); + added_indices = _mm512_add_epi16(added_indices, added_psi); + added_indices = _mm512_add_epi16(added_indices, bucket); + + const __m512i removed_indices0 = _mm512_cvtepi16_epi32(_mm512_castsi512_si256(removed_indices)); + const __m512i removed_indices1 = + _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(removed_indices, 1)); + _mm512_storeu_si512(write_removed, removed_indices0); + _mm512_storeu_si512(write_removed + 16, removed_indices1); + + const __m512i added_indices0 = _mm512_cvtepi16_epi32(_mm512_castsi512_si256(added_indices)); + const __m512i added_indices1 = + _mm512_cvtepi16_epi32(_mm512_extracti64x4_epi64(added_indices, 1)); + _mm512_storeu_si512(write_added, added_indices0); + _mm512_storeu_si512(write_added + 16, added_indices1); +} +#endif + // Index of a feature for a given king position and another piece on some square IndexType HalfKAv2_hm::make_index(Color perspective, Square s, Piece pc, Square ksq) { diff --git a/src/nnue/features/half_ka_v2_hm.h b/src/nnue/features/half_ka_v2_hm.h index 49b0a87a4..312ae306a 100644 --- a/src/nnue/features/half_ka_v2_hm.h +++ b/src/nnue/features/half_ka_v2_hm.h @@ -106,6 +106,18 @@ class HalfKAv2_hm { using IndexList = ValueList; using DiffType = DirtyPiece; +#if defined(USE_AVX512ICL) + // Compute all changed feature indices and write them to the given lists + static void write_indices(const std::array& oldPieces, + const std::array& newPieces, + Bitboard removedBB, + Bitboard addedBB, + Color perspective, + Square ksq, + IndexList& removed, + IndexList& added); +#endif + // Index of a feature for a given king position and another piece on some square static IndexType make_index(Color perspective, Square s, Piece pc, Square ksq); diff --git a/src/nnue/nnue_accumulator.cpp b/src/nnue/nnue_accumulator.cpp index 3f588e37d..10a3dee6e 100644 --- a/src/nnue/nnue_accumulator.cpp +++ b/src/nnue/nnue_accumulator.cpp @@ -710,6 +710,10 @@ void update_accumulator_refresh_cache(Color pers Bitboard removedBB = changedBB & entry.pieceBB; Bitboard addedBB = changedBB & pos.pieces(); +#if defined(USE_AVX512ICL) + PSQFeatureSet::write_indices(entry.pieces, pos.piece_array(), removedBB, addedBB, perspective, + ksq, removed, added); +#else while (removedBB) { Square sq = pop_lsb(removedBB); @@ -720,6 +724,7 @@ void update_accumulator_refresh_cache(Color pers Square sq = pop_lsb(addedBB); added.push_back(PSQFeatureSet::make_index(perspective, sq, pos.piece_on(sq), ksq)); } +#endif entry.pieceBB = pos.pieces(); entry.pieces = pos.piece_array();