# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 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 logging import time from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import torch from executorch.exir import EdgeProgramManager from executorch.exir._warnings import experimental from executorch.exir.program import ExecutorchProgramManager from executorch.exir.schema import Program from executorch.extension.export_util.utils import save_pte_program from tabulate import tabulate from torch import nn from torch.export import ExportedProgram from torch.fx import GraphModule from .recipe import ExportRecipe, LoweringRecipe, QuantizationRecipe from .stages import ( EdgeProgramManagerTransformStage, EdgeTransformAndLowerStage, ExecutorchStage, PipelineArtifact, QuantizeStage, SourceTransformStage, Stage, ToBackendStage, ToEdgeStage, TorchExportStage, ) from .types import StageType @experimental( "This API and all of its related functionality such as ExportSession and ExportRecipe are experimental." ) def export( model: Union[ nn.Module, Dict[str, nn.Module], GraphModule, Dict[str, GraphModule], ExportedProgram, Dict[str, ExportedProgram], str, ], example_inputs: Optional[ Union[ List[tuple[torch.Tensor, ...]], Dict[str, List[tuple[torch.Tensor, ...]]], ] ] = None, export_recipe: ExportRecipe = None, 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, ) -> "ExportSession": """ Create and configure an ExportSession with the given parameters. This function provides a convenient way to create an ExportSession and optionally run the export process in one step. Args: model: The PyTorch model(s) to export. Can be: - nn.Module or Dict[str, nn.Module]: Eager PyTorch model(s) - GraphModule or Dict[str, GraphModule]: Quantized model(s) (e.g., from prepare/convert) - ExportedProgram or Dict[str, ExportedProgram]: Already exported model(s) - str: Path to load an ExportedProgram from disk example_inputs: Example inputs for the model(s), either a list of input tuples or a dictionary mapping method names to lists of input tuples. First sample (index 0) is used for torch.export.export() to export the model. All samples are used as calibration dataset in PT2E Quantize stage. Optional when model is ExportedProgram (not needed). export_recipe: Contains the configuration for the export process name: Optional name for the export dynamic_shapes: Optional dynamic shape specifications constant_methods: Optional dictionary of constant methods artifact_dir: Optional directory to store artifacts generate_etrecord: Optional flag to generate an etrecord Returns: A configured ExportSession instance with the export process completed if requested """ session = ExportSession( model=model, example_inputs=example_inputs, export_recipe=export_recipe, name=name, dynamic_shapes=dynamic_shapes, constant_methods=constant_methods, artifact_dir=artifact_dir, generate_etrecord=generate_etrecord, ) session.export() return session @experimental( "This API and all of its related functionality such as ExportSession and ExportRecipe are experimental." ) class ExportSession: """ Manages the export process for ExecuTorch models. This class handles the export process through a pipeline of stages: 1. (Optional) Quantize - Apply post-training quantization to the model 2. Export - Export PyTorch model to ExportedProgram 3. EdgeTransformAndLower - Transform and lower to EdgeProgramManager 4. Executorch - Convert to ExecutorchProgramManager for final execution """ def __init__( self, model: Union[ nn.Module, Dict[str, nn.Module], GraphModule, Dict[str, GraphModule], ExportedProgram, Dict[str, ExportedProgram], str, ], example_inputs: Optional[ Union[ List[tuple[torch.Tensor, ...]], Dict[str, List[tuple[torch.Tensor, ...]]], ] ] = None, export_recipe: ExportRecipe = None, 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: Optional[bool] = False, ) -> None: """ Initialize the ExportSession with model, inputs, and recipe. Args: model: The PyTorch model(s) to export. Can be: - nn.Module or Dict[str, nn.Module]: Eager PyTorch model(s) - GraphModule or Dict[str, GraphModule]: Quantized model(s) - ExportedProgram or Dict[str, ExportedProgram]: Already exported model(s) - str: Path to load an ExportedProgram from disk example_inputs: Example inputs for the model(s), either a list of input tuples or a dictionary mapping method names to lists of input tuples. First sample (index 0) is used for torch.export.export() to export the model. All samples are used as calibration dataset in PT2E Quantize stage, Optional when model is ExportedProgram (not needed). export_recipe: Contains the configuration for the export process name: Optional name for the export dynamic_shapes: Optional dynamic shape specifications constant_methods: Optional dictionary of constant methods artifact_dir: Optional directory to store artifacts generate_etrecord: Optional flag to generate an etrecord """ # Load model from file if string path provided if isinstance(model, str): model = torch.export.load(model) logging.info(f"Loaded ExportedProgram from {model}") # Detect input model type to determine which stages to skip self._input_model_type = self._detect_model_type(model) # Standardize model to dictionary format self._model = model if isinstance(model, dict) else {"forward": model} # Validate and standardize example_inputs to dictionary format # example_inputs not required for ExportedProgram input if self._input_model_type == "ExportedProgram": self._example_inputs = example_inputs or {} if isinstance(self._example_inputs, list): self._example_inputs = {"forward": self._example_inputs} else: # For nn.Module and GraphModule, example_inputs are required if example_inputs is None: raise ValueError( f"example_inputs are required when model is {self._input_model_type}. " f"Only ExportedProgram inputs can omit example_inputs." ) self._example_inputs = ( example_inputs if isinstance(example_inputs, dict) else {"forward": example_inputs} ) # Standardize dynamic_shapes to dictionary format self._dynamic_shapes = {} if dynamic_shapes is not None: if isinstance(dynamic_shapes, dict): self._dynamic_shapes = dynamic_shapes else: self._dynamic_shapes = {"forward": dynamic_shapes} self._export_recipe = export_recipe self._quant_recipe: Optional[QuantizationRecipe] = ( self._export_recipe.quantization_recipe ) self._lowering_recipe: Optional[LoweringRecipe] = ( self._export_recipe.lowering_recipe ) # Stages to run self._pipeline_stages = ( self._export_recipe.pipeline_stages or self._get_default_pipeline() ) # Stage registry: map of StageType to Stage instances self._stage_registry: Dict[StageType, Stage] = self._build_stages( self._pipeline_stages ) # Intialize run context self._run_context: Dict[str, Any] = { "example_inputs": self._example_inputs, "dynamic_shapes": self._dynamic_shapes, "constant_methods": constant_methods, "export_recipe": self._export_recipe, "session_name": name, "artifact_dir": artifact_dir, "generate_etrecord": generate_etrecord, } self._stage_to_artifacts: Dict[StageType, PipelineArtifact] = {} def _detect_model_type( self, model: Union[nn.Module, GraphModule, ExportedProgram, Dict] ) -> str: """ Detect the type of input model. Args: model: Input model in various formats Returns: String indicating the model type: "nn.Module", "GraphModule", or "ExportedProgram" """ # Handle dict (multi-method) - check first value if isinstance(model, dict): first_value = next(iter(model.values())) return self._detect_model_type(first_value) # Detect single model type if isinstance(model, ExportedProgram): return "ExportedProgram" elif isinstance(model, GraphModule): return "GraphModule" elif isinstance(model, nn.Module): return "nn.Module" else: raise TypeError(f"Unsupported model type: {type(model)}") def _get_default_pipeline(self) -> List[StageType]: """ Get default pipeline stages based on input model type. Returns: List of stages appropriate for the input model type """ stages = [] # Add quantization stages only for eager nn.Module if self._input_model_type == "nn.Module": stages.extend( [ StageType.SOURCE_TRANSFORM, # Optional stage, returns original model if quant recipe is invalid StageType.QUANTIZE, # Optional stage, returns original model if quant recipe is invalid ] ) # Add torch export stage if not already exported if self._input_model_type != "ExportedProgram": stages.append(StageType.TORCH_EXPORT) # Always include edge and executorch stages stages.extend( [ StageType.TO_EDGE_TRANSFORM_AND_LOWER, StageType.TO_EXECUTORCH, ] ) return stages def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]: """Build the stage registry from the given stages.""" stage_registry: Dict[StageType, Stage] = {} stage = None for stage_type in stages or self._get_default_pipeline(): if stage_type == StageType.SOURCE_TRANSFORM: stage = SourceTransformStage(self._quant_recipe) elif stage_type == StageType.QUANTIZE: stage = QuantizeStage(self._quant_recipe) elif stage_type == StageType.TORCH_EXPORT: aten_transform_passes = None if self._export_recipe.aten_transform_passes is not None: aten_transform_passes = list( self._export_recipe.aten_transform_passes ) stage = TorchExportStage( aten_transform_passes, strict=self._export_recipe.strict ) elif stage_type == StageType.TO_EDGE_TRANSFORM_AND_LOWER: stage = EdgeTransformAndLowerStage.from_recipe(self._lowering_recipe) elif stage_type == StageType.TO_EDGE: stage = ToEdgeStage.from_recipe(self._lowering_recipe) elif stage_type == StageType.EDGE_PROGRAM_MANAGER_TRANSFORM: stage = EdgeProgramManagerTransformStage.from_recipe( self._lowering_recipe ) elif stage_type == StageType.TO_BACKEND: stage = ToBackendStage.from_recipe(self._lowering_recipe) elif stage_type == StageType.TO_EXECUTORCH: stage = ExecutorchStage(self._export_recipe.executorch_backend_config) else: logging.info( f"{stage_type} is unknown, you have to register it before executing export()" ) if stage: stage_registry[stage_type] = stage return stage_registry def register_stage(self, stage_type: StageType, stage: Stage) -> None: """ Register a new stage or override an existing stage implementation. Args: stage_type: The type of stage to register stage: The stage instance to register """ self._stage_registry[stage_type] = stage def get_registered_stage(self, stage_type: StageType) -> Optional[Stage]: """ Get a registered stage by its type. Args: stage_type: The type of stage to retrieve Returns: The registered stage instance, or None if not found """ return self._stage_registry.get(stage_type) def get_all_registered_stages(self) -> Dict[StageType, Stage]: """ Get all registered stages. Returns: Dictionary mapping stage types to stage instances """ return self._stage_registry def _validate_pipeline_sequence( self, stages: List[StageType], ) -> None: if not stages: raise ValueError("Pipeline stages cannot be empty") # Validate pipeline compatibility with input model type if self._input_model_type == "GraphModule": # GraphModule input should not run quantization stages incompatible_stages = {StageType.SOURCE_TRANSFORM, StageType.QUANTIZE} found_incompatible = set(stages) & incompatible_stages if found_incompatible: stage_names = ", ".join(s.name for s in found_incompatible) raise ValueError( f"Cannot run {stage_names} stage(s) with GraphModule input. " f"GraphModule is already quantized. " f"Remove {stage_names} from pipeline_stages or use nn.Module input." ) elif self._input_model_type == "ExportedProgram": # ExportedProgram input should not run quantization or torch export stages incompatible_stages = { StageType.SOURCE_TRANSFORM, StageType.QUANTIZE, StageType.TORCH_EXPORT, } found_incompatible = set(stages) & incompatible_stages if found_incompatible: stage_names = ", ".join(s.name for s in found_incompatible) raise ValueError( f"Cannot run {stage_names} stage(s) with ExportedProgram input. " f"ExportedProgram is already exported. " f"Remove {stage_names} from pipeline_stages or use nn.Module/GraphModule input." ) # Validate that the first stage can start a pipeline first_stage = stages[0] first_stage_instance = self._stage_registry.get(first_stage) if first_stage_instance is None: raise ValueError( f"Stage {first_stage} not found in registry, register it using session.register_stage()" ) if not first_stage_instance.can_start_pipeline: raise ValueError(f"Stage {first_stage} cannot start a pipeline. ") # Validate stage transitions for i in range(1, len(stages)): current_stage = stages[i] previous_stage = stages[i - 1] # Get the stage instance to check its valid predecessors stage_instance = self._stage_registry.get(current_stage) if stage_instance is None: raise ValueError( f"Stage {current_stage} not found in registry, , register it using session.register_stage()" ) valid_predecessors = stage_instance.valid_predecessor_stages # Check if the previous stage is valid for the current stage if valid_predecessors and previous_stage not in valid_predecessors: raise ValueError( f"Invalid transition from {previous_stage} to {current_stage}. " f"Valid predecessors for {current_stage}: {valid_predecessors}" ) def _run_pipeline(self) -> None: # Validate if given stage sequence is valid self._validate_pipeline_sequence( stages=self._pipeline_stages, ) current_artifact = PipelineArtifact(data=self._model, context=self._run_context) # Execute stages from registry in the order specified by pipeline_stages for stage_type in self._pipeline_stages: stage = self._stage_registry.get(stage_type) if stage is None: raise ValueError(f"Stage {stage_type} not found in registry") logging.info(f"Executing stage: {stage_type}") start = time.perf_counter() stage.run(current_artifact) elapsed = (time.perf_counter() - start) * 1000 current_artifact = stage.get_artifacts() current_artifact.add_context("duration_ms", int(elapsed)) logging.info(f"Stage {stage_type} execution done") self._stage_to_artifacts[stage_type] = current_artifact def export(self) -> None: """ Execute the full export process. This method orchestrates the export process with optional quantization: 1. (Optional) Apply quantization to the model 2. Export the PyTorch model to ExportedProgram 3. Transform and lower to EdgeProgramManager 4. Convert to ExecutorchProgramManager """ # Run the pipeline from the beginning self._run_pipeline() def get_stage_artifacts(self) -> Dict[StageType, PipelineArtifact]: return self._stage_to_artifacts def get_exported_program(self, method_name: str = "forward") -> ExportedProgram: """ Get the ExportedProgram for a specific method after torch export. Args: method_name: Name of the method to get exported program for, defaults to "forward" Returns: The ExportedProgram for the specified method Raises: RuntimeError: If torch export stage has not been run KeyError: If the method name is not found in exported programs """ artifact = self._stage_to_artifacts.get(StageType.TORCH_EXPORT) if artifact is None or artifact.data is None: raise RuntimeError( "Exported program is not available. Run Torch Export Stage first." ) exported_programs = artifact.data if method_name not in exported_programs: raise KeyError( f"Method name '{method_name}' not found in exported programs. " f"Available methods: {list(exported_programs.keys())}" ) return exported_programs[method_name] def get_edge_program_manager(self) -> "EdgeProgramManager": """ Get the EdgeProgramManager after edge lowering stages. This method checks multiple stages in order of preference: 1. TO_EDGE_TRANSFORM_AND_LOWER (combined stage) 2. TO_BACKEND (separate stage with backend delegation) 3. EDGE_PROGRAM_MANAGER_TRANSFORM (separate stage after TO_EDGE) 4. TO_EDGE (separate stage without backend delegation) Returns: The EdgeProgramManager Raises: RuntimeError: If no edge stage has been run """ # Check stages in order of preference for stage_type in [ StageType.TO_EDGE_TRANSFORM_AND_LOWER, StageType.TO_BACKEND, StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, StageType.TO_EDGE, ]: artifact = self._stage_to_artifacts.get(stage_type) if artifact is not None and artifact.data is not None: logging.info(f"Returning edge program manager from stage {stage_type}") return artifact.data raise RuntimeError( "Edge program manager is not available. " "Run one of the edge stages first: TO_EDGE_TRANSFORM_AND_LOWER, TO_EDGE, EDGE_PROGRAM_MANAGER_TRANSFORM, or TO_BACKEND." ) def get_executorch_program(self) -> Program: """ Get the ExecutorchProgram from the ExecutorchProgramManager. Returns: The ExecutorchProgram Raises: RuntimeError: If the executorch program manager is not initialized """ return self.get_executorch_program_manager().executorch_program def get_executorch_program_manager(self) -> ExecutorchProgramManager: """ Get the ExecutorchProgramManager. Returns: The ExecutorchProgramManager Raises: RuntimeError: If the executorch program manager is not initialized """ artifact = self._stage_to_artifacts.get(StageType.TO_EXECUTORCH) if artifact is None or artifact.data is None: raise RuntimeError( "Executorch program manager is not initialized. Run Executorch Stage first." ) return artifact.data def get_pte_buffer(self) -> bytes: """ Get the PTE buffer as bytes. Returns: The PTE buffer as bytes Raises: RuntimeError: If the executorch program manager is not initialized """ return self.get_executorch_program_manager().buffer def save_to_pte(self, output_name: str) -> None: """ Save the model to a .pte file. Args: output_name (Optional[str]): The name of the .pte file. """ assert output_name, "Need a valid output name" save_pte_program(self.get_executorch_program_manager(), output_name) def get_example_input( self, method_name: str = "forward" ) -> Tuple[torch.Tensor, ...]: """ Get the example input for a specific method. Args: method_name: Name of the method to get example input for, defaults to "forward" Returns: Tuple of tensors representing the example input Raises: KeyError: If the method name is not found in example inputs ValueError: If the example inputs list is empty """ if method_name not in self._example_inputs: raise KeyError(f"Method name '{method_name}' not found in example inputs") # Access the first element of the list for this method example_inputs_list = self._example_inputs[method_name] if not example_inputs_list: raise ValueError(f"Example inputs list for method {method_name} is empty") # The original code expects this to be a tuple of tensors return self._example_inputs[method_name][0] def run_method( self, method_name: str = "forward", example_inputs: Optional[Tuple[torch.Tensor, ...]] = None, ) -> Sequence[Any]: """ Run a specific method with the given inputs. Args: method_name: Name of the method to run, defaults to "forward" example_inputs: Optional inputs to use, defaults to the example inputs Returns: The outputs of the method execution Raises: RuntimeError: If the method cannot be loaded """ # Lazy import to avoid forcing portable_lib dependency at module load time. try: from executorch.runtime import Runtime, Verification except ModuleNotFoundError as e: raise ModuleNotFoundError( "executorch.runtime is not available. " "In OSS: Please ensure executorch is properly installed via pip. " "In fbcode: Please add //executorch/runtime:runtime to your dependencies." ) from e et_runtime = Runtime.get() program = et_runtime.load_program( self.get_pte_buffer(), verification=Verification.Minimal ) forward = program.load_method(method_name) if forward is None: raise RuntimeError( f"Failed to load method '{method_name}' from the program" ) if example_inputs is None: example_inputs = self.get_example_input(method_name) return forward.execute(example_inputs) def print_delegation_info(self) -> None: """ Print delegation information for the exported program. """ lowering_stage = list( set(self._pipeline_stages) & {StageType.TO_EDGE_TRANSFORM_AND_LOWER, StageType.TO_BACKEND} ) if not lowering_stage: RuntimeError( "No delegation info available, atleast one of the lowering stages should be present" ) stage_artifact = self._stage_to_artifacts.get(lowering_stage[0]) if stage_artifact is None: RuntimeError("No delegation info available, run the lowering stage first") # pyre-ignore delegation_info = stage_artifact.get_context("delegation_info", None) if delegation_info: print(delegation_info.get_summary()) df = delegation_info.get_operator_delegation_dataframe() print(tabulate(df, headers="keys", tablefmt="fancy_grid")) else: print("No delegation info available") # Use Any instead of ETRecord as return type to avoid static dependency on etrecord def get_etrecord(self) -> Any: """ Get the etrecord from the ExecuTorchProgramManager. Returns: The etrecord in the ExecuTorchProgramManager Raises: RuntimeError: If the ExecuTorchManager is unavailable, or etrecord is not available in the ExecuTorchProgramManager """ return self.get_executorch_program_manager().get_etrecord()