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 e0f4d2f407c..715a38f18f8 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..6a0053558fd 100644 --- a/devtools/bundled_program/core.py +++ b/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,26 +102,33 @@ 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: @@ -118,12 +136,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( + 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,13 +161,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, @@ -154,31 +168,25 @@ def serialize_to_schema(self) -> bp_schema.BundledProgram: ) 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: 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..0ebb585525e 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. @@ -9,19 +10,235 @@ import tempfile import unittest from typing import List +from unittest.mock import patch 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._serialize._program import _get_extended_header +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.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: + 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_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): + 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 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], 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 50fcd46fa53..938f72393b6 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]: @@ -183,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 diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK index a06a66c4a29..408998c5e27 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", @@ -76,6 +87,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 1b661995501..793b407a75b 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. @@ -19,6 +20,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 @@ -81,6 +88,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)