Implementation of AI algorithm and AI game

Implement the algorith MCTS to perform tic-tac-toe game in kernel space,
in order to achieve the algorithm we also implment fixed-point
arithmetic including logarithm and integer square root.

The main control loop of AI tic-tac-toe game is the function "ai_game()"
which is located in simrupt.c now. However it present a severe problem
which will corrupt too much CPU computation time within one game. Future
improvements are needed to seperate the function into multiple tasklet.

Author: vax-r

Makefile

@@ -1,5 +1,5 @@
 TARGET = kmldrv
-kmldrv-objs = simrupt.o
+kmldrv-objs = simrupt.o game.o mcts.o
 obj-m := $(TARGET).o
 
 KDIR ?= /lib/modules/$(shell uname -r)/build

game.c

@@ -0,0 +1,69 @@
+#include <linux/slab.h>
+
+#include "game.h"
+
+
+const line_t lines[4] = {
+    {1, 0, 0, 0, BOARD_SIZE - GOAL + 1, BOARD_SIZE},             // ROW
+    {0, 1, 0, 0, BOARD_SIZE, BOARD_SIZE - GOAL + 1},             // COL
+    {1, 1, 0, 0, BOARD_SIZE - GOAL + 1, BOARD_SIZE - GOAL + 1},  // PRIMARY
+    {1, -1, 0, GOAL - 1, BOARD_SIZE - GOAL + 1, BOARD_SIZE},     // SECONDARY
+};
+
+static char check_line_segment_win(const char *t, int i, int j, line_t line)
+{
+    char last = t[GET_INDEX(i, j)];
+    if (last == ' ')
+        return ' ';
+    for (int k = 1; k < GOAL; k++) {
+        if (last != t[GET_INDEX(i + k * line.i_shift, j + k * line.j_shift)])
+            return ' ';
+    }
+
+#if !ALLOW_EXCEED
+    if (last == LOOKUP(t, i - line.i_shift, j - line.j_shift, ' ') ||
+        last ==
+            LOOKUP(t, i + GOAL * line.i_shift, j + GOAL * line.j_shift, ' '))
+        return ' ';
+#endif
+    return last;
+}
+
+char check_win(char *t)
+{
+    for (int i_line = 0; i_line < 4; ++i_line) {
+        line_t line = lines[i_line];
+        for (int i = line.i_lower_bound; i < line.i_upper_bound; ++i) {
+            for (int j = line.j_lower_bound; j < line.j_upper_bound; ++j) {
+                char win = check_line_segment_win(t, i, j, line);
+                if (win != ' ')
+                    return win;
+            }
+        }
+    }
+    for (int i = 0; i < N_GRIDS; i++)
+        if (t[i] == ' ')
+            return ' ';
+    return 'D';
+}
+
+fixed_point_t calculate_win_value(char win, char player)
+{
+    if (win == player)
+        return 1U << FIXED_SCALE_BITS;
+    if (win == (player ^ 'O' ^ 'X'))
+        return 0U;
+    return 1U << (FIXED_SCALE_BITS - 1);
+}
+
+int *available_moves(const char *table)
+{
+    int *moves = kzalloc(N_GRIDS * sizeof(int), GFP_KERNEL);
+    int m = 0;
+    for (int i = 0; i < N_GRIDS; i++)
+        if (table[i] == ' ')
+            moves[m++] = i;
+    if (m < N_GRIDS)
+        moves[m] = -1;
+    return moves;
+}

game.h

@@ -17,11 +17,23 @@ typedef struct {
     int i_lower_bound, j_lower_bound, i_upper_bound, j_upper_bound;
 } line_t;
 
+/* Self-defined fixed-point type, using last 10 bits as fractional bits,
+ * starting from lsb */
+#define FIXED_SCALE_BITS 8
+#define FIXED_MAX (~0U)
+#define FIXED_MIN (0U)
+#define GET_SIGN(x) ((x) & (1U << 31))
+#define SET_SIGN(x) ((x) | (1U << 31))
+#define CLR_SIGN(x) ((x) & ((1U << 31) - 1U))
+typedef unsigned fixed_point_t;
+
 #define DRAW_SIZE (N_GRIDS + BOARD_SIZE)
 #define DRAWBUFFER_SIZE                                                 \
     ((BOARD_SIZE * (BOARD_SIZE + 1) << 1) + (BOARD_SIZE * BOARD_SIZE) + \
      ((BOARD_SIZE << 1) + 1))
 
 extern const line_t lines[4];
 
-// void draw_board(const char *t);
+int *available_moves(const char *table);
+char check_win(char *t);
+fixed_point_t calculate_win_value(char win, char player);

mcts.c

