Store external Vela blocks in named data (pytorch#21893)

Differential Revision: D114240131

Pull Request resolved: pytorch#21893

Author: kkch3ng

backends/arm/BUCK

@@ -1,3 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -87,6 +89,7 @@ fbcode_target(
         ":arm_vela",
         "//executorch/backends/arm/tosa:specification",
         "//executorch/backends/arm/tosa:partitioner",
+        "//executorch/exir/_serialize:lib",
     ],
 )
 
@@ -162,6 +165,7 @@ fbcode_target(
     ],
     deps = [
         "fbsource//third-party/pypi/ethos-u-vela:ethos-u-vela",
+        "//executorch/exir/_serialize:lib",
     ],
 )
 fbcode_target(

backends/arm/CMakeLists.txt

@@ -1,3 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
 # Copyright 2023, 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the

backends/arm/__init__.py

@@ -23,6 +23,10 @@
     "EthosUBackend": ("executorch.backends.arm.ethosu", "EthosUBackend"),
     "EthosUCompileSpec": ("executorch.backends.arm.ethosu", "EthosUCompileSpec"),
     "EthosUPartitioner": ("executorch.backends.arm.ethosu", "EthosUPartitioner"),
+    "VelaExternalBlockPlacements": (
+        "executorch.backends.arm.ethosu",
+        "VelaExternalBlockPlacements",
+    ),
     "VgfBackend": ("executorch.backends.arm.vgf", "VgfBackend"),
     "VgfCompileSpec": ("executorch.backends.arm.vgf", "VgfCompileSpec"),
     "VgfPartitioner": ("executorch.backends.arm.vgf", "VgfPartitioner"),

backends/arm/arm_vela.py

@@ -1,13 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
 # Copyright 2023-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 hashlib
 import os
 import struct
 import tempfile
 
+from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import List
 
 import numpy as np
@@ -20,6 +25,33 @@
     has_vela = False
 
 
+_BLOCK_NAME_LENGTH = 16
+_BLOCK_HEADER = struct.Struct(f"<{_BLOCK_NAME_LENGTH}sIB11s")
+_BLOCK_ALIGNMENT = 16
+_STREAM_HEADER = "vela_bin_stream"
+_STREAM_FOOTER = "vela_end_stream"
+_EXTERNAL_REFERENCE = 1
+_RESERVED_BYTES = b"\x00" * 11
+
+
+@dataclass(frozen=True)
+class VelaExternalBlock:
+    """Named-data payload emitted by Vela compilation."""
+
+    key: str
+    payload: bytes
+    alignment: int
+    placement: str
+
+
+@dataclass(frozen=True)
+class VelaCompileResult:
+    """Binary stream and external payloads produced by Vela."""
+
+    processed_bytes: bytes
+    external_blocks: tuple[VelaExternalBlock, ...] = ()
+
+
 def _as_int32(value, name: str) -> int:
     """Convert numpy scalars to signed int32 with a clear error on overflow."""
     arr = np.asarray(value)
@@ -65,16 +97,18 @@ def vela_compile(
     args: List[str],
     verbose: bool = False,
     intermediate_path: str | None = None,
-):
+    block_placements: Mapping[str, str] | None = None,
+) -> VelaCompileResult:
     """Compile a TOSA graph to a binary stream for ArmBackendEthosU using
     Vela.
     """
     if not has_vela:
         raise RuntimeError(
             "ethos-u-vela pip package couldn't be imported. Make sure it's installed!"
         )
+    resolved_block_placements: Mapping[str, str] = block_placements or {}
 
-    def run(dir: str) -> bytes:
+    def run(dir: str) -> VelaCompileResult:
         tosaname = "out.tosa"
         tosa_path = os.path.join(dir, tosaname)
         with open(tosa_path, "wb") as f:
@@ -90,11 +124,10 @@ def run(dir: str) -> bytes:
 
         np_path = os.path.join(dir, "output", "out_vela.npz")
 
-        blocks = b""
         with np.load(np_path, allow_pickle=False) as data:
             # Construct our modified output_blocks with data in a form easily
             # digested on the device side
-            bin_blocks = {"vela_bin_stream": b""}
+            bin_blocks = {_STREAM_HEADER: b""}
 
             # copy command data through unmodified
             bin_blocks["cmd_data"] = data["cmd_data"].tobytes()
@@ -114,28 +147,58 @@ def run(dir: str) -> bytes:
             bin_blocks["inputs"] = vela_bin_pack_io("input", data)
             bin_blocks["outputs"] = vela_bin_pack_io("output", data)
 
-            bin_blocks["vela_end_stream"] = b""
+            bin_blocks[_STREAM_FOOTER] = b""
+
+            unknown_blocks = resolved_block_placements.keys() - bin_blocks.keys()
+            if unknown_blocks:
+                raise ValueError(
+                    "External Vela block placements reference blocks that were not "
+                    f"emitted: {sorted(unknown_blocks)}"
+                )
 
             # Emit the NPZ regions as:
             #  - 16 byte block name null terminated string (padded to 16 if name shorter)
-            #  - 4 bytes of int32 block length and 12 bytes of 0's
+            #  - 4 bytes of int32 block length, 1 byte external flag, and 11 reserved bytes
             #  - block data (padded to 16 byte alignment at end)
             # Repeat for all blocks
