From 1c9e266ae25f7f6067806d139e5971ce8513d2f6 Mon Sep 17 00:00:00 2001 From: tslil Date: Tue, 8 Sep 2026 11:55:28 +0100 Subject: ~25% speedup in selfplay due to: incremental zobrist hashing, smarter road checking, re-using the input computed for the NN --- include/actions.c | 62 ++++++++++++++- include/negamax.c | 3 +- include/nn.c | 26 ++++++- include/nn.h | 2 + include/pv.c | 5 +- include/pv.h | 8 +- include/tak.c | 220 +++++++++++++++++++++++++++++------------------------- include/tak.h | 5 ++ include/zobrist.c | 21 +++--- include/zobrist.h | 6 ++ 10 files changed, 237 insertions(+), 121 deletions(-) (limited to 'include') diff --git a/include/actions.c b/include/actions.c index 541b388..981214f 100644 --- a/include/actions.c +++ b/include/actions.c @@ -17,6 +17,7 @@ #include "actions.h" #include "tak.h" +#include "zobrist.h" #include #include @@ -135,7 +136,7 @@ action_t *actions_generate_at(tak_state_p state, uint32_t *num_actions, // Only try moves after CPS and if the colour is correct if (count) { if (state->ply >= 2 && - ((state->colours[loc] & 1) == state->current_colour)) { + (TOP_COLOUR_AT(state, loc) == state->current_colour)) { // Pre-compute end-stops and crushes uint8_t end_stops[4], crushes[4] = {0, 0, 0, 0}; @@ -245,7 +246,18 @@ action_t *actions_generate_at(tak_state_p state, uint32_t *num_actions, void action_take(tak_state_p state, const action_t action) { const int8_t loc = A_GET_LOC(action); + + // Incremental hash: record the affected cells and their current (pre-move) + // contributions, apply the move, then fold in the change. + uint8_t affected[7]; + uint64_t old_contrib[7]; + int n_aff; + if (A_GET_TYPE(action) == A_PLACE) { + affected[0] = (uint8_t)loc; + n_aff = 1; + old_contrib[0] = zobrist_cell(state, affected[0]); + const uint8_t black = (state->current_colour == C_BLACK); switch (A_GET_DATA0(action)) { case STONE_FLAT: { @@ -293,6 +305,15 @@ void action_take(tak_state_p state, const action_t action) { for (uint8_t _gaps = gaps; _gaps; steps++) _gaps &= _gaps - 1; + // Record the affected cells (source + every drop square) and their + // current contributions + affected[0] = (uint8_t)loc; + for (int8_t k = 1; k <= steps; k++) + affected[k] = (uint8_t)(loc + k * delta); + n_aff = steps + 1; + for (int i = 0; i < n_aff; i++) + old_contrib[i] = zobrist_cell(state, affected[i]); + // Move top stone type to destination state->celldat[loc + steps * delta] &= CLR_STONE; // necessary for crushing state->celldat[loc + steps * delta] |= STONE_AT(state, loc); @@ -327,6 +348,14 @@ void action_take(tak_state_p state, const action_t action) { } // Next ply inline_next_ply(state); + + // Fold the move into the incremental hash: remove each affected cell's old + // contribution, add its new one, and flip the ply-parity term. + for (int i = 0; i < n_aff; i++) { + state->hash ^= old_contrib[i]; + state->hash ^= zobrist_cell(state, affected[i]); + } + state->hash ^= PLY_PARITY_MAGIC; } void action_undo(tak_state_p state, const action_t action) { @@ -334,7 +363,18 @@ void action_undo(tak_state_p state, const action_t action) { inline_prev_ply(state); const int8_t loc = A_GET_LOC(action); + + // Incremental hash: record the affected cells and their current + // (pre-undo) contributions, reverse the move, then fold in the change. + uint8_t affected[7]; + uint64_t old_contrib[7]; + int n_aff; + if (A_GET_TYPE(action) == A_PLACE) { + affected[0] = (uint8_t)loc; + n_aff = 1; + old_contrib[0] = zobrist_cell(state, affected[0]); + const uint8_t black = (state->current_colour == C_BLACK); state->celldat[loc] = 0; state->colours[loc] = 0; @@ -358,6 +398,18 @@ void action_undo(tak_state_p state, const action_t action) { dir = A_GET_DATA1(action) >> 4; const int8_t delta = move_deltas[dir]; + // Record the affected cells (source + every drop square) and their + // current contributions + int8_t n_steps = 1; + for (uint8_t _gaps = gaps; _gaps; n_steps++) + _gaps &= _gaps - 1; + affected[0] = (uint8_t)loc; + for (int8_t k = 1; k <= n_steps; k++) + affected[k] = (uint8_t)(loc + k * delta); + n_aff = n_steps + 1; + for (int i = 0; i < n_aff; i++) + old_contrib[i] = zobrist_cell(state, affected[i]); + int8_t steps = 1; uint8_t gap_bit = 1, total = 1; for (int8_t d = 1; d < num; d++, total++, gap_bit <<= 1) { @@ -389,6 +441,14 @@ void action_undo(tak_state_p state, const action_t action) { STONE_FLAT; // should be optimised out } } + + // Fold the undo into the incremental hash: remove each affected cell's + // current contribution, add its restored one, and flip the ply-parity term. + for (int i = 0; i < n_aff; i++) { + state->hash ^= old_contrib[i]; + state->hash ^= zobrist_cell(state, affected[i]); + } + state->hash ^= PLY_PARITY_MAGIC; } void action_to_ptn(const action_t action, char *out_ptn) { diff --git a/include/negamax.c b/include/negamax.c index 6c29481..c86eeab 100644 --- a/include/negamax.c +++ b/include/negamax.c @@ -100,6 +100,7 @@ float negamax_generate(tak_state_p state) { // this is constant across iterative deepening const uint64_t root_hash = zobrist_compute(state); + state->hash = root_hash; // anchor the incremental hash at the root #ifdef RECORD_PV pv_leaf_t iteration_leaf; #endif @@ -197,7 +198,7 @@ static float negamax(tak_state_p state, const uint8_t cur_depth, action_take(state, actions[idx]); - const uint64_t child_hash = zobrist_compute(state); + const uint64_t child_hash = state->hash; // Compute the value of the node and retain the leaf that produced it. float node_value; diff --git a/include/nn.c b/include/nn.c index c1a806d..f228297 100644 --- a/include/nn.c +++ b/include/nn.c @@ -15,6 +15,8 @@ along with ctak. If not, see . */ +#include + #include "nn.h" #include "weights_5.h" @@ -53,7 +55,7 @@ inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) { default: { // |{flat, standing, cap} * {black, white}| = 6 // index 1 .. 6 - return (uint8_t)(1u + 3u * (state->colours[loc] & 1u) + (d & STONE_MASK)); + return (uint8_t)(1u + 3u * TOP_COLOUR_AT(state, loc) + (d & STONE_MASK)); } } } @@ -72,14 +74,30 @@ inline void nn_prepare_head(float out_head[3], tak_state_p state) { out_head[2] = (float)(state->black_count & 127) * inv_max_flats; } +// The cell indices computed by the most recent nn_prepare_input, and the hash +// of the state that produced them. +static uint8_t prepared_idxs[36]; // max board is 6x6 +static uint64_t prepared_hash = 0; + void nn_prepare_input(float *cur_board, tak_state_p state, const float embedding[NN_N_EMB]) { const unsigned int n_cells = state->board_size * state->board_size; nn_prepare_head(cur_board, state); - uint8_t idxs[36]; - nn_cell_idxs(idxs, state); + nn_cell_idxs(prepared_idxs, state); + prepared_hash = state->hash; for (unsigned int loc = 0; loc < n_cells; loc++) - cur_board[3 + loc] = embedding[idxs[loc]]; + cur_board[3 + loc] = embedding[prepared_idxs[loc]]; +} + +// Copy the cell indices produced by the most recent nn_prepare_input into out +// and return 1, but only if that call was made on exactly this state, otherwise +// return 0 and leave out unfilled. +int nn_last_cell_idxs(uint8_t *out, tak_state_p state) { + const unsigned int n_cells = state->board_size * state->board_size; + if (prepared_hash != state->hash) + return 0; + memcpy(out, prepared_idxs, n_cells); + return 1; } // =================================================================== diff --git a/include/nn.h b/include/nn.h index 9e576c1..871bd0f 100644 --- a/include/nn.h +++ b/include/nn.h @@ -35,6 +35,8 @@ void nn_prepare_head(float out_head[3], tak_state_p state); void nn_prepare_input(float *cur_board, tak_state_p state, const float embedding[NN_N_EMB]); + +int nn_last_cell_idxs(uint8_t *out, tak_state_p state); // =================================================================== // Network // =================================================================== diff --git a/include/pv.c b/include/pv.c index 9c120f5..94c00d1 100644 --- a/include/pv.c +++ b/include/pv.c @@ -38,9 +38,10 @@ static float terminal_value(const enum WIN_TYPE winner) { void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state, const enum WIN_TYPE winner) { - nn_prepare_head(leaf->head, state); - nn_cell_idxs(leaf->cells, state); const uint8_t n_cells = state->board_size * state->board_size; + nn_prepare_head(leaf->head, state); + if (!nn_last_cell_idxs(leaf->cells, state)) + nn_cell_idxs(leaf->cells, state); for (uint8_t loc = n_cells; loc < PV_MAX_CELLS; loc++) leaf->cells[loc] = 0; leaf->current_colour = (uint8_t)state->current_colour; diff --git a/include/pv.h b/include/pv.h index 4be7b24..06455fe 100644 --- a/include/pv.h +++ b/include/pv.h @@ -34,9 +34,11 @@ typedef struct pv_leaf_s { float terminal_value; } pv_leaf_t; -// Capture a horizon leaf (winner == 0xFF) or an exact terminal leaf -void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state, - enum WIN_TYPE winner); +// Capture a horizon leaf (winner == 0xFF) or an exact terminal leaf. +// +// The cell indices are taken from the encoding of the most recent +// nn_prepare_input. +void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state, enum WIN_TYPE winner); // Publish and retrieve the leaf selected by the full-depth root search void pv_publish_leaf(const pv_leaf_t *leaf); diff --git a/include/tak.c b/include/tak.c index 406f8b5..e9f27ff 100644 --- a/include/tak.c +++ b/include/tak.c @@ -50,6 +50,7 @@ void reset_state(tak_state_p state, const uint8_t new_board_size) { state->ply = 0; state->won = 0xFF; // i may live to regret this hack state->current_colour = C_BLACK; + state->hash = 0; // empty board, ply 0 for (uint8_t k = 0; k < NUM_SQUARES(new_board_size); k++) { state->celldat[k] = 0; @@ -154,7 +155,7 @@ enum ACT_RESULT try_move(tak_state_p state, const int8_t location, if (avail == 0) return ACT_ILLEGAL; // Does the current player own the pile? - if ((state->colours[location] & 1) != state->current_colour) + if (TOP_COLOUR_AT(state, location) != state->current_colour) return ACT_ILLEGAL; // Is the desired direction and count on the board? int8_t delta = 0; @@ -239,25 +240,40 @@ enum ACT_RESULT try_move(tak_state_p state, const int8_t location, // Checking win // =================================================================== -// Check for the presence of a road connecting opposite sides -static enum WIN_TYPE check_road_colour(tak_state_p state, - const enum COLOUR colour) { - int component[NUM_SQUARES(state->board_size)], +// Sides of the board a connected component can touch +#define SIDE_BOTTOM 0x1 // row 0 +#define SIDE_TOP 0x2 // last row +#define SIDE_LEFT 0x4 // column 0 +#define SIDE_RIGHT 0x8 // last column +#define SIDES_BOTTOM_TOP (SIDE_BOTTOM | SIDE_TOP) +#define SIDES_LEFT_RIGHT (SIDE_LEFT | SIDE_RIGHT) + +// One bit per colour, set once a road for that colour has been found +#define ROAD_WHITE (1u << C_WHITE) +#define ROAD_BLACK (1u << C_BLACK) +#define ROAD_BOTH (ROAD_WHITE | ROAD_BLACK) + +enum WIN_TYPE check_win(tak_state_p state) { + uint8_t component[NUM_SQUARES(state->board_size)], touching[NUM_SQUARES(state->board_size)]; /* - We're doing a poor version of a disjoint set data structure to - track and merge connected components. The array `component' stores - indices to the representative cells of each connected component. A - cell is representative if component[cell] = cell. We track only - whether representatives are touching sides, and do a 2-dimensional - DP approach to forming these from the board. + We're doing a poor version of a disjoint set data structure to track and + merge connected components. The array `component' stores indices to the + representative cells of each connected component. A cell is representative + if component[cell] = cell. We track only whether representatives are + touching sides, and do a 2-dimensional DP approach to forming these from the + board. Note: we don't track the rank/size of each tree, because we're not - interested in good asymptotic complexity in the size of the board, - merely good performance for a single board size in practice. For - the same reason we also don't do path flattening/halving or - anything. + interested in good asymptotic complexity in the size of the board, merely + good performance for a single board size in practice. For the same reason we + also don't do path flattening/halving or anything. + + A cell is road-eligible only for the colour of its top stone, so black and + white components are disjoint and can share one component/touching pair, + letting us check both colours (and count top-flat stones) in a single pass. + The dragon rule needs the road status of both colours. */ for (int k = 0; k < NUM_SQUARES(state->board_size); k++) { component[k] = k; // every square is in its own connected @@ -265,97 +281,110 @@ static enum WIN_TYPE check_road_colour(tak_state_p state, touching[k] = 0; // and not connected to any sides } - // touching is the bit mask for connectivity, - // bottom | top | left | right - // 1 2 4 8 + uint8_t road = 0; + + // flat count: +1 per top-flat black stone, -1 per top-flat white + int8_t total = 0; + uint8_t board_full = 1; - int touch = 5; + int touch = SIDE_BOTTOM | SIDE_LEFT; for (int row = 0; row < state->board_size; row++) { for (int col = 0; col < state->board_size; col++) { const int cur = THE_COORDS(state->board_size, col, row); - if (COUNT_AT(state, cur) && (state->colours[cur] & 1) == colour && - STONE_AT(state, cur) != STONE_STANDING) { - - // do we have any neighbours to the left and below? - const int left_neighbour = - ((cur % state->board_size > 0) && - (COUNT_AT(state, cur - 1)) // wont ever be out of bounds - && ((state->colours[cur - 1] & 1) == colour) && - (STONE_AT(state, cur - 1) != STONE_STANDING)); - - const int lowr_neighbour = - ((cur >= state->board_size) && - (COUNT_AT(state, cur - state->board_size)) && - ((state->colours[cur - state->board_size] & 1) == colour) && - (STONE_AT(state, cur - state->board_size) != STONE_STANDING)); - - // always take the component of the lower neighbour if - // possible, failing that take the left neighbour, otherwise - // we're not yet connected, so update our own component. - if (lowr_neighbour) { - // look up the root of the lower neighbour - int root = cur - state->board_size; - while (root != component[root]) - root = component[root]; - // join the set - component[cur] = root; - if (touch) { - // something new - touching[root] |= touch; - // are we done? - if ((touching[root] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC) - return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE; - } - // if we also have a left neighbour then we should `merge' - // sets, and here we assume that the left neighbour set is - // always smaller (may not be) for the direction of merge - if (left_neighbour) { - int left_root = cur - 1; - while (left_root != component[left_root]) - left_root = component[left_root]; - // merge - const int left_touch = touching[left_root]; - if (left_touch) { - touching[root] |= left_touch; - if ((touching[root] & 0x3) == 0x3 || - (touching[root] & 0xC) == 0xC) - return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE; + const uint8_t count = COUNT_AT(state, cur); + if (count == 0) { + board_full = 0; + } else if (STONE_AT(state, cur) == STONE_FLAT) { + total += (TOP_COLOUR_AT(state, cur) == C_BLACK) ? +1 : -1; + } + + if (count && STONE_AT(state, cur) != STONE_STANDING) { + const uint8_t colour = TOP_COLOUR_AT(state, cur); + const uint8_t done = 1u << colour; + if ((road & done) == 0) { + + // do we have any neighbours to the left and below? + const int left_neighbour = + ((cur % state->board_size > 0) && + (COUNT_AT(state, cur - 1)) // wont ever be out of bounds + && (TOP_COLOUR_AT(state, cur - 1) == colour) && + (STONE_AT(state, cur - 1) != STONE_STANDING)); + + const int lowr_neighbour = + ((cur >= state->board_size) && + (COUNT_AT(state, cur - state->board_size)) && + (TOP_COLOUR_AT(state, cur - state->board_size) == colour) && + (STONE_AT(state, cur - state->board_size) != STONE_STANDING)); + + // always take the component of the lower neighbour if + // possible, failing that take the left neighbour, otherwise + // we're not yet connected, so update our own component. + if (lowr_neighbour) { + // look up the root of the lower neighbour + int root = cur - state->board_size; + while (root != component[root]) + root = component[root]; + // join the set + component[cur] = root; + if (touch) { + // something new + touching[root] |= touch; + // are we done? + if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP || + (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT) + road |= done; } - component[left_root] = root; - } - } else if (left_neighbour) { - int root = cur - 1; - while (root != component[root]) - root = component[root]; - component[cur] = root; - if (touch) { - touching[root] |= touch; - if ((touching[root] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC) - return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE; + // if we also have a left neighbour then we should `merge' + // sets, and here we assume that the left neighbour set is + // always smaller (may not be) for the direction of merge + if (left_neighbour) { + int left_root = cur - 1; + while (left_root != component[left_root]) + left_root = component[left_root]; + // merge + const int left_touch = touching[left_root]; + if (left_touch) { + touching[root] |= left_touch; + if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP || + (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT) + road |= done; + } + component[left_root] = root; + } + } else if (left_neighbour) { + int root = cur - 1; + while (root != component[root]) + root = component[root]; + component[cur] = root; + if (touch) { + touching[root] |= touch; + if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP || + (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT) + road |= done; + } + } else if (touch) { + // we had no left or lower neighbour, so we're on our own + touching[cur] = touch; } - } else if (touch) { - // we had no left or lower neighbour, so we're on our own - touching[cur] = touch; + if (road == ROAD_BOTH) + break; // both colours decided, the scan is done } } if (col + 2 == state->board_size) - touch |= 8; + touch |= SIDE_RIGHT; else - touch &= 0x3; + touch &= (SIDE_BOTTOM | SIDE_TOP); } if (row + 2 == state->board_size) - touch = 6; + touch = SIDE_TOP | SIDE_LEFT; else - touch = 4; + touch = SIDE_LEFT; + if (road == ROAD_BOTH) + break; } - return 0xFF; -} -enum WIN_TYPE check_win(tak_state_p state) { - // Road? - enum WIN_TYPE rb, rw; - rb = check_road_colour(state, C_BLACK); - rw = check_road_colour(state, C_WHITE); + const enum WIN_TYPE rw = (road & ROAD_WHITE) ? WIN_ROAD_WHITE : 0xFF, + rb = (road & ROAD_BLACK) ? WIN_ROAD_BLACK : 0xFF; if (rb == WIN_ROAD_BLACK && rw == WIN_ROAD_WHITE) { return (state->ply & 1) ? rb : rw; // Dragons @@ -365,15 +394,6 @@ enum WIN_TYPE check_win(tak_state_p state) { return rb; } - // Do we do a flat count? - int8_t total = 0, board_full = 1; - for (uint8_t k = 0; k < NUM_SQUARES(state->board_size); k++) { - if (COUNT_AT(state, k) == 0) { - board_full = 0; - } else if (STONE_AT(state, k) == STONE_FLAT) { - total += ((state->colours[k] & 1) == C_BLACK) ? +1 : -1; - } - } if (state->black_count == 0 || state->white_count == 0 || board_full) { // Decide based on count if (total > 0) diff --git a/include/tak.h b/include/tak.h index d60c221..324c41f 100644 --- a/include/tak.h +++ b/include/tak.h @@ -55,6 +55,9 @@ typedef struct tak_state_s { enum WIN_TYPE won; enum COLOUR current_colour; uint8_t white_count, black_count, ply, board_size; + // incremental zobrist hash; maintained by action_take/action_undo and + // re-anchored at the search root (see zobrist.h) + uint64_t hash; } *tak_state_p; #define NUM_SHIFT 4 @@ -63,6 +66,8 @@ typedef struct tak_state_s { #define STONE_MASK 3 // 0b00000011 #define STONE_AT(state, l) (state->celldat[(l)] & STONE_MASK) #define COUNT_AT(state, l) (state->celldat[(l)] >> NUM_SHIFT) +// the colour stack holds one bit per stone, the top stone in the LSB +#define TOP_COLOUR_AT(state, l) (state->colours[(l)] & 1) #define THE_COORDS(board_size, col, row) ((col) + (row)*board_size) // =================================================================== diff --git a/include/zobrist.c b/include/zobrist.c index 2922abf..5076fe1 100644 --- a/include/zobrist.c +++ b/include/zobrist.c @@ -50,18 +50,19 @@ void zobrist_free(void) { } -// a position's value is from the perspective of the player to move, which flips -// with parity, so we must track this -#define PLY_PARITY_MAGIC 0x9E3779B97F4A7C15ULL +uint64_t zobrist_cell(tak_state_p state, const uint8_t l) { + colour_stack_t c = state->colours[l]; + const uint8_t count = COUNT_AT(state, l); + enum STONE_VARIANT s = STONE_AT(state, l); + uint64_t h = 0; + for (uint8_t hh = 0; hh < count; hh++, c >>= 1) + h ^= zobrist[l * (15 * 2 * 3) + hh * 2 * 3 + (c & 1) * 3 + s]; + return h; +} uint64_t zobrist_compute(tak_state_p state) { uint64_t hash = 0; - for (uint8_t l = 0; l < state->board_size * state->board_size; l++) { - colour_stack_t c = state->colours[l]; - const uint8_t count = COUNT_AT(state, l); - enum STONE_VARIANT s = STONE_AT(state, l); - for (uint8_t h = 0; h < count; h++, c >>= 1) - hash ^= zobrist[l * (15 * 2 * 3) + h * 2 * 3 + (c & 1) * 3 + s]; - } + for (uint8_t l = 0; l < state->board_size * state->board_size; l++) + hash ^= zobrist_cell(state, l); return hash ^ ((state->ply & 1) ? PLY_PARITY_MAGIC : 0); } diff --git a/include/zobrist.h b/include/zobrist.h index 5fd5ace..27053a6 100644 --- a/include/zobrist.h +++ b/include/zobrist.h @@ -19,8 +19,14 @@ #include #include +#define PLY_PARITY_MAGIC 0x9E3779B97F4A7C15ULL + void zobrist_free(void); int zobrist_init(const uint8_t board_size); +// hash for whole state uint64_t zobrist_compute(tak_state_p state); + +// hash contribution from a single board cell at location l +uint64_t zobrist_cell(tak_state_p state, const uint8_t l); -- cgit v1.3.1