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 bc034bc9720..cb82118bb50 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. @@ -341,6 +342,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)