[RFC] Export Stream ETRecord Generation Support

[Gasoonjia]

Overview

Currently, ETRecord generation is separated from the export flow, which increases ramping up costs for users and makes it impossible for high-level export APIs like to_edge_transform_and_lower() and executorch.export() to generate ETRecord due to hiding the essential information under-the-neath. This document proposes integrating ETRecord generation directly into the export flow with configurable support.

Goals

1. Add configurable support for generating ETRecord along with export flow
2. Zero impact to export flow when ETRecord generation is turned off
3. No BC breaking: current ETRecord generation should keep working
4. Seamless integration with existing high-level export APIs

Current State Analysis

Current Export APIs

• executorch.export() - High-level export function returning ExportSession
• to_edge() - Converts ExportedProgram to EdgeProgramManager
• to_edge_transform_and_lower() - Combines edge conversion with partitioning/lowering

Current ETRecord Generation

• Standalone generate_etrecord() function in /data/users/gasoonjia/fbsource/fbcode/executorch/devtools/etrecord/_etrecord.py
• Requires separate calls after export process
• Takes EdgeProgramManager, ExecutorchProgram/ExecutorchProgramManager/BundledProgram, and optional ExportedProgram

Proposed Solution

1. Add generate_etrecord Parameter to Export APIs

Target APIs After Changes:

executorch.export()

def export(
    model: Union[nn.Module, Dict[str, nn.Module]],
    example_inputs: Union[
        List[tuple[torch.Tensor, ...]], Dict[str, List[tuple[torch.Tensor, ...]]]
    ],
    export_recipe: ExportRecipe,
    name: Optional[str] = None,
    dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None,
    constant_methods: Optional[Union[Dict[str, Callable]]] = None,
    artifact_dir: Optional[str] = None,
    generate_etrecord: bool = False,  # NEW PARAMETER
) -> "ExportSession":

to_edge()

def to_edge(
    programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
    constant_methods: Optional[Dict[str, Any]] = None,
    compile_config: Optional[EdgeCompileConfig] = None,
    generate_etrecord: bool = False,  # NEW PARAMETER
) -> "EdgeProgramManager":

to_edge_transform_and_lower()

def to_edge_transform_and_lower(
    programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
    transform_passes: Optional[
        Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
    ] = None,
    partitioner: Optional[...] = None,
    # ... other existing parameters
    generate_etrecord: bool = False,  # NEW PARAMETER
) -> "EdgeProgramManager":

2. Add get_etrecord() Method to Return Objects

ExportSession Enhancement

class ExportSession:
    def __init__(self, ...):
        # ... existing code
        self._etrecord: Optional[ETRecord] = None
        self._generate_etrecord: bool = generate_etrecord
    
    def get_etrecord(self) -> Optional[ETRecord]:
        """
        Get the generated ETRecord if etrecord generation was enabled.
        
        Returns:
            ETRecord object if generation was enabled and export completed,
            
        Raises:
            RuntimeError: If export() hasn't been called yet, or ETRecord object was not generated. 
        """

ExecutorchProgramManager Enhancement

class ExecutorchProgramManager:
    def __init__(self, ...):
        # ... existing code
        self._etrecord: Optional[ETRecord] = None
    
    def get_etrecord(self) -> Optional[ETRecord]:
        """
        Get the generated ETRecord if etrecord generation was enabled.
        
        Returns:
            ETRecord object if generation was enabled, None otherwise
            
        Raises:
            RuntimeError: if ETRecord object was not generated. 
        """

3. ETRecord Dump Functionality

ETRecord Class Enhancement

class ETRecord:
    # ... existing fields
    
    def save(self, target_path: Union[str, os.PathLike]) -> None:
        """
        Dump the ETRecord to disk at the specified path.
        
        Args:
            target_path: Path where the ETRecord file will be saved
            
        Raises:
            RuntimeError: If ETRecord data is incomplete or invalid
        """
        # Implementation will serialize the ETRecord object to the target path
        # Similar to current generate_etrecord() but working with existing ETRecord object

