diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index 60a044158..7e4c669ae 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -704,6 +704,8 @@ class FeatureTransformer { accumulator.computed[Perspective] = true; #ifdef VECTOR + const bool combineLast3 = std::abs((int) removed.size() - (int) added.size()) == 1 + && removed.size() + added.size() > 2; vec_t acc[Tiling::NumRegs]; psqt_vec_t psqt[Tiling::NumPsqtRegs]; @@ -717,7 +719,7 @@ class FeatureTransformer { acc[k] = entryTile[k]; std::size_t i = 0; - for (; i < std::min(removed.size(), added.size()); ++i) + for (; i < std::min(removed.size(), added.size()) - combineLast3; ++i) { IndexType indexR = removed[i]; const IndexType offsetR = HalfDimensions * indexR + j * Tiling::TileHeight; @@ -729,23 +731,56 @@ class FeatureTransformer { for (IndexType k = 0; k < Tiling::NumRegs; ++k) acc[k] = vec_add_16(acc[k], vec_sub_16(columnA[k], columnR[k])); } - for (; i < removed.size(); ++i) + if (combineLast3) { - IndexType index = removed[i]; - const IndexType offset = HalfDimensions * index + j * Tiling::TileHeight; - auto* column = reinterpret_cast(&weights[offset]); + IndexType indexR = removed[i]; + const IndexType offsetR = HalfDimensions * indexR + j * Tiling::TileHeight; + auto* columnR = reinterpret_cast(&weights[offsetR]); + IndexType indexA = added[i]; + const IndexType offsetA = HalfDimensions * indexA + j * Tiling::TileHeight; + auto* columnA = reinterpret_cast(&weights[offsetA]); - for (IndexType k = 0; k < Tiling::NumRegs; ++k) - acc[k] = vec_sub_16(acc[k], column[k]); + if (removed.size() > added.size()) + { + IndexType indexR2 = removed[i + 1]; + const IndexType offsetR2 = HalfDimensions * indexR2 + j * Tiling::TileHeight; + auto* columnR2 = reinterpret_cast(&weights[offsetR2]); + + for (IndexType k = 0; k < Tiling::NumRegs; ++k) + acc[k] = vec_sub_16(vec_add_16(acc[k], columnA[k]), + vec_add_16(columnR[k], columnR2[k])); + } + else + { + IndexType indexA2 = added[i + 1]; + const IndexType offsetA2 = HalfDimensions * indexA2 + j * Tiling::TileHeight; + auto* columnA2 = reinterpret_cast(&weights[offsetA2]); + + for (IndexType k = 0; k < Tiling::NumRegs; ++k) + acc[k] = vec_add_16(vec_sub_16(acc[k], columnR[k]), + vec_add_16(columnA[k], columnA2[k])); + } } - for (; i < added.size(); ++i) + else { - IndexType index = added[i]; - const IndexType offset = HalfDimensions * index + j * Tiling::TileHeight; - auto* column = reinterpret_cast(&weights[offset]); + for (; i < removed.size(); ++i) + { + IndexType index = removed[i]; + const IndexType offset = HalfDimensions * index + j * Tiling::TileHeight; + auto* column = reinterpret_cast(&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) + acc[k] = vec_sub_16(acc[k], column[k]); + } + for (; i < added.size(); ++i) + { + IndexType index = added[i]; + const IndexType offset = HalfDimensions * index + j * Tiling::TileHeight; + auto* column = reinterpret_cast(&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++)