From 9b164d952061213da4fc0f7ac8646e44e8a77cd5 Mon Sep 17 00:00:00 2001
From: Timothy Herchen <timothy.herchen@gmail.com>
Date: Sat, 27 Sep 2025 08:40:42 -0700
Subject: [PATCH] Shave some instructions off a hot loop in affine transform

On x86, GCC generates highly suboptimal code for this loop in its old form,
about 2x as many instructions as necessary. This decreases throughput
especially in an SMT setting. Clang does a better job but this change still has
some improvement. Note that the std::ptrdiff_t type is not optional; using an
unsigned type brings back the bad assembly. (Not sure why, but it seems
reliable on all the GCC versions I tested.)

passed STC:
LLR: 2.93 (-2.94,2.94) <0.00,2.00>
Total: 44672 W: 11841 L: 11527 D: 21304
Ptnml(0-2): 165, 4625, 12415, 4993, 138
https://tests.stockfishchess.org/tests/view/68d8111efa806e2e8393b10e

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

No functional change
---
 AUTHORS                                         |  1 +
 src/nnue/layers/affine_transform_sparse_input.h | 17 ++++++++++++-----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/AUTHORS b/AUTHORS
index 6bd323d2c..1fb91adaf 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -243,6 +243,7 @@ Thanar2
 thaspel
 theo77186
 TierynnB
+Timothy Herchen (anematode)
 Ting-Hsuan Huang (fffelix-huang)
 Tobias Steinmann
 Tomasz Sobczyk (Sopel97)
diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h
index a073c6196..11e460666 100644
--- a/src/nnue/layers/affine_transform_sparse_input.h
+++ b/src/nnue/layers/affine_transform_sparse_input.h
@@ -22,6 +22,7 @@
 #define NNUE_LAYERS_AFFINE_TRANSFORM_SPARSE_INPUT_H_INCLUDED
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <iostream>
 
@@ -287,12 +288,18 @@ class AffineTransformSparseInput {
         for (IndexType k = 0; k < NumRegs; ++k)
             acc[k] = biasvec[k];
 
-        for (IndexType j = 0; j < count; ++j)
+        auto* start = nnz;
+        auto* end   = nnz + count;
+
+        // convince GCC to not do weird pointer arithmetic in the following loop
+        const std::int8_t* weights_cp = weights;
+
+        while (start < end)
         {
-            const auto    i  = nnz[j];
-            const invec_t in = vec_set_32(input32[i]);
-            const auto    col =
-              reinterpret_cast<const invec_t*>(&weights[i * OutputDimensions * ChunkSize]);
+            const std::ptrdiff_t i = *start;
+            start++;
+            const invec_t in  = vec_set_32(input32[i]);
+            const auto    col = (const invec_t*) (&weights_cp[i * OutputDimensions * ChunkSize]);
             for (IndexType k = 0; k < NumRegs; ++k)
                 vec_add_dpbusd_32(acc[k], in, col[k]);
         }