BundleIO: Match bundled tensors to program I/O layouts

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

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,12 +1,12 @@
 # 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
+from typing import Dict, List, NamedTuple, Optional, Sequence, Type, Union
 
 import executorch.devtools.bundled_program.schema as bp_schema
 
@@ -19,7 +19,13 @@
 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._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
 supported_program_type_table: Dict[Type[core_schema.KernelTypes], ConfigValue] = {
@@ -30,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.
@@ -91,39 +102,49 @@ 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 = 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 = 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_val in 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(
-                            TensorSpec.from_tensor(input_val, const=True),
+                            input_val,
+                            plan.inputs[input_idx].dim_order,
                             inputs,
                         )
                     else:
                         self._emit_prim(
                             input_val,
                             inputs,
                         )
-                for expected_output_tensor in cur_plan_test_expected_outputs:
+                for output_idx, expected_output_tensor in enumerate(
+                    test_case.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.outputs[output_idx].dim_order,
                         expected_outputs,
                     )
                 bundled_test_cases.append(
@@ -140,45 +161,32 @@ 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,
             program=program_bytes,
         )
         return self._bundled_program_in_schema
 
+    @staticmethod
     def _emit_bundled_tensor(
-        self, spec: TensorSpec, bundled_values: List[bp_schema.Value]
+        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().resolve_conj().resolve_neg()
+        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),
                 ),
             )
         )
@@ -193,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:

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