+            blocks = b""
+            external_blocks: list[VelaExternalBlock] = []
             for key in bin_blocks.keys():
                 block_name = bytes(key, "utf8")[:15]
                 block_name = block_name + b"\x00" * (16 - len(block_name))
+                block_data = bin_blocks[key]
+                placement = resolved_block_placements.get(key)
+                external_reference = 0
+                if placement is not None:
+                    digest = hashlib.sha256(
+                        placement.encode("ascii") + b"\0" + block_data
+                    ).hexdigest()
+                    external_blocks.append(
+                        VelaExternalBlock(
+                            key=digest,
+                            payload=block_data,
+                            alignment=_BLOCK_ALIGNMENT,
+                            placement=placement,
+                        )
+                    )
+                    block_data = digest.encode("ascii")
+                    external_reference = _EXTERNAL_REFERENCE
 
                 # We need the acual unpadded block lengths for hw setup
-                block_length_bytes = struct.pack("<iiii", len(bin_blocks[key]), 0, 0, 0)
+                block_header = _BLOCK_HEADER.pack(
+                    block_name,
+                    len(block_data),
+                    external_reference,
+                    _RESERVED_BYTES,
+                )
 
                 # Pad block data to multiple of 16 bytes
-                block_data = bin_blocks[key]
                 block_data = block_data + b"\x00" * (15 - (len(block_data) - 1) % 16)
 
-                block = block_name + block_length_bytes + block_data
+                block = block_header + block_data
                 blocks = blocks + block
 
-        return blocks
+            return VelaCompileResult(blocks, tuple(external_blocks))
 
     if intermediate_path is not None:
         return run(intermediate_path)

backends/arm/ethosu/__init__.py

@@ -5,7 +5,12 @@
 #
 
 from .backend import EthosUBackend  # noqa: F401
-from .compile_spec import EthosUCompileSpec  # noqa: F401
+from .compile_spec import EthosUCompileSpec, VelaExternalBlockPlacements  # noqa: F401
 from .partitioner import EthosUPartitioner  # noqa: F401
 
-__all__ = ["EthosUBackend", "EthosUPartitioner", "EthosUCompileSpec"]
+__all__ = [
+    "EthosUBackend",
+    "EthosUPartitioner",
+    "EthosUCompileSpec",
+    "VelaExternalBlockPlacements",
+]

backends/arm/ethosu/backend.py

@@ -1,3 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -14,10 +16,11 @@
 import logging
 from typing import final, List
 
-from executorch.backends.arm.arm_vela import vela_compile
+from executorch.backends.arm.arm_vela import vela_compile, VelaCompileResult
 from executorch.backends.arm.ethosu.compile_spec import EthosUCompileSpec
 
 from executorch.backends.arm.tosa.backend import TOSABackend
+from executorch.exir._serialize._named_data_store import NamedDataStore
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export.exported_program import ExportedProgram
@@ -37,18 +40,18 @@ class EthosUBackend(BackendDetails):
 
     @staticmethod
     def _compile_tosa_flatbuffer(
-        tosa_flatbuffer: bytes, compile_spec: EthosUCompileSpec
-    ) -> bytes:
+        tosa_flatbuffer: bytes,
+        compile_spec: EthosUCompileSpec,
+    ) -> VelaCompileResult:
         """Compile a TOSA flatbuffer into a target-specific binary stream.
 
         Args:
             tosa_flatbuffer (bytes): Serialized TOSA graph produced by
                 ``TOSABackend``.
             compile_spec (EthosUCompileSpec): Compile specification providing
                 Vela flags and intermediate paths.
-
         Returns:
-            bytes: Target-specific binary stream produced by Vela.
+            VelaCompileResult: Binary stream and external payloads from Vela.
 
         """
         compile_flags = compile_spec.compiler_flags
@@ -70,13 +73,15 @@ def _compile_tosa_flatbuffer(
             )
 
         # Pass on the TOSA flatbuffer to the vela compiler.
-        binary = vela_compile(
+        return vela_compile(
             tosa_flatbuffer,
             compile_flags,
             verbose=logger.getEffectiveLevel() <= logging.INFO,
             intermediate_path=compile_spec._get_intermediate_path(),
+            block_placements=(
+                compile_spec.external_block_placements.to_block_placements()
+            ),
         )
-        return binary
 
     @staticmethod
     def preprocess(
@@ -108,8 +113,21 @@ def preprocess(
         # which can be passed on to next compilation step.
         tosa_preprocess = TOSABackend._preprocess(edge_program, tosa_compile_spec)
 
-        binary = EthosUBackend._compile_tosa_flatbuffer(
+        compile_result = EthosUBackend._compile_tosa_flatbuffer(
             tosa_preprocess.processed_bytes, compile_spec
         )
-
-        return PreprocessResult(processed_bytes=binary)
+        if not compile_result.external_blocks:
+            return PreprocessResult(processed_bytes=compile_result.processed_bytes)
+
+        data_store = NamedDataStore()
+        for external_block in compile_result.external_blocks:
+            data_store.add_named_data(
+                external_block.key,
+                external_block.payload,
+                alignment=external_block.alignment,
+                external_tag=external_block.placement,
+            )
+        return PreprocessResult(
+            processed_bytes=compile_result.processed_bytes,
+            data_store_output=data_store.get_named_data_store_output(),
+        )