iree-org
diff --git a/‎README.md‎
Lines changed: 52 additions & 0 deletions b/‎README.md‎
Lines changed: 52 additions & 0 deletions
diff --git a/‎include/fusilli/backend/runtime.h‎
Lines changed: 15 additions & 17 deletions b/‎include/fusilli/backend/runtime.h‎
Lines changed: 15 additions & 17 deletions
diff --git a/‎include/fusilli/graph/graph.h‎
Lines changed: 116 additions & 28 deletions b/‎include/fusilli/graph/graph.h‎
Lines changed: 116 additions & 28 deletions
@@ -30,6 +30,58 @@ software effectively.
 
 ![Fusilli](docs/fusilli.png)
 
+## JIT and AOT support
+
+Fusilli's default API is built for the JIT usecase:
+
+```cpp
+graph.compile(handle);
+graph.execute(handle, variantPack, workspace);
+```
+
+`Graph::compile(handle)` compiles for `handle.getBackend()` and immediately loads
+the resulting artifact into runtime state owned by that `Graph`. The compiled
+artifacts are cached in-process and tied to the lifetime of the `Graph` instance.
+Here `Graph::compile` is just an implementation detail that goes hand-in-hand
+with `Graph::execute` (in the same process and using the same device handle).
+
+For AOT-style callers we recognize that the compilation may happen in a separate
+process, or without access to the specific execution device. To support the AOT
+usecase, Fusilli exposes alternate APIs with explicit artifact save/load steps:
+
+```cpp
+// Compile doesn't require a device handle
+auto vmfbBytes = graph.compileToArtifact(backend);
+// Keep, copy, or serialize the VMFB bytes for later execution.
+
+Graph runtimeGraph;
+// Rebuild and validate the same logical graph.
+runtimeGraph.loadFromArtifact(handle, vmfbBytes);
+runtimeGraph.execute(handle, variantPack, workspace);
+```
+
+Important constraints for multi-backend AOT compilation:
+
+- Compiled artifacts are backend-specific. An artifact produced for one backend
+  must be loaded and executed with a handle for the same backend. `execute()`
+  rejects handles whose backend does not match the currently loaded artifact's
+  backend.
+- `compileToArtifact()` does not load or unload runtime state. If a `Graph`
+  already has an artifact loaded, compiling another artifact leaves the loaded
+  artifact executable until `loadFromArtifact()` replaces it.
+- A `Graph` owns one loaded runtime state at a time. Loading another artifact
+  replaces the previous VM context, function and workspace size.
+- If compiling artifacts for multiple backends from one `Graph`, keep or
+  serialize each returned VMFB byte buffer before compiling the next backend.
+  Fusilli's in-process cache owns only the current compile-side artifact and is
+  not a persistent cache contract.
+- To keep multiple backend artifacts loaded concurrently, use one validated
+  `Graph` instance per backend. To switch backends on the same `Graph`, call
+  `loadFromArtifact(handleForSelectedBackend, artifactForSelectedBackend)`
+  before `execute()`.
+
+AOT samples are under `samples/aot/`, everything else in `samples/` uses the JIT API.
+
 ## Developer Guide
 
 ### Setup
 
@@ -44,7 +44,6 @@
 #include <iree/hal/api.h>
 #include <iree/hal/drivers/hip/api.h>
 #include <iree/hal/drivers/init.h>
-#include <iree/io/file_contents.h>
 #include <iree/modules/hal/module.h>
 #include <iree/vm/api.h>
 #include <iree/vm/bytecode/module.h>