@@ -0,0 +1,206 @@
+#include <linux/slab.h>
+#include <linux/string.h>
+
+#include "game.h"
+#include "mcts.h"
+#include "util.h"
+#include "wyhash.h"
+
+struct node {
+    int move;
+    char player;
+    int n_visits;
+    fixed_point_t score;
+    struct node *parent;
+    struct node *children[N_GRIDS];
+};
+
+static struct node *new_node(int move, char player, struct node *parent)
+{
+    struct node *node = kzalloc(sizeof(struct node), GFP_KERNEL);
+    node->move = move;
+    node->player = player;
+    node->n_visits = 0;
+    node->score = 0;
+    node->parent = parent;
+    memset(node->children, 0, sizeof(node->children));
+    return node;
+}
+
+static void free_node(struct node *node)
+{
+    for (int i = 0; i < N_GRIDS; i++)
+        if (node->children[i])
+            free_node(node->children[i]);
+    kfree(node);
+}
+
+fixed_point_t fixed_sqrt(fixed_point_t x)
+{
+    if (!x || x == (1U << FIXED_SCALE_BITS))
+        return x;
+
+    fixed_point_t s = 0U;
+    for (int i = (31 - __builtin_clz(x | 1)); i >= 0; i--) {
+        fixed_point_t t = (1U << i);
+        if ((((s + t) * (s + t)) >> FIXED_SCALE_BITS) <= x)
+            s += t;
+    }
+    return s;
+}
+
+fixed_point_t fixed_log(fixed_point_t v)
+{
+    if (!v || v == (1U << FIXED_SCALE_BITS))
+        return 0;
+
+    fixed_point_t numerator = (v - (1U << FIXED_SCALE_BITS));
+    int neg = 0;
+    if (GET_SIGN(numerator)) {
+        neg = 1;
+        numerator = CLR_SIGN(numerator);
+        numerator = (1U << 31) - numerator;
+    }
+
+    fixed_point_t y =
+        (numerator << FIXED_SCALE_BITS) / (v + (1U << FIXED_SCALE_BITS));
+
+    fixed_point_t ans = 0U;
+    for (unsigned i = 1; i < 20; i += 2) {
+        fixed_point_t z = (1U << FIXED_SCALE_BITS);
+        for (int j = 0; j < i; j++) {
+            z *= y;
+            z >>= FIXED_SCALE_BITS;
+        }
+        z <<= FIXED_SCALE_BITS;
+        z /= (i << FIXED_SCALE_BITS);
+
+        ans += z;
+    }
+    ans <<= 1;
+    ans = neg ? SET_SIGN(ans) : ans;
+    return ans;
+}
+
+static inline fixed_point_t uct_score(int n_total,
+                                      int n_visits,
+                                      fixed_point_t score)
+{
+    if (n_visits == 0)
+        return FIXED_MAX;
+
+    fixed_point_t result =
+        score << FIXED_SCALE_BITS /
+                     (fixed_point_t) (n_visits << FIXED_SCALE_BITS);
+    fixed_point_t tmp =
+        EXPLORATION_FACTOR *
+        fixed_sqrt(fixed_log(n_total << FIXED_SCALE_BITS) / n_visits);
+    tmp >>= FIXED_SCALE_BITS;
+    return result + tmp;
+}
+
+static struct node *select_move(struct node *node)
+{
+    struct node *best_node = NULL;
+    fixed_point_t best_score = 0U;
+    for (int i = 0; i < N_GRIDS; i++) {
+        if (!node->children[i])
+            continue;
+        fixed_point_t score =
+            uct_score(node->n_visits, node->children[i]->n_visits,
+                      node->children[i]->score);
+        if (score > best_score) {
+            best_score = score;
+            best_node = node->children[i];
+        }
+    }
+    return best_node;
+}
+
+static fixed_point_t simulate(char *table, char player)
+{
+    char current_player = player;
+    char temp_table[N_GRIDS];
+    memcpy(temp_table, table, N_GRIDS);
+    while (1) {
+        int *moves = available_moves(temp_table);
+        if (moves[0] == -1) {
+            kfree(moves);
+            break;
+        }
+        int n_moves = 0;
+        while (n_moves < N_GRIDS && moves[n_moves] != -1)
+            ++n_moves;
+        int move = moves[wyhash64() % n_moves];
+        kfree(moves);
+        temp_table[move] = current_player;
+        char win;
+        if ((win = check_win(temp_table)) != ' ')
+            return calculate_win_value(win, player);
+        current_player ^= 'O' ^ 'X';
+    }
+    return (fixed_point_t) (1UL << (FIXED_SCALE_BITS - 1));
+}
+
+static void backpropagate(struct node *node, fixed_point_t score)
+{
+    while (node) {
+        node->n_visits++;
+        node->score += score;
+        node = node->parent;
+        score = 1 - score;
+    }
+}
+
+static void expand(struct node *node, char *table)
+{
+    int *moves = available_moves(table);
+    int n_moves = 0;
+    while (n_moves < N_GRIDS && moves[n_moves] != -1)
+        ++n_moves;
+    for (int i = 0; i < n_moves; i++) {
+        node->children[i] = new_node(moves[i], node->player ^ 'O' ^ 'X', node);
+    }
+    kfree(moves);
+}
+
+int mcts(char *table, char player)
+{
+    char win;
+    struct node *root = new_node(-1, player, NULL);
+    for (int i = 0; i < ITERATIONS; i++) {
+        struct node *node = root;
+        char temp_table[N_GRIDS];
+        memcpy(temp_table, table, N_GRIDS);
+        while (1) {
+            if ((win = check_win(temp_table)) != ' ') {
+                fixed_point_t score =
+                    calculate_win_value(win, node->player ^ 'O' ^ 'X');
+                backpropagate(node, score);
+                break;
+            }
+            if (node->n_visits == 0) {
+                fixed_point_t score = simulate(temp_table, node->player);
+                backpropagate(node, score);
+                break;
+            }
+            if (node->children[0] == NULL)
+                expand(node, temp_table);
+            node = select_move(node);
+            if (!node)
+                return -1;
+            temp_table[node->move] = node->player ^ 'O' ^ 'X';
+        }
+    }
+    struct node *best_node = root;
+    int most_visits = -1;
+    for (int i = 0; i < N_GRIDS; i++) {
+        if (root->children[i] && root->children[i]->n_visits > most_visits) {
+            most_visits = root->children[i]->n_visits;
+            best_node = root->children[i];
+        }
+    }
+    int best_move = best_node->move;
+    free_node(root);
+    return best_move;
+}

mcts.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#define ITERATIONS 100000
+#define EXPLORATION_FACTOR fixed_sqrt(1U << (FIXED_SCALE_BITS + 1))
+
+int mcts(char *table, char player);