Merge HalfKA and Threats accumulators

Passed STC (https://tests.stockfishchess.org/tests/view/6a2893ad7c758d82accea129):
LLR: 3.20 (-2.94,2.94) <0.00,2.00>
Total: 23328 W: 6145 L: 5838 D: 11345
Ptnml(0-2): 50, 2463, 6346, 2740, 65

Instead of repeatedly doing the sum HalfKA + threats at the end, it's profitable to simply store one accumulator per side that combines them. This also avoids an extra/load store of an accumulator, and halves the cache footprint of the accumulators.

For full refreshes, we always compute both halfka and threats simultaneously. Any threat full refresh is always a halfka refresh because it occurs when the king crosses the center line, while halfka refreshes are required for ANY king move, so we don't need a separate detection path for threats.

I get about a 2.5% speedup locally with this, but I'd appreciate other ppl's measurements.

closes https://github.com/official-stockfish/Stockfish/pull/6890

No functional change

src/nnue/nnue_accumulator.cpp

@@ -20,7 +20,6 @@
 
 #include <cassert>
 #include <new>
-#include <type_traits>
 
 #include "../bitboard.h"
 #include "../misc.h"
@@ -37,78 +36,38 @@ using namespace SIMD;
 
 namespace {
 
-template<bool Forward, typename FeatureSet>
+template<bool Forward>
 void update_accumulator_incremental(Color                     perspective,
                                     const FeatureTransformer& featureTransformer,
                                     const Square              ksq,
-                                    AccumulatorState<FeatureSet>&       target_state,
+                                    AccumulatorState&         target_state,
-                                    const AccumulatorState<FeatureSet>& computed);
+                                    const AccumulatorState&   computed);
 
 void update_accumulator_refresh_cache(Color                     perspective,
                                       const FeatureTransformer& featureTransformer,
                                       const Position&           pos,
-                                      AccumulatorState<PSQFeatureSet>& accumulatorState,
+                                      AccumulatorState&         accumulatorState,
                                       AccumulatorCaches&        cache);
-
-void update_threats_accumulator_full(Color                               perspective,
-                                     const FeatureTransformer&           featureTransformer,
-                                     const Position&                     pos,
-                                     AccumulatorState<ThreatFeatureSet>& accumulatorState);
 }
 
-template<typename T>
+const AccumulatorState& AccumulatorStack::latest() const noexcept { return accumulators[size - 1]; }
-const AccumulatorState<T>& AccumulatorStack::latest() const noexcept {
-    return accumulators<T>()[size - 1];
-}
 
