1/* programs/train's and infer's CPU reference (mesh-ref lib.rs `cubic`, `warp_sum`, `step_sums`,
2 * `apply`, `loss`, `sin_pi`, `grid`, `samples`). */
3#include "mesh/ref/cubic.h"
4
5#include "mesh/abi/view_train.h"
6#include "mesh/ftz.h"
7#include "mesh/ref/block.h"
8
9float mesh_cubic(const float p[4], float x) {
10    float x2 = mesh_ftz_mul(x, x), x3 = mesh_ftz_mul(x2, x);
11    return mesh_ftz_add(mesh_ftz_add(mesh_ftz_add(p[0], mesh_ftz_mul(p[1], x)), mesh_ftz_mul(p[2], x2)),
12                        mesh_ftz_mul(p[3], x3));
13}
14
15float mesh_cubic_warp_sum(const float v[32]) {
16    float w[32];
17    for (int l = 0; l < 32; l++) w[l] = v[l];
18    mesh_block_warp_reduce_f32(w, mesh_block_add_f32);
19    return w[0];
20}
21
22void mesh_cubic_step_sums(const float p[4], const float *xs, const float *ys, uint32_t off, float out[5]) {
23    float lanes[5][32];
24    for (uint32_t l = 0; l < 32; l++) {
25        size_t i = (8 * l + off) % 256;
26        float x = xs[i], y = ys[i];
27        float e = mesh_ftz_sub(mesh_cubic(p, x), y);
28        float g = mesh_ftz_mul(2.0f, e);
29        float x2 = mesh_ftz_mul(x, x);
30        lanes[0][l] = mesh_ftz_mul(e, e);
31        lanes[1][l] = g;
32        lanes[2][l] = mesh_ftz_mul(g, x);
33        lanes[3][l] = mesh_ftz_mul(g, x2);
34        lanes[4][l] = mesh_ftz_mul(g, mesh_ftz_mul(x2, x));
35    }
36    for (int k = 0; k < 5; k++) out[k] = mesh_cubic_warp_sum(lanes[k]);
37}
38
39void mesh_cubic_apply(float p[4], const float s[5], float lr, float inv) {
40    for (int k = 0; k < 4; k++) p[k] = mesh_ftz_sub(p[k], mesh_ftz_mul(mesh_ftz_mul(lr, s[k + 1]), inv));
41}
42
43float mesh_cubic_loss(const float s[5], float inv) { return mesh_ftz_mul(s[0], inv); }
44
45float mesh_cubic_sin_pi(float x) {
46    double t = mesh_ftz_to_f64(x);
47    if (t > 0.5)
48        t = 1.0 - t;
49    else if (t < -0.5)
50        t = -1.0 - t;
51    double a = 3.141592653589793 * t;
52    double a2 = a * a;
53    double s = 1.0;
54    for (uint32_t k = 11; k >= 1; k--) s = 1.0 - a2 / (double)((2 * k) * (2 * k + 1)) * s;
55    return mesh_ftz_from_f64(a * s);
56}
57
58void mesh_cubic_grid(uint32_t n, float *out) {
59    float d = mesh_ftz_from_u32(n - 1);
60    for (uint32_t i = 0; i < n; i++) out[i] = mesh_ftz_add(-1.0f, mesh_ftz_div(mesh_ftz_mul(2.0f, mesh_ftz_from_u32(i)), d));
61}
62
63void mesh_cubic_samples(float *xs, float *ys) {
64    mesh_cubic_grid(MESH_TRAIN_SAMPLES, xs);
65    for (uint32_t i = 0; i < MESH_TRAIN_SAMPLES; i++) ys[i] = mesh_cubic_sin_pi(xs[i]);
66}