Target Usage Flows

Flow 1: High-level export() API with ETRecord

import executorch
from executorch.export import ExportRecipe

# Create model and inputs
model = MyModel()
example_inputs = (torch.randn(1, 3, 224, 224),)
recipe = ExportRecipe(name="my_recipe")

# Export with ETRecord generation enabled
session = executorch.export(
    model=model,
    example_inputs=[example_inputs],
    export_recipe=recipe,
    generate_etrecord=True  # Enable ETRecord generation
)

# Get the ETRecord
etrecord = session.get_etrecord()
etrecord.save("model_debug.etrecord")

# Continue with normal export flow
session.save_pte_file("model.pte")

Flow 2: to_edge_transform_and_lower() with ETRecord

from executorch.exir.program import to_edge_transform_and_lower
from torch.export import export

# Export model first
exported_program = export(model, example_inputs)

# Transform and lower with ETRecord generation
edge_manager = to_edge_transform_and_lower(
    exported_program,
    partitioner=[MyPartitioner()],
    generate_etrecord=True  # Enable ETRecord generation
)

et_manager = edge_manager.to_executorch()

# Get ETRecord from edge manager
etrecord = et_manager.get_etrecord()
etrecord.save("edge_debug.etrecord")

# Continue to ExecutorchProgramManager
executorch_manager = edge_manager.to_executorch()

Flow 3: to_edge() with ETRecord

from executorch.exir.program import to_edge
from torch.export import export

# Export model first
exported_program = export(model, example_inputs)

# Convert to edge with ETRecord generation
edge_manager = to_edge(
    exported_program,
    generate_etrecord=True  # Enable ETRecord generation
)

edge_manager = edge_manager.to_backend()
et_manager = edge_manager.to_executorch()

# Get ETRecord
etrecord = et_manager.get_etrecord()
etrecord.save("edge_only_debug.etrecord")

Flow 4: Backward Compatibility - Existing ETRecord Generation

# Existing code continues to work unchanged
from executorch.devtools import generate_etrecord

# ... existing export flow
aten_dialect = export(model, example_inputs, strict=True)

# Convert to edge
edge_program = to_edge(
    aten_dialect, 
    compile_config=EdgeCompileConfig(_check_ir_validity=True)
)

# Create copy for ETRecord (needed because to_executorch modifies in-place)
edge_program_copy = copy.deepcopy(edge_program)

# Convert to ExecutorchProgramManager
executorch_program = edge_program_copy.to_executorch()

# Generate ETRecord separately (existing approach)
generate_etrecord(
    "debug.etrecord",
    edge_manager,
    executorch_manager,
    exported_program
)

Benefits

1. Reduced Complexity: Users can generate ETRecord in a single export call, without manage separate ETRecord generation calls
2. Better Integration: ETRecord generation becomes an optional part of the standard export workflow
3. Backward Compatibility: Existing code continues to work without changes

This design ensures a smooth transition while providing the requested functionality with zero impact on existing workflows.

[GregoryComer]

Looks great. I have a few small thoughts:

1. On the get_etrecord API, it might make sense to just raise an exception if generate_etrecord wasn't set as opposed to returning None. That can give a clear error message. You could make an argument that the user calling get_etrecord expects to get an ETRecord back, and not getting one is an error. Thoughts?
2. What do you think about renaming ETRecord.dump to ETRecord.save? That might be slightly more intuitive.

[Gasoonjia]

Thanks for sharing your thoughts. Both of them are good for me.

[larryliu0820]

generate_etrecord can it be part of the config?

[Gasoonjia]

I don't want to make it be part of compile_config for two reasons:

1. current compile_config contains the configuration related to export progress. But etrecord does zero impact on export. if we put generate_etrecord into compile_config it may mislead our user.
2. in the future i'd love to expand the generate_etrecord into a config argument (maybe debug_config) or something else to support more debug-related functionality in aot stage, like numerical discrepancy detection, delegate info debugging, etc.