aboutsummaryrefslogtreecommitdiff
path: root/include/nn.c
blob: 909b8706c374edd892844bf83633469c356d6991 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
/*
  This file is part of ctak.

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with ctak. If not, see <https://www.gnu.org/licenses/>.
*/

#include <string.h>

#include "nn.h"

#include "weights_5.h"
#include "weights_6.h"

// ===================================================================
// Helpers
// ===================================================================

static float fast_tanh(const float x) {
  // tanh via a degree-11 odd polynomial, least-squares fit of tanh on [-3, 3]
  //  (max absolute error ~4e-3), saturating to +-1 outside that interval.
  if (x > 3.0f)
    return 1.0f;
  if (x < -3.0f)
    return -1.0f;
  const float x2 = x * x;
  return x * (0.997136032f + x2 * (-0.305215124f + x2 *
                                   (0.0832466371f + x2 *
                                    (-0.0139947924f + x2 *
                                     (0.00123106314f + x2 * -4.29748644e-05f)))));
}

static float fast_silu(const float x) {
  // [4,2] Pade approximant with cutoff
  if (x < -4.6f)
    return 0.0;
  if (x > 4.6f)
    return x;
  return (x * (0.5f + x * (0.25f + x * (0.05f + 0.004167f * x)))) / (1+0.1f*x);
}
// ===================================================================
// Input encoding
// ===================================================================

inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) {
  const uint8_t d = state->celldat[loc];
  switch (d & NUM_MASK) {
    case 0: {
      // index 0
      return 0;
    }
    default: {
      // |{flat, standing, cap} * {black, white}| = 6
      // index 1 .. 6
      const uint8_t top_colour = TOP_COLOUR_AT(state, loc);
      return (uint8_t)(1u + 3u * ((state->ply & 1) ? 1-top_colour : top_colour)  + (d & STONE_MASK));
    }
  }
}

inline void nn_cell_idxs(uint8_t *out, tak_state_p state) {
  const unsigned int n_cells = state->board_size * state->board_size;
  for (unsigned int loc = 0; loc < n_cells; loc++)
    out[loc] = cell_idx_at(state, loc);
}

inline void nn_prepare_head(float out_head[2], tak_state_p state) {
  const float inv_max_flats =
    (state->board_size == 6) ? (1.0f / 30.0f) : (1.0f / 21.0f);
  const uint8_t idx = state->ply & 1;
  out_head[idx] = (float)(state->white_count & 127) * inv_max_flats;
  out_head[1-idx] = (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);
  nn_cell_idxs(prepared_idxs, state);
  prepared_hash = state->hash;
  for (unsigned int loc = 0; loc < n_cells; loc++)
    cur_board[2 + 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;
}

// ===================================================================
// Implementation of small neural networks
// ===================================================================

float eval_net(const float *board, const unsigned int n_inp, const float *w,
               const float *biases, const float *logit_w, const float logit_b,
               const unsigned int n_dense,
              const float parity) {
  float dense1[FIVE_DENSE_NUM]; // FIVE_DENSE_NUM == SIX_DENSE_NUM
  float z = logit_b;

  for (unsigned int d1 = 0; d1 < n_dense; d1++) {
    float h = biases[d1];
    const float *wd = w + d1 * n_inp;
    for (unsigned int fl = 0; fl < n_inp; fl++)
      h += board[fl] * wd[fl];
    dense1[d1] = fast_silu(h);
  }
  for (unsigned int d1 = 0; d1 < n_dense; d1++)
    z += dense1[d1] * logit_w[d1];

  return parity * fast_tanh(z);
}

// ===================================================================
// Interfaces to the two built-in evaluators
// ===================================================================

float five_nn_fn(tak_state_p state) {
  static float cur_board[FIVE_INP_NUM] NN_ALIGNED;

  nn_prepare_input(cur_board, state, five_embedding);
  return eval_net(cur_board, FIVE_INP_NUM, &five_dense_weights[0][0],
                  five_dense_biases, five_logit_weights, five_logit_bias,
                  FIVE_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
}

float six_nn_fn(tak_state_p state) {
  static float cur_board[SIX_INP_NUM] NN_ALIGNED;

  nn_prepare_input(cur_board, state, six_embedding);
  return eval_net(cur_board, SIX_INP_NUM, &six_dense_weights[0][0],
                  six_dense_biases, six_logit_weights, six_logit_bias,
                  SIX_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
}