1/* programs/matmul's exact mode on the CPU (mesh-ref matmul.rs). */
2#include "mesh/ref/matmul.h"
3
4#include "mesh/ftz.h"
5#include "mesh/ref/block.h"
6
7#include <pthread.h>
8#include <stdbool.h>
9
10float mesh_matmul_value(uint32_t seed, uint32_t i) {
11    uint32_t h = mesh_block_mix(((seed * 0x9e3779b9u) + i) * 0x85ebca6bu ^ seed);
12    return mesh_ftz_sub(mesh_ftz_mul(mesh_ftz_from_u32(h >> 8), 1.0f / 8388608.0f), 1.0f);
13}
14
15void mesh_matmul_matrix(uint32_t seed, size_t n, float *out) {
16    for (size_t i = 0; i < n; i++) out[i] = mesh_matmul_value(seed, (uint32_t)i);
17}
18
19typedef struct {
20    const float *a, *b;
21    float *c;
22    size_t k, n, row0, rows;
23} job;
24
25static void *rows(void *arg) {
26    const job *j = arg;
27    for (size_t i = j->row0; i < j->row0 + j->rows; i++)
28        for (size_t col = 0; col < j->n; col++) {
29            float acc = 0.0f;
30            for (size_t kk = 0; kk < j->k; kk++)
31                acc = mesh_ftz_add(acc, mesh_ftz_mul(j->a[i * j->k + kk], j->b[kk * j->n + col]));
32            j->c[i * j->n + col] = acc;
33        }
34    return NULL;
35}
36
37void mesh_matmul(const float *a, const float *b, size_t m, size_t k, size_t n, size_t threads, float *c) {
38    if (threads == 0) threads = 1;
39    size_t per = (m + threads - 1) / threads;
40    if (per == 0) per = 1;
41    size_t chunks = (m + per - 1) / per;
42    if (chunks <= 1) {
43        job one = {a, b, c, k, n, 0, m};
44        rows(&one);
45        return;
46    }
47    job jobs[chunks];
48    pthread_t tid[chunks];
49    bool started[chunks];
50    for (size_t t = 0; t < chunks; t++) {
51        size_t r0 = t * per;
52        jobs[t] = (job){a, b, c, k, n, r0, r0 + per <= m ? per : m - r0};
53        started[t] = pthread_create(&tid[t], NULL, rows, &jobs[t]) == 0;
54        if (!started[t]) rows(&jobs[t]); /* no thread: do it here */
55    }
56    for (size_t t = 0; t < chunks; t++)
57        if (started[t]) pthread_join(tid[t], NULL);
58}
59
60void mesh_matvec(const float *a, const float *x, size_t m, size_t k, float *y) {
61    for (size_t i = 0; i < m; i++) {
62        float lanes[32];
63        for (size_t l = 0; l < 32; l++) {
64            float p = 0.0f;
65            for (size_t kk = l; kk < k; kk += 32) p = mesh_ftz_add(p, mesh_ftz_mul(a[i * k + kk], x[kk]));
66            lanes[l] = p;
67        }
68        mesh_block_warp_reduce_f32(lanes, mesh_block_add_f32);
69        y[i] = lanes[0];
70    }
71}
72
73uint64_t mesh_matmul_digest(const float *c, size_t n) {
74    uint64_t h = 0xcbf29ce484222325ull;
75    for (size_t i = 0; i < n; i++) h = (h ^ mesh_ftz_bits(c[i])) * 0x00000100000001b3ull;
76    return h;
77}