BundleIO: Preserve channels-last input and output layouts

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<int64_t>(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<ScalarType>(bundled_tensor->scalar_type())));
 
   // Validate data buffer exists and has sufficient size

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(

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],

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()

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 = (
             (

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)