up

Author: metascroy

backends/mlx/builder/program_builder.py

@@ -27,7 +27,6 @@
 from typing import Any, DefaultDict, Dict, List, Optional, Set, Tuple, Union
 
 import torch
-
 from executorch.backends.mlx._logging import logger
 from executorch.backends.mlx.builder.op_helpers import torch_dtype_to_scalar_type
 from executorch.backends.mlx.builder.op_registry import (
@@ -132,7 +131,9 @@ class MLXProgramBuilder:
 
     def __init__(self, ep: ExportedProgram, named_data_key_prefix: str = ""):
         self.ep: ExportedProgram = ep
-        self._instrs: List[Instruction] = []
+        self._chains: List[List[Instruction]] = [[]]  # chain 0 = main
+        self._current_chain: int = 0
+        self.init_chain_idx: int = -1
         self.extra_constants: Dict[str, torch.Tensor] = {}
         self.slot_manager = SlotManager()
         self.node_info: DefaultDict[Node, NodeInfo] = defaultdict(NodeInfo)
@@ -163,7 +164,13 @@ def _prefix_key(self, name: str) -> str:
         return name
 
     def emit(self, op: OpNodeUnion) -> None:
-        self._instrs.append(Instruction(op=op))
+        self._chains[self._current_chain].append(Instruction(op=op))
+
+    def emit_init(self, op: OpNodeUnion) -> None:
+        if self.init_chain_idx == -1:
+            self.init_chain_idx = len(self._chains)
+            self._chains.append([])
+        self._chains[self.init_chain_idx].append(Instruction(op=op))
 
     def args(self, node: Node) -> Tuple[Any, ...]:
         return self.slot_map(node.args)
@@ -934,9 +941,11 @@ def _build_mlx_graph(self) -> MLXGraph:
             num_mutable_buffer_tensors=num_tensors[IdSpace.MutableBuffer],
             num_temp_tensors=num_temp_tensors,
             num_values=num_values_count,
-            instruction_chains=[InstructionChain(instructions=self._instrs)],
+            instruction_chains=[
+                InstructionChain(instructions=chain) for chain in self._chains
+            ],
             main_chain_idx=0,
-            init_chain_idx=-1,
+            init_chain_idx=self.init_chain_idx,
             input_map=input_map,
             output_map=output_map,
             mutable_buffer_map=mutable_buffer_map,

backends/mlx/runtime/MLXBackend.cpp

@@ -219,10 +219,24 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
           static_cast<const uint8_t*>(processed->data()), processed->size());
 
       // Validate schema version
-      if (handle->program.version != "1") {
+      int schema_version = 1;
+      if (!handle->program.version.empty()) {
+        try {
+          schema_version = std::stoi(handle->program.version);
+        } catch (...) {
+          throw std::runtime_error(
+              "Invalid MLX schema version '" + handle->program.version +
+              "' (expected integer)");
+        }
+      }
+      constexpr int kMaxSupportedVersion = 1;
+      if (schema_version > kMaxSupportedVersion) {
         throw std::runtime_error(
-            "Unsupported MLX schema version '" + handle->program.version +
-            "' (expected '1'). Rebuild the .pte with a matching SDK version.");
+            "This .pte requires ExecuTorch MLX runtime version " +
+            std::to_string(schema_version) +
+            " but this runtime only supports up to version " +
+            std::to_string(kMaxSupportedVersion) +
+            ". Upgrade ExecuTorch to a newer version.");
       }
 
       // Load constants from named_data_map
