speed up optim

engine/src/main.rs

@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
 fn main() {
     let device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25, 128);
+    let mcts_config = MctsConfig::new(150, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 

engine/src/mcts.rs

@@ -1,9 +1,8 @@
 // use crate::networks::PolicyNetwork;
 use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
-use crate::net::encoding::{encode_board_state_perspective, encode_move};
+use crate::net::encoding::{encode_board_states_batch, encode_move};
 use crate::net::model::ChessModel;
 use burn::prelude::Device;
-use burn::Tensor;
 use chess::BoardStatus::{Checkmate, Stalemate};
 use chess::Color::{Black, White};
 use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
@@ -37,18 +36,34 @@ impl Node {
     }
 
     pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
+        let sqrt_parent = (self.visit_count as f32).sqrt();
+        let c_puct_sqrt = c_puct * sqrt_parent;
+
         self.children
             .iter()
             .copied()
             .max_by(|&a, &b| {
-                ucb_score(self, &arena[a], c_puct)
+                let score_a = ucb_score_fast(&arena[a], c_puct_sqrt);
-                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
+                let score_b = ucb_score_fast(&arena[b], c_puct_sqrt);
-                    .unwrap()
+                score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
             })
             .expect("select_child on leaf")
     }
 }
 
+#[inline]
+fn ucb_score_fast(child: &Node, c_puct_sqrt: f32) -> f32 {
+    let prior_score = c_puct_sqrt * child.prior / (1.0 + child.visit_count as f32);
+
+    let value_score = if child.visit_count > 0 {
+        -child.value()
+    } else {
+        0.0
+    };
+
+    value_score + prior_score
+}
+
 impl Clone for Node {
     fn clone(&self) -> Self {
         Self::new(self.prior, self.board_state.clone(), self.last_move.clone())
@@ -156,6 +171,9 @@ fn softmax(logits: &[f32]) -> Vec<f32> {
     probs
 }
 
+const LEAF_VL_COUNT: u32 = 10;
+const LEAF_VL_WEIGHT: f32 = 10.0;
+
 impl Mcts {
     pub fn search(
         &mut self,
@@ -180,7 +198,6 @@ impl Mcts {
         while sims_done < num_sims {
             let mut leaf_nodes: Vec<usize> = Vec::new();
             let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
-            let mut leaf_states: Vec<Tensor<4>> = Vec::new();
 
             let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
 
@@ -197,14 +214,21 @@ impl Mcts {
                     path.push(current);
                 }
 
+                // Strong virtual loss on unexpanded leaf so batch collector visits diverse leaves across tree
+                nodes[current].visit_count += LEAF_VL_COUNT;
+                nodes[current].value_sum -= LEAF_VL_WEIGHT;
+
                 // Check terminal status of leaf
                 let stm = nodes[current].board_state.board.side_to_move();
                 if let Some(term_val) = terminal_value(nodes[current].board_state.status, stm) {
-                    // Revert virtual loss for path except leaf itself
+                    // Revert virtual loss for path
                     for &idx in &path[..path.len() - 1] {
                         nodes[idx].visit_count -= 1;
                         nodes[idx].value_sum += 1.0;
                     }
+                    nodes[current].visit_count -= LEAF_VL_COUNT;
+                    nodes[current].value_sum += LEAF_VL_WEIGHT;
+
                     self.backpropagate(&mut nodes, &path, term_val, stm);
                     sims_done += 1;
                     continue;
@@ -217,16 +241,13 @@ impl Mcts {
                         nodes[idx].visit_count -= 1;
                         nodes[idx].value_sum += 1.0;
                     }
+                    nodes[current].visit_count -= LEAF_VL_COUNT;
+                    nodes[current].value_sum += LEAF_VL_WEIGHT;
                     break;
                 }
 
                 leaf_nodes.push(current);
                 leaf_paths.push(path);
-
-                let state: Tensor<4> =
-                    encode_board_state_perspective(&nodes[current].board_state, device)
-                        .reshape([1, 18, 8, 8]);
-                leaf_states.push(state);
             }
 
             if leaf_nodes.is_empty() {
@@ -236,8 +257,9 @@ impl Mcts {
                 continue;
             }
 
-            // Batch evaluate non-terminal leaf states
+            // Batch evaluate non-terminal leaf states using single CPU buffer encoding
-            let batch = Tensor::cat(leaf_states, 0);
+            let leaf_refs: Vec<&BoardState> = leaf_nodes.iter().map(|&idx| &nodes[idx].board_state).collect();
+            let batch = encode_board_states_batch(&leaf_refs, device);
             let (policy_batch, value_batch) = model.forward(batch);
 
             let policy_data = policy_batch.into_data().to_vec::<f32>().unwrap();
@@ -253,6 +275,8 @@ impl Mcts {
                     nodes[idx].visit_count -= 1;
                     nodes[idx].value_sum += 1.0;
                 }
+                nodes[node_idx].visit_count -= LEAF_VL_COUNT;
+                nodes[node_idx].value_sum += LEAF_VL_WEIGHT;
 
                 let start = i * num_moves;
                 let end = start + num_moves;
@@ -324,9 +348,7 @@ impl Mcts {
         if let Some(term_val) = terminal_value(status, stm) {
             term_val
         } else {
-            let state: Tensor<4> =
+            let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
-                encode_board_state_perspective(&arena[node_idx].board_state, device)
-                    .reshape([1, 18, 8, 8]);
             let (policy_head, value_head) = model.forward(state);
 
             let legal_moves: Vec<ChessMove> =
@@ -428,19 +450,7 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
     samples
 }
 
-fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
-    let prior_score = c_puct * child.prior * (parent.visit_count as f32).sqrt()
-        / (1.0 + child.visit_count as f32);
 
-    let value_score: f32;
-    if child.visit_count > 0 {
-        value_score = -child.value(); // value from opposing side
-    } else {
-        value_score = 0.0;
-    }
-
-    value_score + prior_score
-}
 
 pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
     // Choose best child of root by visit count (standard MCTS behavior)

engine/src/net/encoding.rs

@@ -1,26 +1,6 @@
 use crate::mcts::BoardState;
 use burn::prelude::*;
-use chess::{ChessMove, Color, File, Piece, Rank, Square, ALL_SQUARES};
+use chess::{ChessMove, Color, File, Piece, Rank, Square};
-/*
-Input planes:
-1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
-7-12: opponent pieces
-13: your kingside
-14: your queenside
-15: opponent kingside
-16: opponent queenside
-17: rep >= 2
-18: En passant square (gap pass over)
-
-
-Output:
-4672 moves represented through planes
-1-56: sliding moves
-57-64: Knight moves
-65-73: underpromotions
-*/
-
-const PLANES_PER_SQUARE: usize = 73;
 
 const SLIDING_DIRS: [(i8, i8); 8] = [
     (0, 1),   // N
@@ -54,99 +34,104 @@ 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];
+    encode_board_state_to_slice(state, &mut planes);
+    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+}
+
+/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
+pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
+    out_slice[..1152].fill(0.0);
 
     let board = &state.board;
     let us = board.side_to_move();
     let them = !us;
+    let flip = us == Color::Black;
 
     let idx = |plane: usize, rank: usize, file: usize| -> usize { plane * 64 + rank * 8 + file };
 
-    let flip = us == Color::Black;
+    let pieces = [
-
+        Piece::Pawn,
-    for &square in ALL_SQUARES.iter() {
+        Piece::Knight,
-        if let Some(piece) = board.piece_on(square) {
+        Piece::Bishop,
-            let color = board.color_on(square).unwrap();
+        Piece::Rook,
-
+        Piece::Queen,
-            let piece_index = match piece {
+        Piece::King,
-                Piece::Pawn => 0,
+    ];
-                Piece::Knight => 1,
-                Piece::Bishop => 2,
-                Piece::Rook => 3,
-                Piece::Queen => 4,
-                Piece::King => 5,
-            };
-
-            // Determine plane based on perspective
-            let plane = if color == us {
-                piece_index
-            } else {
-                piece_index + 6
-            };
 
+    for (piece_index, &piece) in pieces.iter().enumerate() {
+        let us_bb = board.pieces(piece) & board.color_combined(us);
+        for square in us_bb {
             let mut rank = square.get_rank().to_index();
             let mut file = square.get_file().to_index();
-
-            // Flip board if black to move
             if flip {
                 rank = 7 - rank;
                 file = 7 - file;
             }
+            out_slice[idx(piece_index, rank, file)] = 1.0;
+        }
 
-            planes[idx(plane, rank, file)] = 1.0;
+        let them_bb = board.pieces(piece) & board.color_combined(them);
+        for square in them_bb {
+            let mut rank = square.get_rank().to_index();
+            let mut file = square.get_file().to_index();
+            if flip {
+                rank = 7 - rank;
+                file = 7 - file;
+            }
+            out_slice[idx(piece_index + 6, rank, file)] = 1.0;
         }
     }
 
-    // -------------------------
     // Castling Rights (planes 12–15)
-    // From perspective
-    // -------------------------
     let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
 
     if us_castle.has_kingside() {
-        fill_plane(&mut planes, 12);
+        fill_plane(out_slice, 12);
     }
     if us_castle.has_queenside() {
-        fill_plane(&mut planes, 13);
+        fill_plane(out_slice, 13);
     }
     if them_castle.has_kingside() {
-        fill_plane(&mut planes, 14);
+        fill_plane(out_slice, 14);
     }
     if them_castle.has_queenside() {
-        fill_plane(&mut planes, 15);
+        fill_plane(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(&mut planes, 16);
+            fill_plane(out_slice, 16);
         }
     }
 
-    let en_passant = board.en_passant();
+    // En passant plane (17)
-    if let Some(ep) = en_passant {
+    if let Some(ep) = board.en_passant() {
         let mut rank = ep.get_rank().to_index();
         let mut file = ep.get_file().to_index();
-
-        // Flip board if black to move
         if flip {
             rank = 7 - rank;
             file = 7 - file;
         }
-
+        out_slice[idx(17, rank, file)] = 1.0;
-        planes[idx(17, rank, file)] = 1.0;
+    }
 }
 
-    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
+    let b = states.len();
+    let mut buffer = vec![0.0f32; b * 1152];
+    for (i, state) in states.iter().enumerate() {
+        let start = i * 1152;
+        encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
+    }
+    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
 }
 
+#[inline]
 fn fill_plane(buffer: &mut [f32], plane: usize) {
     let start = plane * 64;
-    for i in 0..64 {
+    buffer[start..start + 64].fill(1.0);
-        buffer[start + i] = 1.0;
-    }
 }
 
 /// Rotate a square into the side-to-move perspective.

engine/src/train.rs

@@ -4,9 +4,9 @@ use crate::net::model::{
     ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
 };
 use burn::data::dataloader::batcher::Batcher;
-use burn::module::{AutodiffModule, Module};
+use burn::module::AutodiffModule;
-use burn::optim::{AdamConfig, GradientsParams, Optimizer};
+use burn::optim::{AdamConfig, GradientsParams};
-use burn::prelude::Device;
+use burn::prelude::{Device, Module};
 use burn::store::ModuleRecord;
 use chess::{ChessMove, Color};
 use rand::rngs::SmallRng;
@@ -119,9 +119,9 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
             let mut ply_count: usize = 0;
 
             while board_state.status == BoardStateStatus::Ongoing {
-                // let before = Instant::now();
+                let before = Instant::now();
                 let results = mcts.search(&board_state, &infer_model, &device);
-                // println!("{}", before.elapsed().as_millis());
+                println!("{}", before.elapsed().as_millis());
                 episode_buffer.push(results);
 
                 let temp = if ply_count < 30 { 1.0 } else { 0.0 };
@@ -130,7 +130,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
                 let stm = board_state.board.side_to_move();
                 let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
-                // println!("playing move: {}", mv);
+                println!("playing move: {}", mv);
                 if episode == 0 {
                     game_hist.push_str(mv.to_string().as_str());
                     game_hist.push(' ');
@@ -200,7 +200,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
         let batcher = ChessBatcher {};
 
-        let batch = batcher.batch(samples, &device);
+        let batch = batcher.batch(samples, &autodiff_device);
 
         let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);