From 9ce1314f9d452d0dc56e1f0d87e37ab529d0d1c8 Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Wed, 20 Mar 2024 10:34:29 -0400 Subject: [PATCH] Switch memory from export to import --- crates/wasm/src/lib.rs | 40 ++++++++++++++++++++------------- crates/web/src/lib.rs | 8 ++++++- packages/core/src/impl.ts | 30 ++++++++++++++++++++----- packages/core/src/index.test.ts | 4 ++-- 4 files changed, 58 insertions(+), 24 deletions(-) diff --git a/crates/wasm/src/lib.rs b/crates/wasm/src/lib.rs index 23cbbd5..662444d 100644 --- a/crates/wasm/src/lib.rs +++ b/crates/wasm/src/lib.rs @@ -7,7 +7,7 @@ use std::{ }; use wasm_encoder::{ BlockType, CodeSection, EntityType, ExportSection, Function, FunctionSection, ImportSection, - Instruction, MemArg, MemorySection, MemoryType, Module, TypeSection, ValType, + Instruction, MemArg, MemoryType, Module, TypeSection, ValType, }; /// Resolve `ty` via `generics` and `types`, then return its ID in `typemap`, inserting if need be. @@ -976,18 +976,23 @@ impl<'a, 'b, O: Eq + Hash, T: Refs<'a, Opaque = O>> Codegen<'a, 'b, O, T> { /// A WebAssembly module for a graph of functions. /// -/// The module exports its memory with name `"m"` and its entrypoint function with name `"f"`. The -/// function takes one parameter in addition to its original parameters, which must be an -/// 8-byte-aligned pointer to the start of the memory region it can use for allocation. The memory -/// is the exact number of pages necessary to accommodate the function's own memory allocation as -/// well as memory allocation for all of its parameters, with each node in each parameter's memory -/// allocation tree being 8-byte aligned. That is, the function's last argument should be just large -/// enough to accommodate those allocations for all the parameters; in that case, no memory will be +/// The module exports its entrypoint function with name `"f"`. The function takes one parameter in +/// addition to its original parameters, which must be an 8-byte-aligned pointer to the start of the +/// memory region it can use for allocation. +/// +/// Under module name `"m"`, the module imports a memory whose minimum number of pages is the exact +/// number of pages necessary to accommodate the function's own memory allocation as well as memory +/// allocation for all of its parameters, with each node in each parameter's memory allocation tree +/// being 8-byte aligned. That is, the function's last argument should be just large enough to +/// accommodate those allocations for all the parameters; in that case, no memory will be /// incorrectly overwritten and no out-of-bounds memory accesses will occur. pub struct Wasm { /// The bytes of the WebAssembly module binary. pub bytes: Vec, + /// The minimum number of pages required by the imported memory. + pub pages: u64, + /// All the opaque functions that the WebAssembly module must import, in order. /// /// The module name for each import is the empty string, and the field name is the base-ten @@ -1390,7 +1395,6 @@ pub fn compile<'a, O: Eq + Hash, T: Refs<'a, Opaque = O>>(f: Node<'a, O, T>) -> type_section.function(params.into_vec(), results.into_vec()); } - let mut memory_section = MemorySection::new(); let page_size = 65536; // https://webassembly.github.io/spec/core/exec/runtime.html#page-size let cost = funcs.last().map_or(0, |((def, _), (_, def_types, _))| { def.params @@ -1400,12 +1404,16 @@ pub fn compile<'a, O: Eq + Hash, T: Refs<'a, Opaque = O>>(f: Node<'a, O, T>) -> .sum() }) + costs.last().unwrap_or(&0); let pages = ((cost + page_size - 1) / page_size).into(); // round up to a whole number of pages - memory_section.memory(MemoryType { - minimum: pages, - maximum: Some(pages), - memory64: false, - shared: false, - }); + import_section.import( + "m", + "", + MemoryType { + minimum: pages, + maximum: None, + memory64: false, + shared: false, + }, + ); let mut export_section = ExportSection::new(); export_section.export( @@ -1419,11 +1427,11 @@ pub fn compile<'a, O: Eq + Hash, T: Refs<'a, Opaque = O>>(f: Node<'a, O, T>) -> module.section(&type_section); module.section(&import_section); module.section(&function_section); - module.section(&memory_section); module.section(&export_section); module.section(&code_section); Wasm { bytes: module.finish(), + pages, imports, } } diff --git a/crates/web/src/lib.rs b/crates/web/src/lib.rs index 8c209f0..ccbd532 100644 --- a/crates/web/src/lib.rs +++ b/crates/web/src/lib.rs @@ -330,9 +330,14 @@ impl Func { /// Compile the call graph subtended by this function to WebAssembly. pub fn compile(&self) -> Wasm { - let rose_wasm::Wasm { bytes, imports } = rose_wasm::compile(self.node()); + let rose_wasm::Wasm { + bytes, + pages, + imports, + } = rose_wasm::compile(self.node()); Wasm { bytes: Some(bytes), + pages, imports: Some( imports .into_keys() @@ -488,6 +493,7 @@ impl Func { #[wasm_bindgen] pub struct Wasm { bytes: Option>, + pub pages: u64, imports: Option>, } diff --git a/packages/core/src/impl.ts b/packages/core/src/impl.ts index dc191ef..43ec980 100644 --- a/packages/core/src/impl.ts +++ b/packages/core/src/impl.ts @@ -814,19 +814,39 @@ const getMeta = ( } else return undefined; }; -/** Concretize the abstract function `f` using the compiler. */ +interface CompileOptions { + memory?: WebAssembly.Memory; +} + +/** + * Concretize the abstract function `f` using the compiler. + * + * Creates a new memory if `opts.memory` is not provided, otherwise attempts to + * grow the provided memory to be large enough. + */ export const compile = async ( f: Fn & ((...args: A) => R), + opts?: CompileOptions, ): Promise<(...args: JsArgs) => ToJs> => { const func = f[inner]; const res = func.compile(); const bytes = res.bytes()!; + const pages = Number(res.pages); const imports = res.imports()!; res.free(); - const instance = await WebAssembly.instantiate( - await WebAssembly.compile(bytes), - { "": Object.fromEntries(imports.map((g, i) => [i.toString(), g])) }, - ); + let memory = opts?.memory; + if (memory === undefined) memory = new WebAssembly.Memory({ initial: pages }); + else { + // https://webassembly.github.io/spec/core/exec/runtime.html#page-size + const pageSize = 65536; + const delta = pages - memory.buffer.byteLength / pageSize; + if (delta > 0) memory.grow(delta); + } + const mod = await WebAssembly.compile(bytes); + const instance = await WebAssembly.instantiate(mod, { + m: { "": memory }, + "": Object.fromEntries(imports.map((g, i) => [i.toString(), g])), + }); const { f: g, m } = instance.exports; const metas: (Meta | undefined)[] = []; const n = func.numTypes(); diff --git a/packages/core/src/index.test.ts b/packages/core/src/index.test.ts index 7ce9f51..11d6da1 100644 --- a/packages/core/src/index.test.ts +++ b/packages/core/src/index.test.ts @@ -656,12 +656,12 @@ describe("valid", () => { expect(memory.buffer.byteLength).toBeGreaterThan(pageSize); const a = []; const b = []; - for (let i = 0; i < n; ++i) { + for (let i = 1; i <= n; ++i) { a.push(i); b.push(1 / i); } const c = gCompiled(a, b); - for (let i = 0; i < n; ++i) expect(c[i]).toBe(1); + for (let i = 0; i < n; ++i) expect(c[i]).toBeCloseTo(1); }); test("compile opaque function", async () => {