exir: add flatbuffer Program serialization for performance

.flake8

@@ -74,6 +74,7 @@ exclude =
     ./build,
     ./configurations,
     ./docs,
+    ./exir/_serialize/generated/executorch_flatbuffer,
     ./third_party,
     *.pyi
 

.github/workflows/validate_flatbuffer_gen.yml

@@ -0,0 +1,35 @@
+name: Validate Flatbuffer Generation
+
+on:
+  workflow_dispatch:
+  pull_request:
+    paths:
+      - "schema/**"
+      - "exir/_serialize/generated/executorch_flatbuffer/**"
+
+jobs:
+  exir-flatbuffer:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: setup python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: install executorch  (with flatc)
+        run: ./install_executorch.sh -e
+
+      - name: Generate flatbuffer Python
+        run: python exir/_serialize/generate_program.py
+
+      - name: Validate executorch_flatbuffer is unchanged
+        run: |
+          git add -A exir/_serialize/generated/executorch_flatbuffer
+          if ! git diff --cached --quiet -- exir/_serialize/generated/executorch_flatbuffer; then
+            echo "Error: executorch_flatbuffer has uncommitted changes."
+            echo "Please run 'python exir/_serialize/generate_program.py' to regenerate the files and commit the changes."
+            exit 1
+          fi

.lintrunner.toml

@@ -8,6 +8,7 @@ exclude_patterns = [
     '**/third-party/**',
     '.github/scripts/**',
     'exir/serde/**',
+    'exir/_serialize/generated/executorch_flatbuffer/**',
 ]
 command = [
     'python',
@@ -39,6 +40,7 @@ exclude_patterns = [
     'third-party/**',
     '**/third-party/**',
     'exir/serde/**',
+    'exir/_serialize/generated/executorch_flatbuffer/**',
 ]
 command = [
     'python',

exir/_serialize/BUCK

@@ -31,12 +31,13 @@ fbcode_target(_kind = runtime.python_library,
         "_cord.py",
         "_dataclass.py",
         "_flatbuffer.py",
+        "_flatbuffer_program.py",
         "_named_data_store.py",
         "_program.py",
         "_serialize.py",
         "data_serializer.py",
         "padding.py",
-    ],
+    ] + glob(["generated/**/*.py"]),
     resources = {
         "//executorch/schema:program.fbs": "program.fbs",
         "//executorch/schema:scalar_type.fbs": "scalar_type.fbs",
@@ -47,6 +48,7 @@ fbcode_target(_kind = runtime.python_library,
     # Please ask before changing this.
     visibility = ["PUBLIC"],
     deps = [
+        "fbsource//third-party/pypi/flatbuffers:flatbuffers",
         "//executorch/exir:schema",
         "//executorch/exir:tensor",
         "//executorch/exir:tensor_layout",

exir/_serialize/_flatbuffer.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.
@@ -94,6 +95,28 @@ def __call__(self, schema: bytes) -> bytes:
         return schema
 
 
+class _SchemaFileIdentifierGetter:
+    """Finds the file_identifier value in flatbuffer schemas."""
+
+    def __init__(self) -> None:
+        self.file_identifier: Optional[bytes] = None
+
+    def __call__(self, schema: bytes) -> bytes:
+        identifiers = re.findall(rb'file_identifier\s+"([^"]+)"', schema)
+        for file_identifier in identifiers:
+            if len(file_identifier) != 4:
+                raise ValueError(
+                    f"Invalid file_identifier length {len(file_identifier)} in schema"
+                )
+            if self.file_identifier is None:
+                self.file_identifier = file_identifier
+            elif self.file_identifier != file_identifier:
+                raise ValueError(
+                    f"Mismatched file_identifier {file_identifier} != {self.file_identifier}"
+                )
+        return schema
+
+
 class _ResourceFiles:
     """Manages a collection of python resources that will be written to files."""
 
@@ -112,6 +135,10 @@ def patch_files(self, patch_fn: Callable[[bytes], bytes]) -> None:
         for name in self._files.keys():
             self._files[name] = patch_fn(self._files[name])
 
+    def get(self, name: str) -> bytes:
+        """Returns the current contents of the named file."""
+        return self._files[name]
+
     def write_to(self, out_dir: str) -> None:
         """Writes the files to the specified directory. File names are based on
         the original resource names.
@@ -131,6 +158,15 @@ class _SchemaInfo:
     # schema files.
     max_alignment: int
 
+    # The file identifier declared in the root schema.
+    file_identifier: bytes
+
+    # The alignment for constant tensor data in the schema.
+    tensor_alignment: int
+
+    # The alignment for delegate data in the schema.
+    delegate_alignment: int
+
 
 def _prepare_schema(
     out_dir: str,
@@ -158,13 +194,36 @@ def _prepare_schema(
     # Find the largest alignment used in the patched schema files.
     get_alignments = _SchemaMaxAlignmentGetter()
     schemas.patch_files(get_alignments)
+    get_file_identifier = _SchemaFileIdentifierGetter()
+    schemas.patch_files(get_file_identifier)
+    if get_file_identifier.file_identifier is None:
+        raise ValueError("Missing file_identifier in schema files.")
+
+    def extract_alignment(schema: bytes, marker: bytes) -> int:
+        for line in schema.splitlines():
+            if marker in line:
+                match = re.search(rb"force_align\s*:\s*(\d+)", line)
+                if match:
+                    return int(match.group(1))
+        raise RuntimeError(f"Failed to find marker {marker!r} in program.fbs")
+
+    program_schema_data = schemas.get(program_schema)
+    tensor_alignment = extract_alignment(
+        program_schema_data, b"@executorch-tensor-alignment"
+    )
+    effective_delegate_alignment = extract_alignment(
+        program_schema_data, b"@executorch-delegate-alignment"
+    )
 
     # Write the patched schema files to the filesystem.
     schemas.write_to(out_dir)
 
     return _SchemaInfo(
         root_path=os.path.join(out_dir, program_schema),
         max_alignment=get_alignments.max_alignment,
+        file_identifier=get_file_identifier.file_identifier,
+        tensor_alignment=tensor_alignment,
+        delegate_alignment=effective_delegate_alignment,
     )