From c71053771fbda8983cb6ccd6dc155e54ef8e9da8 Mon Sep 17 00:00:00 2001
From: jw1912 <jamie.whiting.dev@gmail.com>
Date: Wed, 14 Aug 2024 21:37:51 +0100
Subject: [PATCH 1/4] Quantised Value Network Bench: 2241415

---
 src/lib.rs                      |  2 +-
 src/networks.rs                 |  6 +-
 src/networks/accumulator.rs     | 31 +++++++++++
 src/networks/layer.rs           | 56 +++++++++++++++++++
 src/networks/value.rs           | 99 +++++++++++++++------------------
 train/value/src/bin/quantise.rs | 20 +++++++
 6 files changed, 159 insertions(+), 55 deletions(-)
 create mode 100644 src/networks/accumulator.rs
 create mode 100644 src/networks/layer.rs
 create mode 100644 train/value/src/bin/quantise.rs

diff --git a/src/lib.rs b/src/lib.rs
index 812f7b74..c1f37eba 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,7 +7,7 @@ mod uci;
 pub use chess::{Board, Castling, ChessState, GameState, Move};
 pub use mcts::{Limits, MctsParams, Searcher};
 pub use networks::{
-    PolicyFileDefaultName, PolicyNetwork, SubNet, ValueFileDefaultName, ValueNetwork,
+    PolicyFileDefaultName, PolicyNetwork, SubNet, ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork,
 };
 pub use tree::Tree;
 pub use uci::Uci;
diff --git a/src/networks.rs b/src/networks.rs
index bd2c9612..a9120277 100644
--- a/src/networks.rs
+++ b/src/networks.rs
@@ -1,5 +1,9 @@
+mod accumulator;
+mod layer;
 mod policy;
 mod value;
 
 pub use policy::{PolicyFileDefaultName, PolicyNetwork, SubNet};
-pub use value::{ValueFileDefaultName, ValueNetwork};
+pub use value::{ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork};
+
+const QA: i16 = 255;
diff --git a/src/networks/accumulator.rs b/src/networks/accumulator.rs
new file mode 100644
index 00000000..ccb59006
--- /dev/null
+++ b/src/networks/accumulator.rs
@@ -0,0 +1,31 @@
+use std::ops::{AddAssign, Mul};
+
+#[derive(Clone, Copy)]
+#[repr(C)]
+pub struct Accumulator<T: Copy, const N: usize>(pub [T; N]);
+
+impl<T: AddAssign<T> + Copy + Mul<T, Output = T>, const N: usize> Accumulator<T, N> {
+    pub fn add(&mut self, other: &Self) {
+        for (i, &j) in self.0.iter_mut().zip(other.0.iter()) {
+            *i += j;
+        }
+    }
+
+    pub fn madd(&mut self, mul: T, other: &Self) {
+        for (i, &j) in self.0.iter_mut().zip(other.0.iter()) {
+            *i += mul * j;
+        }
+    }
+}
+
+impl<const N: usize> Accumulator<f32, N> {
+    pub fn quantise(&self, qa: i16) -> Accumulator<i16, N> {
+        let mut res = Accumulator([0; N]);
+
+        for (i, j) in res.0.iter_mut().zip(self.0.iter()) {
+            *i = (*j * f32::from(qa)) as i16;
+        }
+
+        res
+    }
+}
\ No newline at end of file
diff --git a/src/networks/layer.rs b/src/networks/layer.rs
new file mode 100644
index 00000000..dcf92fd6
--- /dev/null
+++ b/src/networks/layer.rs
@@ -0,0 +1,56 @@
+use crate::Board;
+
+use super::{accumulator::Accumulator, QA};
+
+#[derive(Clone, Copy)]
+pub struct Layer<T: Copy, const M: usize, const N: usize> {
+    weights: [Accumulator<T, N>; M],
+    biases: Accumulator<T, N>,
+}
+
+impl<const M: usize, const N: usize> Layer<i16, M, N> {
+    pub fn forward(&self, board: &Board) -> Accumulator<i16, N> {
+        let mut out = self.biases;
+
+        board.map_value_features(|feat| out.add(&self.weights[feat]));
+
+        out
+    }
+}
+
+impl<const M: usize, const N: usize> Layer<f32, M, N> {
+    #[inline]
+    fn screlu(x: f32) -> f32 {
+        x.clamp(0.0, 1.0).powi(2)
+    }
+
+    pub fn forward(&self, inputs: &Accumulator<f32, M>) -> Accumulator<f32, N> {
+        let mut fwd = self.biases;
+
+        for (i, d) in inputs.0.iter().zip(self.weights.iter()) {
+            let act = Self::screlu(*i);
+            fwd.madd(act, d);
+        }
+
+        fwd
+    }
+
+    pub fn forward_from_i16(&self, inputs: &Accumulator<i16, M>) -> Accumulator<f32, N> {
+        let mut fwd = self.biases;
+
+        for (i, d) in inputs.0.iter().zip(self.weights.iter()) {
+            let act = Self::screlu(f32::from(*i) / f32::from(QA));
+            fwd.madd(act, d);
+        }
+
+        fwd
+    }
+
+    pub fn quantise_into(&self, dest: &mut Layer<i16, M, N>, qa: i16) {
+        for (acc_i, acc_j) in dest.weights.iter_mut().zip(self.weights.iter()) {
+            *acc_i = acc_j.quantise(qa);
+        }
+
+        dest.biases = self.biases.quantise(qa);
+    }
+}
diff --git a/src/networks/value.rs b/src/networks/value.rs
index bac4a2e0..240a45f1 100644
--- a/src/networks/value.rs
+++ b/src/networks/value.rs
@@ -1,37 +1,32 @@
-use crate::Board;
+use crate::{boxed_and_zeroed, Board};
+
+use super::{layer::Layer, QA};
 
 // DO NOT MOVE
 #[allow(non_upper_case_globals)]
