cache and fixes

Cargo.lock

@@ -149,17 +149,6 @@ dependencies = [
  "portable-atomic-util",
 ]
 
-[[package]]
-name = "async-trait"
-version = "0.1.92"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 3.0.4",
-]
-
 [[package]]
 name = "atomic"
 version = "0.6.1"
@@ -1219,9 +1208,9 @@ dependencies = [
 
 [[package]]
 name = "crossbeam-epoch"
-version = "0.9.18"
+version = "0.9.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
 dependencies = [
  "crossbeam-utils",
 ]
@@ -2177,10 +2166,10 @@ dependencies = [
 name = "engine"
 version = "0.1.0"
 dependencies = [
- "async-trait",
  "burn",
  "chess",
  "ctrlc",
+ "moka",
  "rand 0.10.1",
  "rand_distr",
  "serde",
@@ -3657,6 +3646,23 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b"
 
+[[package]]
+name = "moka"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9"
+dependencies = [
+ "crossbeam-channel",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+ "equivalent",
+ "parking_lot",
+ "portable-atomic",
+ "smallvec",
+ "tagptr",
+ "uuid",
+]
+
 [[package]]
 name = "moxcms"
 version = "0.8.1"
@@ -5609,6 +5615,12 @@ dependencies = [
  "windows",
 ]
 
+[[package]]
+name = "tagptr"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
+
 [[package]]
 name = "tar"
 version = "0.4.46"

engine/Cargo.toml

@@ -12,4 +12,4 @@ ctrlc = "3.5.2"
 serde_json = "1.0.150"
 serde = { version = "1.0.228", features = ["derive"] }
 tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread"] }
-async-trait = "0.1.92"
+moka = { version = "0.12.16", features = ["sync"] }

engine/src/main.rs

@@ -24,8 +24,8 @@ fn main() {
         max_time_s: None,
         num_iters: None,
         max_depth: 100, // unused
-        model_name: String::from("Test1"),
+        model_name: String::from("threadedModel"),
-        load_model: true,
+        load_model: false,
         hidden_channels: 48,
         num_blocks: 4,
         batch_size: 1024,        // num positions sampled to a batch (2048)
@@ -38,6 +38,7 @@ fn main() {
         num_workers: 4,
         nn_batch_size: 64,
         episodes_per_worker: 4,
+        nn_cache_size: 500_000, // 18,736B / item in cache. 100k=1.75GB.
     };
 
     train(training_config, device);

engine/src/mcts.rs

@@ -1,13 +1,13 @@
 // use crate::networks::PolicyNetwork;
 use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
 use crate::net::encoding::encode_move;
+use crate::worker::NNCacheKey;
 use chess::BoardStatus::{Checkmate, Stalemate};
 use chess::Color::{Black, White};
 use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
 use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
 use rand::SeedableRng;
 use std::collections::HashMap;
-
 //TODO retain tree after move selection?
 
 #[derive(Clone)]
@@ -450,6 +450,18 @@ impl BoardState {
         }
     }
 
+    pub fn nn_cache_key(&self) -> NNCacheKey {
+        let board_hash = self.board.get_hash(); // whatever chess crate API you're using
+
+        let repetition_two = self.repetition_table.get(&board_hash).copied().unwrap_or(0) >= 2;
+
+        NNCacheKey {
+            board_hash,
+            halfmove_clock: self.halfmove_clock,
+            repetition_two,
+        }
+    }
+
     pub fn apply_move(&mut self, mv: ChessMove) {
         let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
         let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();

engine/src/net/encoding.rs

@@ -33,14 +33,14 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
 const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
 
 pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
-    let mut planes = vec![0.0f32; 18 * 64];
+    let mut planes = vec![0.0f32; 19 * 64];
     encode_board_state_to_slice(state, &mut planes);
-    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+    Tensor::<3>::from_floats(TensorData::new(planes, [19, 8, 8]), device)
 }
 
-/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
+/// Write 19x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1216 elements long).
 pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
-    out_slice[..1152].fill(0.0);
+    out_slice[..1216].fill(0.0);
 
     let board = &state.board;
     let us = board.side_to_move();
@@ -86,23 +86,23 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
     let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
 
     if us_castle.has_kingside() {
-        fill_plane(out_slice, 12);
+        fill_plane_ones(out_slice, 12);
     }
     if us_castle.has_queenside() {
-        fill_plane(out_slice, 13);
+        fill_plane_ones(out_slice, 13);
     }
     if them_castle.has_kingside() {
-        fill_plane(out_slice, 14);
+        fill_plane_ones(out_slice, 14);
     }
     if them_castle.has_queenside() {
-        fill_plane(out_slice, 15);
+        fill_plane_ones(out_slice, 15);
     }
 
     // Repetition plane (16)
     let current_hash = board.get_hash();
     if let Some(count) = state.repetition_table.get(&current_hash) {
         if *count >= 2 {
-            fill_plane(out_slice, 16);
+            fill_plane_ones(out_slice, 16);
         }
     }
 
@@ -116,22 +116,29 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
         }
         out_slice[idx(17, rank, file)] = 1.0;
     }
+
+    fill_plane(out_slice, 18, Some(state.halfmove_clock as f32 / 100.0));
 }
 
 pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
     let b = states.len();
