From f2d1db4d67c5737dcc23108ca9d5081935b43d06 Mon Sep 17 00:00:00 2001 From: Erik Lundell Date: Wed, 30 Sep 2026 07:39:54 +0000 Subject: [PATCH] BundleIO: Preserve channels-last input and output layouts TensorSpec.from_tensor(..., const=True) makes test tensors contiguous, discarding channels-last strides and reordering their bytes. Bundled inputs can then have incorrect values, while output verification can fail its stride check. Preserve the layout of bundled test tensors. Construct ATen tensors with strides derived from the serialized dimension order before copying their bytes. Cover channels-last inputs and outputs, offset views, and non-dense views. Additionally, a small type hint fix in tensor.py Validation: bundled-program suite (10 tests) Signed-off-by: Erik Lundell Change-Id: I4d05a144aafbb0fbd25bfbbd161c4eebc8dc2b0f --- devtools/bundled_program/bundled_program.cpp | 22 +++-- devtools/bundled_program/core.py | 9 +- .../bundled_program/test/test_bundle_data.py | 89 ++++++++++++++++++- exir/tensor.py | 16 +++- exir/tests/test_tensor.py | 21 +++++ extension/pybindings/test/test_pybindings.py | 40 +++++++++ 6 files changed, 185 insertions(+), 12 deletions(-) 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..1059f9f1c07 100644 --- a/devtools/bundled_program/core.py +++ b/devtools/bundled_program/core.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. @@ -110,7 +111,9 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: for input_val in cur_plan_test_inputs: if type(input_val) is torch.Tensor: self._emit_bundled_tensor( - TensorSpec.from_tensor(input_val, const=True), + TensorSpec.from_tensor( + input_val, const=True, preserve_layout=True + ), inputs, ) else: @@ -123,7 +126,9 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: 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), + TensorSpec.from_tensor( + expected_output_tensor, const=True, preserve_layout=True + ), expected_outputs, ) bundled_test_cases.append( diff --git a/devtools/bundled_program/test/test_bundle_data.py b/devtools/bundled_program/test/test_bundle_data.py index 9fdeb4a776d..17bb768a12c 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,101 @@ 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._serialize import _PTEFile, _serialize_pte_binary +from executorch.exir.tensor import stride_from_dim_order class TestBundle(unittest.TestCase): + def _serialize_tensor_test_case( + self, input_tensor: torch.Tensor, output_tensor: torch.Tensor + ) -> bp_schema.BundledMethodTestCase: + with tempfile.TemporaryDirectory() as tmp_dir: + pte_path = f"{tmp_dir}/model.pte" + with open(pte_path, "wb") as pte_file: + pte_file.write(b"program") + 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 assertIOsetDataEqual( self, program_ioset_data: List[bp_schema.Value], diff --git a/exir/tensor.py b/exir/tensor.py index 50fcd46fa53..9414518f63f 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -50,7 +50,7 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]: def dim_order_from_stride( - stride: Tuple[int], sizes: Optional[Tuple[int]] = None + stride: Tuple[int, ...], sizes: Optional[Tuple[int, ...]] = None ) -> Tuple[bytes]: """ Dimension order represents how dimensions are laid out in memory, @@ -207,10 +207,18 @@ 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, + preserve_layout: bool = False, + ) -> TensorSpec: if const: - # for non-contigous tensors, convert to a contiguous one - tensor = tensor.contiguous() + tensor = ( + tensor.clone(memory_format=torch.preserve_format) + if preserve_layout + else tensor.contiguous() + ) # Weights cannot be views during emission or serialization if tensor.nbytes != tensor.untyped_storage().nbytes(): tensor = tensor.clone() diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py index d6307d716e1..bef95c51803 100644 --- a/exir/tests/test_tensor.py +++ b/exir/tests/test_tensor.py @@ -99,6 +99,27 @@ def test_fp8_tensor_conversion(self) -> None: ) self.compare_tensors(normal_tensor, flatbuffer_tensor) + def test_const_tensor_preserve_layout(self) -> None: + tensor = ( + torch.arange(3 * 3 * 4 * 5, dtype=torch.float32) + .reshape(3, 3, 4, 5) + .contiguous(memory_format=torch.channels_last) + )[1:3] + + default_spec = TensorSpec.from_tensor(tensor, const=True) + self.assertEqual(default_spec.dim_order, (0, 1, 2, 3)) + + spec = TensorSpec.from_tensor(tensor, const=True, preserve_layout=True) + self.assertEqual(spec.dim_order, (0, 2, 3, 1)) + self.assertEqual(spec.stride, (60, 1, 15, 3)) + self.assertTrue(spec.const) + assert spec.storage is not None + self.assertEqual(spec.storage.nbytes(), tensor.nbytes) + self.assertEqual( + bytes(spec.storage), + tensor.permute(0, 2, 3, 1).contiguous().numpy().tobytes(), + ) + def test_allocation_info_succeeds(self) -> None: test_cases = ( ( diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 7e0d3dfce25..9609ad79b2b 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. @@ -193,6 +194,45 @@ 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: + if self.kernel_mode != "aten": + self.skipTest("Requires the ATen bundled-program loader") + + 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, + ) + + 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,))] + ) + ], + ) + bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program) + bundled_module = self.runtime._load_bundled_program_from_buffer(bundled_buffer) + + output = bundled_module.verify_result_with_bundled_expected_output( + "forward", 0 + )[0] + self.assertEqual(output.stride(), expected.stride()) + self.assertTrue(torch.equal(output, expected)) + def test_unsupported_dim_order(self) -> None: model = ModuleChannelsLast() exported_program, inputs = create_program(model)