Introduce CachedAotMachineFactory.

Signed-off-by: Edoardo Vacchi <[email protected]>
extism · Dec 12, 2024 · 3ad144e · 3ad144e
1 parent 4c1d4b6
commit 3ad144e
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 9 deletions.
diff --git a/src/main/java/org/extism/chicory/sdk/CachedAotMachineFactory.java 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<Instance, Machine> {
+    private final Map<WasmModule, AotMachineFactory> 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
@@ -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
@@ -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<HostFunction> 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
@@ -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
@@ -37,17 +37,20 @@ Plugin link() {
 
         Map<String, String> 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
@@ -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);