@@ -240,8 +239,7 @@ inline ErrorObject Handle::createAMDGPUDevice(int deviceId, uintptr_t stream) {
 //===----------------------------------------------------------------------===//
 
 // Create IREE VM context for this graph and load the compiled artifact.
-inline ErrorObject Graph::createVmContext(const Handle &handle,
-                                          const std::string &vmfbPath) {
+inline ErrorObject Graph::createVmContext(const Handle &handle) {
   // Create a context even if one was created earlier, since the handle
   // (hence device) might have changed and we might be re-compiling the graph
   // for the new device.
@@ -268,25 +266,16 @@ inline ErrorObject Graph::createVmContext(const Handle &handle,
     FUSILLI_CHECK_ERROR(status);
   }
 
-  // Read the VMFB file and create a bytecode module from it.
+  // Create a bytecode module from the graph-owned VMFB bytes.
   FUSILLI_LOG_LABEL_ENDL("INFO: Loading bytecode module into IREE VM context");
   {
-    iree_io_file_contents_t *fileContents = nullptr;
-    FUSILLI_CHECK_ERROR(iree_io_file_contents_read(
-        iree_make_cstring_view(vmfbPath.c_str()), allocator, &fileContents));
-
     iree_vm_module_t *bytecodeModule = nullptr;
     iree_status_t status = iree_vm_bytecode_module_create(
         handle.getInstance(), IREE_VM_BYTECODE_MODULE_FLAG_NONE,
-        fileContents->const_buffer,
-        iree_io_file_contents_deallocator(fileContents), allocator,
-        &bytecodeModule);
-    if (!iree_status_is_ok(status)) {
-      iree_io_file_contents_free(fileContents);
-      FUSILLI_CHECK_ERROR(status);
-    }
-    // File contents ownership transferred to bytecode module on success
-    // so there's no `iree_io_file_contents_free` on the success path.
+        iree_make_const_byte_span(loadedArtifactBytes_.data(),
+                                  loadedArtifactBytes_.size()),
+        iree_allocator_null(), allocator, &bytecodeModule);
+    FUSILLI_CHECK_ERROR(status);
 
     status =
         iree_vm_context_register_modules(vmContext_.get(),
@@ -395,6 +384,15 @@ Graph::execute(const Handle &handle,
   if (!kBackendExecuteAsync.contains(handle.getBackend())) // C++20
     return ErrorObject(ErrorCode::InternalError,
                        "Graph::execute got an unknown backend");
+  FUSILLI_RETURN_ERROR_IF(!loadedBackend_.has_value(), ErrorCode::NotCompiled,
+                          "Graph::execute requires a successful compile() first"
+                          " (loaded backend not set)");
+  FUSILLI_RETURN_ERROR_IF(handle.getBackend() != *loadedBackend_,
+                          ErrorCode::InvalidArgument,
+                          "Graph::execute got a handle for backend " +
+                              kBackendToStr.at(handle.getBackend()) +
+                              ", but the loaded artifact uses backend " +
+                              kBackendToStr.at(*loadedBackend_));
   bool executeAsync = kBackendExecuteAsync.at(handle.getBackend());
 
   iree_allocator_t allocator = iree_allocator_system();
 
@@ -52,6 +52,7 @@
 #include <memory>
 #include <optional>
 #include <set>
+#include <span>
 #include <sstream>
 #include <string>
 #include <unordered_map>
@@ -111,29 +112,93 @@ class Graph : public INode {
     return ok();
   }
 
-  // Compiles the graph using IREE compiler and sets up the IREE VM
-  // context for future g->execute calls.
+  // Compiles the graph using IREE compiler and sets up the IREE VM context for
+  // future g->execute calls. This is the default JIT convenience API.
+  //
+  // This is equivalent to:
+  //   compileToArtifact(handle.getBackend(), remove)
+  //   loadFromArtifact(handle, vmfbBytes)
+  //
+  // The compiled artifact is backend-specific and is immediately loaded using
+  // the same backend from `handle`.
   //
   // Set `remove = true` to remove compilation artifacts (cache files) when
   // this `Graph` instance goes out of scope.
   ErrorObject compile(const Handle &handle, bool remove = false) {
     FUSILLI_LOG_LABEL_ENDL("INFO: Compiling Graph");
+    FUSILLI_ASSIGN_OR_RETURN(auto vmfbBytes,
+                             compileToArtifact(handle.getBackend(), remove));
+    return loadFromArtifact(handle, vmfbBytes);
+  }
+
+  // Compiles the graph using IREE compiler to produce a backend-specific VMFB
+  // artifact. This does not create any runtime state; call
+  // `loadFromArtifact()` before executing.
+  //
+  // If callers need artifacts for multiple backends, they should keep or
+  // serialize each returned VMFB byte buffer before compiling the next backend.
+  // Fusilli owns only the in-process `cache_` for the most recent compile-side
+  // artifact; it does not provide persistent artifact storage or invalidation.
+  //
+  // Compiling a new artifact does not load or unload runtime state. If this
+  // `Graph` already has an artifact loaded, that artifact remains executable
+  // until a later `loadFromArtifact()` replaces it.
+  //
+  // Set `remove = true` to remove compilation artifacts (cache files) when
+  // this `Graph` instance goes out of scope.
+  ErrorOr<std::vector<uint8_t>> compileToArtifact(Backend backend,
+                                                  bool remove = false) {
+    FUSILLI_LOG_LABEL_ENDL("INFO: Compiling Graph to artifact");
     FUSILLI_RETURN_ERROR_IF(!isValidated_, ErrorCode::NotValidated,
                             "Graph must be validated before being compiled");
 
     // Generate MLIR assembly for this graph.
     FUSILLI_ASSIGN_OR_RETURN(std::string generatedAsm, emitAsm());
 
     // Compile using IREE compiler or reuse cached artifact.
-    FUSILLI_ASSIGN_OR_RETURN(auto vmfbPath,
-                             getCompiledArtifact(handle, generatedAsm, remove));
+    FUSILLI_ASSIGN_OR_RETURN(
+        auto vmfbPath, getCompiledArtifact(backend, generatedAsm, remove));
 
     FUSILLI_LOG_LABEL_ENDL("INFO: Compiled Graph cached at \"" +
                            vmfbPath.string() + "\"");
 
-    // Create per-graph IREE VM context and load the compiled artifact.
-    FUSILLI_CHECK_ERROR(createVmContext(handle, vmfbPath.string()));
+    return readFileBytes(vmfbPath);
+  }
 
+  // Loads a compiled VMFB artifact onto the device owned by `handle`, creating
+  // per-graph runtime state and querying workspace size. This does not require
+  // the artifact to have been produced by this `Graph` instance.
+  //
+  // The artifact must be compatible with `handle.getBackend()`. A VMFB compiled
+  // for one backend is not a portable artifact for another backend.
+  //
+  // A `Graph` owns one loaded runtime state at a time. Loading a new artifact
+  // replaces any previously loaded VM context, function, workspace size, and VM
+  // input list capacity. If loading fails, the `Graph` is left without loaded
+  // runtime state. To keep multiple backends loaded concurrently, use one
+  // `Graph` instance per backend artifact. To switch backends on one `Graph`,
+  // call `loadFromArtifact()` with the selected backend artifact before
+  // `execute()`.
+  ErrorObject loadFromArtifact(const Handle &handle,
+                               std::span<const uint8_t> vmfbBytes) {
+    FUSILLI_LOG_LABEL_ENDL("INFO: Loading compiled artifact into VM context");
+    FUSILLI_RETURN_ERROR_IF(
+        !isValidated_, ErrorCode::NotValidated,
+        "Graph must be validated before loading a compiled artifact");
+    std::vector<uint8_t> artifactBytes(vmfbBytes.begin(), vmfbBytes.end());
+    // Loading replaces the currently executable artifact. Clear first so a
+    // failed load cannot leave execute() using stale runtime state from an
+    // older artifact.
+    clearRuntimeState();
+    loadedArtifactBytes_ = std::move(artifactBytes);
+    ErrorObject status = createVmContext(handle);
+    if (isError(status)) {
+      // createVmContext may have partially populated runtime state before
+      // failing; leave the graph in a clean "not loaded" state.
+      clearRuntimeState();
+      return status;
+    }
+    loadedBackend_ = handle.getBackend();
     return ok();
   }
 
@@ -235,7 +300,7 @@ class Graph : public INode {
                                    std::shared_ptr<Buffer>> &variantPack,
           const std::shared_ptr<Buffer> &workspace) const;
 
-  // Delete copy constructors, keep default move constructor and destructor.
+  // Delete copy constructors and keep default move operations.
   Graph(const Graph &) = delete;
   Graph &operator=(const Graph &) = delete;
   Graph(Graph &&) noexcept = default;
@@ -354,11 +419,11 @@ class Graph : public INode {
   // TODO(#13): Make this private. It is public for now to aid testing and
   // debuggability, however the intended user facing API is `Graph::compile()`.
   ErrorOr<std::filesystem::path>
-  getCompiledArtifact(const Handle &handle, const std::string &generatedAsm,
+  getCompiledArtifact(Backend backend, const std::string &generatedAsm,
                       bool remove, std::optional<bool> *reCompiled = nullptr) {
     // Check for cache hit.
     FUSILLI_ASSIGN_OR_RETURN(bool cacheValid,
-                             validateCache(handle, generatedAsm));
+                             validateCache(backend, generatedAsm));
     if (cacheValid) {
       if (reCompiled)
         *reCompiled = false;
@@ -367,7 +432,7 @@ class Graph : public INode {
     // (Re)generate cache.
     FUSILLI_ASSIGN_OR_RETURN(
         auto generatedCache,
-        generateCompiledArtifact(handle, generatedAsm, remove));
+        generateCompiledArtifact(backend, generatedAsm, remove));
     cache_ = std::move(generatedCache);
     if (reCompiled)
       *reCompiled = true;
@@ -397,8 +462,16 @@ class Graph : public INode {
 
 private:
   // Definition in `fusilli/backend/runtime.h`.
-  ErrorObject createVmContext(const Handle &handle,
-                              const std::string &vmfbPath);
+  ErrorObject createVmContext(const Handle &handle);
+
+  void clearRuntimeState() {
+    vmFunction_.reset();
+    vmContext_.reset();
+    workspaceSize_.reset();
+    loadedBackend_.reset();
+    loadedArtifactBytes_.clear();
+    vmInputListCapacity_ = 0;
+  }
 
   // Queries the required transient/workspace buffer size from the compiled
   // module. Returns the size in bytes, or 0 if no transients are needed.
@@ -410,8 +483,8 @@ class Graph : public INode {
   // `remove = true` to remove cache files when returned `CachedAssets` lifetime
   // ends.
   ErrorOr<CachedAssets>
-  generateCompiledArtifact(const Handle &handle,
-                           const std::string &generatedAsm, bool remove) {
+  generateCompiledArtifact(Backend backend, const std::string &generatedAsm,
+                           bool remove) {
     FUSILLI_LOG_LABEL_ENDL("INFO: Generating compiled artifacts");
 
     // Create cache files.
@@ -448,16 +521,16 @@ class Graph : public INode {
     if (checkCompileBackendEnv()) {
       // Use CompileCommand (CLI).
       CompileCommand cmd = CompileCommand::build(
-          handle.getBackend(), cache.input, cache.output, cache.statistics);
+          backend, cache.input, cache.output, cache.statistics);
       FUSILLI_CHECK_ERROR(cmd.writeTo(cache.command));
       FUSILLI_LOG_LABEL_ENDL("INFO: iree-compile command (CLI)");
       FUSILLI_LOG_ENDL(cmd.toString());
       FUSILLI_CHECK_ERROR(cmd.execute());
     } else {
       // Use CompileSession (C API) - DEFAULT.
       FUSILLI_ASSIGN_OR_RETURN(CompileSession session,
-                               CompileSession::build(handle.getBackend(),
-                                                     cache.input, cache.output,
+                               CompileSession::build(backend, cache.input,
+                                                     cache.output,
                                                      cache.statistics));
       FUSILLI_CHECK_ERROR(session.writeTo(cache.command));
       FUSILLI_LOG_LABEL_ENDL("INFO: iree-compile command (C API)");
@@ -473,8 +546,8 @@ class Graph : public INode {
   //  - Graph name (and therefore cache path) has changed
   //  - Generated assembly differs
   //  - Compile commands have changed
-  //  - Handle/backend (and therefore compile command) has changed
-  ErrorOr<bool> validateCache(const Handle &handle,
+  //  - Backend (and therefore compile command) has changed
+  ErrorOr<bool> validateCache(Backend backend,
                               const std::string &generatedAsm) {
     FUSILLI_LOG_LABEL_ENDL("INFO: Validating cache");
 
@@ -544,13 +617,13 @@ class Graph : public INode {
     if (checkCompileBackendEnv()) {
       // Use CompileCommand (CLI).
       CompileCommand cmd =
-          CompileCommand::build(handle.getBackend(), input, output, statistics);
+          CompileCommand::build(backend, input, output, statistics);
       cmdString = cmd.toString();
     } else {
       // Use CompileSession (C API) - DEFAULT.
-      FUSILLI_ASSIGN_OR_RETURN(auto session,
-                               CompileSession::build(handle.getBackend(), input,
-                                                     output, statistics));
+      FUSILLI_ASSIGN_OR_RETURN(
+          auto session,
+          CompileSession::build(backend, input, output, statistics));
       cmdString = session.toString();
     }
 
@@ -618,10 +691,14 @@ class Graph : public INode {
   // This is set after `validate()` is run at least once successfully.
   bool isValidated_ = false;
 
-  // Required workspace buffer size in bytes. Set during createVmContext()
-  // by querying the iree.abi.transients.size.constant attribute.
-  // std::nullopt indicates the graph has not been compiled yet.
-  std::optional<size_t> workspaceSize_;
+  // Bytes backing the currently loaded VMFB artifact. IREE's bytecode module
+  // may retain references to this archive, so keep it alive for as long as the
+  // VM context is alive.
+  //
+  // Keep this declared before vmContext_: members destruct in reverse
+  // declaration order, and the context must be released before the backing VMFB
+  // bytes.
+  std::vector<uint8_t> loadedArtifactBytes_;
 
   // IREE VM context lifetime managed by the `Graph` object
   // (deleted when the `Graph` object goes out of scope).
@@ -631,11 +708,22 @@ class Graph : public INode {
   // Avoids repeated function lookup on every execute() call.
   std::optional<iree_vm_function_t> vmFunction_;
 
+  // Required workspace buffer size in bytes. Set during createVmContext()
+  // by querying the iree.abi.transients.size.constant attribute.
+  // std::nullopt indicates the graph has not been compiled yet.
+  std::optional<size_t> workspaceSize_;
+
   // Pre-computed VM input list capacity for iree_vm_list_create().
   // Set during createVmContext() to avoid recomputing on every execute().
   iree_host_size_t vmInputListCapacity_ = 0;
 
-  // Cache set by `getCompiledArtifact()`.
+  // Backend for the currently loaded runtime state. `execute()` requires a
+  // handle for this backend because VMFB artifacts are backend-specific.
+  std::optional<Backend> loadedBackend_;
+
+  // Compile-side cache set by `getCompiledArtifact()`. The AOT
+  // `loadFromArtifact()` API uses caller-provided VMFB bytes directly and does
+  // not consult this cache.
   //
   // Note: new instances should always re-generate cache even if the results
   // could be read from the file system. Old results may have been generated