-pub const ValueFileDefaultName: &str = "nn-2eeff9457b79.network";
+pub const ValueFileDefaultName: &str = "nn-f3bd88cdf0ac.network";
 
 const SCALE: i32 = 400;
 
 #[repr(C)]
 pub struct ValueNetwork {
-    l1: Layer<{ 768 * 4 }, 2048>,
-    l2: Layer<2048, 16>,
-    l3: Layer<16, 16>,
-    l4: Layer<16, 16>,
-    l5: Layer<16, 16>,
-    l6: Layer<16, 16>,
-    l7: Layer<16, 16>,
-    l8: Layer<16, 16>,
-    l9: Layer<16, 16>,
-    l10: Layer<16, 16>,
-    l11: Layer<16, 1>,
+    l1: Layer<i16, { 768 * 4 }, 2048>,
+    l2: Layer<f32, 2048, 16>,
+    l3: Layer<f32, 16, 16>,
+    l4: Layer<f32, 16, 16>,
+    l5: Layer<f32, 16, 16>,
+    l6: Layer<f32, 16, 16>,
+    l7: Layer<f32, 16, 16>,
+    l8: Layer<f32, 16, 16>,
+    l9: Layer<f32, 16, 16>,
+    l10: Layer<f32, 16, 16>,
+    l11: Layer<f32, 16, 1>,
 }
 
 impl ValueNetwork {
     pub fn eval(&self, board: &Board) -> i32 {
-        let mut l2 = self.l1.biases;
-
-        board.map_value_features(|feat| {
-            for (i, d) in l2.vals.iter_mut().zip(&self.l1.weights[feat].vals) {
-                *i += *d;
-            }
-        });
-
-        let l3 = self.l2.forward(&l2);
+        let l2 = self.l1.forward(board);
+        let l3 = self.l2.forward_from_i16(&l2);
         let l4 = self.l3.forward(&l3);
         let l5 = self.l4.forward(&l4);
         let l6 = self.l5.forward(&l5);
@@ -42,44 +37,42 @@ impl ValueNetwork {
         let l11 = self.l10.forward(&l10);
         let out = self.l11.forward(&l11);
 
-        (out.vals[0] * SCALE as f32) as i32
+        (out.0[0] * SCALE as f32) as i32
     }
 }
 
-#[derive(Clone, Copy)]
-struct Layer<const M: usize, const N: usize> {
-    weights: [Accumulator<N>; M],
-    biases: Accumulator<N>,
+#[repr(C)]
+pub struct UnquantisedValueNetwork {
+    l1: Layer<f32, { 768 * 4 }, 2048>,
+    l2: Layer<f32, 2048, 16>,
+    l3: Layer<f32, 16, 16>,
+    l4: Layer<f32, 16, 16>,
+    l5: Layer<f32, 16, 16>,
+    l6: Layer<f32, 16, 16>,
+    l7: Layer<f32, 16, 16>,
+    l8: Layer<f32, 16, 16>,
+    l9: Layer<f32, 16, 16>,
+    l10: Layer<f32, 16, 16>,
+    l11: Layer<f32, 16, 1>,
 }
 
-impl<const M: usize, const N: usize> Layer<M, N> {
-    fn forward(&self, inputs: &Accumulator<M>) -> Accumulator<N> {
-        let mut fwd = self.biases;
+impl UnquantisedValueNetwork {
+    pub fn quantise(&self) -> Box<ValueNetwork> {
+        let mut quantised: Box<ValueNetwork> = unsafe { boxed_and_zeroed() };
 
-        for (i, d) in inputs.vals.iter().zip(self.weights.iter()) {
-            let act = screlu(*i);
-            fwd.madd(act, d);
-        }
+        self.l1.quantise_into(&mut quantised.l1, QA);
 
-        fwd
-    }
-}
-
-#[inline]
-fn screlu(x: f32) -> f32 {
-    x.clamp(0.0, 1.0).powi(2)
-}
-
-#[derive(Clone, Copy)]
-#[repr(C)]
-struct Accumulator<const HIDDEN: usize> {
-    vals: [f32; HIDDEN],
-}
+        quantised.l2 = self.l2;
+        quantised.l3 = self.l3;
+        quantised.l4 = self.l4;
+        quantised.l5 = self.l5;
+        quantised.l6 = self.l6;
+        quantised.l7 = self.l7;
+        quantised.l8 = self.l8;
+        quantised.l9 = self.l9;
+        quantised.l10 = self.l10;
+        quantised.l11 = self.l11;
 
-impl<const HIDDEN: usize> Accumulator<HIDDEN> {
-    fn madd(&mut self, mul: f32, other: &Self) {
-        for (i, &j) in self.vals.iter_mut().zip(other.vals.iter()) {
-            *i += mul * j;
-        }
+        quantised
     }
 }
diff --git a/train/value/src/bin/quantise.rs b/train/value/src/bin/quantise.rs
new file mode 100644
index 00000000..c4ab06f0
--- /dev/null
+++ b/train/value/src/bin/quantise.rs
@@ -0,0 +1,20 @@
+use std::io::Write;
+
+use monty::{read_into_struct_unchecked, UnquantisedValueNetwork, ValueNetwork};
+
+fn main() {
+    let unquantised: Box<UnquantisedValueNetwork> = unsafe {
+        read_into_struct_unchecked("nn-2eeff9457b79.network")
+    };
+
+    let quantised = unquantised.quantise();
+
+    let mut file = std::fs::File::create("quantised.network").unwrap();
+
+    unsafe {
+        let ptr: *const ValueNetwork = quantised.as_ref();
+        let slice_ptr: *const u8 = std::mem::transmute(ptr);
+        let slice = std::slice::from_raw_parts(slice_ptr, std::mem::size_of::<ValueNetwork>());
+        file.write_all(slice).unwrap();
+    }
+}
\ No newline at end of file

From d3f4dfcd8a6ed088041741e7c0ad34f62d139649 Mon Sep 17 00:00:00 2001
From: jw1912 <jamie.whiting.dev@gmail.com>
Date: Wed, 14 Aug 2024 22:21:48 +0100
Subject: [PATCH 2/4] Better Quantisation Bench: 2269584

---
 src/networks.rs       | 2 +-
 src/networks/value.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/networks.rs b/src/networks.rs
index a9120277..3889580c 100644
--- a/src/networks.rs
+++ b/src/networks.rs
@@ -6,4 +6,4 @@ mod value;
 pub use policy::{PolicyFileDefaultName, PolicyNetwork, SubNet};
 pub use value::{ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork};
 
-const QA: i16 = 255;
+const QA: i16 = 512;
diff --git a/src/networks/value.rs b/src/networks/value.rs
index 240a45f1..86708973 100644
--- a/src/networks/value.rs
+++ b/src/networks/value.rs
@@ -4,7 +4,7 @@ use super::{layer::Layer, QA};
 
 // DO NOT MOVE
 #[allow(non_upper_case_globals)]
-pub const ValueFileDefaultName: &str = "nn-f3bd88cdf0ac.network";
+pub const ValueFileDefaultName: &str = "nn-c3e7b78c4f09.network";
 
 const SCALE: i32 = 400;
 

From a83a87847e4be9c3a31cf1e16faa34ad2dfa4f12 Mon Sep 17 00:00:00 2001
From: jw1912 <jamie.whiting.dev@gmail.com>
Date: Wed, 14 Aug 2024 23:04:27 +0100
Subject: [PATCH 3/4] Bench: 2233112

---
 Cargo.lock                       | 11 ++--
 Cargo.toml                       |  3 -
 src/chess.rs                     |  6 +-
 src/lib.rs                       |  2 +-
 src/networks.rs                  |  3 +-
 src/networks/accumulator.rs      | 22 +++++++-
 src/networks/activation.rs       | 19 +++++++
 src/networks/layer.rs            | 36 ++++++++----
 src/networks/policy.rs           | 97 +++++++++++++++++++-------------
 src/networks/value.rs            | 22 ++++----
 train/policy/src/bin/quantise.rs | 20 +++++++
 train/policy/src/chess.rs        | 76 ++++++++++++++++++++++++-
 train/policy/src/main.rs         |  2 +-
 13 files changed, 237 insertions(+), 82 deletions(-)
 create mode 100644 src/networks/activation.rs
 create mode 100644 train/policy/src/bin/quantise.rs

diff --git a/Cargo.lock b/Cargo.lock
index 2ca4e08c..8abfde48 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -140,7 +140,7 @@ checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
 [[package]]
 name = "goober"
 version = "0.1.0"
-source = "git+https://github.com/jw1912/goober.git#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
+source = "git+https://github.com/jw1912/goober#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
 dependencies = [
  "goober-core",
  "goober-derive",
@@ -150,12 +150,12 @@ dependencies = [
 [[package]]
 name = "goober-core"
 version = "0.1.0"
-source = "git+https://github.com/jw1912/goober.git#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
+source = "git+https://github.com/jw1912/goober#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
 
 [[package]]
 name = "goober-derive"
 version = "0.1.0"
-source = "git+https://github.com/jw1912/goober.git#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
+source = "git+https://github.com/jw1912/goober#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -165,7 +165,7 @@ dependencies = [
 [[package]]
 name = "goober-layer"
 version = "0.1.0"
-source = "git+https://github.com/jw1912/goober.git#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
+source = "git+https://github.com/jw1912/goober#32b9b52e68ef03d9d706548fb21bb3c4535c4dd2"
 dependencies = [
  "goober-core",
 ]
@@ -240,9 +240,6 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
 [[package]]
 name = "monty"
 version = "1.0.0"
-dependencies = [
- "goober",
-]
 
 [[package]]
 name = "montyformat"
diff --git a/Cargo.toml b/Cargo.toml
index 253ef440..ab1b14b7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,9 +10,6 @@ strip = true
 lto = true
 codegen-units = 1
 
-[dependencies]
-goober = { git = 'https://github.com/jw1912/goober.git' }
-
 [features]
 embed = []
 datagen = []
diff --git a/src/chess.rs b/src/chess.rs
index d355307a..c679de5b 100644
--- a/src/chess.rs
+++ b/src/chess.rs
@@ -140,8 +140,8 @@ impl ChessState {
         self.stm()
     }
 
-    pub fn get_policy_feats(&self) -> (goober::SparseVector, u64) {
-        let mut feats = goober::SparseVector::with_capacity(32);
+    pub fn get_policy_feats(&self) -> (Vec<usize>, u64) {
+        let mut feats = Vec::with_capacity(32);
         self.board.map_policy_features(|feat| feats.push(feat));
         (feats, self.board.threats())
     }
@@ -149,7 +149,7 @@ impl ChessState {
     pub fn get_policy(
         &self,
         mov: Move,
-        (feats, threats): &(goober::SparseVector, u64),
+        (feats, threats): &(Vec<usize>, u64),
         policy: &PolicyNetwork,
     ) -> f32 {
         policy.get(&self.board, &mov, feats, *threats)
diff --git a/src/lib.rs b/src/lib.rs
index c1f37eba..991ba777 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,7 +7,7 @@ mod uci;
 pub use chess::{Board, Castling, ChessState, GameState, Move};
 pub use mcts::{Limits, MctsParams, Searcher};
 pub use networks::{
-    PolicyFileDefaultName, PolicyNetwork, SubNet, ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork,
+    PolicyFileDefaultName, PolicyNetwork, ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork, UnquantisedPolicyNetwork,
 };
 pub use tree::Tree;
 pub use uci::Uci;
diff --git a/src/networks.rs b/src/networks.rs
index 3889580c..b1dacf8c 100644
--- a/src/networks.rs
+++ b/src/networks.rs
@@ -1,9 +1,10 @@
 mod accumulator;
+mod activation;
 mod layer;
 mod policy;
 mod value;
 
-pub use policy::{PolicyFileDefaultName, PolicyNetwork, SubNet};
+pub use policy::{PolicyFileDefaultName, PolicyNetwork, UnquantisedPolicyNetwork};
 pub use value::{ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork};
 
 const QA: i16 = 512;
diff --git a/src/networks/accumulator.rs b/src/networks/accumulator.rs
index ccb59006..c09ae3c3 100644
--- a/src/networks/accumulator.rs
+++ b/src/networks/accumulator.rs
@@ -1,7 +1,9 @@
 use std::ops::{AddAssign, Mul};
 
-#[derive(Clone, Copy)]
+use super::activation::Activation;
+
 #[repr(C)]
+#[derive(Clone, Copy)]
 pub struct Accumulator<T: Copy, const N: usize>(pub [T; N]);
 
 impl<T: AddAssign<T> + Copy + Mul<T, Output = T>, const N: usize> Accumulator<T, N> {
@@ -22,8 +24,22 @@ impl<const N: usize> Accumulator<f32, N> {
     pub fn quantise(&self, qa: i16) -> Accumulator<i16, N> {
         let mut res = Accumulator([0; N]);
 
-        for (i, j) in res.0.iter_mut().zip(self.0.iter()) {
-            *i = (*j * f32::from(qa)) as i16;
+        for (i, &j) in res.0.iter_mut().zip(self.0.iter()) {
+            if j > 1.98 {
+                println!("{j}")
+            }
+
+            *i = (j * f32::from(qa)) as i16;
+        }
+
+        res
+    }
+
+    pub fn dot<T: Activation>(&self, other: &Self) -> f32 {
+        let mut res = 0.0;
+
+        for (i, j) in self.0.iter().zip(other.0.iter()) {
+            res += T::activate(*i) * T::activate(*j);
         }
 
         res
diff --git a/src/networks/activation.rs b/src/networks/activation.rs
new file mode 100644
index 00000000..af0a4295
--- /dev/null
+++ b/src/networks/activation.rs
@@ -0,0 +1,19 @@
+pub trait Activation {
+    fn activate(x: f32) -> f32;
+}
+
+pub struct ReLU;
+impl Activation for ReLU {
+    #[inline]
+    fn activate(x: f32) -> f32 {
+        x.max(0.0)
+    }
+}
+
+pub struct SCReLU;
+impl Activation for SCReLU {
+    #[inline]
+    fn activate(x: f32) -> f32 {
+        x.clamp(0.0, 1.0).powi(2)
+    }
+}
diff --git a/src/networks/layer.rs b/src/networks/layer.rs
index dcf92fd6..318a3b4b 100644
--- a/src/networks/layer.rs
+++ b/src/networks/layer.rs
@@ -1,6 +1,6 @@
 use crate::Board;
 
-use super::{accumulator::Accumulator, QA};
+use super::{accumulator::Accumulator, activation::Activation, QA};
 
 #[derive(Clone, Copy)]
 pub struct Layer<T: Copy, const M: usize, const N: usize> {
@@ -16,30 +16,35 @@ impl<const M: usize, const N: usize> Layer<i16, M, N> {
 
         out
     }
-}
 
-impl<const M: usize, const N: usize> Layer<f32, M, N> {
-    #[inline]
-    fn screlu(x: f32) -> f32 {
-        x.clamp(0.0, 1.0).powi(2)
+    pub fn forward_from_slice(&self, feats: &[usize]) -> Accumulator<i16, N> {
+        let mut out = self.biases;
+
+        for &feat in feats {
+            out.add(&self.weights[feat])
+        }
+
+        out
     }
+}
 
-    pub fn forward(&self, inputs: &Accumulator<f32, M>) -> Accumulator<f32, N> {
+impl<const M: usize, const N: usize> Layer<f32, M, N> {
+    pub fn forward<T: Activation>(&self, inputs: &Accumulator<f32, M>) -> Accumulator<f32, N> {
         let mut fwd = self.biases;
 
         for (i, d) in inputs.0.iter().zip(self.weights.iter()) {
-            let act = Self::screlu(*i);
+            let act = T::activate(*i);
             fwd.madd(act, d);
         }
 
         fwd
     }
 
-    pub fn forward_from_i16(&self, inputs: &Accumulator<i16, M>) -> Accumulator<f32, N> {
+    pub fn forward_from_i16<T: Activation>(&self, inputs: &Accumulator<i16, M>) -> Accumulator<f32, N> {
         let mut fwd = self.biases;
 
         for (i, d) in inputs.0.iter().zip(self.weights.iter()) {
-            let act = Self::screlu(f32::from(*i) / f32::from(QA));
+            let act = T::activate(f32::from(*i) / f32::from(QA));
             fwd.madd(act, d);
         }
 
@@ -53,4 +58,15 @@ impl<const M: usize, const N: usize> Layer<f32, M, N> {
 
         dest.biases = self.biases.quantise(qa);
     }
+
+    pub fn quantise(&self, qa: i16) -> Layer<i16, M, N> {
+        let mut res = Layer {
+            weights: [Accumulator([0; N]); M],
+            biases: Accumulator([0; N]),
+        };
+
+        self.quantise_into(&mut res, qa);
+
+        res
+    }
 }
diff --git a/src/networks/policy.rs b/src/networks/policy.rs
index e7df5adf..0f937b36 100644
--- a/src/networks/policy.rs
+++ b/src/networks/policy.rs
@@ -1,56 +1,34 @@
-use crate::chess::{Board, Move};
+use crate::{boxed_and_zeroed, chess::{Board, Move}};
 
-use goober::{activation, layer, FeedForwardNetwork, Matrix, SparseVector, Vector};
+use super::{accumulator::Accumulator, activation::ReLU, layer::Layer, QA};
 
 // DO NOT MOVE
 #[allow(non_upper_case_globals)]
-pub const PolicyFileDefaultName: &str = "nn-6b5dc1d7fff9.network";
+pub const PolicyFileDefaultName: &str = "nn-e2a03baa505c.network";
 
 #[repr(C)]
-#[derive(Clone, Copy, FeedForwardNetwork)]
-pub struct SubNet {
-    ft: layer::SparseConnected<activation::ReLU, 768, 16>,
-    l2: layer::DenseConnected<activation::ReLU, 16, 16>,
+#[derive(Clone, Copy)]
+struct SubNet {
+    ft: Layer<i16, 768, 16>,
+    l2: Layer<f32, 16, 16>,
 }
 
 impl SubNet {
-    pub const fn zeroed() -> Self {
-        Self {
-            ft: layer::SparseConnected::zeroed(),
-            l2: layer::DenseConnected::zeroed(),
-        }
-    }
-
-    pub fn from_fn<F: FnMut() -> f32>(mut f: F) -> Self {
-        let matrix = Matrix::from_fn(|_, _| f());
-        let vector = Vector::from_fn(|_| f());
-
-        let matrix2 = Matrix::from_fn(|_, _| f());
-        let vector2 = Vector::from_fn(|_| f());
-
-        Self {
-            ft: layer::SparseConnected::from_raw(matrix, vector),
-            l2: layer::DenseConnected::from_raw(matrix2, vector2),
-        }
+    fn out(&self, feats: &[usize]) -> Accumulator<f32, 16> {
+        let l2 = self.ft.forward_from_slice(feats);
+        self.l2.forward_from_i16::<ReLU>(&l2)
     }
 }
 
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct PolicyNetwork {
-    pub subnets: [[SubNet; 2]; 448],
-    pub hce: layer::DenseConnected<activation::Identity, 4, 1>,
+    subnets: [[SubNet; 2]; 448],
+    hce: Layer<f32, 4, 1>,
 }
 
 impl PolicyNetwork {
-    pub const fn zeroed() -> Self {
-        Self {
-            subnets: [[SubNet::zeroed(); 2]; 448],
-            hce: layer::DenseConnected::zeroed(),
-        }
-    }
-
-    pub fn get(&self, pos: &Board, mov: &Move, feats: &SparseVector, threats: u64) -> f32 {
+    pub fn get(&self, pos: &Board, mov: &Move, feats: &[usize], threats: u64) -> f32 {
         let flip = pos.flip_val();
         let pc = pos.get_pc(1 << mov.src()) - 1;
 
@@ -62,18 +40,57 @@ impl PolicyNetwork {
         let to_subnet = &self.subnets[64 * pc + usize::from(mov.to() ^ flip)][good_see];
         let to_vec = to_subnet.out(feats);
 
-        let hce = self.hce.out(&Self::get_hce_feats(pos, mov))[0];
+        let hce = self.hce.forward::<ReLU>(&Self::get_hce_feats(pos, mov)).0[0];
 
-        from_vec.dot(&to_vec) + hce
+        from_vec.dot::<ReLU>(&to_vec) + hce
     }
 
-    pub fn get_hce_feats(_: &Board, mov: &Move) -> Vector<4> {
-        let mut feats = Vector::zeroed();
+    pub fn get_hce_feats(_: &Board, mov: &Move) -> Accumulator<f32, 4> {
+        let mut feats = [0.0; 4];
 
         if mov.is_promo() {
             feats[mov.promo_pc() - 3] = 1.0;
         }
 
-        feats
+        Accumulator(feats)
+    }
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+struct UnquantisedSubNet {
+    ft: Layer<f32, 768, 16>,
+    l2: Layer<f32, 16, 16>,
+}
+
+impl UnquantisedSubNet {
+    fn quantise(&self, qa: i16) -> SubNet {
+        SubNet {
+            ft: self.ft.quantise(qa),
+            l2: self.l2,
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct UnquantisedPolicyNetwork {
+    subnets: [[UnquantisedSubNet; 2]; 448],
+    hce: Layer<f32, 4, 1>,
+}
+
+impl UnquantisedPolicyNetwork {
+    pub fn quantise(&self) -> Box<PolicyNetwork> {
+        let mut quant: Box<PolicyNetwork> = unsafe { boxed_and_zeroed() };
+
+        for (qpair, unqpair) in quant.subnets.iter_mut().zip(self.subnets.iter()) {
+            for (qsubnet, unqsubnet) in qpair.iter_mut().zip(unqpair.iter()) {
+                *qsubnet = unqsubnet.quantise(QA);
+            }
+        }
+
+        quant.hce = self.hce;
+
+        quant
     }
 }
diff --git a/src/networks/value.rs b/src/networks/value.rs
index 86708973..70546d09 100644
--- a/src/networks/value.rs
+++ b/src/networks/value.rs
@@ -1,6 +1,6 @@
 use crate::{boxed_and_zeroed, Board};
 
-use super::{layer::Layer, QA};
+use super::{activation::SCReLU, layer::Layer, QA};
 
 // DO NOT MOVE
 #[allow(non_upper_case_globals)]
@@ -26,16 +26,16 @@ pub struct ValueNetwork {
 impl ValueNetwork {
     pub fn eval(&self, board: &Board) -> i32 {
         let l2 = self.l1.forward(board);
-        let l3 = self.l2.forward_from_i16(&l2);
-        let l4 = self.l3.forward(&l3);
-        let l5 = self.l4.forward(&l4);
-        let l6 = self.l5.forward(&l5);
-        let l7 = self.l6.forward(&l6);
-        let l8 = self.l7.forward(&l7);
-        let l9 = self.l8.forward(&l8);
-        let l10 = self.l9.forward(&l9);
-        let l11 = self.l10.forward(&l10);
-        let out = self.l11.forward(&l11);
+        let l3 = self.l2.forward_from_i16::<SCReLU>(&l2);
+        let l4 = self.l3.forward::<SCReLU>(&l3);
+        let l5 = self.l4.forward::<SCReLU>(&l4);
+        let l6 = self.l5.forward::<SCReLU>(&l5);
+        let l7 = self.l6.forward::<SCReLU>(&l6);
+        let l8 = self.l7.forward::<SCReLU>(&l7);
+        let l9 = self.l8.forward::<SCReLU>(&l8);
+        let l10 = self.l9.forward::<SCReLU>(&l9);
+        let l11 = self.l10.forward::<SCReLU>(&l10);
+        let out = self.l11.forward::<SCReLU>(&l11);
 
         (out.0[0] * SCALE as f32) as i32
     }
diff --git a/train/policy/src/bin/quantise.rs b/train/policy/src/bin/quantise.rs
new file mode 100644
index 00000000..eb220d6a
--- /dev/null
+++ b/train/policy/src/bin/quantise.rs
@@ -0,0 +1,20 @@
+use std::io::Write;
+
+use monty::{read_into_struct_unchecked, UnquantisedPolicyNetwork, PolicyNetwork};
+
+fn main() {
+    let unquantised: Box<UnquantisedPolicyNetwork> = unsafe {
+        read_into_struct_unchecked("nn-6b5dc1d7fff9.network")
+    };
+
+    let quantised = unquantised.quantise();
+
+    let mut file = std::fs::File::create("quantised.network").unwrap();
+
+    unsafe {
+        let ptr: *const PolicyNetwork = quantised.as_ref();
+        let slice_ptr: *const u8 = std::mem::transmute(ptr);
+        let slice = std::slice::from_raw_parts(slice_ptr, std::mem::size_of::<PolicyNetwork>());
+        file.write_all(slice).unwrap();
+    }
+}
diff --git a/train/policy/src/chess.rs b/train/policy/src/chess.rs
index 784bb9a7..396af31d 100644
--- a/train/policy/src/chess.rs
+++ b/train/policy/src/chess.rs
@@ -1,9 +1,81 @@
 use datagen::{PolicyData, Rand};
-use goober::{FeedForwardNetwork, OutputLayer, SparseVector, Vector};
-use monty::{Board, Move, PolicyNetwork, SubNet};
+use goober::{FeedForwardNetwork, OutputLayer, SparseVector, Vector, layer, activation, Matrix};
+use monty::{Board, Move};
 
 use crate::TrainablePolicy;
 
+#[repr(C)]
+#[derive(Clone, Copy, FeedForwardNetwork)]
+pub struct SubNet {
+    ft: layer::SparseConnected<activation::ReLU, 768, 16>,
+    l2: layer::DenseConnected<activation::ReLU, 16, 16>,
+}
+
+impl SubNet {
+    pub const fn zeroed() -> Self {
+        Self {
+            ft: layer::SparseConnected::zeroed(),
+            l2: layer::DenseConnected::zeroed(),
+        }
+    }
+
+    pub fn from_fn<F: FnMut() -> f32>(mut f: F) -> Self {
+        let matrix = Matrix::from_fn(|_, _| f());
+        let vector = Vector::from_fn(|_| f());
+
+        let matrix2 = Matrix::from_fn(|_, _| f());
+        let vector2 = Vector::from_fn(|_| f());
+
+        Self {
+            ft: layer::SparseConnected::from_raw(matrix, vector),
+            l2: layer::DenseConnected::from_raw(matrix2, vector2),
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct PolicyNetwork {
+    pub subnets: [[SubNet; 2]; 448],
+    pub hce: layer::DenseConnected<activation::Identity, 4, 1>,
+}
+
+impl PolicyNetwork {
+    pub const fn zeroed() -> Self {
+        Self {
+            subnets: [[SubNet::zeroed(); 2]; 448],
+            hce: layer::DenseConnected::zeroed(),
+        }
+    }
+
+    pub fn get(&self, pos: &Board, mov: &Move, feats: &SparseVector, threats: u64) -> f32 {
+        let flip = pos.flip_val();
+        let pc = pos.get_pc(1 << mov.src()) - 1;
+
+        let from_threat = usize::from(threats & (1 << mov.src()) > 0);
+        let from_subnet = &self.subnets[usize::from(mov.src() ^ flip)][from_threat];
+        let from_vec = from_subnet.out(feats);
+
+        let good_see = usize::from(pos.see(mov, -108));
+        let to_subnet = &self.subnets[64 * pc + usize::from(mov.to() ^ flip)][good_see];
+        let to_vec = to_subnet.out(feats);
+
+        let hce = self.hce.out(&Self::get_hce_feats(pos, mov))[0];
+
+        from_vec.dot(&to_vec) + hce
+    }
+
+    pub fn get_hce_feats(_: &Board, mov: &Move) -> Vector<4> {
+        let mut feats = Vector::zeroed();
+
+        if mov.is_promo() {
+            feats[mov.promo_pc() - 3] = 1.0;
+        }
+
+        feats
+    }
+}
+
 impl TrainablePolicy for PolicyNetwork {
     type Data = PolicyData;
 
diff --git a/train/policy/src/main.rs b/train/policy/src/main.rs
index 6966d0db..28ff32f4 100644
--- a/train/policy/src/main.rs
+++ b/train/policy/src/main.rs
@@ -1,4 +1,4 @@
-use monty::PolicyNetwork;
+use policy::chess::PolicyNetwork;
 
 fn main() {
     let mut args = std::env::args();

From e0240019a6e84817ca9e936f152d7bf252e272e7 Mon Sep 17 00:00:00 2001
From: jw1912 <jamie.whiting.dev@gmail.com>
Date: Thu, 15 Aug 2024 02:22:59 +0100
Subject: [PATCH 4/4] Format Bench: 2233112

---
 src/lib.rs                       | 3 ++-
 src/networks.rs                  | 2 +-
 src/networks/accumulator.rs      | 2 +-
 src/networks/layer.rs            | 5 ++++-
 src/networks/policy.rs           | 5 ++++-
 train/policy/src/bin/quantise.rs | 7 +++----
 train/policy/src/chess.rs        | 2 +-
 train/value/src/bin/quantise.rs  | 7 +++----
 8 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 991ba777..73600c7c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,7 +7,8 @@ mod uci;
 pub use chess::{Board, Castling, ChessState, GameState, Move};
 pub use mcts::{Limits, MctsParams, Searcher};
 pub use networks::{
-    PolicyFileDefaultName, PolicyNetwork, ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork, UnquantisedPolicyNetwork,
+    PolicyFileDefaultName, PolicyNetwork, UnquantisedPolicyNetwork, UnquantisedValueNetwork,
+    ValueFileDefaultName, ValueNetwork,
 };
 pub use tree::Tree;
 pub use uci::Uci;
diff --git a/src/networks.rs b/src/networks.rs
index b1dacf8c..0afea34e 100644
--- a/src/networks.rs
+++ b/src/networks.rs
@@ -5,6 +5,6 @@ mod policy;
 mod value;
 
 pub use policy::{PolicyFileDefaultName, PolicyNetwork, UnquantisedPolicyNetwork};
-pub use value::{ValueFileDefaultName, ValueNetwork, UnquantisedValueNetwork};
+pub use value::{UnquantisedValueNetwork, ValueFileDefaultName, ValueNetwork};
 
 const QA: i16 = 512;
diff --git a/src/networks/accumulator.rs b/src/networks/accumulator.rs
index c09ae3c3..616c4337 100644
--- a/src/networks/accumulator.rs
+++ b/src/networks/accumulator.rs
@@ -44,4 +44,4 @@ impl<const N: usize> Accumulator<f32, N> {
 
         res
     }
-}
\ No newline at end of file
+}
diff --git a/src/networks/layer.rs b/src/networks/layer.rs
index 318a3b4b..8a470d20 100644
--- a/src/networks/layer.rs
+++ b/src/networks/layer.rs
@@ -40,7 +40,10 @@ impl<const M: usize, const N: usize> Layer<f32, M, N> {
         fwd
     }
 
-    pub fn forward_from_i16<T: Activation>(&self, inputs: &Accumulator<i16, M>) -> Accumulator<f32, N> {
+    pub fn forward_from_i16<T: Activation>(
+        &self,
+        inputs: &Accumulator<i16, M>,
+    ) -> Accumulator<f32, N> {
         let mut fwd = self.biases;
 
         for (i, d) in inputs.0.iter().zip(self.weights.iter()) {
diff --git a/src/networks/policy.rs b/src/networks/policy.rs
index 0f937b36..d3659a85 100644
--- a/src/networks/policy.rs
+++ b/src/networks/policy.rs
@@ -1,4 +1,7 @@
-use crate::{boxed_and_zeroed, chess::{Board, Move}};
+use crate::{
+    boxed_and_zeroed,
+    chess::{Board, Move},
+};
 
 use super::{accumulator::Accumulator, activation::ReLU, layer::Layer, QA};
 
diff --git a/train/policy/src/bin/quantise.rs b/train/policy/src/bin/quantise.rs
index eb220d6a..728b0161 100644
--- a/train/policy/src/bin/quantise.rs
+++ b/train/policy/src/bin/quantise.rs
@@ -1,11 +1,10 @@
 use std::io::Write;
 
-use monty::{read_into_struct_unchecked, UnquantisedPolicyNetwork, PolicyNetwork};
+use monty::{read_into_struct_unchecked, PolicyNetwork, UnquantisedPolicyNetwork};
 
 fn main() {
-    let unquantised: Box<UnquantisedPolicyNetwork> = unsafe {
-        read_into_struct_unchecked("nn-6b5dc1d7fff9.network")
-    };
+    let unquantised: Box<UnquantisedPolicyNetwork> =
+        unsafe { read_into_struct_unchecked("nn-6b5dc1d7fff9.network") };
 
     let quantised = unquantised.quantise();
 
diff --git a/train/policy/src/chess.rs b/train/policy/src/chess.rs
index 396af31d..c9f603bd 100644
--- a/train/policy/src/chess.rs
+++ b/train/policy/src/chess.rs
@@ -1,5 +1,5 @@
 use datagen::{PolicyData, Rand};
-use goober::{FeedForwardNetwork, OutputLayer, SparseVector, Vector, layer, activation, Matrix};
+use goober::{activation, layer, FeedForwardNetwork, Matrix, OutputLayer, SparseVector, Vector};
 use monty::{Board, Move};
 
 use crate::TrainablePolicy;
diff --git a/train/value/src/bin/quantise.rs b/train/value/src/bin/quantise.rs
index c4ab06f0..a041e5fc 100644
--- a/train/value/src/bin/quantise.rs
+++ b/train/value/src/bin/quantise.rs
@@ -3,9 +3,8 @@ use std::io::Write;
 use monty::{read_into_struct_unchecked, UnquantisedValueNetwork, ValueNetwork};
 
 fn main() {
-    let unquantised: Box<UnquantisedValueNetwork> = unsafe {
-        read_into_struct_unchecked("nn-2eeff9457b79.network")
-    };
+    let unquantised: Box<UnquantisedValueNetwork> =
+        unsafe { read_into_struct_unchecked("nn-2eeff9457b79.network") };
 
     let quantised = unquantised.quantise();
 
@@ -17,4 +16,4 @@ fn main() {
         let slice = std::slice::from_raw_parts(slice_ptr, std::mem::size_of::<ValueNetwork>());
         file.write_all(slice).unwrap();
     }
-}
\ No newline at end of file
+}