Shared memory for multiple entry points with XNNPACK delegate

[Olgacmt]

🚀 The feature, motivation and pitch

The feature, motivation and pitch

We have a set of production models written in PyTorch and exported to mobile using TorchScript with the XNNPACK delegate. These models:

• Have multiple entry points (e.g., forward, decode, infer)
• Share internal state (e.g., buffers like self.hidden)
• Use custom operators
• Are expected to share memory between entry points at runtime

With Executorch 0.6.0, it is currently not possible to export these models cleanly with support for shared state and multiple entry points when using the XNNPACK delegate.

✅ What currently works (badly):

We can sort-of simulate shared memory by:

• Lifting state tensors out of the module using memory planning,
• Giving them fixed buffer IDs,
• And manually ensuring the runtime binds those buffers to the same address.

This is error-prone, brittle, and not user-friendly at all. It works, but it’s not scalable or robust — especially with models that have delegate compatibility requirements or custom ops.

💣 What’s broken:

• No first-class support for multiple entry points during export.
• No clean UX for declaring shared memory/state across entry points.
• Delegates like XNNPACK or CoreML can’t currently share tensor state across entry points without hacks.
• Overriding forward() as a workaround for multi-entry export is extremely fragile. If the aliased function calls forward internally, it breaks in dumb ways.

🧠 Why this matters:

This breaks real deployment scenarios. Without proper shared state handling:

• Delegates reallocate or reset state per entry point
• You can’t maintain persistent buffers across calls (e.g., RNN hidden states)
• Exported models can't reuse memory, killing performance on mobile
• You can’t export realistic multi-stage or streaming inference pipelines

💡 What we propose:

At minimum:

• Provide a default memory planner utility that allows you to declare shared buffers by name (e.g. "hidden_state", "cache"), and automatically lifts + aligns them to shared IDs
• Longer term: Add first-class multi-entry point support in export, and expose shared buffer semantics to delegates in a clean, non-hacky way
• Explore delegate-level shared context support so operators that span entry points (e.g., XNNPACK-conv using the same weights) don’t get re-initialized unnecessarily

This is currently being actively investigated on the delegate side, but the lack of upstream support in export/memory planning makes it hard to build consistent tooling around this.

Alternatives

No response

Additional context

Existing work and conversations regarding these issues:
#9012
#8030
#10144
#8870
#7458

RFC (Optional)

No response

cc @digantdesai @mcr229 @cbilgin

[Olgacmt]

Example of the current code that does not work but would be nice to have it operational:

import copy
import inspect
import logging
import os
import sys
from contextlib import contextmanager
from typing import Any, Dict, Final, List, Optional, Tuple, Union
import torch
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import (
    EdgeCompileConfig,
    EdgeProgramManager,
    ExecutorchProgramManager,
    ExportedProgram,
    control_flow,
    to_edge,
    to_edge_transform_and_lower,
)
from executorch.exir.capture._config import ExecutorchBackendConfig
from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass
from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass
from executorch.runtime import Runtime
from torch._dynamo import explain


CUSTOM_OPS_LIB = os.getenv(
    "CUSTOM_OPS_LIB",
    "/shared_mem_test_etx/cmake-build/libmyops_ops_aot.dylib",
)
print(f"Loading custom ops library {CUSTOM_OPS_LIB}")
custom_ops_lib = torch.ops.load_library(CUSTOM_OPS_LIB)

NATIVE_LIB_PATH = os.getenv(
    "NATIVE_LIB_PATH",
    "/shared_mem_test_etx/cmake-install",
)
print(f"Adding path to native lib to a Python path {NATIVE_LIB_PATH}")
sys.path.insert(0, NATIVE_LIB_PATH)

# Now import from the local my_lib module
from my_lib._native import create_context_parameters, create_module_context

# torch._logging.set_logs(dynamo=logging.DEBUG)
torch._dynamo.config.verbose = True
torch._dynamo.config.capture_dynamic_output_shape_ops = True
torch._dynamo.config.capture_scalar_outputs = True
torch._dynamo.config.dynamic_shapes = True


class TestTorchModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._ctx = 0

    def _init(
        self,
        ctx: int,
        params_ptr: int = 0,
    ):
        if ctx:
            self._ctx = ctx
            if params_ptr:
                self.init(torch.ops.my_ops.get_dict(params_ptr))
        else:
            self.flush()

    def init(self, params_ptr: int):
        """Override to initialize/reset internal state and extract config parameters into module variables."""
        pass

    def algo_version(self) -> str:
        return torch.ops.my_ops.algo_version(self._ctx)

    def get_timeseries_window(
        self,
        name: str,
        start_time: torch.Tensor,
        end_time: torch.Tensor,
        transposed: bool = False,
    ) -> torch.Tensor:
        return torch.ops.my_ops.get_timeseries_window(
            self._ctx, name, start_time, end_time, not transposed
        )


class TestModel(TestTorchModule):
    speed_bands: List[float]
    thresholds: List[float]

    def __init__(
        self,
        max_batch_size: int,
        max_seq_len: int,
    ):
        super().__init__()
        context = create_module_context()
        super()._init(context)
        # State tracking variables - just use 1D tensors for simplicity
        self.register_buffer(
            "params",
            torch.tensor(create_module_context(), dtype=torch.int64),
            persistent=True,
        )
        self.register_buffer(
            "prev_time",
            torch.tensor(float("nan"), dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "prev_speed",
            torch.tensor(float("nan"), dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "prediction_count",
            torch.tensor(float("nan"), dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "cache",
            torch.zeros((3, 3), dtype=torch.float32),
            persistent=True,
        )

    def init(self, params: torch.Tensor) -> torch.Tensor:
        self.params.copy_(params)
        self.prev_time.copy_(torch.tensor(5.0))
        self.prev_speed.copy_(torch.tensor(1.0))
        return self.params.clone()

    def set_cache(self, data: torch.Tensor) -> None:
        self.cache.copy_(data)

    # need narrow here:
    def get_cache(self) -> torch.Tensor:
        return self.cache.clone()

    # fmt: off
    def forward(self, curr_time: torch.Tensor, curr_speed: torch.Tensor) -> torch.Tensor:
        # result = self._compute(curr_time, curr_speed)
        prev_time = self.prev_time
        prev_speed = self.prev_speed
        self.prev_time.copy_(curr_time)
        self.prev_speed.copy_(curr_speed)
        print(self.algo_version())
        gps = self.get_timeseries_window("gps", prev_time, curr_time, True)
        self.prev_time.copy_(curr_time)
        self.prev_speed.copy_(curr_speed)
        return self.params.clone()

    def init_example_inputs(self) -> Dict[str, dict[str, torch.Tensor] | None]:
        return {
            "inputs": {"params": self.params},
            "dynamic_shapes": {"params": None},
        }

    def forward_example_inputs(self) -> Dict[str, dict[str, torch.Tensor] | None]:
        return {
            "inputs": {
                "curr_time": torch.tensor(10.0, dtype=torch.float32),
                "curr_speed": torch.tensor(2.0, dtype=torch.float32),
            },
            "dynamic_shapes": {"curr_time": None, "curr_speed": None},
        }

    def set_cache_example_inputs(self) -> Dict[str, dict[str, torch.Tensor] | None]:
        data = torch.zeros((3,3), dtype=torch.float32, device="cpu")
        return {
            "inputs": {"data": data},
            "dynamic_shapes": {"data": None},
        }

    def get_cache_example_inputs(self) -> dict[str, dict[str, int]]:
        return {
            "inputs": {},
            "dynamic_shapes": None,
        }

    def get_shared_buffer_names(self) -> List[str]:
        return ["params", "prev_time", "prev_speed", "prediction_count", "cache"]

    def get_entrypoint_names(self) -> List[str]:
        return [
            "init",
            "forward",
            "set_cache",
            "get_cache",
        ]

    def get_shared_buffer_memory_plan(self) -> Dict[str, torch.Tensor]:
        return {
            "params": self.params,
            "prev_time": self.prev_time,
            "prev_speed": self.prev_speed,
            "prediction_count": self.prediction_count,
            "cache": self.cache,
        }


# fmt: on


@contextmanager
def patch_forward(obj: torch.nn.Module, new_method):
    """
    Patches the forward method of a PyTorch module, needs to patch the class not
    just the instance b/c otherwise export will error out when trying to produce
    items that reference self.
    """
    original_method = obj.__class__.forward
    obj.__class__.forward = new_method
    try:
        yield
    finally:
        # Restore the original method
        obj.__class__.forward = original_method


def main():
    max_batch_size = 10
    max_seq_len = 20

    module = TestModel(max_batch_size, max_seq_len)
    module.eval()

    shared_context_ptr = torch.tensor(create_context_parameters())
    module.init(shared_context_ptr)
    # import ipdb; ipdb.set_trace()
    for entry_point in module.get_entrypoint_names():
        example_inputs = getattr(module, f"{entry_point}_example_inputs")()["inputs"]
        result = getattr(module, entry_point)(**example_inputs)
        print(f"Model output: {result}. Example inputs for {entry_point}: {example_inputs}")

    explanation = explain(module)(**module.forward_example_inputs()["inputs"])
    print(explanation)

    method_graphs: dict[str, ExportedProgram] = {}
    with torch.no_grad():
        for i, entry_point in enumerate(module.get_entrypoint_names()):
            method = getattr(module, entry_point)
            with patch_forward(module, method):
                sig = inspect.signature(getattr(module, entry_point))
                param_names = [param.name for param in sig.parameters.values()]
                print(param_names)
                example_args = {}
                dynamic_shapes = {}

                for param in param_names:
                    print(f"Parameter: {param}")
                    example_args[param] = getattr(module, f"{entry_point}_example_inputs")()[
                        "inputs"
                    ][param]
                    dynamic_shapes[param] = getattr(module, f"{entry_point}_example_inputs")()[
                        "dynamic_shapes"
                    ][param]
                print("Example args:", example_args)
                print("Dynamic shapes:", dynamic_shapes)
                method_graph: torch.nn.Module = torch.export.export_for_training(
                    module,
                    (),
                    kwargs=example_args,
                    dynamic_shapes=dynamic_shapes,
                    strict=False,
                ).module()
                method_graph: ExportedProgram = torch.export.export(
                    method_graph,
                    (),
                    kwargs=example_args,
                    dynamic_shapes=dynamic_shapes,
                    strict=False,
                )
                print("--------------------------------")
                print(f"Method Graph: {method}")
                print("Exported program graph (shows original operations):")
                method_graph.graph.print_tabular()
                method_graphs[method.__name__] = method_graph
                torch.export.save(method_graph, f"test_model_aten_graph_{method.__name__}.pte")
            for key, val in method_graphs.items():
                print("  ------------------------------")
                print(f"  method: {key}")
                print("  " + str(val.graph))

    print("Partition/lower to Edge IR for XNNPACK")
    xnnpack_program: EdgeProgramManager = to_edge_transform_and_lower(
        method_graphs,
        compile_config=EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=True),
        partitioner=[XnnpackPartitioner()],
    )
    for entry_point in module.get_entrypoint_names():
        xnnpack_program._edge_programs[entry_point].graph.print_tabular()

    passes = [
        InitializedMutableBufferPass(module.get_shared_buffer_names()),
    ]
    memory_pass = MemoryPlanningPass(
        allow_lifetime_and_storage_overlap=True,
        alloc_graph_input=True,
        alloc_graph_output=True,
    )
    et_config = ExecutorchBackendConfig(
        passes=passes,
        memory_planning_pass=memory_pass,
    )
    executorch_program: ExecutorchProgramManager = xnnpack_program.to_executorch(config=et_config)
    with open("test_model_xnnpack_program.pte", "wb") as file:
        executorch_program.write_to_file(file)

    runtime = Runtime.get()
    program = runtime.load_program(executorch_program.buffer)
    for method_name in program.method_names:
        _ = program.load_method(method_name)

    print("done")


if __name__ == "__main__":
    main()