Implement direct serialization of .bpte:s (#20217)

This patch is heavily inspired by
#17333

Similarly to the old serialization of pte:s,
bundled program serialization went through a
program -> string json -> flatbuffer path, which
is heavily inefficient. As in the PR above,
generate python flatbuffer bindings from the fbs schema and use them to
cut out the json step.

Checked in flatbuffer bindings and fbs schema coherence is enforced by
the validate_flatbuffer_gen job.

e6586533af

Testing:
Tested by devtools/bundled_program/test/test_end2end.py, Additionally,
did the following benchmark locally:

/usr/bin/time -v -- aot_arm_compiler.py --delegate --bundleio -t vgf
[args]

Time metric: Elapsed (wall clock) time
Memory metric: Maximum resident set size

```
Model       Baseline (s / MiB)    With-patch (s / MiB)    Speed-up   Relative memory   same.bpte  bpte size (MiB)
---------   -------------------   ---------------------   --------   ---------------   --------   ---------
mv2         31.36 / 1349.7        22.37 / 845.5           1.40x      0.63x (-37.4%)    yes        13.91
resnet50    128.57 / 10380.9      28.24 / 2003.4          4.55x      0.19x (-80.7%)    yes        97.99
w2l         730.23 / 12102.9      15.85 / 2376.2          46.07x     0.20x (-80.4%)    yes        131.45
deit_tiny   87.57 / 2669.5        70.23 / 1037.6          1.25x      0.39x (-61.1%)    yes        22.88
```


cc @Gasoonjia @digantdesai @freddan80 @per @zingo @oscarandersson8218
@mansnils @Sebastian-Larsson @robell @rascani

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

.flake8

@@ -75,6 +75,7 @@ exclude =
     ./configurations,
     ./docs,
     ./exir/_serialize/generated/executorch_flatbuffer,
+    ./devtools/bundled_program/serialize/generated,
     ./third_party,
     *.pyi
 

.github/workflows/validate_flatbuffer_gen.yml

@@ -5,7 +5,9 @@ on:
   pull_request:
     paths:
       - "schema/**"
-      - "exir/_serialize/generated/executorch_flatbuffer/**"
+      - "devtools/bundled_program/schema/**"
+      - "exir/_serialize/generated/**"
+      - "devtools/bundled_program/serialize/generated/**"
 
 jobs:
   exir-flatbuffer:
@@ -33,3 +35,15 @@ jobs:
             echo "Please run 'python exir/_serialize/generate_program.py' to regenerate the files and commit the changes."
             exit 1
           fi
+
+      - name: Generate bundled program flatbuffer Python
+        run: python devtools/bundled_program/serialize/generate_bundled_program.py
+
+      - name: Validate bundled_program_flatbuffer is unchanged
+        run: |
+          git add -A devtools/bundled_program/serialize/generated
+          if ! git diff --cached --quiet -- devtools/bundled_program/serialize/generated; then
+            echo "Error: bundled_program_flatbuffer has uncommitted changes."
+            echo "Please run 'python devtools/bundled_program/serialize/generate_bundled_program.py' to regenerate the files and commit the changes."
+            exit 1
+          fi

.lintrunner.toml

@@ -9,6 +9,7 @@ exclude_patterns = [
     '.github/scripts/**',
     'exir/serde/**',
     'exir/_serialize/generated/executorch_flatbuffer/**',
+    'devtools/bundled_program/serialize/generated/**',
 ]
 command = [
     'python',
@@ -41,6 +42,7 @@ exclude_patterns = [
     '**/third-party/**',
     'exir/serde/**',
     'exir/_serialize/generated/executorch_flatbuffer/**',
+    'devtools/bundled_program/serialize/generated/**',
 ]
 command = [
     'python',

backends/arm/scripts/pre-push

@@ -177,7 +177,7 @@ for COMMIT in ${COMMITS}; do
     for committed_file in "${license_files[@]}"; do
         # Skip files with certain extensions
         case "$committed_file" in
-            *.md|*.md.in|*.json|*.yml|*.yaml|*.cmake|*.patch|.gitignore|*.bzl|BUCK|*/BUCK|TARGETS|*/TARGETS)
+            *.md|*.md.in|*.json|*.yml|*.yaml|*.cmake|*.patch|.gitignore|*.bzl|BUCK|*/BUCK|TARGETS|*/TARGETS|*/generated/*)
                 echo -e "${INFO} Skipping license check for ${committed_file} (excluded extension)"
                 continue
                 ;;

devtools/bundled_program/schema/README.md

@@ -4,3 +4,13 @@ and other useful info together for verifying the correctness of ExecuTorch progr
 
 ## Rules to ensure forward/backward compatibility
 Please check the rules in [here](../../../schema/README.md) for more info.
+
+
+## Regenerating generated code
+
+Schema changes require regenerating the Python bindings in
+`devtools/bundled_program/serialize/generated` and committing the updated files. From the repo root:
+
+```sh
+python devtools/bundled_program/serialize/generate_bundled_program.py
+```

devtools/bundled_program/serialize/BUCK

@@ -9,7 +9,7 @@ fbcode_target(_kind = runtime.python_library,
     name = "lib",
     srcs = [
         "__init__.py",
-    ],
+    ] + glob(["generated/**/*.py"]),
     resources = {
         "//executorch/devtools/bundled_program/schema:bundled_program_schema.fbs": "bundled_program_schema.fbs",
         "//executorch/devtools/bundled_program/schema:scalar_type.fbs": "scalar_type.fbs",
@@ -19,6 +19,7 @@ fbcode_target(_kind = runtime.python_library,
     # Please ask before changing this.
     visibility = ["PUBLIC"],
     deps = [
+        "fbsource//third-party/pypi/flatbuffers:flatbuffers",
         "fbsource//third-party/pypi/setuptools:setuptools",
         "//executorch/devtools/bundled_program/schema:bundled_program_schema_py",
         "//executorch/exir/_serialize:lib",

devtools/bundled_program/serialize/__init__.py

@@ -1,6 +1,6 @@
 # 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.
@@ -9,23 +9,62 @@
 
 # TODO(T138924864): Refactor to unify the serialization for bundled program and executorch program.
 
+import functools
 import importlib.resources as _resources
 import json
 import os
+import re
 import tempfile
+from typing import Any
 
 import executorch.devtools.bundled_program.schema as bp_schema
 
 import executorch.devtools.bundled_program.serialize as serialization_package
+
+import flatbuffers  # pyre-ignore[21]
 from executorch.devtools.bundled_program.core import BundledProgram
+from executorch.devtools.bundled_program.serialize.generated.bundled_program_flatbuffer import (
+    Bool as _Bool,
+    BundledMethodTestCase as _BundledMethodTestCase,
+    BundledMethodTestSuite as _BundledMethodTestSuite,
+    BundledProgram as _BundledProgram,
+    Double as _Double,
+    Int as _Int,
+    Tensor as _Tensor,
+    Value as _Value,
+    ValueUnion as _ValueUnion,
+)
 from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
 from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile
+from executorch.exir._serialize._flatbuffer_program import (
+    _coerce_bytes,
+    _create_aligned_byte_vector,
+)
 
 # The prefix of schema files used for bundled program
 BUNDLED_PROGRAM_SCHEMA_NAME = "bundled_program_schema"
 SCALAR_TYPE_SCHEMA_NAME = "scalar_type"
 
 
+@functools.lru_cache(maxsize=1)
+def _bundled_program_file_identifier() -> bytes:
+    schema = _resources.read_binary(
+        serialization_package, f"{BUNDLED_PROGRAM_SCHEMA_NAME}.fbs"
+    )
+    match = re.search(rb'file_identifier\s+"([^"]+)"', schema)
+    if match is None:
+        raise ValueError(
+            f"Missing file_identifier in {BUNDLED_PROGRAM_SCHEMA_NAME}.fbs"
+        )
+    file_identifier = match.group(1)
+    if len(file_identifier) != 4:
+        raise ValueError(
+            f"Invalid file_identifier length {len(file_identifier)} "
+            f"in {BUNDLED_PROGRAM_SCHEMA_NAME}.fbs"
+        )
+    return file_identifier
+
+
 def write_schema(d: str, schema_name: str) -> None:
     schema_path = os.path.join(d, "{}.fbs".format(schema_name))
     with open(schema_path, "wb") as schema_file:
@@ -78,6 +117,145 @@ def convert_from_flatbuffer(program_flatbuffer: bytes) -> bytes:
             return output_file.read()
 
 
+def _pack_tensor(self: Any, builder: Any) -> int:
+    if self.sizes is not None:
+        _Tensor.TensorStartSizesVector(builder, len(self.sizes))
+        for i in reversed(range(len(self.sizes))):
+            builder.PrependInt32(self.sizes[i])
+        sizes = builder.EndVector()
+    if self.data is not None:
+        data = _create_aligned_byte_vector(builder, _coerce_bytes(self.data), 16)
+    if self.dimOrder is not None:
+        dim_order = _create_aligned_byte_vector(
+            builder, _coerce_bytes(self.dimOrder), 1
+        )
+
+    _Tensor.TensorStart(builder)
+    _Tensor.TensorAddScalarType(builder, self.scalarType)
+    if self.sizes is not None:
+        _Tensor.TensorAddSizes(builder, sizes)
+    if self.data is not None:
+        _Tensor.TensorAddData(builder, data)
+    if self.dimOrder is not None:
+        _Tensor.TensorAddDimOrder(builder, dim_order)
+    return _Tensor.TensorEnd(builder)
+
+
+def _pack_bundled_program(self: Any, builder: Any) -> int:
+    if self.methodTestSuites is not None:
+        method_test_suites_list = [
+            method_test_suite.Pack(builder)
+            for method_test_suite in self.methodTestSuites
+        ]
+        _BundledProgram.BundledProgramStartMethodTestSuitesVector(
+            builder, len(self.methodTestSuites)
+        )
+        for i in reversed(range(len(self.methodTestSuites))):
+            builder.PrependUOffsetTRelative(method_test_suites_list[i])
+        method_test_suites = builder.EndVector()
+    if self.program is not None:
+        program = _create_aligned_byte_vector(builder, _coerce_bytes(self.program), 32)
+
+    _BundledProgram.BundledProgramStart(builder)
+    _BundledProgram.BundledProgramAddVersion(builder, self.version)
+    if self.methodTestSuites is not None:
+        _BundledProgram.BundledProgramAddMethodTestSuites(builder, method_test_suites)
+    if self.program is not None:
+        _BundledProgram.BundledProgramAddProgram(builder, program)
+    return _BundledProgram.BundledProgramEnd(builder)
+
+
+@functools.lru_cache(maxsize=1)
+def _install_fast_packers() -> None:
+    _Tensor.TensorT.Pack = _pack_tensor
+    _BundledProgram.BundledProgramT.Pack = _pack_bundled_program
+
+
+def _convert_tensor(val: bp_schema.Tensor) -> Any:
+    result = _Tensor.TensorT()
+    result.scalarType = int(val.scalar_type)
+    result.sizes = list(val.sizes)
+    result.data = _coerce_bytes(val.data)
+    result.dimOrder = _coerce_bytes(val.dim_order)
+    return result
+
+
+def _convert_int(val: bp_schema.Int) -> Any:
+    result = _Int.IntT()
+    result.intVal = val.int_val
+    return result
+
+
+def _convert_bool(val: bp_schema.Bool) -> Any:
+    result = _Bool.BoolT()
+    result.boolVal = val.bool_val
+    return result
+
+
+def _convert_double(val: bp_schema.Double) -> Any:
+    result = _Double.DoubleT()
+    result.doubleVal = val.double_val
+    return result
+
+
+def _convert_value_union(val: bp_schema.ValueUnion) -> tuple[int, Any]:
+    if isinstance(val, bp_schema.Tensor):
+        return _ValueUnion.ValueUnion.Tensor, _convert_tensor(val)
+    if isinstance(val, bp_schema.Int):
+        return _ValueUnion.ValueUnion.Int, _convert_int(val)
+    if isinstance(val, bp_schema.Bool):
+        return _ValueUnion.ValueUnion.Bool, _convert_bool(val)
+    if isinstance(val, bp_schema.Double):
+        return _ValueUnion.ValueUnion.Double, _convert_double(val)
+    return _ValueUnion.ValueUnion.NONE, None
+
+
+def _convert_value(val: bp_schema.Value) -> Any:
+    result = _Value.ValueT()
+    result.valType, result.val = _convert_value_union(val.val)
+    return result
+
+
+def _convert_method_test_case(val: bp_schema.BundledMethodTestCase) -> Any:
+    result = _BundledMethodTestCase.BundledMethodTestCaseT()
+    result.inputs = [_convert_value(value) for value in val.inputs]
+    result.expectedOutputs = [_convert_value(value) for value in val.expected_outputs]
+    return result
+
+
+def _convert_method_test_suite(val: bp_schema.BundledMethodTestSuite) -> Any:
+    result = _BundledMethodTestSuite.BundledMethodTestSuiteT()
+    result.methodName = val.method_name
+    result.testCases = [
+        _convert_method_test_case(test_case) for test_case in val.test_cases
+    ]
+    return result
+
+
+def _convert_bundled_program(val: bp_schema.BundledProgram) -> Any:
+    result = _BundledProgram.BundledProgramT()
+    result.version = val.version
+    result.methodTestSuites = [
+        _convert_method_test_suite(suite) for suite in val.method_test_suites
+    ]
+    result.program = _coerce_bytes(val.program)
+    return result
+
+
+def _bundled_program_schema_to_flatbuffer(
+    bundled_program: bp_schema.BundledProgram,
+) -> bytes:
+    _install_fast_packers()
+    bundled_program_t = _convert_bundled_program(bundled_program)
+    builder = flatbuffers.Builder()
+    bundled_program_offset = bundled_program_t.Pack(builder)
+    builder.Finish(
+        bundled_program_offset,
+        file_identifier=_bundled_program_file_identifier(),
+    )
+    return bytes(builder.Output())
+
+
 # from bundled program to flatbuffer
 def serialize_from_bundled_program_to_flatbuffer(
     bundled_program: BundledProgram,
@@ -94,9 +272,7 @@ def serialize_from_bundled_program_to_flatbuffer(
 
     bundled_program_in_schema = bundled_program.serialize_to_schema()
 
-    return convert_to_flatbuffer(
-        serialize_from_bundled_program_to_json(bundled_program_in_schema)
-    )
+    return _bundled_program_schema_to_flatbuffer(bundled_program_in_schema)
 
 
 # From flatbuffer to bundled program in schema.