1/* infer for the wasm executor (alpha-rust-mesh programs/infer/wasm): the trained cubic over a
2 * 64-point grid, reading train's params through a peer view (manifest reads = ["train"]), in the
3 * cluster float mode (Z): the same bits as the GPUs' image. Layout: mesh-abi views::infer. */
4#include "mesh/abi/view_infer.h"
5#include "mesh/abi/view_train.h"
6#include "mesh/ref/cubic.h"
7#include "mesh/rt_wasm.h"
8
9static float xs[MESH_INFER_POINTS], ys[MESH_INFER_POINTS];
10
11/* The grid: x_i = -1 + 2 i / (POINTS - 1) (mesh-ref's, operation for operation the GPU's). */
12static void init(const mesh_ctx *ctx) {
13    mesh_cubic_grid(MESH_INFER_POINTS, xs);
14    mesh_view_write_f32s(mesh_ctx_view(ctx), MESH_INFER_I_X, xs, MESH_INFER_POINTS);
15}
16
17/* Returns POINTS, or R_NO_TRAIN when no train is loaded. */
18static uint32_t run(const mesh_ctx *ctx) {
19    uint32_t s;
20    mesh_view tv;
21    if (!mesh_ctx_find(ctx, "train", 5, &s) || !mesh_ctx_peer_view(ctx, s, &tv)) return MESH_INFER_R_NO_TRAIN;
22    float p[4];
23    mesh_view_read_f32s(tv, MESH_TRAIN_T_PARAMS, p, 4);
24    mesh_view v = mesh_ctx_view(ctx);
25    mesh_view_read_f32s(v, MESH_INFER_I_X, xs, MESH_INFER_POINTS);
26    for (uint32_t i = 0; i < MESH_INFER_POINTS; i++) ys[i] = mesh_cubic(p, xs[i]);
27    mesh_view_write_f32s(v, MESH_INFER_I_Y, ys, MESH_INFER_POINTS);
28    mesh_view_write_f32s(v, MESH_INFER_I_PARAMS, p, 4);
29    mesh_view_set_u32(v, MESH_INFER_I_RUNS, mesh_view_u32(v, MESH_INFER_I_RUNS) + 1u);
30    mesh_view_set_u32(v, MESH_INFER_I_TRAIN_STEP, mesh_view_u32(tv, MESH_TRAIN_T_NEXT_STEP));
31    mesh_ctx_count(ctx, 0, 1);
32    return MESH_INFER_POINTS;
33}
34
35static uint32_t task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
36    switch (op) {
37    case MESH_OP_TASK: return run(ctx);
38    case MESH_INFER_OP_PROBE_READ: {
39        mesh_view pv;
40        return mesh_ctx_peer_view(ctx, arg, &pv) ? mesh_view_len(pv) : MESH_INFER_R_NO_VIEW;
41    }
42    default: return MESH_R_BAD_OP;
43    }
44}
45
46MESH_PROGRAM(init, task)