From 3ad144e6465d51c9334d61397f53c09c5b77f891 Mon Sep 17 00:00:00 2001 From: Edoardo Vacchi Date: Thu, 12 Dec 2024 16:25:51 +0100 Subject: [PATCH] Introduce CachedAotMachineFactory. Signed-off-by: Edoardo Vacchi --- .../chicory/sdk/CachedAotMachineFactory.java | 27 +++++++++++++++++++ .../org/extism/chicory/sdk/ChicoryModule.java | 4 +-- .../extism/chicory/sdk/DependencyGraph.java | 10 ++++++- .../java/org/extism/chicory/sdk/Kernel.java | 17 +++++++++--- .../java/org/extism/chicory/sdk/Linker.java | 6 +++-- .../org/extism/chicory/sdk/PluginTest.java | 2 +- 6 files changed, 57 insertions(+), 9 deletions(-) create mode 100644 src/main/java/org/extism/chicory/sdk/CachedAotMachineFactory.java diff --git a/src/main/java/org/extism/chicory/sdk/CachedAotMachineFactory.java b/src/main/java/org/extism/chicory/sdk/CachedAotMachineFactory.java new file mode 100644 index 0000000..057dc96 --- /dev/null +++ b/src/main/java/org/extism/chicory/sdk/CachedAotMachineFactory.java @@ -0,0 +1,27 @@ +package org.extism.chicory.sdk; + +import com.dylibso.chicory.experimental.aot.AotMachineFactory; +import com.dylibso.chicory.runtime.Instance; +import com.dylibso.chicory.runtime.Machine; +import com.dylibso.chicory.wasm.WasmModule; + +import java.util.HashMap; +import java.util.Map; +import java.util.function.Function; + +class CachedAotMachineFactory implements Function { + private final Map factories = new HashMap<>(); + + public CachedAotMachineFactory compile(WasmModule wasmModule) { + factories.put(wasmModule, new AotMachineFactory(wasmModule)); + return this; + } + + public Machine apply(Instance instance) { + if (!factories.containsKey(instance.module())) { + throw new IllegalArgumentException("Instance module is not cached"); + } + var factory = factories.get(instance.module()); + return factory.apply(instance); + } +} diff --git a/src/main/java/org/extism/chicory/sdk/ChicoryModule.java b/src/main/java/org/extism/chicory/sdk/ChicoryModule.java index 3c68d33..2f2f1ff 100644 --- a/src/main/java/org/extism/chicory/sdk/ChicoryModule.java +++ b/src/main/java/org/extism/chicory/sdk/ChicoryModule.java @@ -29,13 +29,13 @@ static WasmModule fromWasm(ManifestWasm m) { } } - static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts) { + static Instance.Builder instanceWithOptions(Instance.Builder m, Manifest.Options opts, CachedAotMachineFactory aotMachineFactory) { if (opts == null) { return m; } // This feature is not compatibly with the native-image builder. if (opts.aot && !IS_NATIVE_IMAGE_AOT) { - m.withMachineFactory(AotMachine::new); + m.withMachineFactory(aotMachineFactory); } if (!opts.validationFlags.isEmpty()) { throw new UnsupportedOperationException("Validation flags are not supported yet"); diff --git a/src/main/java/org/extism/chicory/sdk/DependencyGraph.java b/src/main/java/org/extism/chicory/sdk/DependencyGraph.java index a84032b..c63336d 100644 --- a/src/main/java/org/extism/chicory/sdk/DependencyGraph.java +++ b/src/main/java/org/extism/chicory/sdk/DependencyGraph.java @@ -1,5 +1,6 @@ package org.extism.chicory.sdk; +import com.dylibso.chicory.experimental.aot.AotMachineFactory; import com.dylibso.chicory.log.Logger; import com.dylibso.chicory.runtime.ExportFunction; import com.dylibso.chicory.runtime.HostFunction; @@ -41,6 +42,7 @@ class DependencyGraph { private final Store store = new Store(); private Manifest.Options options; + private CachedAotMachineFactory aotMachineFactory; public DependencyGraph(Logger logger) { this.logger = logger; @@ -51,6 +53,9 @@ public DependencyGraph(Logger logger) { */ public void setOptions(Manifest.Options options) { this.options = options; + if (options != null && options.aot) { + this.aotMachineFactory = new CachedAotMachineFactory(); + } } /** @@ -93,6 +98,9 @@ private void checkCollision(String moduleName, String symbol) { */ public void registerModule(String name, WasmModule m) { checkCollision(name, null); + if (aotMachineFactory != null) { + aotMachineFactory.compile(m); + } ExportSection exportSection = m.exportSection(); for (int i = 0; i < exportSection.exportCount(); i++) { @@ -223,7 +231,7 @@ private Instance instantiate(String moduleId, List moreHostFunctio Instance instance = ChicoryModule.instanceWithOptions( - Instance.builder(m), this.options) + Instance.builder(m), this.options, aotMachineFactory) .withImportValues(importValues) .withStart(false) .build(); diff --git a/src/main/java/org/extism/chicory/sdk/Kernel.java b/src/main/java/org/extism/chicory/sdk/Kernel.java index 759ba49..cb77774 100644 --- a/src/main/java/org/extism/chicory/sdk/Kernel.java +++ b/src/main/java/org/extism/chicory/sdk/Kernel.java @@ -3,10 +3,13 @@ import com.dylibso.chicory.runtime.ExportFunction; import com.dylibso.chicory.runtime.HostFunction; import com.dylibso.chicory.runtime.Instance; +import com.dylibso.chicory.runtime.Machine; import com.dylibso.chicory.wasm.Parser; +import com.dylibso.chicory.wasm.WasmModule; import com.dylibso.chicory.wasm.types.ValueType; import java.util.List; +import java.util.function.Function; public class Kernel { static final String IMPORT_MODULE_NAME = "extism:host/env"; @@ -33,7 +36,11 @@ public class Kernel { final ExportFunction memoryBytes; public Kernel() { - Instance kernel = instance(); + this(null); + } + + Kernel(CachedAotMachineFactory machineFactory) { + Instance kernel = instance(machineFactory); instanceMemory = kernel.memory(); alloc = kernel.export("alloc"); free = kernel.export("free"); @@ -57,9 +64,13 @@ public Kernel() { memoryBytes = kernel.export("memory_bytes"); } - public static Instance instance() { + private static Instance instance(CachedAotMachineFactory machineFactory) { var kernelStream = Kernel.class.getClassLoader().getResourceAsStream("extism-runtime.wasm"); - return Instance.builder(Parser.parse(kernelStream)).build(); + WasmModule module = Parser.parse(kernelStream); + if (machineFactory != null) { + machineFactory.compile(module); + } + return Instance.builder(module).withMachineFactory(machineFactory).build(); } public void setInput(byte[] input) { diff --git a/src/main/java/org/extism/chicory/sdk/Linker.java b/src/main/java/org/extism/chicory/sdk/Linker.java index b70f4da..8c464aa 100644 --- a/src/main/java/org/extism/chicory/sdk/Linker.java +++ b/src/main/java/org/extism/chicory/sdk/Linker.java @@ -37,17 +37,20 @@ Plugin link() { Map config; WasiOptions wasiOptions; + CachedAotMachineFactory aotMachineFactory; if (manifest.options == null) { config = Map.of(); wasiOptions = null; + aotMachineFactory = null; } else { dg.setOptions(manifest.options); config = manifest.options.config; wasiOptions = manifest.options.wasiOptions; + aotMachineFactory = manifest.options.aot? new CachedAotMachineFactory() : null; } // Register the HostEnv exports. - var hostEnv = new HostEnv(new Kernel(), config, logger); + var hostEnv = new HostEnv(new Kernel(aotMachineFactory), config, logger); dg.registerFunctions(hostEnv.toHostFunctions()); // Register the WASI host functions. @@ -57,7 +60,6 @@ Plugin link() { logger, wasiOptions).toHostFunctions()); } - // Register the user-provided host functions. dg.registerFunctions(Arrays.stream(this.hostFunctions) .map(ExtismHostFunction::asHostFunction) diff --git a/src/test/java/org/extism/chicory/sdk/PluginTest.java b/src/test/java/org/extism/chicory/sdk/PluginTest.java index 958ac3a..ae7a948 100644 --- a/src/test/java/org/extism/chicory/sdk/PluginTest.java +++ b/src/test/java/org/extism/chicory/sdk/PluginTest.java @@ -29,7 +29,7 @@ public void testGreet() { public void testGreetAoT() { var url = "https://github.com/extism/plugins/releases/download/v1.1.1/greet.wasm"; var wasm = ManifestWasm.fromUrl(url).build(); - var manifest = Manifest.ofWasms(wasm).build(); + var manifest = Manifest.ofWasms(wasm).withOptions(new Manifest.Options().withAoT()).build(); var plugin = Plugin.ofManifest(manifest).build(); var input = "Benjamin"; var result = new String(plugin.call("greet", input.getBytes(StandardCharsets.UTF_8)), StandardCharsets.UTF_8);