From d3bc5fbd533c4da23cfe290c27aab066d0892e5c Mon Sep 17 00:00:00 2001 From: DragonDuck24 Date: Tue, 25 Aug 2026 13:46:00 -0400 Subject: [PATCH] speed up optim --- engine/src/main.rs | 2 +- engine/src/mcts.rs | 68 ++++++++++++--------- engine/src/net/encoding.rs | 117 ++++++++++++++++--------------------- engine/src/train.rs | 14 ++--- 4 files changed, 98 insertions(+), 103 deletions(-) diff --git a/engine/src/main.rs b/engine/src/main.rs index 0000800..aedc872 100644 --- a/engine/src/main.rs +++ b/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(); diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index 7d49c82..8e062f3 100644 --- a/engine/src/mcts.rs +++ b/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 { 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 = Vec::new(); let mut leaf_paths: Vec> = Vec::new(); - let mut leaf_states: Vec> = 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::().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 = @@ -428,19 +450,7 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec { 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) diff --git a/engine/src/net/encoding.rs b/engine/src/net/encoding.rs index 7593c8c..13e829a 100644 --- a/engine/src/net/encoding.rs +++ b/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}; -/* -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; +use chess::{ChessMove, Color, File, Piece, Rank, Square}; 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; - - for &square in ALL_SQUARES.iter() { - if let Some(piece) = board.piece_on(square) { - let color = board.color_on(square).unwrap(); - - let piece_index = match piece { - 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 - }; + let pieces = [ + Piece::Pawn, + Piece::Knight, + Piece::Bishop, + Piece::Rook, + Piece::Queen, + Piece::King, + ]; + 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(¤t_hash) { if *count >= 2 { - fill_plane(&mut planes, 16); + fill_plane(out_slice, 16); } } - let en_passant = board.en_passant(); - if let Some(ep) = en_passant { + // En passant plane (17) + 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; } - - planes[idx(17, rank, file)] = 1.0; + out_slice[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 + i] = 1.0; - } + buffer[start..start + 64].fill(1.0); } /// Rotate a square into the side-to-move perspective. diff --git a/engine/src/train.rs b/engine/src/train.rs index d41c354..61264a6 100644 --- a/engine/src/train.rs +++ b/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::optim::{AdamConfig, GradientsParams, Optimizer}; -use burn::prelude::Device; +use burn::module::AutodiffModule; +use burn::optim::{AdamConfig, GradientsParams}; +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) { 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) { 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) { 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);