From 210a94d6a8bc3b5aff562c3811db6e363a130d84 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Fri, 24 Jul 2026 01:22:34 -0700 Subject: [PATCH] ETRecord: capture delegate `_delegate_info_meta` into ETRecord Summary: Backends can attach arbitrary per-delegate info during lowering via `PreprocessResult._delegate_info_meta`, which `to_backend` copies onto `lowered_module.meta["_delegate_info_meta"]`. Until now that info dead-ended on the lowered module and never reached the ETRecord. This diff carries it the last mile so it can be inspected after an ETRecord is parsed back. What changed: - `ETRecord` gains a `delegate_info: Dict[str, Any]` field, stored as a single JSON blob (reserved name `delegate_info`), keyed by `f"{method}/{lowered_name}"`. - Extraction lives entirely inside ETRecord (no delegate-specific logic in `_program.py`). `_maybe_extract_delegate_info` is called from `add_edge_dialect_program` and `add_executorch_program`, gated so it runs only when (a) the ETRecord has no delegate_info yet and (b) the incoming graph actually carries it. All recorded graphs are the same model, so it is captured exactly once. It walks lowered submodules (recursing into control-flow submodules) and skips non-JSON-serializable values (with a warning) so a backend cannot break ETRecord saving. - Covers both flows: `to_edge_transform_and_lower(...)` (captured at `add_edge_dialect_program`) and `to_edge(...).to_backend(...)` (captured at `add_executorch_program`, since the delegated graph is only available there). - Example backend now logs a minimal, non-static fingerprint (`sha256(processed_bytes)`) via `_delegate_info_meta` as the reference pattern for other backends (Arm/XNNPACK) to follow. Note: `exir/program/_program.py` is intentionally untouched. Differential Revision: D113189648 --- backends/example/example_backend.py | 13 +- devtools/etrecord/BUCK | 2 + devtools/etrecord/_etrecord.py | 178 ++++++++++++++++++++++- devtools/etrecord/tests/BUCK | 2 + devtools/etrecord/tests/etrecord_test.py | 73 ++++++++++ 5 files changed, 265 insertions(+), 3 deletions(-) diff --git a/backends/example/example_backend.py b/backends/example/example_backend.py index 9f1e17dedd6..8ebe4a45cb6 100644 --- a/backends/example/example_backend.py +++ b/backends/example/example_backend.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import copy +import hashlib from typing import final, List from executorch.backends.example.example_backend_delegate_passes.merge_to_dim_pass import ( @@ -37,5 +38,13 @@ def preprocess( assert graph_module_res is not None graph_module_res = MergeToDimPass()(graph_module_res.graph_module) assert graph_module_res is not None - processed_bytes = str(graph_module_res.graph_module.graph) - return PreprocessResult(bytes(processed_bytes, encoding="utf8")) + processed_bytes = bytes( + str(graph_module_res.graph_module.graph), encoding="utf8" + ) + return PreprocessResult( + processed_bytes, + # Log a small, non-static fingerprint of the compiled blob into the + # ETRecord (via lowered_module.meta) so it can be inspected after the + # ETRecord is parsed back. + _delegate_info_meta=hashlib.sha256(processed_bytes).hexdigest(), + ) diff --git a/devtools/etrecord/BUCK b/devtools/etrecord/BUCK index 5197c39d3e0..16e08a7f00a 100644 --- a/devtools/etrecord/BUCK +++ b/devtools/etrecord/BUCK @@ -12,7 +12,9 @@ fbcode_target(_kind = runtime.python_library, deps = [ "//executorch/devtools/bundled_program:core", "//executorch/devtools/bundled_program/schema:bundled_program_schema_py", + "//executorch/exir:graph_module", "//executorch/exir:lib", + "//executorch/exir:lowered_backend_module", "//executorch/exir/emit:emit", "//executorch/exir/serde:serialize", ], diff --git a/devtools/etrecord/_etrecord.py b/devtools/etrecord/_etrecord.py index 71bdc7804c4..fea90fca992 100644 --- a/devtools/etrecord/_etrecord.py +++ b/devtools/etrecord/_etrecord.py @@ -8,8 +8,9 @@ import io import json +import logging import os -from typing import BinaryIO, Dict, IO, List, Optional, Union +from typing import Any, BinaryIO, Dict, IO, List, Optional, Union from zipfile import BadZipFile, ZipFile import torch @@ -26,6 +27,8 @@ ExportedProgram, ) from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap +from executorch.exir.graph_module import get_control_flow_submodules +from executorch.exir.lowered_backend_module import get_lowered_submodules from executorch.exir.serde.export_serialize import SerializedArtifact from executorch.exir.serde.serialize import deserialize, serialize @@ -55,6 +58,10 @@ class ETRecordReservedFileNames(StrEnum): INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME = "instruction_id_to_num_outs_map" REFERENCE_OUTPUTS = "reference_outputs" REPRESENTATIVE_INPUTS = "representative_inputs" + # Single JSON blob mapping each lowered module's key to the JSON-able value the + # backend attached via PreprocessResult._delegate_info_meta. Key format: + # f"{method}/{lowered_name}" (e.g. "forward/lowered_module_0"). + DELEGATE_INFO = "delegate_info" class ETRecord: @@ -75,6 +82,7 @@ def __init__( ] = None, _reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None, _representative_inputs: Optional[List[ProgramInput]] = None, + delegate_info: Optional[Dict[str, Any]] = None, ): """ Please do not construct an ETRecord object directly. @@ -126,6 +134,11 @@ def __init__( self._instruction_id_to_num_outs_map = _instruction_id_to_num_outs_map self._reference_outputs = _reference_outputs self._representative_inputs = _representative_inputs + # JSON-able delegate info collected from lowered modules' + # meta["_delegate_info_meta"]. Keyed by f"{method}/{lowered_name}" + # (e.g. "forward/lowered_module_0"); each value is whatever JSON-able + # object the backend attached (string, int, dict, ...). + self.delegate_info = delegate_info @property def _debug_handle_map(self): @@ -170,6 +183,7 @@ def save(self, path: Union[str, os.PathLike, BinaryIO, IO[bytes]]) -> None: self._write_identifier(etrecord_zip) self._save_programs(etrecord_zip) self._save_graph_map(etrecord_zip) + self._save_delegate_info(etrecord_zip) self._save_metadata(etrecord_zip) finally: etrecord_zip.close() @@ -207,6 +221,20 @@ def _save_graph_map(self, etrecord_zip: ZipFile) -> None: etrecord_zip, module_name, "forward", export_module ) + def _save_delegate_info(self, etrecord_zip: ZipFile) -> None: + """Save the delegate info map as a single JSON blob. + + ``delegate_info`` maps ``f"{method}/{lowered_name}"`` to whatever JSON-able + value the backend attached via ``PreprocessResult._delegate_info_meta``. + """ + if not self.delegate_info: + return + + etrecord_zip.writestr( + ETRecordReservedFileNames.DELEGATE_INFO, + json.dumps(self.delegate_info), + ) + def _save_metadata(self, etrecord_zip: ZipFile) -> None: """Save debug maps, reference outputs, and other metadata.""" if self._debug_handle_map is not None: @@ -265,6 +293,7 @@ def copy(self) -> "ETRecord": export_graph_id=self.export_graph_id, edge_dialect_program=self.edge_dialect_program, graph_map=self.graph_map, + delegate_info=self.delegate_info, # Explicitly exclude executorch-specific fields for clean separation _debug_handle_map=None, _delegate_map=None, @@ -273,6 +302,91 @@ def copy(self) -> "ETRecord": _representative_inputs=None, ) + def _maybe_extract_delegate_info( + self, + exported_programs: Optional[Dict[str, ExportedProgram]], + ) -> None: + """ + Extract delegate info from the lowered modules of the given exported + programs into this ETRecord, if it has not been extracted already. + + All programs recorded in an ETRecord correspond to the same model, so the + delegate info only needs to be captured once. This method is a no-op when + (a) this ETRecord already has ``delegate_info``, or (b) none of the given + programs carry any delegate info. It is called from the ``add_*`` methods + so extraction happens wherever the delegated graph first becomes available + (``add_edge_dialect_program`` for the ``to_edge_transform_and_lower`` flow, + ``add_executorch_program`` for the ``to_edge(...).to_backend(...)`` flow). + + Backends may attach an arbitrary JSON-able value to a lowered module via + ``PreprocessResult._delegate_info_meta`` (see + ``exir/backend/backend_details.py``). Every lowered submodule is walked + (recursing through control-flow submodules) and its value recorded under + the key ``f"{method}/{lowered_name}"`` (e.g. ``forward/lowered_module_0``) + so it can be serialized as JSON and reloaded from the ETRecord. Values + that are not JSON-serializable are skipped (with a warning) so a backend + attaching a non-JSON object cannot break ETRecord saving. + + Args: + exported_programs: Mapping of method name to the ``ExportedProgram`` + whose graph module may still carry the ``LoweredBackendModule``s. + """ + # (a) Already captured: nothing to do (all programs are one model). + if self.delegate_info is not None: + return + if not exported_programs: + return + + delegate_info: Dict[str, Any] = {} + for method_name, exported_program in exported_programs.items(): + self._collect_delegate_info( + exported_program.graph_module, method_name, "", delegate_info + ) + + # (b) None of the graphs carried delegate info: leave it unset so a later + # add_* call (with the delegated graph) can still capture it. + if delegate_info: + self.delegate_info = delegate_info + + def _collect_delegate_info( + self, + graph_module: torch.fx.GraphModule, + method_name: str, + name_prefix: str, + delegate_info: Dict[str, Any], + ) -> None: + """Recursively collect ``_delegate_info_meta`` values from lowered submodules.""" + for lowered_name, lowered_module, _ in get_lowered_submodules(graph_module): + meta = getattr(lowered_module, "meta", None) + if not isinstance(meta, dict) or "_delegate_info_meta" not in meta: + continue + value = meta["_delegate_info_meta"] + # Guard: only keep JSON-serializable values so ETRecord.save() (which + # json.dumps the whole delegate_info map) cannot fail. + try: + json.dumps(value) + except (TypeError, ValueError): + logging.warning( + "Skipping non-JSON-serializable _delegate_info_meta on lowered " + "module %s%s (method %s).", + name_prefix, + lowered_name, + method_name, + ) + continue + key = f"{method_name}/{name_prefix}{lowered_name}" + delegate_info[key] = value + + # Recurse into control-flow submodules (cond/map/scan), prefixing the + # submodule name so nested lowered-module names cannot collide. + for submodule_name, submodule, _ in get_control_flow_submodules(graph_module): + self._collect_delegate_info( + submodule, + method_name, + f"{name_prefix}{submodule_name}/", + delegate_info, + ) + def _save_exported_program( self, etrecord_zip: ZipFile, @@ -398,6 +512,13 @@ def add_executorch_program( self._reference_outputs = reference_outputs self._representative_inputs = representative_inputs + # For the to_edge(...).to_backend(...) flow the delegated graph is not + # available until here, so capture any delegate info from the executorch + # program's graphs (no-op if already captured at the edge stage). + self._maybe_extract_delegate_info( + _get_exported_programs_from_executorch_program(executorch_program) + ) + def add_exported_program( self, exported_program: Optional[Union[ExportedProgram, Dict[str, ExportedProgram]]], @@ -467,6 +588,12 @@ def add_edge_dialect_program( # Set the extracted data self.edge_dialect_program = processed_edge_dialect_program + # If this edge program is already delegated (e.g. the + # to_edge_transform_and_lower flow), capture any delegate info now. + self._maybe_extract_delegate_info( + _as_exported_program_dict(processed_edge_dialect_program) + ) + def update_representative_inputs( self, representative_inputs: Union[List[ProgramInput], BundledProgram], @@ -668,6 +795,48 @@ def _add_module_to_graph_map( raise RuntimeError(f"Unsupported graph module type. {type(export_module)}") +def _as_exported_program_dict( + program: Union[ExportedProgram, Dict[str, ExportedProgram], None] +) -> Optional[Dict[str, ExportedProgram]]: + """Normalize a single ExportedProgram or a method->EP dict to a dict. + + A bare ExportedProgram is treated as the "forward" method. + """ + if program is None: + return None + if isinstance(program, ExportedProgram): + return {"forward": program} + return program + + +def _get_exported_programs_from_executorch_program( + executorch_program: Union[ + ExecutorchProgram, ExecutorchProgramManager, BundledProgram + ] +) -> Optional[Dict[str, ExportedProgram]]: + """Return the method->ExportedProgram map backing an executorch program. + + These graphs still carry the ``LoweredBackendModule``s (and their ``.meta``), + so they are what delegate-info extraction walks. Supports the manager, the + single-program, and the bundled-program forms; returns None if the graphs + cannot be reached. + """ + if isinstance(executorch_program, BundledProgram): + return _get_exported_programs_from_executorch_program( + executorch_program.executorch_program + ) + if isinstance(executorch_program, ExecutorchProgramManager): + return { + method: executorch_program.exported_program(method) + for method in executorch_program.methods + } + # ExecutorchProgram exposes a single (forward) exported program. + exported_program = getattr(executorch_program, "exported_program", None) + if isinstance(exported_program, ExportedProgram): + return {"forward": exported_program} + return None + + def _process_edge_dialect_program( edge_dialect_program: Union[EdgeProgramManager, ExirExportedProgram] ) -> Union[ExportedProgram, Dict[str, ExportedProgram]]: @@ -788,6 +957,8 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 serialized_constants_files = set() serialized_example_inputs_files = set() + delegate_info: Optional[Dict[str, Any]] = None + edge_dialect_prefix = f"{ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM}/" for entry in file_list: @@ -823,6 +994,10 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 ): # New format: edge_dialect_exported_program/method_name serialized_edge_dialect_program_files.add(entry) + elif entry == ETRecordReservedFileNames.DELEGATE_INFO: + delegate_info = json.loads( + etrecord_zip.read(ETRecordReservedFileNames.DELEGATE_INFO) + ) elif entry == ETRecordReservedFileNames.EXPORTED_PROGRAM: serialized_artifact = SerializedArtifact( etrecord_zip.read(ETRecordReservedFileNames.EXPORTED_PROGRAM), @@ -920,6 +1095,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 exported_program=exported_program, edge_dialect_program=edge_dialect_program, graph_map=graph_map, + delegate_info=delegate_info, _debug_handle_map=debug_handle_map, _delegate_map=delegate_map, _instruction_id_to_num_outs_map=instruction_id_to_num_outs_map, diff --git a/devtools/etrecord/tests/BUCK b/devtools/etrecord/tests/BUCK index f5d031d40fb..0ab86cdd86b 100644 --- a/devtools/etrecord/tests/BUCK +++ b/devtools/etrecord/tests/BUCK @@ -17,6 +17,8 @@ fbcode_target(_kind = runtime.python_library, srcs = ["etrecord_test.py"], deps = [ "//caffe2:torch", + "//executorch/backends/example:example_partitioner", + "//executorch/backends/example:example_quantizer", "//executorch/devtools/bundled_program:config", "//executorch/devtools/bundled_program:core", "//executorch/devtools/etrecord:etrecord", diff --git a/devtools/etrecord/tests/etrecord_test.py b/devtools/etrecord/tests/etrecord_test.py index 127f2cd1d00..f0065a01b36 100644 --- a/devtools/etrecord/tests/etrecord_test.py +++ b/devtools/etrecord/tests/etrecord_test.py @@ -14,6 +14,8 @@ import executorch.exir.tests.models as models import torch from executorch import exir +from executorch.backends.example.example_partitioner import ExamplePartitioner +from executorch.backends.example.example_quantizer import ExampleQuantizer from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite from executorch.devtools.bundled_program.core import BundledProgram @@ -30,6 +32,8 @@ from executorch.export import export as etexport, ExportRecipe, StageType from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + # TODO : T154728484 Add test cases to cover multiple entry points class TestETRecord(unittest.TestCase): @@ -1816,6 +1820,75 @@ def test_multi_method_etrecord_generation(self): self.assertIsNotNone(parsed_etrecord._debug_handle_map) self.assertIsNotNone(parsed_etrecord._delegate_map) + def test_delegate_info_save_and_parse(self): + """Test that a backend's _delegate_info_meta is logged into the ETRecord and + can be decoded after parse. + + Lowers a quantized Conv2d model to the Example backend (which attaches a + sha256 fingerprint of its processed_bytes via + PreprocessResult._delegate_info_meta), then saves + parses the ETRecord and + checks the delegate_info round-trips. + """ + + class Conv2dModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv2d = torch.nn.Conv2d(16, 33, 3) + + def forward(self, arg): + return self.conv2d(arg) + + example_inputs = (torch.randn(20, 16, 50, 100),) + edge_compile_config = exir.EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ) + + m = Conv2dModule().eval() + m = torch.export.export(m, example_inputs, strict=True).module() + m = prepare_pt2e(m, ExampleQuantizer()) + m(*example_inputs) + m = convert_pt2e(m) + + edge_manager = to_edge( + export(m, example_inputs, strict=True), + compile_config=edge_compile_config, + generate_etrecord=True, + ) + lowered = edge_manager.to_backend(ExamplePartitioner()) + et_manager = lowered.to_executorch() + + etrecord = et_manager.get_etrecord() + # delegate_info is populated on the ExecuTorch-stage ETRecord. + self.assertIsNotNone(etrecord.delegate_info) + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_delegate_info.bin" + etrecord.save(etrecord_path) + parsed_etrecord = parse_etrecord(etrecord_path) + + # delegate_info decodes back as plain JSON. + self.assertIsNotNone(parsed_etrecord.delegate_info) + self.assertEqual( + parsed_etrecord.delegate_info, + json.loads(json.dumps(etrecord.delegate_info)), + ) + + # There is exactly one lowered module, keyed by "/". + self.assertEqual(len(parsed_etrecord.delegate_info), 1) + key = next(iter(parsed_etrecord.delegate_info)) + self.assertTrue( + key.startswith("forward/"), + f"unexpected delegate_info key: {key}", + ) + + # The logged value is the Example backend's sha256 fingerprint (non-static + # hex string of the processed blob). + value = parsed_etrecord.delegate_info[key] + self.assertIsInstance(value, str) + self.assertEqual(len(value), 64) + int(value, 16) # valid hex, raises ValueError otherwise + def test_edge_after_transform_graph_capture(self): """Test that to_edge_transform_and_lower with transform_passes captures the after-transform graph.