speed up optim

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)
-                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
-                    .unwrap()
+                let score_a = ucb_score_fast(&arena[a], c_puct_sqrt);
+                let score_b = ucb_score_fast(&arena[b], c_puct_sqrt);
+                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
-            let batch = Tensor::cat(leaf_states, 0);
+            // Batch evaluate non-terminal leaf states using single CPU buffer encoding
+            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> =
-                encode_board_state_perspective(&arena[node_idx].board_state, device)
-                    .reshape([1, 18, 8, 8]);
+            let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
             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)