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

#![allow(clippy::multiple_bound_locations)]

//!Module for layer types.

use std::ops::AddAssign;
use std::usize;

use crate::nnue_rs::ops::*;

use bytemuck::Zeroable;

///A dense layer.
#[derive(Debug, Clone, Zeroable)]
pub struct Dense<W: Zeroable, B: Zeroable, const INPUTS: usize, const OUTPUTS: usize> {
    pub weights: [[W; INPUTS]; OUTPUTS],
    pub biases: [B; OUTPUTS],
}

impl<
        W: Copy + Zeroable,
        B: Copy + Zeroable + AddAssign + From<<[W; INPUTS] as Dot>::Output>,
        const INPUTS: usize,
        const OUTPUTS: usize,
    > Dense<W, B, INPUTS, OUTPUTS>
where
    [W; INPUTS]: Dot,
{
    pub fn activate(&self, inputs: &[W; INPUTS], outputs: &mut [B; OUTPUTS]) {
        *outputs = self.biases;
        for (o, w) in outputs.iter_mut().zip(&self.weights) {
            *o += inputs.dot(w).into();
        }
    }
}

///A specialized [`Dense`] layer that operates on boolean inputs
///and can incrementally update the output accumulator.
#[derive(Debug, Clone, Zeroable)]
pub struct BitDense<WB: Zeroable, const INPUTS: usize, const OUTPUTS: usize> {
    pub weights: [[WB; OUTPUTS]; INPUTS],
    pub biases: [WB; OUTPUTS],
}

impl<WB: Zeroable + Clone, const INPUTS: usize, const OUTPUTS: usize> BitDense<WB, INPUTS, OUTPUTS>
where
    [WB; OUTPUTS]: VecAdd + VecSub,
{
    ///Clear an accumulator to a default state.
    pub fn empty(&self, outputs: &mut [WB; OUTPUTS]) {
        outputs.clone_from(&self.biases);
    }

    ///Add an input feature to an accumulator.
    pub fn add(&self, index: usize, outputs: &mut [WB; OUTPUTS]) {
        outputs.vec_add(&self.weights[index]);
    }

    ///Remove an input feature from an accumulator.
    pub fn sub(&self, index: usize, outputs: &mut [WB; OUTPUTS]) {
        outputs.vec_sub(&self.weights[index]);
    }

    ///Debug function for testing; Recommended to use the [`add`](Self::add) and [`sub`](Self::sub)
    ///functions instead to incrementally add and remove input features from an accumulator.
    pub fn activate(&self, inputs: &[bool; INPUTS], outputs: &mut [WB; OUTPUTS]) {
        self.empty(outputs);
        for (i, &input) in inputs.iter().enumerate() {
            if input {
                self.add(i, outputs);
            }
        }
    }
}