1/* infer: the trained cubic over a 64-point grid. It reads train's params
2 * through mesh_ctx_peer_view (manifest `reads = ["train"]`, so the runtime
3 * refuses any other view). Layout: @ajax/mesh-abi-view-infer. The loops
4 * mirror mesh-kcheck's CPU reference index by index (bitwise checks).
5 * Ported from alpha-rust-mesh programs/infer/src/lib.rs. */
6#include "mesh/rt.h"
7
8static float infer_cubic(const float p[4], float x) {
9    float x2 = x * x;
10    float x3 = x2 * x;
11    return p[0] + p[1] * x + p[2] * x2 + p[3] * x3;
12}
13
14/* The grid: x_i = -1 + 2 i / (POINTS - 1); lane l writes points l and l + 32. */
15static void infer_init(const mesh_ctx *ctx) {
16    uint32_t l = mesh_ctx_lane(ctx);
17    uint32_t is[2] = {l, l + 32};
18    for (int k = 0; k < 2; k++) {
19        uint32_t i = is[k];
20        mesh_ctx_set_f32(ctx, MESH_INFER_I_X + 4 * i, -1.0f + 2.0f * (float)i / (float)(MESH_INFER_POINTS - 1));
21    }
22}
23
24/* Returns POINTS, or R_NO_TRAIN when no train is loaded. */
25static uint32_t infer_run(const mesh_ctx *ctx) {
26    uint32_t s;
27    mesh_view tv;
28    if (!mesh_ctx_find(ctx, "train", 5, &s) || !mesh_ctx_peer_view(ctx, s, &tv)) return MESH_INFER_R_NO_TRAIN;
29    float p[4] = {mesh_view_f32(tv, MESH_TRAIN_T_PARAMS), mesh_view_f32(tv, MESH_TRAIN_T_PARAMS + 4),
30                  mesh_view_f32(tv, MESH_TRAIN_T_PARAMS + 8), mesh_view_f32(tv, MESH_TRAIN_T_PARAMS + 12)};
31    uint32_t l = mesh_ctx_lane(ctx);
32    uint32_t is[2] = {l, l + 32};
33    for (int k = 0; k < 2; k++) {
34        uint32_t i = is[k];
35        mesh_ctx_set_f32(ctx, MESH_INFER_I_Y + 4 * i, infer_cubic(p, mesh_ctx_view_f32(ctx, MESH_INFER_I_X + 4 * i)));
36    }
37    if (mesh_ctx_leader(ctx)) {
38        for (uint32_t k = 0; k < 4; k++) mesh_ctx_set_f32(ctx, MESH_INFER_I_PARAMS + 4 * k, p[k]);
39        mesh_ctx_set_u32(ctx, MESH_INFER_I_RUNS, mesh_ctx_view_u32(ctx, MESH_INFER_I_RUNS) + 1);
40        mesh_ctx_set_u32(ctx, MESH_INFER_I_TRAIN_STEP, mesh_view_u32(tv, MESH_TRAIN_T_NEXT_STEP));
41        mesh_ctx_count(ctx, 0, 1);
42    }
43    return MESH_INFER_POINTS;
44}
45
46static uint32_t infer_task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
47    if (op == MESH_OP_TASK) return infer_run(ctx);
48    if (op == MESH_INFER_OP_PROBE_READ) {
49        mesh_view v;
50        return mesh_ctx_peer_view(ctx, arg, &v) ? mesh_view_len(v) : MESH_INFER_R_NO_VIEW;
51    }
52    return MESH_R_BAD_OP;
53}
54
55MESH_PROGRAM(infer_init, infer_task);