-// Explicit template instantiations
+AccumulatorState& AccumulatorStack::mut_latest() noexcept { return accumulators[size - 1]; }
-template const AccumulatorState<PSQFeatureSet>&    AccumulatorStack::latest() const noexcept;
-template const AccumulatorState<ThreatFeatureSet>& AccumulatorStack::latest() const noexcept;
-
-template<typename T>
-AccumulatorState<T>& AccumulatorStack::mut_latest() noexcept {
-    return mut_accumulators<T>()[size - 1];
-}
-
-template<typename T>
-const std::array<AccumulatorState<T>, AccumulatorStack::MaxSize>&
-AccumulatorStack::accumulators() const noexcept {
-    static_assert(std::is_same_v<T, PSQFeatureSet> || std::is_same_v<T, ThreatFeatureSet>,
-                  "Invalid Feature Set Type");
-
-    if constexpr (std::is_same_v<T, PSQFeatureSet>)
-        return psq_accumulators;
-
-    if constexpr (std::is_same_v<T, ThreatFeatureSet>)
-        return threat_accumulators;
-}
-
-template<typename T>
-std::array<AccumulatorState<T>, AccumulatorStack::MaxSize>&
-AccumulatorStack::mut_accumulators() noexcept {
-    static_assert(std::is_same_v<T, PSQFeatureSet> || std::is_same_v<T, ThreatFeatureSet>,
-                  "Invalid Feature Set Type");
-
-    if constexpr (std::is_same_v<T, PSQFeatureSet>)
-        return psq_accumulators;
-
-    if constexpr (std::is_same_v<T, ThreatFeatureSet>)
-        return threat_accumulators;
-}
 
 void AccumulatorStack::reset() noexcept {
-    psq_accumulators[0].reset({});
+    accumulators[0].dirtyPiece = {};
-    threat_accumulators[0].reset({});
+    new (&accumulators[0].dirtyThreats) DirtyThreats;
+    accumulators[0].computed.fill(false);
     size = 1;
 }
 
 std::pair<DirtyPiece&, DirtyThreats&> AccumulatorStack::push() noexcept {
     assert(size < MaxSize);
-    auto& dp  = psq_accumulators[size].reset();
+    auto& st = accumulators[size];
-    auto& dts = threat_accumulators[size].reset();
+    st.computed.fill(false);
-    new (&dts) DirtyThreats;
+    new (&st.dirtyThreats) DirtyThreats;
     size++;
-    return {dp, dts};
+    return {st.dirtyPiece, st.dirtyThreats};
 }
 
 void AccumulatorStack::pop() noexcept {
@@ -120,158 +79,90 @@ void AccumulatorStack::evaluate(const Position&           pos,
                                 const FeatureTransformer& featureTransformer,
                                 // Silence spurious warning on GCC 10
                                 [[maybe_unused]] AccumulatorCaches& cache) noexcept {
-    evaluate_side<PSQFeatureSet>(WHITE, pos, featureTransformer, cache);
+    evaluate_side(WHITE, pos, featureTransformer, cache);
-    evaluate_side<PSQFeatureSet>(BLACK, pos, featureTransformer, cache);
+    evaluate_side(BLACK, pos, featureTransformer, cache);
-
-    evaluate_side<ThreatFeatureSet>(WHITE, pos, featureTransformer, cache);
-    evaluate_side<ThreatFeatureSet>(BLACK, pos, featureTransformer, cache);
 }
 
-template<typename FeatureSet>
 void AccumulatorStack::evaluate_side(Color                     perspective,
                                      const Position&           pos,
                                      const FeatureTransformer& featureTransformer,
                                      AccumulatorCaches&        cache) noexcept {
 
-    const auto last_usable_accum = find_last_usable_accumulator<FeatureSet>(perspective);
+    const auto last_usable_accum = find_last_usable_accumulator(perspective);
 
-    if (accumulators<FeatureSet>()[last_usable_accum].computed[perspective])
+    if (accumulators[last_usable_accum].computed[perspective])
-        forward_update_incremental<FeatureSet>(perspective, pos, featureTransformer,
+        forward_update_incremental(perspective, pos, featureTransformer, last_usable_accum);
-                                               last_usable_accum);
 
     else
     {
-        if constexpr (std::is_same_v<FeatureSet, PSQFeatureSet>)
+        update_accumulator_refresh_cache(perspective, featureTransformer, pos, mut_latest(), cache);
-            update_accumulator_refresh_cache(perspective, featureTransformer, pos,
+        backward_update_incremental(perspective, pos, featureTransformer, last_usable_accum);
-                                             mut_latest<PSQFeatureSet>(), cache);
-        else
-            update_threats_accumulator_full(perspective, featureTransformer, pos,
-                                            mut_latest<ThreatFeatureSet>());
-
-        backward_update_incremental<FeatureSet>(perspective, pos, featureTransformer,
-                                                last_usable_accum);
     }
 }
 
 // Find the earliest usable accumulator, this can either be a computed accumulator or the accumulator
 // state just before a change that requires full refresh.
-template<typename FeatureSet>
 usize AccumulatorStack::find_last_usable_accumulator(Color perspective) const noexcept {
 
     for (usize curr_idx = size - 1; curr_idx > 0; curr_idx--)
     {
-        if (accumulators<FeatureSet>()[curr_idx].computed[perspective])
+        if (accumulators[curr_idx].computed[perspective])
             return curr_idx;
 
-        if (FeatureSet::requires_refresh(accumulators<FeatureSet>()[curr_idx].diff, perspective))
+        // Threat feature set refreshes require a king move across the center, i.e.,
+        // a subset of halfka refreshes
+        if (PSQFeatureSet::requires_refresh(accumulators[curr_idx].dirtyPiece, perspective))
             return curr_idx;
     }
 
     return 0;
 }
 
-template<typename FeatureSet>
 void AccumulatorStack::forward_update_incremental(Color                     perspective,
                                                   const Position&           pos,
                                                   const FeatureTransformer& featureTransformer,
                                                   const usize               begin) noexcept {
 
-    assert(begin < accumulators<FeatureSet>().size());
+    assert(begin < accumulators.size());
-    assert(accumulators<FeatureSet>()[begin].computed[perspective]);
+    assert(accumulators[begin].computed[perspective]);
 
     const Square ksq = pos.square<KING>(perspective);
 
     for (usize next = begin + 1; next < size; next++)
-    {
         update_accumulator_incremental<true>(perspective, featureTransformer, ksq,
-                                             mut_accumulators<FeatureSet>()[next],
+                                             accumulators[next], accumulators[next - 1]);
-                                             accumulators<FeatureSet>()[next - 1]);
-    }
 
-    assert(latest<FeatureSet>().computed[perspective]);
+    assert(latest().computed[perspective]);
 }
 
-template<typename FeatureSet>
 void AccumulatorStack::backward_update_incremental(Color                     perspective,
-
                                                    const Position&           pos,
                                                    const FeatureTransformer& featureTransformer,
                                                    const usize               end) noexcept {
 
-    assert(end < accumulators<FeatureSet>().size());
+    assert(end < accumulators.size());
     assert(end < size);
-    assert(latest<FeatureSet>().computed[perspective]);
+    assert(latest().computed[perspective]);
 
     const Square ksq = pos.square<KING>(perspective);
 
     for (i64 next = i64(size) - 2; next >= i64(end); next--)
         update_accumulator_incremental<false>(perspective, featureTransformer, ksq,
-                                              mut_accumulators<FeatureSet>()[next],
+                                              accumulators[next], accumulators[next + 1]);
-                                              accumulators<FeatureSet>()[next + 1]);
 
