From 5edf6e341b5bc7266ec34c726933ec8ceaae1b4a Mon Sep 17 00:00:00 2001 From: Erik Lundell Date: Wed, 30 Sep 2026 07:39:54 +0000 Subject: [PATCH 1/6] BundleIO: Match bundled tensors to program I/O layouts Serialize bundled inputs and expected outputs in the dimension order declared by the program. This preserves channels-last values and avoids layout mismatches when eager test tensors have permuted strides. Reconstruct ATen tensors using the serialized dimension order. Keep layout conversion within BundledProgram instead of adding a preserve_layout flag to TensorSpec.from_tensor. Cover channels-last tensors, offset and non-dense views, and permuted outputs, including portable runtime verification. Correct the dim_order_from_stride return annotation and update Buck dependencies. Assisted by OpenAI Codex. Co-authored-by: Zingo Andersen Signed-off-by: Erik Lundell Change-Id: I4d05a144aafbb0fbd25bfbbd161c4eebc8dc2b0f --- devtools/bundled_program/BUCK | 6 + devtools/bundled_program/bundled_program.cpp | 22 ++- devtools/bundled_program/core.py | 67 ++++----- devtools/bundled_program/test/BUCK | 7 + .../bundled_program/test/test_bundle_data.py | 134 +++++++++++++++++- exir/tensor.py | 14 +- extension/pybindings/test/BUCK | 14 ++ extension/pybindings/test/test_pybindings.py | 67 +++++++++ 8 files changed, 287 insertions(+), 44 deletions(-) diff --git a/devtools/bundled_program/BUCK b/devtools/bundled_program/BUCK index 06951d344fb..1a509c9a994 100644 --- a/devtools/bundled_program/BUCK +++ b/devtools/bundled_program/BUCK @@ -1,3 +1,8 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load(":targets.bzl", "define_common_targets") @@ -20,6 +25,7 @@ fbcode_target( "//executorch/devtools/bundled_program/schema:bundled_program_schema_py", "//executorch/exir:schema", "//executorch/exir:tensor", + "//executorch/exir/_serialize:lib", ], ) diff --git a/devtools/bundled_program/bundled_program.cpp b/devtools/bundled_program/bundled_program.cpp index 7fe8bfd5f5d..24028d25ff8 100644 --- a/devtools/bundled_program/bundled_program.cpp +++ b/devtools/bundled_program/bundled_program.cpp @@ -1,7 +1,7 @@ /* * Copyright (c) Meta Platforms, Inc. and affiliates. * All rights reserved. - * Copyright 2025 Arm Limited and/or its affiliates. + * Copyright 2025-2026 Arm Limited and/or its affiliates. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. @@ -46,13 +46,25 @@ constexpr size_t kMaxDim = 16; at::Tensor tensor_like(bundled_program_flatbuffer::Tensor* bundled_tensor) { ET_CHECK(bundled_tensor->sizes()->size() <= kMaxDim); int64_t ret_t_sizes[kMaxDim]; + int64_t ret_t_strides[kMaxDim]; + const size_t dim = bundled_tensor->sizes()->size(); - for (size_t i = 0; i < bundled_tensor->sizes()->size(); i++) { + ET_CHECK(bundled_tensor->dim_order() != nullptr); + ET_CHECK(bundled_tensor->dim_order()->size() == dim); + + for (size_t i = 0; i < dim; i++) { ret_t_sizes[i] = static_cast(bundled_tensor->sizes()->data()[i]); } - - at::Tensor ret_tensor = at::zeros( - {ret_t_sizes, bundled_tensor->sizes()->size()}, + ET_CHECK( + executorch::runtime::dim_order_to_stride( + ret_t_sizes, + bundled_tensor->dim_order()->data(), + dim, + ret_t_strides) == Error::Ok); + + at::Tensor ret_tensor = at::empty_strided( + {ret_t_sizes, dim}, + {ret_t_strides, dim}, at::dtype(static_cast(bundled_tensor->scalar_type()))); // Validate data buffer exists and has sufficient size diff --git a/devtools/bundled_program/core.py b/devtools/bundled_program/core.py index a72b7fa962f..00375a6f60f 100644 --- a/devtools/bundled_program/core.py +++ b/devtools/bundled_program/core.py @@ -1,11 +1,11 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import ctypes -import typing from typing import Dict, List, Optional, Sequence, Type, Union import executorch.devtools.bundled_program.schema as bp_schema @@ -19,7 +19,8 @@ from executorch.devtools.bundled_program.version import BUNDLED_PROGRAM_SCHEMA_VERSION from executorch.exir import ExecutorchProgram, ExecutorchProgramManager -from executorch.exir.tensor import get_scalar_type, scalar_type_enum, TensorSpec +from executorch.exir._serialize import _deserialize_pte_binary +from executorch.exir.tensor import get_scalar_type, scalar_type_enum # pyre-ignore supported_program_type_table: Dict[Type[core_schema.KernelTypes], ConfigValue] = { @@ -91,10 +92,21 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: if self._bundled_program_in_schema is not None: return self._bundled_program_in_schema + if self.executorch_program: + program = self._extract_program(self.executorch_program) + program_bytes = self.executorch_program.buffer + else: + assert self._pte_file_path is not None + with open(self._pte_file_path, "rb") as f: + program_bytes = f.read() + program = _deserialize_pte_binary(program_bytes).program + plans = {plan.name: plan for plan in program.execution_plan} + bundled_method_test_suites: List[bp_schema.BundledMethodTestSuite] = [] # Emit data and metadata of bundled tensor for method_test_suite in self.method_test_suites: + plan = plans[method_test_suite.method_name] bundled_test_cases: List[bp_schema.BundledMethodTestCase] = [] # emit I/O sets for each method test case @@ -107,10 +119,11 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: i ].expected_outputs - for input_val in cur_plan_test_inputs: + for input_idx, input_val in enumerate(cur_plan_test_inputs): if type(input_val) is torch.Tensor: self._emit_bundled_tensor( - TensorSpec.from_tensor(input_val, const=True), + input_val, + plan.values[plan.inputs[input_idx]].val.dim_order, inputs, ) else: @@ -118,12 +131,15 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: input_val, inputs, ) - for expected_output_tensor in cur_plan_test_expected_outputs: + for output_idx, expected_output_tensor in enumerate( + cur_plan_test_expected_outputs + ): assert ( type(expected_output_tensor) is torch.Tensor ), "Only tensor outputs are currently supported." self._emit_bundled_tensor( - TensorSpec.from_tensor(expected_output_tensor, const=True), + expected_output_tensor, + plan.values[plan.outputs[output_idx]].val.dim_order, expected_outputs, ) bundled_test_cases.append( @@ -140,13 +156,6 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: ) ) - if self.executorch_program: - program_bytes = self.executorch_program.buffer - else: - assert self._pte_file_path is not None - with open(self._pte_file_path, "rb") as f: - program_bytes = f.read() - self._bundled_program_in_schema = bp_schema.BundledProgram( version=BUNDLED_PROGRAM_SCHEMA_VERSION, method_test_suites=bundled_method_test_suites, @@ -155,30 +164,24 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: return self._bundled_program_in_schema def _emit_bundled_tensor( - self, spec: TensorSpec, bundled_values: List[bp_schema.Value] + self, + tensor: torch.Tensor, + dim_order: List[int], + bundled_values: List[bp_schema.Value], ) -> None: - # QuantizedSchema in tensor has deprecated and may not be used anymore. - # So here we don't emit it. - - if spec.allocated_memory == 0: - tensor_data: bytes = b"" - else: - array_type = ( - ctypes.c_char * typing.cast(torch.UntypedStorage, spec.storage).nbytes() - ) - spec_array = ctypes.cast( - typing.cast(torch.UntypedStorage, spec.storage).data_ptr(), - ctypes.POINTER(array_type), - ).contents - tensor_data: bytes = bytes(spec_array) - + sizes = list(tensor.shape) + # Store bytes in the physical order declared by the method's I/O tensor. + tensor = tensor.permute(dim_order).contiguous() + tensor_data = ( + ctypes.string_at(tensor.data_ptr(), tensor.nbytes) if tensor.nbytes else b"" + ) bundled_values.append( bp_schema.Value( val=bp_schema.Tensor( - scalar_type=scalar_type_enum(spec.dtype), - sizes=spec.shape, + scalar_type=scalar_type_enum(tensor.dtype), + sizes=sizes, data=tensor_data, - dim_order=list(spec.dim_order), + dim_order=list(dim_order), ), ) ) diff --git a/devtools/bundled_program/test/BUCK b/devtools/bundled_program/test/BUCK index 7809f382bb5..a12b43e4ea2 100644 --- a/devtools/bundled_program/test/BUCK +++ b/devtools/bundled_program/test/BUCK @@ -1,3 +1,8 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") oncall("executorch") # add this empty BUCK file to unblock landing. Without this, we get land error: @@ -22,6 +27,8 @@ fbcode_target(_kind = python_unittest, "//executorch/devtools/bundled_program:core", "//executorch/devtools/bundled_program/schema:bundled_program_schema_py", "//executorch/devtools/bundled_program/util:test_util", + "//executorch/exir:lib", + "//executorch/exir:tensor", "//executorch/exir/_serialize:lib", ], ) diff --git a/devtools/bundled_program/test/test_bundle_data.py b/devtools/bundled_program/test/test_bundle_data.py index 9fdeb4a776d..eb5b9efd21b 100644 --- a/devtools/bundled_program/test/test_bundle_data.py +++ b/devtools/bundled_program/test/test_bundle_data.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -13,15 +14,146 @@ import executorch.devtools.bundled_program.schema as bp_schema import torch -from executorch.devtools.bundled_program.config import ConfigValue +from executorch.devtools.bundled_program.config import ( + ConfigValue, + MethodTestCase, + MethodTestSuite, +) from executorch.devtools.bundled_program.core import BundledProgram from executorch.devtools.bundled_program.util.test_util import ( get_common_executorch_program, ) +from executorch.exir import to_edge from executorch.exir._serialize import _PTEFile, _serialize_pte_binary +from executorch.exir.tensor import stride_from_dim_order +from torch.export import export class TestBundle(unittest.TestCase): + def _serialize_tensor_test_case( + self, input_tensor: torch.Tensor, output_tensor: torch.Tensor + ) -> bp_schema.BundledMethodTestCase: + class Model(torch.nn.Module): + def forward(self, x): + return x * 2 + + program = to_edge(export(Model(), (input_tensor,))).to_executorch() + with tempfile.TemporaryDirectory() as tmp_dir: + pte_path = f"{tmp_dir}/model.pte" + with open(pte_path, "wb") as pte_file: + pte_file.write(program.buffer) + bundled_program = BundledProgram( + executorch_program=None, + method_test_suites=[ + MethodTestSuite( + "forward", [MethodTestCase((input_tensor,), (output_tensor,))] + ) + ], + pte_file_path=pte_path, + ) + return ( + bundled_program.serialize_to_schema() + .method_test_suites[0] + .test_cases[0] + ) + + def test_channels_last_input_and_output(self) -> None: + input_tensor = ( + torch.arange(2 * 3 * 4 * 5, dtype=torch.float32) + .reshape(2, 3, 4, 5) + .contiguous(memory_format=torch.channels_last) + ) + output_tensor = (input_tensor * 2).contiguous(memory_format=torch.channels_last) + test_case = self._serialize_tensor_test_case(input_tensor, output_tensor) + + for bundled_value, source in zip( + test_case.inputs + test_case.expected_outputs, + (input_tensor, output_tensor), + ): + bundled_tensor = bundled_value.val + assert isinstance(bundled_tensor, bp_schema.Tensor) + self.assertEqual(bundled_tensor.dim_order, [0, 2, 3, 1]) + self.assertEqual(bundled_tensor.sizes, list(source.shape)) + self.assertEqual( + bundled_tensor.data, + source.permute(0, 2, 3, 1).contiguous().numpy().tobytes(), + ) + + def test_channels_last_view_with_storage_offset(self) -> None: + base = ( + torch.arange(3 * 3 * 4 * 5, dtype=torch.float32) + .reshape(3, 3, 4, 5) + .contiguous(memory_format=torch.channels_last) + ) + view = base[1:3] + self.assertGreater(view.storage_offset(), 0) + test_case = self._serialize_tensor_test_case(view, view) + + for bundled_value in test_case.inputs + test_case.expected_outputs: + bundled_tensor = bundled_value.val + assert isinstance(bundled_tensor, bp_schema.Tensor) + self.assertEqual(bundled_tensor.dim_order, [0, 2, 3, 1]) + self.assertEqual(len(bundled_tensor.data), view.nbytes) + self.assertEqual( + bundled_tensor.data, + view.permute(0, 2, 3, 1).contiguous().numpy().tobytes(), + ) + + def test_non_dense_view_uses_only_its_own_elements(self) -> None: + base = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5) + view = base[:, :, ::2, :] + test_case = self._serialize_tensor_test_case(view, view) + + for bundled_value in test_case.inputs + test_case.expected_outputs: + bundled_tensor = bundled_value.val + assert isinstance(bundled_tensor, bp_schema.Tensor) + self.assertEqual(len(bundled_tensor.data), view.nbytes) + decoded = torch.frombuffer( + bytearray(bundled_tensor.data), dtype=view.dtype + ).as_strided( + bundled_tensor.sizes, + stride_from_dim_order(bundled_tensor.sizes, bundled_tensor.dim_order), + ) + self.assertTrue(torch.equal(decoded, view)) + + def test_bundle_uses_program_dim_order(self) -> None: + class Permute(torch.nn.Module): + def forward(self, x): + return x.permute(0, 2, 3, 1) + + model = Permute() + export_input = torch.zeros(2, 3, 4, 5) + program = to_edge(export(model, (export_input,))).to_executorch() + input_tensor = torch.arange(120, dtype=torch.float32).reshape(2, 4, 3, 5) + input_tensor = input_tensor.transpose(1, 2) + expected = model(input_tensor) + self.assertFalse(input_tensor.is_contiguous()) + self.assertFalse(expected.is_contiguous()) + suites = [ + MethodTestSuite("forward", [MethodTestCase((input_tensor,), (expected,))]) + ] + with tempfile.TemporaryDirectory() as tmp_dir: + pte_path = f"{tmp_dir}/model.pte" + with open(pte_path, "wb") as pte_file: + pte_file.write(program.buffer) + for from_file in (False, True): + with self.subTest(from_file=from_file): + bundle = ( + BundledProgram(None, suites, pte_file_path=pte_path) + if from_file + else BundledProgram(program, suites) + ) + case = ( + bundle.serialize_to_schema().method_test_suites[0].test_cases[0] + ) + for value, source in zip( + case.inputs + case.expected_outputs, (input_tensor, expected) + ): + self.assertEqual(value.val.dim_order, [0, 1, 2, 3]) + self.assertEqual( + value.val.data, source.contiguous().numpy().tobytes() + ) + def assertIOsetDataEqual( self, program_ioset_data: List[bp_schema.Value], diff --git a/exir/tensor.py b/exir/tensor.py index 50fcd46fa53..09ad208994f 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -11,7 +11,6 @@ from __future__ import annotations import math -import typing from typing import Dict, List, NamedTuple, Optional, Tuple, Union import executorch.exir.schema as schema @@ -50,8 +49,8 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]: def dim_order_from_stride( - stride: Tuple[int], sizes: Optional[Tuple[int]] = None -) -> Tuple[bytes]: + stride: Tuple[int, ...], sizes: Optional[Tuple[int, ...]] = None +) -> Tuple[int, ...]: """ Dimension order represents how dimensions are laid out in memory, starting from the outer-most to the inner-most dimension. @@ -115,7 +114,7 @@ def __lt__(self, other): expected = make_channels_last_strides_for(sizes) if all(statically_known_true(s == e) for s, e in zip(stride, expected)): sorted_dims = [0, *range(2, ndim), 1] - return tuple(typing.cast(Tuple[bytes], sorted_dims)) + return tuple(sorted_dims) def stride_from_dim_order(sizes: List[int], dim_order: List[int]) -> List[int]: @@ -207,9 +206,12 @@ def nbytes(self) -> int: return num_bytes_from_shape_and_dtype(self.shape, self.dtype) @classmethod - def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec: + def from_tensor( + cls, + tensor: torch.Tensor, + const: bool = False, + ) -> TensorSpec: if const: - # for non-contigous tensors, convert to a contiguous one tensor = tensor.contiguous() # Weights cannot be views during emission or serialization if tensor.nbytes != tensor.untyped_storage().nbytes(): diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK index e90e837732d..b07390cdeba 100644 --- a/extension/pybindings/test/BUCK +++ b/extension/pybindings/test/BUCK @@ -1,3 +1,8 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") @@ -40,6 +45,9 @@ fbcode_target( preload_deps = ["//executorch/kernels/quantized:aot_lib"], deps = [ ":make_test", + "//executorch/devtools/bundled_program:config", + "//executorch/devtools/bundled_program:core", + "//executorch/devtools/bundled_program/serialize:lib", "//executorch/exir:schema", "//executorch/exir/backend/test:device_util", "//executorch/extension/pybindings:portable_lib", @@ -53,6 +61,9 @@ fbcode_target( preload_deps = ["//executorch/kernels/quantized:aot_lib"], deps = [ ":make_test", + "//executorch/devtools/bundled_program:config", + "//executorch/devtools/bundled_program:core", + "//executorch/devtools/bundled_program/serialize:lib", "//executorch/exir:schema", "//executorch/exir/backend/test:device_util", "//executorch/extension/pybindings:aten_lib", @@ -66,6 +77,9 @@ fbcode_target( srcs = ["test_pybindings.py"], deps = [ ":make_test", + "//executorch/devtools/bundled_program:config", + "//executorch/devtools/bundled_program:core", + "//executorch/devtools/bundled_program/serialize:lib", "//executorch/exir:schema", "//executorch/exir/backend/test:device_util", ], diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index bc034bc9720..3eb4e3c5178 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -15,6 +16,12 @@ import numpy as np import torch +from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite +from executorch.devtools.bundled_program.core import BundledProgram +from executorch.devtools.bundled_program.serialize import ( + serialize_from_bundled_program_to_flatbuffer, +) + from executorch.exir import ExecutorchBackendConfig, to_edge from executorch.exir.backend.test.device_util import DeviceAwarePartitioner from executorch.exir.passes import MemoryPlanningPass @@ -341,6 +348,66 @@ def test_channels_last(self) -> None: expected = model(inputs[0]) self.assertTrue(torch.allclose(expected, executorch_output)) + def test_bundled_channels_last_input_and_output(self) -> None: + input_tensor = ( + torch.arange(2 * 3 * 4 * 5, dtype=torch.float32) + .reshape(2, 3, 4, 5) + .contiguous(memory_format=torch.channels_last) + ) + model = ModuleChannelsLast() + expected = model(input_tensor) + self.assertTrue(expected.is_contiguous(memory_format=torch.channels_last)) + program = to_edge(export(model, (input_tensor,))).to_executorch() + bundled_program = BundledProgram( + program, + [ + MethodTestSuite( + "forward", + [ + MethodTestCase((input_tensor,), (expected,)), + MethodTestCase( + (input_tensor.contiguous(),), (expected.contiguous(),) + ), + ], + ) + ], + ) + bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program) + bundled_module = self.runtime._load_bundled_program_from_buffer(bundled_buffer) + + for test_index in range(2): + output = bundled_module.verify_result_with_bundled_expected_output( + "forward", test_index + )[0] + self.assertEqual(output.stride(), expected.stride()) + self.assertTrue(torch.equal(output, expected)) + + def test_bundled_permute_output(self) -> None: + class Permute(torch.nn.Module): + def forward(self, x): + return x.permute(0, 2, 3, 1) + + model = Permute() + input_tensor = torch.arange(120, dtype=torch.float32).reshape(2, 3, 4, 5) + expected = model(input_tensor) + self.assertFalse(expected.is_contiguous()) + program = to_edge(export(model, (input_tensor,))).to_executorch() + bundle = BundledProgram( + program, + [ + MethodTestSuite( + "forward", [MethodTestCase((input_tensor,), (expected,))] + ) + ], + ) + bundled_module = self.runtime._load_bundled_program_from_buffer( + serialize_from_bundled_program_to_flatbuffer(bundle) + ) + output = bundled_module.verify_result_with_bundled_expected_output( + "forward", 0 + )[0] + self.assertTrue(torch.equal(output, expected)) + def test_unsupported_dim_order(self) -> None: model = ModuleChannelsLast() exported_program, inputs = create_program(model) From c4daafa294490144cb76042b606b24c73811b5b2 Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Mon, 26 Oct 2026 15:28:02 +0100 Subject: [PATCH 2/6] BundleIO: Move bundled tests away from upstream layout additions Keep the two bundled-program tests separate from the ordinary input layout tests added upstream. Move the existing tests before test_e2e without changing their bodies. Keep this adjustment on top of the original patch for separate code and IP review. Squash the two commits before the upstream merge. Authored with assistance from OpenAI Codex. Change-Id: I4e170bbbe3c3dbbe2c4e5c907a8de7c7acbf02eb Signed-off-by: Zingo Andersen --- extension/pybindings/test/test_pybindings.py | 120 +++++++++---------- 1 file changed, 60 insertions(+), 60 deletions(-) diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 3eb4e3c5178..dfbc36504c1 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -72,6 +72,66 @@ def setUp(self): self.load_prog_fn = runtime._load_program_from_buffer self.runtime = runtime + def test_bundled_channels_last_input_and_output(self) -> None: + input_tensor = ( + torch.arange(2 * 3 * 4 * 5, dtype=torch.float32) + .reshape(2, 3, 4, 5) + .contiguous(memory_format=torch.channels_last) + ) + model = ModuleChannelsLast() + expected = model(input_tensor) + self.assertTrue(expected.is_contiguous(memory_format=torch.channels_last)) + program = to_edge(export(model, (input_tensor,))).to_executorch() + bundled_program = BundledProgram( + program, + [ + MethodTestSuite( + "forward", + [ + MethodTestCase((input_tensor,), (expected,)), + MethodTestCase( + (input_tensor.contiguous(),), (expected.contiguous(),) + ), + ], + ) + ], + ) + bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program) + bundled_module = self.runtime._load_bundled_program_from_buffer(bundled_buffer) + + for test_index in range(2): + output = bundled_module.verify_result_with_bundled_expected_output( + "forward", test_index + )[0] + self.assertEqual(output.stride(), expected.stride()) + self.assertTrue(torch.equal(output, expected)) + + def test_bundled_permute_output(self) -> None: + class Permute(torch.nn.Module): + def forward(self, x): + return x.permute(0, 2, 3, 1) + + model = Permute() + input_tensor = torch.arange(120, dtype=torch.float32).reshape(2, 3, 4, 5) + expected = model(input_tensor) + self.assertFalse(expected.is_contiguous()) + program = to_edge(export(model, (input_tensor,))).to_executorch() + bundle = BundledProgram( + program, + [ + MethodTestSuite( + "forward", [MethodTestCase((input_tensor,), (expected,))] + ) + ], + ) + bundled_module = self.runtime._load_bundled_program_from_buffer( + serialize_from_bundled_program_to_flatbuffer(bundle) + ) + output = bundled_module.verify_result_with_bundled_expected_output( + "forward", 0 + )[0] + self.assertTrue(torch.equal(output, expected)) + def test_e2e(self): exported_program, inputs = create_program(ModuleAdd()) executorch_module = self.load_fn(exported_program.buffer) @@ -348,66 +408,6 @@ def test_channels_last(self) -> None: expected = model(inputs[0]) self.assertTrue(torch.allclose(expected, executorch_output)) - def test_bundled_channels_last_input_and_output(self) -> None: - input_tensor = ( - torch.arange(2 * 3 * 4 * 5, dtype=torch.float32) - .reshape(2, 3, 4, 5) - .contiguous(memory_format=torch.channels_last) - ) - model = ModuleChannelsLast() - expected = model(input_tensor) - self.assertTrue(expected.is_contiguous(memory_format=torch.channels_last)) - program = to_edge(export(model, (input_tensor,))).to_executorch() - bundled_program = BundledProgram( - program, - [ - MethodTestSuite( - "forward", - [ - MethodTestCase((input_tensor,), (expected,)), - MethodTestCase( - (input_tensor.contiguous(),), (expected.contiguous(),) - ), - ], - ) - ], - ) - bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program) - bundled_module = self.runtime._load_bundled_program_from_buffer(bundled_buffer) - - for test_index in range(2): - output = bundled_module.verify_result_with_bundled_expected_output( - "forward", test_index - )[0] - self.assertEqual(output.stride(), expected.stride()) - self.assertTrue(torch.equal(output, expected)) - - def test_bundled_permute_output(self) -> None: - class Permute(torch.nn.Module): - def forward(self, x): - return x.permute(0, 2, 3, 1) - - model = Permute() - input_tensor = torch.arange(120, dtype=torch.float32).reshape(2, 3, 4, 5) - expected = model(input_tensor) - self.assertFalse(expected.is_contiguous()) - program = to_edge(export(model, (input_tensor,))).to_executorch() - bundle = BundledProgram( - program, - [ - MethodTestSuite( - "forward", [MethodTestCase((input_tensor,), (expected,))] - ) - ], - ) - bundled_module = self.runtime._load_bundled_program_from_buffer( - serialize_from_bundled_program_to_flatbuffer(bundle) - ) - output = bundled_module.verify_result_with_bundled_expected_output( - "forward", 0 - )[0] - self.assertTrue(torch.equal(output, expected)) - def test_unsupported_dim_order(self) -> None: model = ModuleChannelsLast() exported_program, inputs = create_program(model) From 953f1d9938bf2b47085eff073bde0e43478157f6 Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Fri, 16 Oct 2026 16:43:24 +0200 Subject: [PATCH 3/6] BundleIO: Revert unrelated TensorSpec formatting Restore the TensorSpec.from_tensor signature and comment to keep the original layout fix focused. Authored with assistance from OpenAI Codex. Change-Id: I00b17a0cf26675edda350694a37e241605727bd4 Signed-off-by: Zingo Andersen --- exir/tensor.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/exir/tensor.py b/exir/tensor.py index 09ad208994f..0bb808aa02e 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -206,12 +206,9 @@ def nbytes(self) -> int: return num_bytes_from_shape_and_dtype(self.shape, self.dtype) @classmethod - def from_tensor( - cls, - tensor: torch.Tensor, - const: bool = False, - ) -> TensorSpec: + def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec: if const: + # for non-contigous tensors, convert to a contiguous one tensor = tensor.contiguous() # Weights cannot be views during emission or serialization if tensor.nbytes != tensor.untyped_storage().nbytes(): From d8dd4620572a900e07168299ef77f3aa7049ccca Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Fri, 16 Oct 2026 16:43:37 +0200 Subject: [PATCH 4/6] BundleIO: Read only I/O metadata from saved programs Read method I/O metadata through FlatBuffers accessors instead of unpacking the entire program. Avoid copying weights and delegate payloads just to obtain tensor dimension orders. Add a regression test that bundles a saved program without unpacking the program or restoring its payload segments. Check file and memory bundles agree for multiple methods and primitive inputs. Authored with assistance from OpenAI Codex. Change-Id: Ic966f59f261a96501abcd9f37a26a4494aa2354f Signed-off-by: Zingo Andersen --- devtools/bundled_program/core.py | 71 +++++++++++++++---- .../bundled_program/test/test_bundle_data.py | 58 +++++++++++++++ 2 files changed, 114 insertions(+), 15 deletions(-) diff --git a/devtools/bundled_program/core.py b/devtools/bundled_program/core.py index 00375a6f60f..c8851741af0 100644 --- a/devtools/bundled_program/core.py +++ b/devtools/bundled_program/core.py @@ -6,7 +6,7 @@ # LICENSE file in the root directory of this source tree. import ctypes -from typing import Dict, List, Optional, Sequence, Type, Union +from typing import Dict, List, NamedTuple, Optional, Sequence, Type, Union import executorch.devtools.bundled_program.schema as bp_schema @@ -19,7 +19,12 @@ from executorch.devtools.bundled_program.version import BUNDLED_PROGRAM_SCHEMA_VERSION from executorch.exir import ExecutorchProgram, ExecutorchProgramManager -from executorch.exir._serialize import _deserialize_pte_binary +from executorch.exir._serialize._flatbuffer_program import ( + _convert_from_flatbuffer_dataclass, + _patch_generated_flatbuffer_aliases, +) +from executorch.exir._serialize.generated.executorch_flatbuffer.EValue import EValueT +from executorch.exir._serialize.generated.executorch_flatbuffer.Program import Program from executorch.exir.tensor import get_scalar_type, scalar_type_enum # pyre-ignore @@ -31,6 +36,11 @@ } +class _MethodIO(NamedTuple): + inputs: List[core_schema.KernelTypes] + outputs: List[core_schema.KernelTypes] + + class BundledProgram: """ Bundled program contains all information needed to execute and verify the program on device. @@ -99,31 +109,26 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: assert self._pte_file_path is not None with open(self._pte_file_path, "rb") as f: program_bytes = f.read() - program = _deserialize_pte_binary(program_bytes).program - plans = {plan.name: plan for plan in program.execution_plan} + program = program_bytes + method_io = self._get_method_io(program) bundled_method_test_suites: List[bp_schema.BundledMethodTestSuite] = [] # Emit data and metadata of bundled tensor for method_test_suite in self.method_test_suites: - plan = plans[method_test_suite.method_name] + plan = method_io[method_test_suite.method_name] bundled_test_cases: List[bp_schema.BundledMethodTestCase] = [] # emit I/O sets for each method test case - for i in range(len(method_test_suite.test_cases)): + for test_case in method_test_suite.test_cases: inputs: List[bp_schema.Value] = [] expected_outputs: List[bp_schema.Value] = [] - cur_plan_test_inputs = method_test_suite.test_cases[i].inputs - cur_plan_test_expected_outputs = method_test_suite.test_cases[ - i - ].expected_outputs - - for input_idx, input_val in enumerate(cur_plan_test_inputs): + for input_idx, input_val in enumerate(test_case.inputs): if type(input_val) is torch.Tensor: self._emit_bundled_tensor( input_val, - plan.values[plan.inputs[input_idx]].val.dim_order, + plan.inputs[input_idx].dim_order, inputs, ) else: @@ -132,14 +137,14 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: inputs, ) for output_idx, expected_output_tensor in enumerate( - cur_plan_test_expected_outputs + test_case.expected_outputs ): assert ( type(expected_output_tensor) is torch.Tensor ), "Only tensor outputs are currently supported." self._emit_bundled_tensor( expected_output_tensor, - plan.values[plan.outputs[output_idx]].val.dim_order, + plan.outputs[output_idx].dim_order, expected_outputs, ) bundled_test_cases.append( @@ -196,6 +201,42 @@ def _emit_prim(self, val: ConfigValue, bundled_values: List[bp_schema.Value]): else: assert 0, "Unsupported primitive type received." + def _get_method_io( + self, program: Union[core_schema.Program, bytes] + ) -> Dict[str, _MethodIO]: + if isinstance(program, core_schema.Program): + return { + plan.name: _MethodIO( + inputs=[plan.values[index].val for index in plan.inputs], + outputs=[plan.values[index].val for index in plan.outputs], + ) + for plan in program.execution_plan + } + + # Read only I/O metadata, without unpacking constants or delegate payloads. + _patch_generated_flatbuffer_aliases() + flatbuffer_program = Program.GetRootAs(program) + method_io = {} + for plan_index in range(flatbuffer_program.ExecutionPlanLength()): + plan = flatbuffer_program.ExecutionPlan(plan_index) + method_io[plan.Name().decode("utf-8")] = _MethodIO( + inputs=[ + _convert_from_flatbuffer_dataclass( + EValueT.InitFromObj(plan.Values(plan.Inputs(index))), + core_schema.EValue, + ).val + for index in range(plan.InputsLength()) + ], + outputs=[ + _convert_from_flatbuffer_dataclass( + EValueT.InitFromObj(plan.Values(plan.Outputs(index))), + core_schema.EValue, + ).val + for index in range(plan.OutputsLength()) + ], + ) + return method_io + def _get_program_input( self, program: core_schema.Program, plan_idx: int, input_idx: int ) -> core_schema.KernelTypes: diff --git a/devtools/bundled_program/test/test_bundle_data.py b/devtools/bundled_program/test/test_bundle_data.py index eb5b9efd21b..0fa9b8f6668 100644 --- a/devtools/bundled_program/test/test_bundle_data.py +++ b/devtools/bundled_program/test/test_bundle_data.py @@ -10,6 +10,7 @@ import tempfile import unittest from typing import List +from unittest.mock import patch import executorch.devtools.bundled_program.schema as bp_schema @@ -25,6 +26,7 @@ ) from executorch.exir import to_edge from executorch.exir._serialize import _PTEFile, _serialize_pte_binary +from executorch.exir._serialize._program import _get_extended_header from executorch.exir.tensor import stride_from_dim_order from torch.export import export @@ -154,6 +156,62 @@ def forward(self, x): value.val.data, source.contiguous().numpy().tobytes() ) + def test_pte_reader_matches_in_memory_program(self) -> None: + program, suites = get_common_executorch_program() + expected = BundledProgram(program, suites).serialize_to_schema() + with tempfile.TemporaryDirectory() as tmp_dir: + pte_path = f"{tmp_dir}/model.pte" + with open(pte_path, "wb") as pte_file: + pte_file.write(program.buffer) + actual = BundledProgram( + None, suites, pte_file_path=pte_path + ).serialize_to_schema() + self.assertEqual(actual, expected) + + def test_pte_reader_does_not_unpack_program_payloads(self) -> None: + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.arange(6).reshape(2, 3).float()) + + def forward(self, x): + return x + self.weight + + model = Model() + tensor = torch.arange(6).reshape(3, 2).t().float() + expected = model(tensor) + program = to_edge(export(model, (torch.zeros(2, 3),))).to_executorch() + program_bytes = program.buffer + header = _get_extended_header(program_bytes) + assert header is not None + self.assertGreater(header.segment_data_size, 0) + suites = [MethodTestSuite("forward", [MethodTestCase((tensor,), (expected,))])] + with tempfile.TemporaryDirectory() as tmp_dir: + pte_path = f"{tmp_dir}/model.pte" + with open(pte_path, "wb") as pte_file: + pte_file.write(program_bytes) + with ( + patch( + "executorch.exir._serialize._flatbuffer_program.ProgramT.InitFromPackedBuf", + side_effect=AssertionError("Full program must not be unpacked"), + ), + patch( + "executorch.exir._serialize._program._restore_segments", + side_effect=AssertionError("Program payloads must not be restored"), + ), + ): + bundle = BundledProgram(None, suites, pte_file_path=pte_path) + schema = bundle.serialize_to_schema() + + self.assertEqual(schema.program, program_bytes) + case = schema.method_test_suites[0].test_cases[0] + for value, source in zip( + case.inputs + case.expected_outputs, (tensor, expected) + ): + self.assertEqual(value.val.dim_order, [0, 1]) + self.assertEqual(value.val.sizes, [2, 3]) + self.assertEqual(value.val.data, source.contiguous().numpy().tobytes()) + def assertIOsetDataEqual( self, program_ioset_data: List[bp_schema.Value], From 03b9c46dbe5401f66b768b65b7edccb9ab22f31e Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Fri, 16 Oct 2026 16:44:26 +0200 Subject: [PATCH 5/6] BundleIO: Resolve tensor view values before copying bytes Materialize pending conjugation and negation before copying tensor bytes. Contiguous views can retain these flags and otherwise serialize the underlying storage values instead of their logical values. Keep tensor byte packing independent of a program instance. Cover a negative view created from a conjugate and a tensor with pending conjugation. The former previously serialized correctly through TensorSpec, which cloned the view before copying storage. Authored with assistance from OpenAI Codex. Change-Id: Ia7c77e1b55ecbaa3ce3578e36aca4690cd400ccb Signed-off-by: Zingo Andersen --- devtools/bundled_program/core.py | 4 +-- .../bundled_program/test/test_bundle_data.py | 29 ++++++++++++++++++- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/devtools/bundled_program/core.py b/devtools/bundled_program/core.py index c8851741af0..6a0053558fd 100644 --- a/devtools/bundled_program/core.py +++ b/devtools/bundled_program/core.py @@ -168,15 +168,15 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: ) return self._bundled_program_in_schema + @staticmethod def _emit_bundled_tensor( - self, tensor: torch.Tensor, dim_order: List[int], bundled_values: List[bp_schema.Value], ) -> None: sizes = list(tensor.shape) # Store bytes in the physical order declared by the method's I/O tensor. - tensor = tensor.permute(dim_order).contiguous() + tensor = tensor.permute(dim_order).contiguous().resolve_conj().resolve_neg() tensor_data = ( ctypes.string_at(tensor.data_ptr(), tensor.nbytes) if tensor.nbytes else b"" ) diff --git a/devtools/bundled_program/test/test_bundle_data.py b/devtools/bundled_program/test/test_bundle_data.py index 0fa9b8f6668..0ebb585525e 100644 --- a/devtools/bundled_program/test/test_bundle_data.py +++ b/devtools/bundled_program/test/test_bundle_data.py @@ -39,7 +39,9 @@ class Model(torch.nn.Module): def forward(self, x): return x * 2 - program = to_edge(export(Model(), (input_tensor,))).to_executorch() + program = to_edge( + export(Model(), (input_tensor.resolve_neg().resolve_conj(),)) + ).to_executorch() with tempfile.TemporaryDirectory() as tmp_dir: pte_path = f"{tmp_dir}/model.pte" with open(pte_path, "wb") as pte_file: @@ -118,6 +120,31 @@ def test_non_dense_view_uses_only_its_own_elements(self) -> None: ) self.assertTrue(torch.equal(decoded, view)) + def test_pending_negative_input_and_output(self) -> None: + tensor = torch.tensor([1 + 2j]).conj().imag + self.assertTrue(tensor.is_contiguous()) + self.assertTrue(tensor.is_neg()) + expected = torch.tensor([1 + 4j]).conj().imag + test_case = self._serialize_tensor_test_case(tensor, expected) + + for value, expected_value in zip( + test_case.inputs + test_case.expected_outputs, (-2.0, -4.0) + ): + self.assertEqual(value.val.dim_order, [0]) + self.assertEqual(value.val.sizes, [1]) + self.assertEqual( + value.val.data, torch.tensor([expected_value]).numpy().tobytes() + ) + + def test_pending_conjugate_tensor(self) -> None: + tensor = torch.tensor([1 + 2j]).conj() + self.assertTrue(tensor.is_conj()) + values: List[bp_schema.Value] = [] + BundledProgram._emit_bundled_tensor(tensor, [0], values) + self.assertEqual(values[0].val.dim_order, [0]) + self.assertEqual(values[0].val.sizes, [1]) + self.assertEqual(values[0].val.data, tensor.resolve_conj().numpy().tobytes()) + def test_bundle_uses_program_dim_order(self) -> None: class Permute(torch.nn.Module): def forward(self, x): From 123a6de774477316f2a70e578428f303f71baa71 Mon Sep 17 00:00:00 2001 From: Zingo Andersen Date: Wed, 28 Oct 2026 10:06:22 +0100 Subject: [PATCH 6/6] Fix dimension-order type annotations in tensor specs Use Tuple[int, ...] for dim_order in TensorSpec and _ViewSpec to match the return type of dim_order_from_stride. Runtime behavior is unchanged. Authored with assistance from OpenAI Codex. Signed-off-by: Zingo Andersen Change-Id: I049655d7d62ec276082c7dca3e3453e5d7643638 --- exir/passes/replace_view_copy_with_view_pass.py | 2 +- exir/tensor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/exir/passes/replace_view_copy_with_view_pass.py b/exir/passes/replace_view_copy_with_view_pass.py index a6fd449e3db..860a1351664 100644 --- a/exir/passes/replace_view_copy_with_view_pass.py +++ b/exir/passes/replace_view_copy_with_view_pass.py @@ -315,7 +315,7 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None: self._base = base self.shape: List[int] = shape self.stride: Tuple[int] = contiguous_stride_from_shape(torch.Size(self.shape)) - self.dim_order: Tuple[bytes] = dim_order_from_stride(self.stride) + self.dim_order: Tuple[int, ...] = dim_order_from_stride(self.stride) self.shape_dynamism: TensorShapeDynamism = determine_tensor_dynanism( torch.Size(self.shape) ) diff --git a/exir/tensor.py b/exir/tensor.py index 0bb808aa02e..938f72393b6 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -182,7 +182,7 @@ def __init__( # convert to list making it easier to handle type checking self.shape: List[int] = list(shape) self.stride: Tuple[int] = contiguous_stride_from_shape(shape) - self.dim_order: Tuple[bytes] = dim_order_from_stride(self.stride) + self.dim_order: Tuple[int, ...] = dim_order_from_stride(self.stride) self.requires_grad = requires_grad self.layout = layout self.is_sparse = is_sparse