From fb1d777250c77532cf9a7d0f11598ed309b49c34 Mon Sep 17 00:00:00 2001
From: Linmiao Xu <linmiao.xu@gmail.com>
Date: Fri, 10 Jul 2026 20:18:25 +0200
Subject: [PATCH] Split AVX-VNNI sparse affine accumulation two ways

Speedup measured by @anematode
```
Result of 100 runs
==================
base (...h.master.exe) = 1504448 +/- 15119
test (./stockfish.exe) = 1519239 +/- 15980
diff = +14791 +/- 4113

speedup = +0.0098
P(speedup > 0) = 1.0000

CPU: Intel Core Ultra 9 285H
```

Passed STC:
https://tests.stockfishchess.org/tests/view/6a4dc4c65529b8472df7fd56
LLR: 2.94 (-2.94,2.94) <0.00,2.00>
Total: 48256 W: 12746 L: 12420 D: 23090
Ptnml(0-2): 144, 5310, 12913, 5598, 163

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

No functional change
---
 .../layers/affine_transform_sparse_input.h    | 42 +++++++++++++++++--
 1 file changed, 38 insertions(+), 4 deletions(-)

diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h
index 168bf2f32..2b3503959 100644
--- a/src/nnue/layers/affine_transform_sparse_input.h
+++ b/src/nnue/layers/affine_transform_sparse_input.h
@@ -167,10 +167,12 @@ class AffineTransformSparseInput {
         constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType);
         constexpr IndexType NumAccums       = OutputDimensions / OutputSimdWidth;
         // If we're using high-latency dot product instructions, split the accumulators
-        // to create 3 separate dependency chains and merge at the end
+        // into separate dependency chains and merge at the end
         constexpr IndexType NumRegs =
     #if (defined(USE_VNNI) && defined(USE_AVX512)) || defined(USE_NEON_DOTPROD)
           3 * NumAccums;
+    #elif defined(USE_AVXVNNI)
+          2 * NumAccums;
     #else
           NumAccums;
     #endif
@@ -180,7 +182,10 @@ class AffineTransformSparseInput {
         for (IndexType k = 0; k < NumAccums; ++k)
             acc[k] = biasvec[k];
 
-    #if defined(USE_NEON_DOTPROD)
+    #if defined(USE_AVXVNNI)
+        for (IndexType k = NumAccums; k < NumRegs; ++k)
+            acc[k] = vec_set_32(0);
+    #elif defined(USE_NEON_DOTPROD)
         for (IndexType k = NumAccums; k < NumRegs; ++k)
             acc[k] = vdupq_n_s32(0);
     #endif
@@ -255,7 +260,33 @@ class AffineTransformSparseInput {
             asm("" : "+r"(base_addr), "+r"(weights_base));  // opt barrier
         #endif
 
-        #if defined(USE_NEON_DOTPROD)
+        #if defined(USE_AVXVNNI)
+            while (bits)
+            {
+                const isize   i0   = pop_lsb(bits);
+                const invec_t in0  = vec_set_32(load_as<i32>(base_addr + i0 * sizeof(i32)));
+                const auto    col0 = reinterpret_cast<const invec_t*>(
+                  &weights_base[i0 * OutputDimensions * ChunkSize]);
+
+                if (!bits)
+                {
+                    for (IndexType l = 0; l < NumAccums; ++l)
+                        vec_add_dpbusd_32(acc[l], in0, col0[l]);
+                    break;
+                }
+
+                const isize   i1   = pop_lsb(bits);
+                const invec_t in1  = vec_set_32(load_as<i32>(base_addr + i1 * sizeof(i32)));
+                const auto    col1 = reinterpret_cast<const invec_t*>(
+                  &weights_base[i1 * OutputDimensions * ChunkSize]);
+
+                for (IndexType l = 0; l < NumAccums; ++l)
+                {
+                    vec_add_dpbusd_32(acc[l], in0, col0[l]);
+                    vec_add_dpbusd_32(acc[l + NumAccums], in1, col1[l]);
+                }
+            }
+        #elif defined(USE_NEON_DOTPROD)
             while (bits)
             {
                 const isize i0 = pop_lsb(bits);
@@ -323,7 +354,10 @@ class AffineTransformSparseInput {
         #endif
         }
 
-        #if defined(USE_NEON_DOTPROD)
+        #if defined(USE_AVXVNNI)
+        for (IndexType l = 0; l < NumAccums; ++l)
+            acc[l] = vec_add_32(acc[l], acc[l + NumAccums]);
+        #elif defined(USE_NEON_DOTPROD)
         for (IndexType l = 0; l < NumAccums; ++l)
             acc[l] = vaddq_s32(vaddq_s32(acc[l], acc[l + NumAccums]), acc[l + 2 * NumAccums]);
         #endif