-    assert(accumulators<FeatureSet>()[end].computed[perspective]);
+    assert(accumulators[end].computed[perspective]);
 }
 
 namespace {
 
-template<typename VectorWrapper,
+void apply_combined(Color                              perspective,
-         IndexType Width,
+                    const FeatureTransformer&          featureTransformer,
-         UpdateOperation... ops,
+                    const AccumulatorState&            from,
-         typename ElementType,
+                    AccumulatorState&                  to,
-         typename... Ts,
+                    const PSQFeatureSet::IndexList&    psqAdded,
-         std::enable_if_t<is_all_same_v<ElementType, Ts...>, bool> = true>
+                    const PSQFeatureSet::IndexList&    psqRemoved,
-void fused_row_reduce(const ElementType* in, ElementType* out, const Ts* const... rows) {
+                    const ThreatFeatureSet::IndexList& thrAdded,
-    constexpr IndexType size = Width * sizeof(ElementType) / sizeof(typename VectorWrapper::type);
+                    const ThreatFeatureSet::IndexList& thrRemoved) {
-
-    auto* vecIn  = reinterpret_cast<const typename VectorWrapper::type*>(in);
-    auto* vecOut = reinterpret_cast<typename VectorWrapper::type*>(out);
-
-    for (IndexType i = 0; i < size; ++i)
-        vecOut[i] = fused<VectorWrapper, ops...>(
-          vecIn[i], reinterpret_cast<const typename VectorWrapper::type*>(rows)[i]...);
-}
-
-template<typename FeatureSet>
-struct AccumulatorUpdateContext {
-    Color                               perspective;
-    const FeatureTransformer&           featureTransformer;
-    const AccumulatorState<FeatureSet>& from;
-    AccumulatorState<FeatureSet>&       to;
-
-    AccumulatorUpdateContext(Color                               persp,
-                             const FeatureTransformer&           ft,
-                             const AccumulatorState<FeatureSet>& accF,
-                             AccumulatorState<FeatureSet>&       accT) noexcept :
-        perspective{persp},
-        featureTransformer{ft},
-        from{accF},
-        to{accT} {}
-
-    template<UpdateOperation... ops,
-             typename... Ts,
-             std::enable_if_t<is_all_same_v<IndexType, Ts...>, bool> = true>
-    void apply(const Ts... indices) {
-        constexpr IndexType Dimensions = FeatureTransformer::OutputDimensions;
-
-        auto to_weight_vector = [&](const IndexType index) {
-            return &featureTransformer.weights[index * Dimensions];
-        };
-
-        auto to_psqt_weight_vector = [&](const IndexType index) {
-            return &featureTransformer.psqtWeights[index * PSQTBuckets];
-        };
-
-        fused_row_reduce<Vec16Wrapper, Dimensions, ops...>(from.accumulation[perspective].data(),
-                                                           to.accumulation[perspective].data(),
-                                                           to_weight_vector(indices)...);
-
-        fused_row_reduce<Vec32Wrapper, PSQTBuckets, ops...>(
-          from.psqtAccumulation[perspective].data(), to.psqtAccumulation[perspective].data(),
-          to_psqt_weight_vector(indices)...);
-    }
-
-    void apply(const typename FeatureSet::IndexList& added,
-               const typename FeatureSet::IndexList& removed) {
     constexpr IndexType Dimensions = FeatureTransformer::OutputDimensions;
 
     const auto& fromAcc = from.accumulation[perspective];
@@ -286,21 +177,38 @@ struct AccumulatorUpdateContext {
     vec_t      acc[Tiling::NumRegs];
     psqt_vec_t psqt[Tiling::NumPsqtRegs];
 
+    const auto* psqWeights    = &featureTransformer.weights[0];
     const auto* threatWeights = &featureTransformer.threatWeights[0];
 
     for (IndexType j = 0; j < Dimensions / Tiling::TileHeight; ++j)
     {
-            auto* fromTile = reinterpret_cast<const vec_t*>(&fromAcc[j * Tiling::TileHeight]);
+        const usize tileOff  = j * Tiling::TileHeight;
-            auto* toTile   = reinterpret_cast<vec_t*>(&toAcc[j * Tiling::TileHeight]);
+        auto*       fromTile = reinterpret_cast<const vec_t*>(&fromAcc[tileOff]);
+        auto*       toTile   = reinterpret_cast<vec_t*>(&toAcc[tileOff]);
 
         for (IndexType k = 0; k < Tiling::NumRegs; ++k)
             acc[k] = fromTile[k];
 
-            for (int i = 0; i < removed.ssize(); ++i)
+        for (int i = 0; i < psqRemoved.ssize(); ++i)
         {
-                usize       index  = removed[i];
+            auto* row =
-                const usize offset = Dimensions * index;
+              reinterpret_cast<const vec_t*>(&psqWeights[psqRemoved[i] * Dimensions + tileOff]);
-                auto*       column = reinterpret_cast<const vec_i8_t*>(&threatWeights[offset]);
+            for (IndexType k = 0; k < Tiling::NumRegs; ++k)
+                acc[k] = vec_sub_16(acc[k], row[k]);
+        }
+
+        for (int i = 0; i < psqAdded.ssize(); ++i)
+        {
+            auto* row =
+              reinterpret_cast<const vec_t*>(&psqWeights[psqAdded[i] * Dimensions + tileOff]);
+            for (IndexType k = 0; k < Tiling::NumRegs; ++k)
+                acc[k] = vec_add_16(acc[k], row[k]);
+        }
+
+        for (int i = 0; i < thrRemoved.ssize(); ++i)
+        {
+            auto* column = reinterpret_cast<const vec_i8_t*>(
+              &threatWeights[thrRemoved[i] * Dimensions + tileOff]);
 
     #ifdef USE_NEON
             for (IndexType k = 0; k < Tiling::NumRegs; k += 2)
@@ -314,11 +222,10 @@ struct AccumulatorUpdateContext {
     #endif
         }
 
-            for (int i = 0; i < added.ssize(); ++i)
+        for (int i = 0; i < thrAdded.ssize(); ++i)
         {
-                usize       index  = added[i];
+            auto* column =
-                const usize offset = Dimensions * index;
+              reinterpret_cast<const vec_i8_t*>(&threatWeights[thrAdded[i] * Dimensions + tileOff]);
-                auto*       column = reinterpret_cast<const vec_i8_t*>(&threatWeights[offset]);
 
     #ifdef USE_NEON
             for (IndexType k = 0; k < Tiling::NumRegs; k += 2)
@@ -334,38 +241,45 @@ struct AccumulatorUpdateContext {
 
         for (IndexType k = 0; k < Tiling::NumRegs; k++)
             vec_store(&toTile[k], acc[k]);
-
-            threatWeights += Tiling::TileHeight;
     }
 
     for (IndexType j = 0; j < PSQTBuckets / Tiling::PsqtTileHeight; ++j)
     {
-            auto* fromTilePsqt =
+        const usize psqtTileOff  = j * Tiling::PsqtTileHeight;
-              reinterpret_cast<const psqt_vec_t*>(&fromPsqtAcc[j * Tiling::PsqtTileHeight]);
+        auto*       fromTilePsqt = reinterpret_cast<const psqt_vec_t*>(&fromPsqtAcc[psqtTileOff]);
-            auto* toTilePsqt =
+        auto*       toTilePsqt   = reinterpret_cast<psqt_vec_t*>(&toPsqtAcc[psqtTileOff]);
-              reinterpret_cast<psqt_vec_t*>(&toPsqtAcc[j * Tiling::PsqtTileHeight]);
 
         for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
             psqt[k] = fromTilePsqt[k];
 
-            for (int i = 0; i < removed.ssize(); ++i)
+        for (int i = 0; i < psqRemoved.ssize(); ++i)
         {
-                usize       index      = removed[i];
-                const usize offset     = PSQTBuckets * index + j * Tiling::PsqtTileHeight;
             auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
-                  &featureTransformer.threatPsqtWeights[offset]);
+              &featureTransformer.psqtWeights[psqRemoved[i] * PSQTBuckets + psqtTileOff]);
-
             for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
                 psqt[k] = vec_sub_psqt_32(psqt[k], columnPsqt[k]);
         }
 
-            for (int i = 0; i < added.ssize(); ++i)
+        for (int i = 0; i < psqAdded.ssize(); ++i)
         {
-                usize       index      = added[i];
-                const usize offset     = PSQTBuckets * index + j * Tiling::PsqtTileHeight;
             auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
-                  &featureTransformer.threatPsqtWeights[offset]);
+              &featureTransformer.psqtWeights[psqAdded[i] * PSQTBuckets + psqtTileOff]);
+            for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
+                psqt[k] = vec_add_psqt_32(psqt[k], columnPsqt[k]);
+        }
 
+        for (int i = 0; i < thrRemoved.ssize(); ++i)
+        {
+            auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
+              &featureTransformer.threatPsqtWeights[thrRemoved[i] * PSQTBuckets + psqtTileOff]);
+            for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
+                psqt[k] = vec_sub_psqt_32(psqt[k], columnPsqt[k]);
+        }
+
+        for (int i = 0; i < thrAdded.ssize(); ++i)
+        {
+            auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
+              &featureTransformer.threatPsqtWeights[thrAdded[i] * PSQTBuckets + psqtTileOff]);
             for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
                 psqt[k] = vec_add_psqt_32(psqt[k], columnPsqt[k]);
         }
@@ -379,47 +293,51 @@ struct AccumulatorUpdateContext {
     toAcc     = fromAcc;
     toPsqtAcc = fromPsqtAcc;
 
-        for (const auto index : removed)
+    for (const auto index : psqRemoved)
     {
         const IndexType offset = Dimensions * index;
+        for (IndexType j = 0; j < Dimensions; ++j)
+            toAcc[j] -= featureTransformer.weights[offset + j];
+        for (usize k = 0; k < PSQTBuckets; ++k)
+            toPsqtAcc[k] -= featureTransformer.psqtWeights[index * PSQTBuckets + k];
+    }
 
+    for (const auto index : psqAdded)
+    {
+        const IndexType offset = Dimensions * index;
+        for (IndexType j = 0; j < Dimensions; ++j)
+            toAcc[j] += featureTransformer.weights[offset + j];
+        for (usize k = 0; k < PSQTBuckets; ++k)
+            toPsqtAcc[k] += featureTransformer.psqtWeights[index * PSQTBuckets + k];
+    }
+
+    for (const auto index : thrRemoved)
+    {
+        const IndexType offset = Dimensions * index;
         for (IndexType j = 0; j < Dimensions; ++j)
             toAcc[j] -= featureTransformer.threatWeights[offset + j];
-
         for (usize k = 0; k < PSQTBuckets; ++k)
             toPsqtAcc[k] -= featureTransformer.threatPsqtWeights[index * PSQTBuckets + k];
     }
 
-        for (const auto index : added)
+    for (const auto index : thrAdded)
     {
         const IndexType offset = Dimensions * index;
-
         for (IndexType j = 0; j < Dimensions; ++j)
             toAcc[j] += featureTransformer.threatWeights[offset + j];
-
         for (usize k = 0; k < PSQTBuckets; ++k)
             toPsqtAcc[k] += featureTransformer.threatPsqtWeights[index * PSQTBuckets + k];
     }
 
 #endif
-    }
-};
-
-template<typename FeatureSet>
-auto make_accumulator_update_context(Color                               perspective,
-                                     const FeatureTransformer&           featureTransformer,
-                                     const AccumulatorState<FeatureSet>& accumulatorFrom,
-                                     AccumulatorState<FeatureSet>&       accumulatorTo) noexcept {
-    return AccumulatorUpdateContext<FeatureSet>{perspective, featureTransformer, accumulatorFrom,
-                                                accumulatorTo};
 }
 
-template<bool Forward, typename FeatureSet>
+template<bool Forward>
 void update_accumulator_incremental(Color                     perspective,
                                     const FeatureTransformer& featureTransformer,
                                     const Square              ksq,
-                                    AccumulatorState<FeatureSet>&       target_state,
+                                    AccumulatorState&         target_state,
-                                    const AccumulatorState<FeatureSet>& computed) {
+                                    const AccumulatorState&   computed) {
 
     assert(computed.computed[perspective]);
     assert(!target_state.computed[perspective]);
@@ -428,74 +346,30 @@ void update_accumulator_incremental(Color                               perspect
     // That might depend on the feature set and generally relies on the
     // feature set's update cost calculation to be correct and never allow
     // updates with more added/removed features than MaxActiveDimensions.
-    // In this case, the maximum size of both feature addition and removal
+    PSQFeatureSet::IndexList    psqRemoved, psqAdded;
-    // is 2, since we are incrementally updating one move at a time.
+    ThreatFeatureSet::IndexList thrRemoved, thrAdded;
-    typename FeatureSet::IndexList removed, added;
+
-    if constexpr (std::is_same_v<FeatureSet, ThreatFeatureSet>)
+    const auto& dirtyPiece   = Forward ? target_state.dirtyPiece : computed.dirtyPiece;
-    {
+    const auto& dirtyThreats = Forward ? target_state.dirtyThreats : computed.dirtyThreats;
+
     const auto* pfBase   = &featureTransformer.threatWeights[0];
     IndexType   pfStride = FeatureTransformer::OutputDimensions;
+
     if constexpr (Forward)
-            FeatureSet::append_changed_indices(perspective, ksq, target_state.diff, removed, added,
+    {
-                                               nullptr, false, pfBase, pfStride);
+        ThreatFeatureSet::append_changed_indices(perspective, ksq, dirtyThreats, thrRemoved,
-        else
+                                                 thrAdded, nullptr, false, pfBase, pfStride);
-            FeatureSet::append_changed_indices(perspective, ksq, computed.diff, added, removed,
+        PSQFeatureSet::append_changed_indices(perspective, ksq, dirtyPiece, psqRemoved, psqAdded);
-                                               nullptr, false, pfBase, pfStride);
     }
     else
     {
-        if constexpr (Forward)
+        ThreatFeatureSet::append_changed_indices(perspective, ksq, dirtyThreats, thrAdded,
-            FeatureSet::append_changed_indices(perspective, ksq, target_state.diff, removed, added);
+                                                 thrRemoved, nullptr, false, pfBase, pfStride);
-        else
+        PSQFeatureSet::append_changed_indices(perspective, ksq, dirtyPiece, psqAdded, psqRemoved);
-            FeatureSet::append_changed_indices(perspective, ksq, computed.diff, added, removed);
     }
 
-    auto updateContext =
+    apply_combined(perspective, featureTransformer, computed, target_state, psqAdded, psqRemoved,
-      make_accumulator_update_context(perspective, featureTransformer, computed, target_state);
+                   thrAdded, thrRemoved);
-
-    if constexpr (std::is_same_v<FeatureSet, ThreatFeatureSet>)
-        updateContext.apply(added, removed);
-    else
-    {
-        [[maybe_unused]] const int addedSize   = added.ssize();
-        [[maybe_unused]] const int removedSize = removed.ssize();
-
-        assert(addedSize == 1 || addedSize == 2);
-        assert(removedSize == 1 || removedSize == 2);
-        assert((Forward && addedSize <= removedSize) || (!Forward && addedSize >= removedSize));
-
-        // Workaround compiler warning for uninitialized variables, replicated
-        // on profile builds on windows with gcc 14.2.0.
-        // Also helps with optimizations on some compilers.
-
-        sf_assume(addedSize == 1 || addedSize == 2);
-        sf_assume(removedSize == 1 || removedSize == 2);
-
-        if (!(removedSize == 1 || removedSize == 2) || !(addedSize == 1 || addedSize == 2))
-            sf_unreachable();
-
-        if ((Forward && removedSize == 1) || (!Forward && addedSize == 1))
-        {
-            assert(addedSize == 1 && removedSize == 1);
-            updateContext.template apply<Add, Sub>(added[0], removed[0]);
-        }
-        else if (Forward && addedSize == 1)
-        {
-            assert(removedSize == 2);
-            updateContext.template apply<Add, Sub, Sub>(added[0], removed[0], removed[1]);
-        }
-        else if (!Forward && removedSize == 1)
-        {
-            assert(addedSize == 2);
-            updateContext.template apply<Add, Add, Sub>(added[0], added[1], removed[0]);
-        }
-        else
-        {
-            assert(addedSize == 2 && removedSize == 2);
-            updateContext.template apply<Add, Add, Sub, Sub>(added[0], added[1], removed[0],
-                                                             removed[1]);
-        }
-    }
 
     target_state.computed[perspective] = true;
 }
@@ -584,10 +458,12 @@ Bitboard get_changed_pieces(const std::array<Piece, SQUARE_NB>& oldPieces,
 #endif
 }
 
+// HalfKA data comes from the Finny table entry, while the threats are built
+// from the active threat features
 void update_accumulator_refresh_cache(Color                     perspective,
                                       const FeatureTransformer& featureTransformer,
                                       const Position&           pos,
-                                      AccumulatorState<PSQFeatureSet>& accumulator,
+                                      AccumulatorState&         accumulator,
                                       AccumulatorCaches&        cache) {
     constexpr auto Dimensions = FeatureTransformer::OutputDimensions;
 
@@ -620,6 +496,9 @@ void update_accumulator_refresh_cache(Color                            perspecti
     entry.pieceBB = pos.pieces();
     entry.pieces  = pos.piece_array();
 
+    ThreatFeatureSet::IndexList active;
+    ThreatFeatureSet::append_active_indices(perspective, pos, active);
+
     accumulator.computed[perspective] = true;
 
 #ifdef VECTOR
@@ -627,76 +506,92 @@ void update_accumulator_refresh_cache(Color                            perspecti
     psqt_vec_t psqt[Tiling::NumPsqtRegs];
 
     const auto* weights       = &featureTransformer.weights[0];
+    const auto* threatWeights = &featureTransformer.threatWeights[0];
 
     for (IndexType j = 0; j < Dimensions / Tiling::TileHeight; ++j)
     {
-        auto* accTile =
+        const usize tileOff = j * Tiling::TileHeight;
-          reinterpret_cast<vec_t*>(&accumulator.accumulation[perspective][j * Tiling::TileHeight]);
+        auto* accTile   = reinterpret_cast<vec_t*>(&accumulator.accumulation[perspective][tileOff]);
-        auto* entryTile = reinterpret_cast<vec_t*>(&entry.accumulation[j * Tiling::TileHeight]);
+        auto* entryTile = reinterpret_cast<vec_t*>(&entry.accumulation[tileOff]);
 
         for (IndexType k = 0; k < Tiling::NumRegs; ++k)
             acc[k] = entryTile[k];
 
         for (int i = 0; i < removed.ssize(); ++i)
         {
-            usize       index  = removed[i];
+            auto* column =
-            const usize offset = Dimensions * index;
+              reinterpret_cast<const vec_t*>(&weights[removed[i] * Dimensions + tileOff]);
-            auto*       column = reinterpret_cast<const vec_t*>(&weights[offset]);
-
             for (IndexType k = 0; k < Tiling::NumRegs; ++k)
                 acc[k] = vec_sub_16(acc[k], column[k]);
         }
         for (int i = 0; i < added.ssize(); ++i)
         {
-            usize       index  = added[i];
+            auto* column =
-            const usize offset = Dimensions * index;
+              reinterpret_cast<const vec_t*>(&weights[added[i] * Dimensions + tileOff]);
-            auto*       column = reinterpret_cast<const vec_t*>(&weights[offset]);
-
             for (IndexType k = 0; k < Tiling::NumRegs; ++k)
                 acc[k] = vec_add_16(acc[k], column[k]);
         }
 
         for (IndexType k = 0; k < Tiling::NumRegs; k++)
             vec_store(&entryTile[k], acc[k]);
+
+        for (int i = 0; i < active.ssize(); ++i)
+        {
+            auto* column =
+              reinterpret_cast<const vec_i8_t*>(&threatWeights[active[i] * Dimensions + tileOff]);
+
+    #ifdef USE_NEON
+            for (IndexType k = 0; k < Tiling::NumRegs; k += 2)
+            {
+                acc[k]     = vaddw_s8(acc[k], vget_low_s8(column[k / 2]));
+                acc[k + 1] = vaddw_high_s8(acc[k + 1], column[k / 2]);
+            }
+    #else
+            for (IndexType k = 0; k < Tiling::NumRegs; ++k)
+                acc[k] = vec_add_16(acc[k], vec_convert_8_16(column[k]));
+    #endif
+        }
+
         for (IndexType k = 0; k < Tiling::NumRegs; k++)
             vec_store(&accTile[k], acc[k]);
-
-        weights += Tiling::TileHeight;
     }
 
     for (IndexType j = 0; j < PSQTBuckets / Tiling::PsqtTileHeight; ++j)
     {
-        auto* accTilePsqt = reinterpret_cast<psqt_vec_t*>(
+        const usize psqtTileOff = j * Tiling::PsqtTileHeight;
-          &accumulator.psqtAccumulation[perspective][j * Tiling::PsqtTileHeight]);
+        auto*       accTilePsqt =
-        auto* entryTilePsqt =
+          reinterpret_cast<psqt_vec_t*>(&accumulator.psqtAccumulation[perspective][psqtTileOff]);
-          reinterpret_cast<psqt_vec_t*>(&entry.psqtAccumulation[j * Tiling::PsqtTileHeight]);
+        auto* entryTilePsqt = reinterpret_cast<psqt_vec_t*>(&entry.psqtAccumulation[psqtTileOff]);
 
         for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
             psqt[k] = entryTilePsqt[k];
 
         for (int i = 0; i < removed.ssize(); ++i)
         {
-            usize       index  = removed[i];
+            auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
-            const usize offset = PSQTBuckets * index + j * Tiling::PsqtTileHeight;
+              &featureTransformer.psqtWeights[removed[i] * PSQTBuckets + psqtTileOff]);
-            auto*       columnPsqt =
-              reinterpret_cast<const psqt_vec_t*>(&featureTransformer.psqtWeights[offset]);
-
             for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
                 psqt[k] = vec_sub_psqt_32(psqt[k], columnPsqt[k]);
         }
         for (int i = 0; i < added.ssize(); ++i)
         {
-            usize       index  = added[i];
+            auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
-            const usize offset = PSQTBuckets * index + j * Tiling::PsqtTileHeight;
+              &featureTransformer.psqtWeights[added[i] * PSQTBuckets + psqtTileOff]);
-            auto*       columnPsqt =
-              reinterpret_cast<const psqt_vec_t*>(&featureTransformer.psqtWeights[offset]);
-
             for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
                 psqt[k] = vec_add_psqt_32(psqt[k], columnPsqt[k]);
         }
 
         for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
             vec_store_psqt(&entryTilePsqt[k], psqt[k]);
+
+        for (int i = 0; i < active.ssize(); ++i)
+        {
+            auto* columnPsqt = reinterpret_cast<const psqt_vec_t*>(
+              &featureTransformer.threatPsqtWeights[active[i] * PSQTBuckets + psqtTileOff]);
+            for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
+                psqt[k] = vec_add_psqt_32(psqt[k], columnPsqt[k]);
+        }
+
         for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
             vec_store_psqt(&accTilePsqt[k], psqt[k]);
     }
@@ -726,91 +621,6 @@ void update_accumulator_refresh_cache(Color                            perspecti
     // Now copy its content to the actual accumulator we were refreshing.
     accumulator.accumulation[perspective]     = entry.accumulation;
     accumulator.psqtAccumulation[perspective] = entry.psqtAccumulation;
-#endif
-}
-
-void update_threats_accumulator_full(Color                               perspective,
-                                     const FeatureTransformer&           featureTransformer,
-                                     const Position&                     pos,
-                                     AccumulatorState<ThreatFeatureSet>& accumulator) {
-    constexpr IndexType Dimensions = FeatureTransformer::OutputDimensions;
-    using Tiling [[maybe_unused]]  = SIMDTiling<Dimensions, Dimensions, PSQTBuckets>;
-
-    ThreatFeatureSet::IndexList active;
-    ThreatFeatureSet::append_active_indices(perspective, pos, active);
-
-    accumulator.computed[perspective] = true;
-
-#ifdef VECTOR
-    vec_t      acc[Tiling::NumRegs];
-    psqt_vec_t psqt[Tiling::NumPsqtRegs];
-
-    const auto* threatWeights = &featureTransformer.threatWeights[0];
-
-    for (IndexType j = 0; j < Dimensions / Tiling::TileHeight; ++j)
-    {
-        auto* accTile =
-          reinterpret_cast<vec_t*>(&accumulator.accumulation[perspective][j * Tiling::TileHeight]);
-
-        for (IndexType k = 0; k < Tiling::NumRegs; ++k)
-            acc[k] = vec_zero();
-
-        int i = 0;
-
-        for (; i < active.ssize(); ++i)
-        {
-            usize       index  = active[i];
-            const usize offset = Dimensions * index;
-            auto*       column = reinterpret_cast<const vec_i8_t*>(&threatWeights[offset]);
-
-    #ifdef USE_NEON
-            for (IndexType k = 0; k < Tiling::NumRegs; k += 2)
-            {
-                acc[k]     = vaddw_s8(acc[k], vget_low_s8(column[k / 2]));
-                acc[k + 1] = vaddw_high_s8(acc[k + 1], column[k / 2]);
-            }
-    #else
-            for (IndexType k = 0; k < Tiling::NumRegs; ++k)
-                acc[k] = vec_add_16(acc[k], vec_convert_8_16(column[k]));
-    #endif
-        }
-
-        for (IndexType k = 0; k < Tiling::NumRegs; k++)
-            vec_store(&accTile[k], acc[k]);
-
-        threatWeights += Tiling::TileHeight;
-    }
-
-    for (IndexType j = 0; j < PSQTBuckets / Tiling::PsqtTileHeight; ++j)
-    {
-        auto* accTilePsqt = reinterpret_cast<psqt_vec_t*>(
-          &accumulator.psqtAccumulation[perspective][j * Tiling::PsqtTileHeight]);
-
-        for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
-            psqt[k] = vec_zero_psqt();
-
-        for (int i = 0; i < active.ssize(); ++i)
-        {
-            usize       index  = active[i];
-            const usize offset = PSQTBuckets * index + j * Tiling::PsqtTileHeight;
-            auto*       columnPsqt =
-              reinterpret_cast<const psqt_vec_t*>(&featureTransformer.threatPsqtWeights[offset]);
-
-            for (usize k = 0; k < Tiling::NumPsqtRegs; ++k)
-                psqt[k] = vec_add_psqt_32(psqt[k], columnPsqt[k]);
-        }
-
-        for (IndexType k = 0; k < Tiling::NumPsqtRegs; ++k)
-            vec_store_psqt(&accTilePsqt[k], psqt[k]);
-    }
-
-#else
-
-    for (IndexType j = 0; j < Dimensions; ++j)
-        accumulator.accumulation[perspective][j] = 0;
-
-    for (usize k = 0; k < PSQTBuckets; ++k)
-        accumulator.psqtAccumulation[perspective][k] = 0;
 
     for (const auto index : active)
     {

src/nnue/nnue_accumulator.h

@@ -41,7 +41,8 @@ struct alignas(CacheLineSize) Accumulator;
 
 class FeatureTransformer;
 
-// Class that holds the result of affine transformation of input features
+// Class that holds the result of affine transformation of input features,
+// combined HalfKA + Threats
 struct alignas(CacheLineSize) Accumulator {
     std::array<std::array<i16, L1>, COLOR_NB>          accumulation;
     std::array<std::array<i32, PSQTBuckets>, COLOR_NB> psqtAccumulation;
@@ -89,27 +90,16 @@ struct AccumulatorCaches {
 };
 
 
-template<typename FeatureSet>
 struct AccumulatorState: public Accumulator {
-    typename FeatureSet::DiffType diff;
+    DirtyPiece   dirtyPiece;
-
+    DirtyThreats dirtyThreats;
-    void reset(const typename FeatureSet::DiffType& dp) noexcept {
-        diff = dp;
-        computed.fill(false);
-    }
-
-    typename FeatureSet::DiffType& reset() noexcept {
-        computed.fill(false);
-        return diff;
-    }
 };
 
 class AccumulatorStack {
    public:
     static constexpr usize MaxSize = MAX_PLY + 1;
 
-    template<typename T>
+    [[nodiscard]] const AccumulatorState& latest() const noexcept;
-    [[nodiscard]] const AccumulatorState<T>& latest() const noexcept;
 
     void                                  reset() noexcept;
     std::pair<DirtyPiece&, DirtyThreats&> push() noexcept;
@@ -121,39 +111,27 @@ class AccumulatorStack {
                   [[maybe_unused]] AccumulatorCaches& cache) noexcept;
 
    private:
-    template<typename T>
+    [[nodiscard]] AccumulatorState& mut_latest() noexcept;
-    [[nodiscard]] AccumulatorState<T>& mut_latest() noexcept;
 
-    template<typename T>
-    [[nodiscard]] const std::array<AccumulatorState<T>, MaxSize>& accumulators() const noexcept;
-
-    template<typename T>
-    [[nodiscard]] std::array<AccumulatorState<T>, MaxSize>& mut_accumulators() noexcept;
-
-    template<typename FeatureSet>
     void evaluate_side(Color                     perspective,
                        const Position&           pos,
                        const FeatureTransformer& featureTransformer,
                        // Silence spurious warning on GCC 10
                        [[maybe_unused]] AccumulatorCaches& cache) noexcept;
 
-    template<typename FeatureSet>
     [[nodiscard]] usize find_last_usable_accumulator(Color perspective) const noexcept;
 
-    template<typename FeatureSet>
     void forward_update_incremental(Color                     perspective,
                                     const Position&           pos,
                                     const FeatureTransformer& featureTransformer,
                                     const usize               begin) noexcept;
 
-    template<typename FeatureSet>
     void backward_update_incremental(Color                     perspective,
                                      const Position&           pos,
                                      const FeatureTransformer& featureTransformer,
                                      const usize               end) noexcept;
 
-    std::array<AccumulatorState<PSQFeatureSet>, MaxSize>    psq_accumulators;
+    std::array<AccumulatorState, MaxSize> accumulators;
-    std::array<AccumulatorState<ThreatFeatureSet>, MaxSize> threat_accumulators;
     usize                                 size = 1;
 };
 

src/nnue/nnue_feature_transformer.h

@@ -205,21 +205,15 @@ class FeatureTransformer {
 
         using namespace SIMD;
         accumulatorStack.evaluate(pos, *this, cache);
-        const auto& accumulatorState       = accumulatorStack.latest<PSQFeatureSet>();
+        const auto& accumulatorState = accumulatorStack.latest();
-        const auto& threatAccumulatorState = accumulatorStack.latest<ThreatFeatureSet>();
 
         const Color perspectives[2]  = {pos.side_to_move(), ~pos.side_to_move()};
         const auto& psqtAccumulation = accumulatorState.psqtAccumulation;
-        auto        psqt =
+        const auto  psqt =
-          (psqtAccumulation[perspectives[0]][bucket] - psqtAccumulation[perspectives[1]][bucket]);
+          (psqtAccumulation[perspectives[0]][bucket] - psqtAccumulation[perspectives[1]][bucket])
-
-        const auto& threatPsqtAccumulation = threatAccumulatorState.psqtAccumulation;
-        psqt                               = (psqt + threatPsqtAccumulation[perspectives[0]][bucket]
-                - threatPsqtAccumulation[perspectives[1]][bucket])
           / 2;
 
         const auto& accumulation = accumulatorState.accumulation;
-        const auto& threatAccumulation = threatAccumulatorState.accumulation;
 
         for (IndexType p = 0; p < 2; ++p)
         {
@@ -289,10 +283,6 @@ class FeatureTransformer {
             // 8 bits. Shifting it by 7 bits left will no longer occupy the
             // signed bit, so we are safe.
 
-            const vec_t* tin0 =
-              reinterpret_cast<const vec_t*>(&(threatAccumulation[perspectives[p]][0]));
-            const vec_t* tin1 = reinterpret_cast<const vec_t*>(
-              &(threatAccumulation[perspectives[p]][HalfDimensions / 2]));
             for (IndexType j = 0; j < NumOutputChunks; j += 2)
             {
                 vec_t packed[2];
@@ -300,10 +290,10 @@ class FeatureTransformer {
                 {
                     const IndexType i = (j + k) * 2;
 
-                    vec_t acc0a = vec_add_16(in0[i + 0], tin0[i + 0]);
+                    vec_t acc0a = in0[i + 0];
-                    vec_t acc0b = vec_add_16(in0[i + 1], tin0[i + 1]);
+                    vec_t acc0b = in0[i + 1];
-                    vec_t acc1a = vec_add_16(in1[i + 0], tin1[i + 0]);
+                    vec_t acc1a = in1[i + 0];
-                    vec_t acc1b = vec_add_16(in1[i + 1], tin1[i + 1]);
+                    vec_t acc1b = in1[i + 1];
 
                     static_assert(FtMaxVal == 255);
 
@@ -360,10 +350,6 @@ class FeatureTransformer {
                 BiasType sum1 =
                   accumulation[static_cast<int>(perspectives[p])][j + HalfDimensions / 2];
 
-                sum0 += threatAccumulation[static_cast<int>(perspectives[p])][j + 0];
-                sum1 +=
-                  threatAccumulation[static_cast<int>(perspectives[p])][j + HalfDimensions / 2];
-
                 sum0 = std::clamp<BiasType>(sum0, 0, FtMaxVal);
                 sum1 = std::clamp<BiasType>(sum1, 0, FtMaxVal);