@@ -251,11 +265,17 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       // SAFETY: The >= 0 check ensures init_chain_idx is non-negative, so the
       // static_cast<uint32_t> cannot produce UINT32_MAX from a -1 sentinel.
       if (handle->program.init_chain_idx >= 0) {
+        handle->state.is_init_chain = true;
         handle->interpreter.run_chain(
             handle->program,
             static_cast<uint32_t>(handle->program.init_chain_idx),
             handle->state,
             handle->stream);
+        handle->state.is_init_chain = false;
+
+        // Evaluate any constants written by the init chain so the first
+        // execute() doesn't pay the cost of materializing them.
+        eval(handle->constants.tensors);
       }
 
     } catch (const std::exception& e) {

backends/mlx/runtime/MLXExecutor.h

@@ -97,6 +97,13 @@ struct ConstantData {
     return tensors[id.idx];
   }
 
+  inline void set(Tid id, Tensor t) {
+    if (id.idx >= tensors.size()) {
+      throw std::out_of_range("ConstantData::set: id out of range");
+    }
+    tensors[id.idx] = std::move(t);
+  }
+
   inline void add(Tensor t) {
     tensors.push_back(std::move(t));
   }
@@ -153,6 +160,9 @@ struct ExecutionState {
   // Non-constant values (SymInt, etc.)
   std::vector<std::optional<Value>> values;
 
+  // Init chain flag: when true, set_tensor allows writing to constants
+  bool is_init_chain{false};
+
   // Logging context
   size_t current_op_idx{0};
   const char* current_op_name{nullptr};
@@ -478,7 +488,15 @@ struct ExecutionState {
       throw std::runtime_error("set_tensor: Program not bound");
     }
     if (id.idx < program->num_constant_tensors) {
-      throw std::runtime_error("set_tensor: cannot write to constant tensor");
+      if (!is_init_chain) {
+        throw std::runtime_error("set_tensor: cannot write to constant tensor");
+      }
+      // Init chain can write over constants
+      if (!constants) {
+        throw std::runtime_error("set_tensor: constants not bound");
+      }
+      const_cast<ConstantData*>(constants)->set(id, std::move(arr));
+      return;
     }
     // Route to mutable buffers or per-execution tensors
     if (is_mutable_buffer(id)) {

backends/mlx/test/test_utils.py

@@ -23,6 +23,7 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
+import executorch.exir as exir
 import numpy as np
 import torch
 
@@ -268,6 +269,7 @@ def export_model_to_pte(
     output_path: Union[str, Path],
     dynamic_shapes: Optional[Dict] = None,
     verbose: bool = False,
+    edge_compile_config: Optional[exir.EdgeCompileConfig] = None,
 ) -> None:
     """
     Export a PyTorch model to a .pte file using the MLX delegate.
@@ -281,7 +283,6 @@ def export_model_to_pte(
             Example: {0: {0: Dim("batch", min=1, max=32)}} for dynamic batch on first input.
         verbose: Whether to print the exported program for debugging.
     """
-    import executorch.exir as exir
     from executorch.backends.mlx import MLXPartitioner
     from executorch.exir.capture._config import ExecutorchBackendConfig
     from torch.export import export
@@ -301,9 +302,11 @@ def export_model_to_pte(
         print(exported_program)
 
     # Lower to edge and delegate to MLX
+    compile_config = edge_compile_config or exir.EdgeCompileConfig()
     edge_program = exir.to_edge_transform_and_lower(
         exported_program,
         partitioner=[MLXPartitioner()],
+        compile_config=compile_config,
     )
 
     # Print edge program if verbose
@@ -865,6 +868,10 @@ def get_dynamic_shapes(self) -> Optional[Dict]:
         """Return dynamic shapes specification for torch.export, or None for static shapes."""
         return None
 
+    def get_edge_compile_config(self) -> Optional[exir.EdgeCompileConfig]:
+        """Return EdgeCompileConfig for export, or None for default."""
+        return None
+
     def get_test_dir(self) -> Path:
         """Get the directory for this test's files."""
         test_dir = Path(__file__).parent / "op_tests" / self.name
@@ -924,6 +931,7 @@ def generate_test_files(self, verbose: bool = False) -> Tuple[Path, Path, Path]:
             pte_path,
             dynamic_shapes=dynamic_shapes,
             verbose=verbose,
+            edge_compile_config=self.get_edge_compile_config(),
         )
 
         # Save test inputs