-    let mut buffer = vec![0.0f32; b * 1152];
+    let mut buffer = vec![0.0f32; b * 1216];
     for (i, state) in states.iter().enumerate() {
-        let start = i * 1152;
+        let start = i * 1216;
-        encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
+        encode_board_state_to_slice(state, &mut buffer[start..start + 1216]);
     }
-    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
+    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 19, 8, 8]), device)
 }
 
 #[inline]
-fn fill_plane(buffer: &mut [f32], plane: usize) {
+fn fill_plane(buffer: &mut [f32], plane: usize, value: Option<f32>) {
     let start = plane * 64;
-    buffer[start..start + 64].fill(1.0);
+    buffer[start..start + 64].fill(value.unwrap_or(1.0));
+}
+
+#[inline]
+fn fill_plane_ones(buffer: &mut [f32], plane: usize) {
+    fill_plane(buffer, plane, Some(1.0));
 }
 
 /// Rotate a square into the side-to-move perspective.

engine/src/net/model.rs

@@ -21,6 +21,7 @@ Input planes:
 16: opponent queenside
 17: rep >= 2
 18: En passant square (gap pass over)
+19: No-progress plys / 100
 
 
 Output:
@@ -45,10 +46,10 @@ pub struct ResidualBlock {
     activation: Relu,
 }
 
-#[derive(Config, Debug)]
+// #[derive(Config, Debug)]
-pub struct ResidualBlockConfig {
+// pub struct ResidualBlockConfig {
-    channels: usize,
+//     channels: usize,
-}
+// }
 
 impl ResidualBlock {
     pub fn new(channels: usize, device: &Device) -> Self {
@@ -110,7 +111,7 @@ pub struct ChessModelConfig {
 impl ChessModelConfig {
     pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
         ChessModel {
-            conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
+            conv: Conv2dConfig::new([19, channels], [3, 3]) // 19 plane input
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
                 .init(device),
             bn: BatchNormConfig::new(channels).init(device),
@@ -210,8 +211,15 @@ pub struct ChessOutput {
 
 impl ItemLazy for ChessOutput {
     fn sync(self) -> Self {
-        let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
+        let [
-            Transaction::default()
+            policy_logits,
+            value,
+            policy_targets,
+            value_targets,
+            policy_loss,
+            value_loss,
+            loss,
+        ] = Transaction::default()
             .register(self.policy_logits)
             .register(self.value)
             .register(self.policy_targets)
@@ -299,7 +307,7 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
         let state_tensors = items
             .iter()
             .map(|item| {
-                encode_board_state_perspective(&item.board_state, device).reshape([1, 18, 8, 8])
+                encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 8, 8])
             })
             .collect::<Vec<_>>();
 
@@ -331,7 +339,7 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
             })
             .collect::<Vec<_>>();
 
-        let states = Tensor::cat(state_tensors, 0); // [B, 18, 8, 8]
+        let states = Tensor::cat(state_tensors, 0); // [B, 19, 8, 8]
         let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672]
         let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1]
 

engine/src/train.rs

@@ -42,6 +42,7 @@ pub struct TrainingConfig {
     pub num_workers: usize,
     pub nn_batch_size: usize,
     pub episodes_per_worker: usize,
+    pub nn_cache_size: u64,
 }
 
 pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
@@ -140,6 +141,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                 infer_model,
                 server_device,
                 training_config.nn_batch_size,
+                training_config.nn_cache_size,
             );
 
             server.run_server()
@@ -211,6 +213,9 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
             }
         }
 
