I asked Claude to fix the bug:Spunc595 wrote: ↑Wed Jul 15, 2026 1:57 pm Hi everyone,
I edited the nnue.rs file as suggested by Jim and it resulted in a significant boost in both nodes and NPS. The engine is now much faster and more stable during search.
However, I’ve hit a critical issue with the NNUE evaluation. Despite the performance gains, the evaluation itself doesn’t seem to be correct. I’ve already performed an in-depth check of the NNUE’s Layers 1, 2, and 3, and the calculation logic seems fine; I just can’t isolate the exact point in my code where the error or inconsistency lies.
I am attaching the command-line logs from a few test positions (including the starting position and K vs K) below so you can see what’s happening:
NNUE: Weights loaded (Safe mode, AVX2 backend).
Book: File 'book.bin' non trovato. Si userà solo la rete.
Luna CE v1.2.0
uci
id name Luna_CE v1.2.0
id author Daniele Marpino
option name Hash type spin default 256 min 1 max 1024
uciok
isready
readyok
position startpos
go depth 10
info depth 1 score cp 29 nodes 41 nps 0 time 0 pv d2d3
info depth 2 score cp 29 nodes 229 nps 0 time 0 pv g1f3 d7d5
info depth 3 score cp 29 nodes 1219 nps 1219000 time 1 pv g1f3 d7d5 d2d3
info depth 4 score cp 29 nodes 2529 nps 632250 time 4 pv g1f3 d7d5 b2b3 h7h6
info depth 5 score cp 29 nodes 3845 nps 640833 time 6 pv g1f3 d7d5 b2b3 h7h6 c2c3
info depth 6 score cp 29 nodes 11646 nps 776400 time 15 pv g1f3 d7d5 b1c3 h7h6 c3d5 d8d5
info depth 7 score cp 29 nodes 17200 nps 781818 time 22 pv g1f3 d7d5 b1c3 h7h6 c3d5 d8d5 d2d4
info depth 8 score cp 29 nodes 43793 nps 826283 time 53 pv g1f3 d7d5 b1c3 h7h6 c3d5 d8d5 b2b3 d5f3
info depth 9 score cp 29 nodes 165278 nps 888591 time 186 pv g1f3 d7d5 b2b3 b8c6 b1c3 h7h5 c3d5 d8d5 c1b2
info depth 10 score cp 29 nodes 859558 nps 810139 time 1061 pv g1f3 d7d5 b1c3 h7h5 c3b5 b8c6 b5c7 d8c7 b2b3 f7f6 d2d3
bestmove g1f3
position fen rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBN1 w KQkq - 0 1
go depth 10
info depth 1 score cp 29 nodes 43 nps 0 time 0 pv b2b3
info depth 2 score cp 28 nodes 150 nps 0 time 0 pv b2b3 e7e6
info depth 3 score cp 29 nodes 1154 nps 1154000 time 1 pv b2b3 e7e6 d2d4
info depth 4 score cp 28 nodes 2400 nps 800000 time 3 pv b2b3 e7e6 d2d4 d7d5
info depth 5 score cp 29 nodes 6670 nps 741111 time 9 pv g1f3 d7d5 g2g3 h7h5 b2b4
info depth 6 score cp 28 nodes 14356 nps 683619 time 21 pv g1f3 d7d5 g2g3 h7h5 b2b4 b8c6
info depth 7 score cp 28 nodes 40330 nps 876739 time 46 pv g1f3 d7d5 g2g3 h7h5 b2b4 b8c6 d2d3
info depth 8 score cp 28 nodes 241421 nps 807428 time 299 pv g1f3 d7d5 g2g3 h7h5 b2b4 b8c6 d2d3 c6b4
info depth 9 score cp 28 nodes 360386 nps 828473 time 435 pv g1f3 d7d5 g2g3 h7h5 b2b4 b8c6 d2d3 c6b4 e2e3
info depth 10 score cp 28 nodes 2489841 nps 786679 time 3165 pv g1f3 d7d5 g2g3 h7h5 b2b4 b8c6 d2d3 c6b4 e2e3 b4a2
bestmove g1f3
position fen 8/8/8/8/4k3/8/8/2K5 w - - 0 1
go depth 10
info depth 1 score cp 28 nodes 11 nps 0 time 0 pv c1b1
info depth 2 score cp 28 nodes 41 nps 0 time 0 pv c1b1 e4d3
info depth 3 score cp 28 nodes 131 nps 0 time 0 pv c1b1 e4d3 b1a1
info depth 4 score cp 28 nodes 512 nps 512000 time 1 pv c1b2 e4d3 b2b1 d3e2
info depth 5 score cp 28 nodes 713 nps 713000 time 1 pv c1b2 e4d3 b2b1 d3e2 b1a2
info depth 6 score cp 28 nodes 1617 nps 808500 time 2 pv c1b2 e4d3 b2b3 d3d2 b3a2 d2c1
info depth 7 score cp 28 nodes 2113 nps 704333 time 3 pv c1b2 e4d3 b2b3 d3d2 b3a2 d2c1 a2a3
info depth 8 score cp 27 nodes 3189 nps 637800 time 5 pv c1b2 e4d3 b2b3 d3d2 b3a2 d2e3 a2a3 e3e2
info depth 9 score cp 28 nodes 4582 nps 763666 time 6 pv c1b2 e4d3 b2b3 d3d2 b3a4 d2e2 a4a5 e2f3 a5a6
info depth 10 score cp 28 nodes 6312 nps 789000 time 8 pv c1b2 e4d3 b2b3 d3d2 b3a4 d2e2 a4a5 e2f3 a5a6 f3e2
bestmove c1b2
Has anyone encountered similar issues during NNUE pipeline integration/optimization? Could this be an alignment issue, a feature input problem, or some detail in the accumulator that gets missed after the performance optimization?
Any suggestions on where to focus my debugging efforts would be invaluable. Thanks in advance!
NNUE: Weights loaded (Safe mode, AVX2 backend).
Book: File 'book.bin' non trovato. Si userà solo la rete.The bug: in nnue.rs, Layer 1's CReLU clamps to a ceiling of 2048:
rustlet mut l1_out = [0i16; L1_SIZE];
simd::crelu(&l1_acc, &mut l1_out, 2048);
But the comment right above the accumulator explicitly documents the quantization scale as QA=255:
rust// --- LAYER 1: FEATURE TRANSFORMER (QA=255)
255 and 2048 are off by ~8x. With a real, QA=255-scaled weight file, the accumulator (bias + summed active-feature weights, saturating i16 add) very quickly runs past 255 for most neurons — but the code doesn't clip there, it lets values run all the way up to 2048 (or floor at 0) before clipping. In practice this means most neurons end up pinned at one of the two extremes of [0, 2048] almost regardless of which pieces are on the board, which is exactly the flat, material-blind score you're seeing. The L2 clamp (clamp(0, 128), documented as QB=64) is internally consistent with its own comment — it's only L1's ceiling that contradicts its documentation.
Code: Select all
// nnue.rs (fixed)
use std::fs::File;
use std::io::Read;
use crate::board::{Scacchiera, Colore};
// --- NNUE ARCHITECTURE CONFIGURATION ---
// 768 input features (12 piece types * 64 squares).
const INPUT_SIZE: usize = 768;
// Size of the first hidden layer (Feature Transformer).
const L1_SIZE: usize = 256;
// Size of the second hidden layer.
const L2_SIZE: usize = 32;
// ============================================================================
// SIMD BACKEND SELECTION
// ============================================================================
// The hot path of NNUE inference is:
// 1) L1 accumulation: for every occupied square, add a 256-wide i16 row
// from the weight matrix into the accumulator (saturating).
// 2) CReLU clamp of the 256-wide accumulator.
// 3) L2: a 256x32 mat-vec product (i16 x i16 -> i32).
// 4) L3: a 32-wide dot product.
//
// All four steps are data-parallel over contiguous i16 arrays, which is the
// textbook case for SIMD. We provide:
// - an AVX2 implementation (x86_64, runtime-detected via CPUID, cached),
// - a NEON implementation (aarch64, always available on that target),
// - a portable scalar fallback used anywhere else, or if AVX2 is missing.
//
// Every SIMD routine has a scalar twin computing the identical arithmetic
// (same saturation points, same integer division/clamping), so switching
// backends never changes the evaluation output.
// ============================================================================
#[cfg(target_arch = "x86_64")]
fn has_avx2() -> bool {
use std::sync::OnceLock;
static AVX2: OnceLock<bool> = OnceLock::new();
*AVX2.get_or_init(|| is_x86_feature_detected!("avx2"))
}
/// Main structure for the quantized LunaNNUE neural network.
/// Stores weights and biases as 16-bit integers (`i16`) to optimize performance on CPUs without an FPU.
pub struct LunaNNUE {
l1_weights: Vec<i16>,
l1_bias: Vec<i16>,
l2_weights: Vec<i16>,
l2_bias: Vec<i16>,
l3_weights: Vec<i16>,
l3_bias: Vec<i16>,
}
impl LunaNNUE {
/// Loads the network's binary weights from an external file in Little-Endian format.
/// Returns `None` in case of I/O error or corrupted file.
pub fn load(path: &str) -> Option<Self> {
let mut file = File::open(path).ok()?;
let mut buffer = Vec::new();
file.read_to_end(&mut buffer).ok()?;
let mut net = LunaNNUE {
l1_weights: vec![0; INPUT_SIZE * L1_SIZE],
l1_bias: vec![0; L1_SIZE],
l2_weights: vec![0; L1_SIZE * L2_SIZE],
l2_bias: vec![0; L2_SIZE],
l3_weights: vec![0; L2_SIZE],
l3_bias: vec![0; 1],
};
let mut offset = 0;
// Helper closure to read paired bytes sequentially and convert them into i16.
let mut read_i16 = |count: usize| -> Vec<i16> {
let mut v = Vec::with_capacity(count);
for _ in 0..count {
if offset + 2 <= buffer.len() {
v.push(i16::from_le_bytes([buffer[offset], buffer[offset + 1]]));
offset += 2;
}
}
v
};
// Ordered mapping of weight and bias vectors.
net.l1_weights = read_i16(INPUT_SIZE * L1_SIZE);
net.l1_bias = read_i16(L1_SIZE);
net.l2_weights = read_i16(L1_SIZE * L2_SIZE);
net.l2_bias = read_i16(L2_SIZE);
net.l3_weights = read_i16(L2_SIZE);
net.l3_bias = read_i16(1);
// Initial diagnostic check on the binary file status.
if net.l1_weights.iter().take(100).all(|&x| x == 0) {
println!("⚠️ WARNING: The NNUE file seems empty!");
} else {
#[cfg(target_arch = "x86_64")]
let backend = if has_avx2() { "AVX2" } else { "scalar" };
#[cfg(target_arch = "aarch64")]
let backend = "NEON";
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
let backend = "scalar";
println!("✅ NNUE: Weights loaded (Safe mode, {} backend).", backend);
}
Some(net)
}
/// Performs the feedforward (inference) step of the network for the current position.
/// Extracts active features via bitboards and propagates values through quantized layers.
pub fn evaluate(&self, s: &Scacchiera) -> i32 {
// ==========================================
// --- LAYER 1: FEATURE TRANSFORMER (QA=255)
// ==========================================
// Initialize the L1 accumulator by cloning the initial bias values.
let mut l1_acc = [0i16; L1_SIZE];
l1_acc.copy_from_slice(&self.l1_bias);
// Iterate over the 6 piece types.
for p_idx in 0..6 {
// Sub-step for White (Color Index 0).
let mut bb_w = s.pezzi[p_idx] & s.colori[0];
while bb_w != 0 {
let sq = bb_w.trailing_zeros() as usize;
let offset = (p_idx * 64 + sq) * L1_SIZE;
simd::accumulate(&mut l1_acc, &self.l1_weights[offset..offset + L1_SIZE]);
bb_w &= bb_w - 1; // Clear the processed bit (Bit-pop)
}
// Sub-step for Black (Color Index 1).
// Black pieces are shifted by 6 positions in the feature index (+6).
let mut bb_b = s.pezzi[p_idx] & s.colori[1];
while bb_b != 0 {
let sq = bb_b.trailing_zeros() as usize;
let offset = ((p_idx + 6) * 64 + sq) * L1_SIZE;
simd::accumulate(&mut l1_acc, &self.l1_weights[offset..offset + L1_SIZE]);
bb_b &= bb_b - 1;
}
}
// Apply CReLU activation function on Layer 1.
// Ceiling must match the QA=255 quantization scale documented above.
// (Was hardcoded to 2048 - ~8x too high - which let most L1 neurons
// run unclipped into saturation regardless of the actual position,
// making the network output nearly independent of material/board state.)
let mut l1_out = [0i16; L1_SIZE];
simd::crelu(&l1_acc, &mut l1_out, 255);
// ==========================================
// --- LAYER 2: HIDDEN LAYER (QB=64)
// ==========================================
let mut l2_acc = [0i32; L2_SIZE];
simd::forward_l2(&l1_out, &self.l2_weights, &self.l2_bias, &mut l2_acc);
// CReLU activation function on Layer 2.
// The ceiling is set to 128 to respect the QB=64 quantization scale.
let mut l2_out = [0i16; L2_SIZE];
for i in 0..L2_SIZE {
l2_out[i] = l2_acc[i].clamp(0, 128) as i16;
}
// ==========================================
// --- LAYER 3: LINEAR OUTPUT
// ==========================================
let score = self.l3_bias[0] as i32
+ simd::weighted_sum_div_pow2(&l2_out, &self.l3_weights, 8); // /256 = >>8
// Synchronize output with the Negamax search perspective.
if s.turno == Colore::Bianco { score } else { -score }
}
}
// ============================================================================
// SIMD implementations
// ============================================================================
mod simd {
use super::{L1_SIZE, L2_SIZE};
/// acc[i] = saturating_add(acc[i], weights[i]) for i in 0..L1_SIZE
#[inline(always)]
pub fn accumulate(acc: &mut [i16; L1_SIZE], weights: &[i16]) {
debug_assert_eq!(weights.len(), L1_SIZE);
#[cfg(target_arch = "x86_64")]
{
if super::has_avx2() {
unsafe { x86::accumulate_avx2(acc, weights) };
return;
}
}
#[cfg(target_arch = "aarch64")]
{
unsafe { arm::accumulate_neon(acc, weights) };
return;
}
#[allow(unreachable_code)]
accumulate_scalar(acc, weights);
}
#[inline(always)]
#[allow(dead_code)]
fn accumulate_scalar(acc: &mut [i16; L1_SIZE], weights: &[i16]) {
for i in 0..L1_SIZE {
acc[i] = acc[i].saturating_add(weights[i]);
}
}
/// out[i] = clamp(acc[i], 0, ceiling)
#[inline(always)]
pub fn crelu(acc: &[i16; L1_SIZE], out: &mut [i16; L1_SIZE], ceiling: i16) {
#[cfg(target_arch = "x86_64")]
{
if super::has_avx2() {
unsafe { x86::crelu_avx2(acc, out, ceiling) };
return;
}
}
#[cfg(target_arch = "aarch64")]
{
unsafe { arm::crelu_neon(acc, out, ceiling) };
return;
}
#[allow(unreachable_code)]
crelu_scalar(acc, out, ceiling);
}
#[inline(always)]
#[allow(dead_code)]
fn crelu_scalar(acc: &[i16; L1_SIZE], out: &mut [i16; L1_SIZE], ceiling: i16) {
for i in 0..L1_SIZE {
out[i] = acc[i].clamp(0, ceiling);
}
}
/// acc2[o] = bias[o] + sum_j( l1_out[j] * weights[j * L2_SIZE + o] ) / 2048
/// (weights are laid out row-major: row j, L2_SIZE contiguous columns)
#[inline(always)]
pub fn forward_l2(
l1_out: &[i16; L1_SIZE],
weights: &[i16],
bias: &[i16],
acc2: &mut [i32; L2_SIZE],
) {
debug_assert_eq!(weights.len(), L1_SIZE * L2_SIZE);
debug_assert_eq!(bias.len(), L2_SIZE);
#[cfg(target_arch = "x86_64")]
{
if super::has_avx2() && L2_SIZE == 32 {
unsafe { x86::forward_l2_avx2(l1_out, weights, bias, acc2) };
return;
}
}
#[cfg(target_arch = "aarch64")]
{
unsafe { arm::forward_l2_neon(l1_out, weights, bias, acc2) };
return;
}
#[allow(unreachable_code)]
forward_l2_scalar(l1_out, weights, bias, acc2);
}
#[inline(always)]
#[allow(dead_code)]
fn forward_l2_scalar(
l1_out: &[i16; L1_SIZE],
weights: &[i16],
bias: &[i16],
acc2: &mut [i32; L2_SIZE],
) {
for i in 0..L2_SIZE {
acc2[i] = bias[i] as i32;
}
for j in 0..L1_SIZE {
let lj = l1_out[j] as i32;
if lj == 0 {
continue; // CReLU sparsity: skip zeroed activations
}
let row = &weights[j * L2_SIZE..j * L2_SIZE + L2_SIZE];
for i in 0..L2_SIZE {
acc2[i] += (lj * row[i] as i32) / 2048;
}
}
}
/// sum_i( truncating_div(a[i] * b[i], 2^shift) )
///
/// NOTE: this deliberately divides *each term* before summing, matching
/// the reference scalar loop `score += (a[i]*b[i]) / divisor`. Integer
/// truncating division does not distribute over addition for negative
/// operands (e.g. (-1)/2 + (-1)/2 == 0 but (-2)/2 == -1), so summing the
/// raw products first and dividing once would silently change the score.
#[inline(always)]
pub fn weighted_sum_div_pow2(a: &[i16], b: &[i16], shift: u32) -> i32 {
debug_assert_eq!(a.len(), b.len());
#[cfg(target_arch = "x86_64")]
{
if super::has_avx2() && a.len() % 16 == 0 {
return unsafe { x86::weighted_sum_div_pow2_avx2(a, b, shift) };
}
}
#[cfg(target_arch = "aarch64")]
{
return unsafe { arm::weighted_sum_div_pow2_neon(a, b, shift) };
}
#[allow(unreachable_code)]
weighted_sum_div_pow2_scalar(a, b, shift)
}
#[inline(always)]
#[allow(dead_code)]
fn weighted_sum_div_pow2_scalar(a: &[i16], b: &[i16], shift: u32) -> i32 {
let divisor = 1i32 << shift;
let mut sum = 0i32;
for i in 0..a.len() {
sum += (a[i] as i32 * b[i] as i32) / divisor;
}
sum
}
// ------------------------------------------------------------------
// x86_64 / AVX2 backend
// ------------------------------------------------------------------
#[cfg(target_arch = "x86_64")]
mod x86 {
use super::{L1_SIZE, L2_SIZE};
use std::arch::x86_64::*;
#[target_feature(enable = "avx2")]
pub unsafe fn accumulate_avx2(acc: &mut [i16; L1_SIZE], weights: &[i16]) {
// 16 x i16 lanes per __m256i -> L1_SIZE / 16 iterations.
let mut i = 0;
while i < L1_SIZE {
let a = _mm256_loadu_si256(acc.as_ptr().add(i) as *const __m256i);
let w = _mm256_loadu_si256(weights.as_ptr().add(i) as *const __m256i);
let sum = _mm256_adds_epi16(a, w); // saturating add, matches i16::saturating_add
_mm256_storeu_si256(acc.as_mut_ptr().add(i) as *mut __m256i, sum);
i += 16;
}
}
#[target_feature(enable = "avx2")]
pub unsafe fn crelu_avx2(acc: &[i16; L1_SIZE], out: &mut [i16; L1_SIZE], ceiling: i16) {
let zero = _mm256_setzero_si256();
let cap = _mm256_set1_epi16(ceiling);
let mut i = 0;
while i < L1_SIZE {
let a = _mm256_loadu_si256(acc.as_ptr().add(i) as *const __m256i);
let clamped = _mm256_min_epi16(_mm256_max_epi16(a, zero), cap);
_mm256_storeu_si256(out.as_mut_ptr().add(i) as *mut __m256i, clamped);
i += 16;
}
}
/// Truncating (round-toward-zero) division of each i32 lane by 2^shift,
/// matching Rust's `/` operator for i32 (as opposed to an arithmetic
/// shift, which rounds toward negative infinity for negative inputs).
/// Uses the variable-count shift intrinsics since `shift` is a runtime
/// value here (the immediate-shift intrinsics require a compile-time
/// constant).
#[target_feature(enable = "avx2")]
unsafe fn div_pow2_trunc(x: __m256i, shift: u32) -> __m256i {
let sign = _mm256_srai_epi32(x, 31); // all-1s if negative, else all-0s (31 is a literal immediate)
let bias = _mm256_srl_epi32(sign, _mm_cvtsi32_si128((32 - shift) as i32));
let biased = _mm256_add_epi32(x, bias);
_mm256_sra_epi32(biased, _mm_cvtsi32_si128(shift as i32))
}
/// Outer-product accumulation: for each active (non-zero) input feature j,
/// broadcast l1_out[j], multiply by the corresponding 32-wide weight row,
/// divide *each* resulting term by 2048 (matching the scalar reference's
/// per-term division), then accumulate into a 32-lane i32 accumulator
/// (4 x __m256i of 8 lanes).
#[target_feature(enable = "avx2")]
pub unsafe fn forward_l2_avx2(
l1_out: &[i16; L1_SIZE],
weights: &[i16],
bias: &[i16],
acc2: &mut [i32; L2_SIZE],
) {
let mut acc0 = _mm256_setzero_si256();
let mut acc1 = _mm256_setzero_si256();
let mut acc2v = _mm256_setzero_si256();
let mut acc3 = _mm256_setzero_si256();
for j in 0..L1_SIZE {
let lj = l1_out[j];
if lj == 0 {
continue; // CReLU sparsity: skip zeroed activations
}
let broadcast = _mm256_set1_epi32(lj as i32);
let row_ptr = weights.as_ptr().add(j * L2_SIZE);
// Widen 4 groups of 8 x i16 -> 8 x i32, multiply, divide, accumulate.
let w0 = _mm256_cvtepi16_epi32(_mm_loadu_si128(row_ptr as *const __m128i));
let w1 = _mm256_cvtepi16_epi32(_mm_loadu_si128(row_ptr.add(8) as *const __m128i));
let w2 = _mm256_cvtepi16_epi32(_mm_loadu_si128(row_ptr.add(16) as *const __m128i));
let w3 = _mm256_cvtepi16_epi32(_mm_loadu_si128(row_ptr.add(24) as *const __m128i));
let p0 = div_pow2_trunc(_mm256_mullo_epi32(broadcast, w0), 11);
let p1 = div_pow2_trunc(_mm256_mullo_epi32(broadcast, w1), 11);
let p2 = div_pow2_trunc(_mm256_mullo_epi32(broadcast, w2), 11);
let p3 = div_pow2_trunc(_mm256_mullo_epi32(broadcast, w3), 11);
acc0 = _mm256_add_epi32(acc0, p0);
acc1 = _mm256_add_epi32(acc1, p1);
acc2v = _mm256_add_epi32(acc2v, p2);
acc3 = _mm256_add_epi32(acc3, p3);
}
let mut raw = [0i32; L2_SIZE];
_mm256_storeu_si256(raw.as_mut_ptr().add(0) as *mut __m256i, acc0);
_mm256_storeu_si256(raw.as_mut_ptr().add(8) as *mut __m256i, acc1);
_mm256_storeu_si256(raw.as_mut_ptr().add(16) as *mut __m256i, acc2v);
_mm256_storeu_si256(raw.as_mut_ptr().add(24) as *mut __m256i, acc3);
for i in 0..L2_SIZE {
acc2[i] = bias[i] as i32 + raw[i];
}
}
/// sum_i( truncating_div(a[i]*b[i], 2^shift) ), see the scalar twin
/// `weighted_sum_div_pow2_scalar` for the exact semantics preserved here.
#[target_feature(enable = "avx2")]
pub unsafe fn weighted_sum_div_pow2_avx2(a: &[i16], b: &[i16], shift: u32) -> i32 {
let mut acc = _mm256_setzero_si256();
let mut i = 0;
while i + 16 <= a.len() {
let va = _mm256_loadu_si256(a.as_ptr().add(i) as *const __m256i);
let vb = _mm256_loadu_si256(b.as_ptr().add(i) as *const __m256i);
// Widen the low/high 8 x i16 lanes to i32 so each product can be
// divided individually before summation (matches scalar order).
let va_lo = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(va));
let va_hi = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(va, 1));
let vb_lo = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vb));
let vb_hi = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vb, 1));
let prod_lo = div_pow2_trunc(_mm256_mullo_epi32(va_lo, vb_lo), shift);
let prod_hi = div_pow2_trunc(_mm256_mullo_epi32(va_hi, vb_hi), shift);
acc = _mm256_add_epi32(acc, prod_lo);
acc = _mm256_add_epi32(acc, prod_hi);
i += 16;
}
let mut sum = hsum_epi32(acc);
let divisor = 1i32 << shift;
while i < a.len() {
sum += (a[i] as i32 * b[i] as i32) / divisor;
i += 1;
}
sum
}
#[target_feature(enable = "avx2")]
unsafe fn hsum_epi32(v: __m256i) -> i32 {
let hi = _mm256_extracti128_si256(v, 1);
let lo = _mm256_castsi256_si128(v);
let sum128 = _mm_add_epi32(hi, lo);
let hi64 = _mm_unpackhi_epi64(sum128, sum128);
let sum64 = _mm_add_epi32(sum128, hi64);
let hi32 = _mm_shuffle_epi32(sum64, 0b01);
let sum32 = _mm_add_epi32(sum64, hi32);
_mm_cvtsi128_si32(sum32)
}
}
// ------------------------------------------------------------------
// aarch64 / NEON backend
// ------------------------------------------------------------------
#[cfg(target_arch = "aarch64")]
mod arm {
use super::{L1_SIZE, L2_SIZE};
use std::arch::aarch64::*;
#[target_feature(enable = "neon")]
pub unsafe fn accumulate_neon(acc: &mut [i16; L1_SIZE], weights: &[i16]) {
// 8 x i16 lanes per 128-bit vector -> L1_SIZE / 8 iterations.
let mut i = 0;
while i < L1_SIZE {
let a = vld1q_s16(acc.as_ptr().add(i));
let w = vld1q_s16(weights.as_ptr().add(i));
let sum = vqaddq_s16(a, w); // saturating add
vst1q_s16(acc.as_mut_ptr().add(i), sum);
i += 8;
}
}
#[target_feature(enable = "neon")]
pub unsafe fn crelu_neon(acc: &[i16; L1_SIZE], out: &mut [i16; L1_SIZE], ceiling: i16) {
let zero = vdupq_n_s16(0);
let cap = vdupq_n_s16(ceiling);
let mut i = 0;
while i < L1_SIZE {
let a = vld1q_s16(acc.as_ptr().add(i));
let clamped = vminq_s16(vmaxq_s16(a, zero), cap);
vst1q_s16(out.as_mut_ptr().add(i), clamped);
i += 8;
}
}
/// Truncating (round-toward-zero) division of each i32 lane by 2^shift,
/// matching Rust's `/` operator for i32. Uses variable-count shifts
/// (vshlq_* with a negative amount shifts right) since `shift` is a
/// runtime value; the `vshrq_n_*` intrinsics require a compile-time
/// constant immediate.
#[target_feature(enable = "neon")]
unsafe fn div_pow2_trunc(x: int32x4_t, shift: i32) -> int32x4_t {
let sign = vshrq_n_s32(x, 31); // 31 is a literal immediate: all-1s if negative, else all-0s
let sign_u = vreinterpretq_u32_s32(sign);
let neg_amt = vdupq_n_s32(-(32 - shift));
let bias_u = vshlq_u32(sign_u, neg_amt); // negative amount => logical right shift
let bias = vreinterpretq_s32_u32(bias_u);
let biased = vaddq_s32(x, bias);
vshlq_s32(biased, vdupq_n_s32(-shift)) // negative amount => arithmetic right shift
}
#[target_feature(enable = "neon")]
pub unsafe fn forward_l2_neon(
l1_out: &[i16; L1_SIZE],
weights: &[i16],
bias: &[i16],
acc2: &mut [i32; L2_SIZE],
) {
// 4 lanes of i32 per vector -> L2_SIZE / 4 accumulators.
let mut accs = [vdupq_n_s32(0); L2_SIZE / 4];
for j in 0..L1_SIZE {
let lj = l1_out[j];
if lj == 0 {
continue;
}
let broadcast = vdupq_n_s32(lj as i32);
let row_ptr = weights.as_ptr().add(j * L2_SIZE);
for k in 0..(L2_SIZE / 4) {
let w16 = vld1_s16(row_ptr.add(k * 4));
let w32 = vmovl_s16(w16); // widen 4 x i16 -> 4 x i32
let prod = vmulq_s32(broadcast, w32);
accs[k] = vaddq_s32(accs[k], div_pow2_trunc(prod, 11));
}
}
let mut raw = [0i32; L2_SIZE];
for k in 0..(L2_SIZE / 4) {
vst1q_s32(raw.as_mut_ptr().add(k * 4), accs[k]);
}
for i in 0..L2_SIZE {
acc2[i] = bias[i] as i32 + raw[i];
}
}
/// sum_i( truncating_div(a[i]*b[i], 2^shift) ), see the scalar twin
/// `weighted_sum_div_pow2_scalar` for the exact semantics preserved here.
#[target_feature(enable = "neon")]
pub unsafe fn weighted_sum_div_pow2_neon(a: &[i16], b: &[i16], shift: u32) -> i32 {
let mut acc = vdupq_n_s32(0);
let mut i = 0;
let len = a.len();
while i + 8 <= len {
let va = vld1q_s16(a.as_ptr().add(i));
let vb = vld1q_s16(b.as_ptr().add(i));
let lo = vmull_s16(vget_low_s16(va), vget_low_s16(vb));
let hi = vmull_s16(vget_high_s16(va), vget_high_s16(vb));
acc = vaddq_s32(acc, div_pow2_trunc(lo, shift as i32));
acc = vaddq_s32(acc, div_pow2_trunc(hi, shift as i32));
i += 8;
}
let mut sum = vaddvq_s32(acc);
let divisor = 1i32 << shift;
while i < len {
sum += (a[i] as i32 * b[i] as i32) / divisor;
i += 1;
}
sum
}
}
}
// ============================================================================
// Tests: every SIMD backend must match the scalar reference exactly.
// ============================================================================
#[cfg(test)]
mod tests {
use super::simd::*;
use super::{L1_SIZE, L2_SIZE};
fn rand_i16_vec(len: usize, seed: &mut u64) -> Vec<i16> {
let mut v = Vec::with_capacity(len);
for _ in 0..len {
*seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
v.push(((*seed >> 48) as i16) % 2000);
}
v
}
#[test]
fn accumulate_matches_scalar() {
let mut seed = 12345u64;
let base = rand_i16_vec(L1_SIZE, &mut seed);
let weights = rand_i16_vec(L1_SIZE, &mut seed);
let mut acc_simd = [0i16; L1_SIZE];
acc_simd.copy_from_slice(&base);
accumulate(&mut acc_simd, &weights);
let mut acc_scalar = [0i16; L1_SIZE];
acc_scalar.copy_from_slice(&base);
for i in 0..L1_SIZE {
acc_scalar[i] = acc_scalar[i].saturating_add(weights[i]);
}
assert_eq!(acc_simd, acc_scalar);
}
#[test]
fn crelu_matches_scalar() {
let mut seed = 999u64;
let acc: Vec<i16> = rand_i16_vec(L1_SIZE, &mut seed)
.iter()
.map(|&x| x.wrapping_sub(1000))
.collect();
let mut acc_arr = [0i16; L1_SIZE];
acc_arr.copy_from_slice(&acc);
let mut out_simd = [0i16; L1_SIZE];
crelu(&acc_arr, &mut out_simd, 2048);
let mut out_scalar = [0i16; L1_SIZE];
for i in 0..L1_SIZE {
out_scalar[i] = acc_arr[i].clamp(0, 2048);
}
assert_eq!(out_simd, out_scalar);
}
#[test]
fn forward_l2_matches_scalar() {
let mut seed = 42u64;
let l1_out_v: Vec<i16> = rand_i16_vec(L1_SIZE, &mut seed)
.iter()
.map(|&x| x.unsigned_abs() as i16 % 2048) // valid CReLU output range
.collect();
let mut l1_out = [0i16; L1_SIZE];
l1_out.copy_from_slice(&l1_out_v);
let weights = rand_i16_vec(L1_SIZE * L2_SIZE, &mut seed);
let bias = rand_i16_vec(L2_SIZE, &mut seed);
let mut acc_simd = [0i32; L2_SIZE];
forward_l2(&l1_out, &weights, &bias, &mut acc_simd);
let mut acc_scalar = [0i32; L2_SIZE];
for i in 0..L2_SIZE {
acc_scalar[i] = bias[i] as i32;
}
for j in 0..L1_SIZE {
let lj = l1_out[j] as i32;
if lj == 0 { continue; }
let row = &weights[j * L2_SIZE..j * L2_SIZE + L2_SIZE];
for i in 0..L2_SIZE {
acc_scalar[i] += (lj * row[i] as i32) / 2048;
}
}
assert_eq!(acc_simd, acc_scalar);
}
#[test]
fn weighted_sum_div_pow2_matches_scalar() {
let mut seed = 777u64;
// Include plenty of negative values: this is exactly the case where
// sum-then-divide would diverge from divide-then-sum.
let a = rand_i16_vec(L2_SIZE, &mut seed);
let b = rand_i16_vec(L2_SIZE, &mut seed);
let simd_result = weighted_sum_div_pow2(&a, &b, 8);
let scalar_result: i32 = a
.iter()
.zip(b.iter())
.map(|(&x, &y)| (x as i32 * y as i32) / 256)
.sum();
assert_eq!(simd_result, scalar_result);
}
#[test]
fn weighted_sum_div_pow2_matches_scalar_l1_sized() {
// Also exercise the L1_SIZE-length case (used nowhere directly today,
// but confirms the AVX2/NEON remainder-handling loop is correct for
// lengths that aren't a clean multiple of 16).
let mut seed = 2024u64;
let a = rand_i16_vec(37, &mut seed);
let b = rand_i16_vec(37, &mut seed);
let simd_result = weighted_sum_div_pow2(&a, &b, 11);
let scalar_result: i32 = a
.iter()
.zip(b.iter())
.map(|(&x, &y)| (x as i32 * y as i32) / 2048)
.sum();
assert_eq!(simd_result, scalar_result);
}
#[test]
fn stress_full_i16_range_many_trials() {
// Real quantized weights can span the full i16 range and be negative;
// hammer every primitive with 500 trials to catch rare lane/rounding bugs.
let mut seed = 0xC0FFEEu64;
let mut next = |lo: i32, hi: i32| -> i16 {
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let r = (seed >> 33) as i64;
(lo as i64 + (r % (hi - lo + 1) as i64)) as i16
};
for _ in 0..500 {
// accumulate
let base: Vec<i16> = (0..L1_SIZE).map(|_| next(-32000, 32000)).collect();
let weights: Vec<i16> = (0..L1_SIZE).map(|_| next(i16::MIN as i32, i16::MAX as i32)).collect();
let mut a1 = [0i16; L1_SIZE];
a1.copy_from_slice(&base);
accumulate(&mut a1, &weights);
let mut a2 = [0i16; L1_SIZE];
a2.copy_from_slice(&base);
for i in 0..L1_SIZE { a2[i] = a2[i].saturating_add(weights[i]); }
assert_eq!(a1, a2, "accumulate mismatch");
// crelu
let mut c1 = [0i16; L1_SIZE];
crelu(&a1, &mut c1, 2048);
let mut c2 = [0i16; L1_SIZE];
for i in 0..L1_SIZE { c2[i] = a1[i].clamp(0, 2048); }
assert_eq!(c1, c2, "crelu mismatch");
// forward_l2 (c1 is a valid 0..=2048 CReLU output)
let l2w: Vec<i16> = (0..L1_SIZE * L2_SIZE).map(|_| next(i16::MIN as i32, i16::MAX as i32)).collect();
let l2b: Vec<i16> = (0..L2_SIZE).map(|_| next(i16::MIN as i32, i16::MAX as i32)).collect();
let mut l2acc1 = [0i32; L2_SIZE];
forward_l2(&c1, &l2w, &l2b, &mut l2acc1);
let mut l2acc2 = [0i32; L2_SIZE];
for i in 0..L2_SIZE { l2acc2[i] = l2b[i] as i32; }
for j in 0..L1_SIZE {
let lj = c1[j] as i32;
if lj == 0 { continue; }
for i in 0..L2_SIZE {
l2acc2[i] += (lj * l2w[j * L2_SIZE + i] as i32) / 2048;
}
}
assert_eq!(l2acc1, l2acc2, "forward_l2 mismatch");
// weighted_sum_div_pow2 (L3-style)
let l2out: Vec<i16> = (0..L2_SIZE).map(|_| next(-32000, 32000)).collect();
let l3w: Vec<i16> = (0..L2_SIZE).map(|_| next(i16::MIN as i32, i16::MAX as i32)).collect();
let s1 = weighted_sum_div_pow2(&l2out, &l3w, 8);
let s2: i32 = l2out.iter().zip(l3w.iter()).map(|(&x, &y)| (x as i32 * y as i32) / 256).sum();
assert_eq!(s1, s2, "weighted_sum_div_pow2 mismatch");
}
}
#[test]
#[ignore]
fn bench_evaluate_throughput() {
use crate::board::Scacchiera;
use crate::zobrist::get_zobrist_keys;
use std::time::Instant;
let net = super::LunaNNUE::load("luna.nnue").expect("luna.nnue must be present for this bench");
let z = get_zobrist_keys();
let board = Scacchiera::new_iniziale(z);
let iters = 2_000_000u64;
let start = Instant::now();
let mut sink = 0i64;
for _ in 0..iters {
sink += net.evaluate(&board) as i64;
}
let elapsed = start.elapsed();
println!(
"evaluate(): {} calls in {:?} ({:.1} ns/call, {:.2} M evals/sec) [sink={}]",
iters,
elapsed,
elapsed.as_nanos() as f64 / iters as f64,
iters as f64 / elapsed.as_secs_f64() / 1e6,
sink
);
}
#[test]
fn full_forward_pass_is_deterministic() {
let mut seed = 314159u64;
let l1_out_v: Vec<i16> = rand_i16_vec(L1_SIZE, &mut seed)
.iter()
.map(|&x| x.unsigned_abs() as i16 % 2048)
.collect();
let mut l1_out = [0i16; L1_SIZE];
l1_out.copy_from_slice(&l1_out_v);
let weights = rand_i16_vec(L1_SIZE * L2_SIZE, &mut seed);
let bias = rand_i16_vec(L2_SIZE, &mut seed);
let mut a = [0i32; L2_SIZE];
let mut b = [0i32; L2_SIZE];
forward_l2(&l1_out, &weights, &bias, &mut a);
forward_l2(&l1_out, &weights, &bias, &mut b);
assert_eq!(a, b);
}
}