==> src/compile.c <==1/* Images (mesh-host/src/wasm.rs `Program`, `instantiate`, `WasmCompiler::prepare`). */
2#include "mesh/exec_wasm.h"
3
4#include "mesh/abi/backend.h"
5#include "mesh/abi/base.h"
6#include "mesh/abi/fnv.h"
7#include "mesh/abi/wasm.h"
8#include "mesh/log.h"
9
10#include <errno.h>
11#include <stdio.h>
12#include <stdlib.h>
13#include <string.h>
14#include <time.h>
15
16#include "exec_internal.h"
17#include "rt_internal.h"
18
19wasm_instance *mesh_wasm_instantiate(const mesh_wasm_program *p, mesh_mailbox mb, uint32_t slot, uint32_t node,
20 uint32_t nodes, mesh_wasm_lane_data *data, char **err) {
21 mesh_wasm_lane_data_init(data, mb, slot, node, nodes, p->task_ms, &p->rt->epoch);
22 size_t n;
23 const wasm_host_func *imports = mesh_wasm_imports(&n);
24 /* StoreLimits: memory at mem_bytes, one instance, 4 tables, 1 memory. */
25 wasm_store_limits lim = {p->mem_bytes, 4, 1, 0, 0};
26 char msg[512];
27 wasm_instance *inst = wasm_instantiate(p->module, imports, n, &lim, data, msg, sizeof msg);
28 if (!inst) {
29 if (err) {
30 bool link = strncmp(msg, "unknown import", 14) == 0 || strncmp(msg, "incompatible import", 19) == 0;
31 *err = link ? mesh_wasm_fmt("load: %s: %s", p->name, msg) : mesh_wasm_fmt("load: %s: instantiate: %s", p->name, msg);
32 }
33 return NULL;
34 }
35 /* No deadline until a task. */
36 wasm_set_epoch(inst, &p->rt->epoch);
37 wasm_set_deadline(inst, UINT64_MAX / 2);
38 return inst;
39}
40
41bool mesh_wasm_entry_points(const mesh_wasm_program *p, wasm_instance *inst, uint32_t *init, uint32_t *task, char **err) {
42 char msg[256];
43 if (!wasm_module_export(wasm_compiled_module(p->module), MESH_WASM_EXPORT_MEMORY, WASM_EXTERN_MEMORY) ||
44 !wasm_memory(inst, NULL)) {
45 if (err) *err = mesh_wasm_fmt("load: %s: no `%s` export", p->name, MESH_WASM_EXPORT_MEMORY);
46 return false;
47 }
48 int64_t f = wasm_instance_func(inst, MESH_WASM_EXPORT_INIT, ">", msg, sizeof msg);
49 if (f < 0) {
50 if (err) *err = mesh_wasm_fmt("load: %s: %s: %s", p->name, MESH_WASM_EXPORT_INIT, msg);
51 return false;
52 }
53 *init = (uint32_t)f;
54 f = wasm_instance_func(inst, MESH_WASM_EXPORT_TASK, "ii>i", msg, sizeof msg);
55 if (f < 0) {
56 if (err) *err = mesh_wasm_fmt("load: %s: %s: %s", p->name, MESH_WASM_EXPORT_TASK, msg);
57 return false;
58 }
59 *task = (uint32_t)f;
60 return true;
61}
62
63void mesh_wasm_program_free(mesh_wasm_program *p) {
64 if (!p || atomic_fetch_sub(&p->refs, 1) > 1) return;
65 wasm_compiled_free(p->module);
66 free(p->name);
67 free(p);
68}
69
70/* ---- the kernel of a wasm slot ---------------------------------------------------------- */
71
72static bool k_zero(const mesh_kernel *k, uint32_t *out, char **err) {
73 (void)k, (void)err;
74 *out = 0;
75 return true;
76}
77/* A lane is a thread: nothing is resident (reported as 1, never 0 = "does not fit"). */
78static bool k_one(const mesh_kernel *k, uint32_t *out, char **err) {
79 (void)k, (void)err;
80 *out = 1;
81 return true;
82}
83static uint64_t k_scratch(const mesh_kernel *k) {
84 (void)k;
85 return 0;
86}
87static bool k_global(const mesh_kernel *k, const char *name, uint32_t *out, char **err) {
88 const mesh_wasm_program *p = k->self;
89 if (strcmp(name, MESH_RT_ABI_SYMBOL) == 0) {
90 *out = p->rt_abi;
91 return true;
92 }
93 if (err) *err = mesh_wasm_fmt("global %s: a wasm image reports only %s (its mesh_rt_abi export)", name, MESH_RT_ABI_SYMBOL);
94 return false;
95}
96static void k_drop(mesh_kernel *k) { mesh_wasm_program_free(k->self); }
97
98const mesh_kernel_ops mesh_wasm_kernel_ops = {MESH_BACKEND_WASM, k_zero, k_zero, k_one, k_scratch, k_global, k_drop};
99
100const mesh_wasm_program *mesh_wasm_kernel_program(const mesh_kernel *k) {
101 return k && k->ops == &mesh_wasm_kernel_ops ? k->self : NULL;
102}
103
104/* ---- the compiler ----------------------------------------------------------------------- */
105
106static bool c_needs_ptx(const mesh_compiler *c) {
107 (void)c;
108 return false;
109}
110
111static bool read_file(const char *path, uint8_t **out, size_t *n, char **why) {
112 FILE *f = fopen(path, "rb");
113 if (!f) {
114 *why = mesh_wasm_fmt("%s", strerror(errno));
115 return false;
116 }
117 uint8_t *b = NULL;
118 size_t len = 0, cap = 0;
119 for (;;) {
120 if (len == cap) {
121 cap = cap ? cap * 2 : 65536;
122 uint8_t *q = realloc(b, cap);
123 if (!q) break;
124 b = q;
125 }
126 size_t k = fread(b + len, 1, cap - len, f);
127 len += k;
128 if (k == 0) break;
129 }
130 bool ok = !ferror(f) && b;
131 if (!ok) *why = mesh_wasm_fmt("%s", ferror(f) ? strerror(errno) : "out of memory");
132 fclose(f);
133 if (!ok) {
134 free(b);
135 return false;
136 }
137 *out = b;
138 *n = len;
139 return true;
140}
141
142static bool c_prepare(const mesh_compiler *c, mesh_payload *src, const uint64_t *expect, const char *root, mesh_prepared *out,
143 char **err) {
144 mesh_wasm_rt *rt = c->self;
145 mesh_manifest m;
146 memset(out, 0, sizeof *out);
147 if (!mesh_check_manifest(src, MESH_BACKEND_WASM, &m, err)) {
148 mesh_payload_free(src);
149 return false;
150 }
151 bool ok = false;
152 char *dir = NULL, *path = NULL, *why = NULL;
153 uint8_t *wasm = NULL;
154 size_t wlen = 0;
155 mesh_wasm_program *p = NULL;
156 wasm_instance *inst = NULL;
157 mesh_wasm_lane_data *data = NULL;
158 uint8_t *scratch = NULL;
159 if (!(dir = mesh_loader_resolve(root, src->dir, err))) goto out;
160 path = mesh_wasm_fmt("%s/%s.wasm", dir, m.name);
161 if (!read_file(path, &wasm, &wlen, &why)) {
162 if (err) *err = mesh_wasm_fmt("manifest: %s declares backend wasm, but %s: %s", m.name, path, why);
163 goto out;
164 }
165 const char *what = src->ptx_len == 0 ? "wasm" : "PTX";
166 uint64_t hash = src->ptx_len == 0 ? mesh_image_hash(wasm, wlen, m.param) : mesh_image_hash(src->ptx, src->ptx_len, m.param);
167 if (!mesh_check_hash(m.name, what, hash, expect, err)) goto out;
168 wasm_float_report fr;
169 char msg[1024];
170 if (!wasm_float_check(wasm, wlen, m.ftz, &fr, msg, sizeof msg)) {
171 if (err) *err = mesh_wasm_fmt("%s (%s)", msg, path);
172 goto out;
173 }
174 double compile_ms = 0;
175 const char *how = "compiled";
176 wasm_compiled *module = mesh_wasm_rt_module(rt, m.name, wasm, wlen, &compile_ms, &how, err);
177 if (!module) goto out;
178 if (!(p = calloc(1, sizeof *p))) {
179 wasm_compiled_free(module);
180 goto out;
181 }
182 atomic_init(&p->refs, 1);
183 p->name = strdup(m.name);
184 p->module = module;
185 p->rt = rt;
186 p->wasm_hash = mesh_fnv1a64(wasm, wlen);
187 p->mem_bytes = (uint64_t)m.mem_mb << 20;
188 p->task_ms = m.task_ms;
189 p->float_report = fr;
190 /* MESH_RT_ABI: the image's mesh_rt_abi export, in a scratch instance (which also proves the
191 * exports and the memory cap fit). */
192 if (!(scratch = calloc(1, MESH_MAILBOX_BYTES)) || !(data = malloc(sizeof *data))) {
193 if (err) *err = mesh_wasm_fmt("load: %s: out of memory", m.name);
194 goto out;
195 }
196 if (!(inst = mesh_wasm_instantiate(p, mesh_mailbox_new(scratch, 0), 0, 0, 1, data, err))) goto out;
197 uint32_t init, task;
198 if (!mesh_wasm_entry_points(p, inst, &init, &task, err)) goto out;
199 int64_t g = wasm_instance_func(inst, MESH_WASM_EXPORT_RT_ABI, ">i", msg, sizeof msg);
200 if (g < 0) {
201 if (err) *err = mesh_wasm_fmt("abi: %s: no %s export: %s", m.name, MESH_WASM_EXPORT_RT_ABI, msg);
202 goto out;
203 }
204 uint64_t io[1] = {0};
205 int trap = wasm_call(inst, (uint32_t)g, io);
206 if (trap != WASM_NUM_OK) {
207 if (err) *err = mesh_wasm_fmt("abi: %s: %s: wasm trap: %s", m.name, MESH_WASM_EXPORT_RT_ABI, wasm_trap_message(trap));
208 goto out;
209 }
210 p->rt_abi = (uint32_t)io[0];
211 mesh_info("wasm: %s %s in %.1f ms (%zu bytes of wasm, wasm hash %016llx, image hash %016llx from its %s; float %s: %u f32 "
212 "ops, all in %u %s* helpers; lanes capped at %u MiB, %u ms a task)",
213 m.name, how, compile_ms, wlen, (unsigned long long)p->wasm_hash, (unsigned long long)hash, what,
214 m.ftz ? "ftz" : "ieee", fr.f32_ops, fr.helpers, MESH_WASM_FTZ_PREFIX, m.mem_mb, m.task_ms);
215 mesh_kernel *k = malloc(sizeof *k);
216 if (!k) goto out;
217 k->ops = &mesh_wasm_kernel_ops;
218 k->self = p;
219 out->manifest = m;
220 out->hash = hash;
221 out->code = wasm;
222 out->code_len = wlen;
223 out->compile_ms = compile_ms;
224 out->src = *src;
225 memset(src, 0, sizeof *src);
226 out->kernel = k;
227 out->exec = MESH_BACKEND_WASM;
228 out->reg_cap = 0;
229 wasm = NULL;
230 p = NULL;
231 ok = true;
232out:
233 wasm_instance_free(inst);
234 free(data);
235 free(scratch);
236 mesh_wasm_program_free(p);
237 free(wasm);
238 free(path);
239 free(dir);
240 free(why);
241 if (!ok) {
242 mesh_manifest_free(&m);
243 mesh_payload_free(src);
244 }
245 return ok;
246}
247
248const mesh_compiler_ops mesh_wasm_compiler_ops = {c_needs_ptx, c_prepare};