+        server_cmd_tx
+            .send(NNServerCmd::Shutdown)
+            .expect("Expected shutdown command to be sent");
         // train
 
         if replay_buffer.is_empty() {

engine/src/worker.rs

@@ -5,10 +5,13 @@ use crate::train::{apply_temperature, sample_move};
 use burn::tensor::Device;
 use burn::Tensor;
 use chess::Color;
+use moka::sync::Cache;
 use rand::prelude::SmallRng;
 use rand::SeedableRng;
 use std::collections::VecDeque;
-use std::sync::mpsc::{Receiver, Sender};
+use std::hash::Hash;
+use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
+use std::time::Duration;
 
 pub struct Worker {
     id: usize,
@@ -134,33 +137,33 @@ impl Worker {
         data_tx.send(worker_data).unwrap();
     }
 
-    async fn model_inference(&self, board_state: BoardState) -> (Vec<f32>, f32) {
+    // async fn model_inference(&self, board_state: BoardState) -> (Vec<f32>, f32) {
-        let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>();
+    //     let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>();
-        let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
+    //     let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
-
+    //
-        self.request_tx
+    //     self.request_tx
-            .send(NNRequest {
+    //         .send(NNRequest {
-                request_tensor: state,
+    //             request_tensor: state,
-                response_tx: tx,
+    //             response_tx: tx,
-            })
+    //         })
-            .unwrap();
+    //         .unwrap();
-
+    //
-        let response = rx.await.expect("worker dropped inference response");
+    //     let response = rx.await.expect("worker dropped inference response");
-        let p = response
+    //     let p = response
-            .response_tensors
+    //         .response_tensors
-            .0
+    //         .0
-            .into_data()
+    //         .into_data()
-            .to_vec::<f32>()
+    //         .to_vec::<f32>()
-            .unwrap();
+    //         .unwrap();
-        let v = response
+    //     let v = response
-            .response_tensors
+    //         .response_tensors
-            .1
+    //         .1
-            .into_data()
+    //         .into_data()
-            .to_vec::<f32>()
+    //         .to_vec::<f32>()
-            .unwrap()[0];
+    //         .unwrap()[0];
-
+    //
-        (p, v)
+    //     (p, v)
-    }
+    // }
 
     pub async fn run_worker(self, cmd_rx: Receiver<WorkerCmd>) {
         println!("Worker started");
@@ -223,12 +226,14 @@ async fn model_inference_static(
     request_tx: Sender<NNRequest>,
     device: Device,
 ) -> (Vec<f32>, f32) {
+    let cache_key = board_state.nn_cache_key();
     let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>();
 
     let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
 
     request_tx
         .send(NNRequest {
+            cache_key,
             request_tensor: state,
             response_tx: tx,
         })
@@ -236,21 +241,7 @@ async fn model_inference_static(
 
     let response = rx.await.expect("worker dropped inference response");
 
-    let p = response
+    (response.policy, response.value)
-        .response_tensors
-        .0
-        .into_data()
-        .to_vec::<f32>()
-        .unwrap();
-
-    let v = response
-        .response_tensors
-        .1
-        .into_data()
-        .to_vec::<f32>()
-        .unwrap()[0];
-
-    (p, v)
 }
 
 pub enum NNServerCmd {
@@ -268,6 +259,7 @@ pub struct NNServer {
     model: ChessModel,
     device: Device,
     batch_size: usize,
+    cache_size: u64,
 }
 
 impl NNServer {
@@ -277,6 +269,7 @@ impl NNServer {
         model: ChessModel,
         device: Device,
         batch_size: usize,
+        cache_size: u64,
     ) -> NNServer {
         NNServer {
             command_rx,
@@ -284,6 +277,7 @@ impl NNServer {
             model,
             device,
             batch_size,
+            cache_size,
         }
     }
 }
@@ -294,31 +288,56 @@ struct NNJobData {
     // game_id: u32,
 }
 pub struct NNRequest {
+    cache_key: NNCacheKey,
     request_tensor: Tensor<4>,
     response_tx: tokio::sync::oneshot::Sender<NNResponse>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct NNResponse {
-    // nn_job_data: NNJobData,
+    policy: Vec<f32>,
-    response_tensors: (Tensor<2>, Tensor<2>),
+    value: f32,
 }
 
 impl NNServer {
     pub fn run_server(&self) {
+        let mut cache: Cache<NNCacheKey, NNResponse> = Cache::new(self.cache_size);
         let mut queue: VecDeque<NNRequest> = VecDeque::new();
 
         loop {
+            // Check for shutdown before doing any work.
+            if let Ok(cmd) = self.command_rx.try_recv() {
+                match cmd {
+                    NNServerCmd::Shutdown => {
+                        println!("NN server shutting down");
+                        break;
+                    } // other commands...
+                }
+            }
+
+            // If we have no queued requests, wait for either a request
+            // or periodically wake up to check for shutdown.
             if queue.is_empty() {
-                let first = self
+                match self.request_rx.recv_timeout(Duration::from_millis(10)) {
-                    .request_rx
+                    Ok(first) => {
-                    .recv()
+                        if try_use_cache(first, &mut cache, &mut queue) {
-                    .expect("failed to receive NN request");
+                            continue;
-                queue.push_back(first);
+                        }
+                    }
+
+                    Err(RecvTimeoutError::Timeout) => {
+                        continue;
+                    }
+
+                    Err(RecvTimeoutError::Disconnected) => {
+                        println!("Request channel disconnected");
+                        break;
+                    }
+                }
             }
 
             while let Ok(req) = self.request_rx.try_recv() {
-                queue.push_back(req);
+                try_use_cache(req, &mut cache, &mut queue);
             }
 
             let mut jobs_to_send: Vec<NNRequest> = Vec::new();
@@ -339,18 +358,49 @@ impl NNServer {
                 0,
             );
 
-            println!("Sending {} jobs to gpu", jobs_to_send.len());
+            // println!("Sending {} jobs to gpu", jobs_to_send.len());
             let (p, v) = self.model.forward(batched_tensor);
             let p_split = p.chunk(jobs_to_send.len(), 0);
             let v_split = v.chunk(jobs_to_send.len(), 0);
 
             for (i, job) in jobs_to_send.into_iter().enumerate() {
+                let p = p_split[i].clone().into_data().to_vec::<f32>().unwrap();
+
+                let v = v_split[i].clone().into_data().to_vec::<f32>().unwrap()[0];
+
+                let res = NNResponse {
+                    policy: p,
+                    value: v,
+                };
+
+                cache.insert(job.cache_key.clone(), res.clone());
+
                 job.response_tx
-                    .send(NNResponse {
+                    .send(res)
-                        response_tensors: (p_split[i].clone(), v_split[i].clone()),
-                    })
                     .expect("worker dropped inference request");
             }
         }
     }
 }
+
+fn try_use_cache(
+    req: NNRequest,
+    cache: &mut Cache<NNCacheKey, NNResponse>,
+    queue: &mut VecDeque<NNRequest>,
+) -> bool {
+    if let Some(cached) = cache.get(&req.cache_key) {
+        // println!("Cache hit");
+        req.response_tx.send(cached).unwrap();
+        true
+    } else {
+        queue.push_back(req);
+        false
+    }
+}
+
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct NNCacheKey {
+    pub board_hash: u64,
+    pub halfmove_clock: u8,
+    pub repetition_two: bool,
+}