Higher Order Op Output Memory Planning

[gtyukasz]

🐛 Describe the bug

Submodules for consecutive higher order ops are memory planned to a shared memory region. The runtime uses MoveCall to make parent-graph value slots reference output tensors that reside in the shared submodule memory arena. It may cause aliasing if a subsequent submodule is called while such an output is still live.

import executorch.runtime
import torch

from executorch.exir import to_edge, ExecutorchBackendConfig
from executorch.exir.passes import MemoryPlanningPass
from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass
from torch.export import export
from typing import Final, Tuple


class TestModule(torch.nn.Module):
    FLOAT64_NAN: Final = torch.tensor([torch.nan], dtype=torch.float64)

    def __init__(self) -> None:
        super().__init__()
        self.register_buffer('state', self.FLOAT64_NAN.clone())

    def forward(self, data: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        val = data[0].unsqueeze(0)
        pred = (val >= 0.0)

        # First cond: produce a value
        result1 = torch.cond(
            pred,
            lambda v: v.clone(),
            lambda v: self.state.clone(),
            [val]
        )

        # Second cond with NESTED cond, using result1
        result2 = torch.cond(
            pred,
            lambda r: torch.cond(
                torch.isnan(self.state),
                lambda r2: r2.clone(),
                lambda r2: r2 + 1.0,
                [r]
            ),
            lambda r: self.state.clone(),
            [result1]
        )

        self.state.copy_(result1)

        return result1, result2


INPUTS: Final = [
    torch.tensor([5.0], dtype=torch.float64),
    torch.tensor([6.0], dtype=torch.float64),
]

module = TestModule().eval()
forward = export(module, (INPUTS[0],))
edge = to_edge(forward)
et = edge.to_executorch(config=ExecutorchBackendConfig(
    passes=[InitializedMutableBufferPass(list(module._buffers.keys()))],
    memory_planning_pass=MemoryPlanningPass(
        alloc_graph_input=False,
        alloc_graph_output=True,
    )
))

rt_prog = executorch.runtime.Runtime.get().load_program(et.buffer)
rt_forward = rt_prog.load_method('forward')

module = TestModule().eval()
for i in INPUTS:
    result = tuple(rt_forward.execute([i]))
    expect = module.forward(i)
    print(f'\nforward({i})')
    print(f'  result: {result}')
    print(f'  expect: {expect}')
    match = all(torch.allclose(r, e, atol=0, rtol=0, equal_nan=True) for r, e in zip(result, expect))
    print(f'  match: {match}')

Output:

forward(tensor([5.], dtype=torch.float64))
  result: (tensor([5.], dtype=torch.float64), tensor([5.], dtype=torch.float64))
  expect: (tensor([5.], dtype=torch.float64), tensor([5.], dtype=torch.float64))
  match: True

forward(tensor([6.], dtype=torch.float64))
  result: (tensor([7.], dtype=torch.float64), tensor([7.], dtype=torch.float64))
  expect: (tensor([6.], dtype=torch.float64), tensor([7.], dtype=torch.float64))
  match: False

Versions

PyTorch version: 2.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 26.6 (arm64)
GCC version: Could not collect
Clang version: 21.0.0 (clang-2100.1.1.101)
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.13 (main, Mar  3 2026, 12:39:30) [Clang 21.0.0 (clang-2100.0.123.102)] (64-bit runtime)
Python platform: macOS-26.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Apple M4 Pro

Versions of relevant libraries:
[pip3] executorch==1.3.1
[pip3] numpy==2.5.1
[pip3] pytorch_tokenizers==1.3.0
[pip3] torch==2.12.1
[pip3] torchao==0.17.0
[conda] Could not collect

cc @JacobSzwejbka @angelayi

[SS-JIA]

@JacobSzwejbka assigning to you, as the one I think is most knowledgable about HOPs at the moment. Feel free to reassign!