From 93efd54d0b430fd91e1cfae352498e75f5c5826a Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Tue, 28 Jul 2026 13:54:37 -0700 Subject: [PATCH] feat(executorch): expose composable Edge export API ## What this adds `torch_tensorrt.save(..., output_format="executorch")` turns a compiled model into an ExecuTorch `.pte` file in one step. That is the easiest path and it stays the recommendation. Some workflows need to stop earlier. You may want to look at the delegated Edge graph, run your own Edge transforms, send operations TensorRT cannot handle to another backend, keep separate methods such as `prefill` and `decode`, add constant methods, or pick the final ExecuTorch configuration yourself. This adds `torch_tensorrt.executorch.export()`. It returns ExecuTorch's standard `EdgeProgramManager`, which is the supported place to inspect and customize an Edge program before you call `to_executorch()`. ```python import torch_tensorrt.executorch edge = torch_tensorrt.executorch.export(trt_module) # Inspect or transform the Edge program here. print(edge.exported_program().graph) program = edge.to_executorch() with open("model.pte", "wb") as output: program.write_to_file(output) # Advanced callers must also persist external tensor data when present. program.write_tensor_data_to_file(".") ``` `save()` now calls the same implementation, so the simple and advanced paths share one lowering path instead of two. ## Which one should I use? Use `save()` when you just want a `.pte`. Use `torch_tensorrt.executorch.export()` when you need control before the `.pte` is created. ## Example: TensorRT with a fallback backend TensorRT always gets the first chance to claim its prebuilt engine nodes. Any partitioners you pass run afterward and can claim operations TensorRT does not support. ```python edge = torch_tensorrt.executorch.export( trt_module, partitioners=[cuda_partitioner], ) program = edge.to_executorch() ``` ## Example: multiple methods ```python edge = torch_tensorrt.executorch.export( { "prefill": prefill_program, "decode": decode_program, }, compile_specs={ "prefill": prefill_specs, "decode": decode_specs, }, ) ``` Method names are preserved, and each method can have its own partitioner and compile-spec pipeline. This does not make mutable state shared between methods. ## Accepted inputs - a TensorRT-compiled `torch.fx.GraphModule` - an engine-bearing `torch.export.ExportedProgram` - a mapping from method names to independent `ExportedProgram` objects A plain `torch.nn.Module` must be compiled with Torch-TensorRT first. ## Safety and memory behavior Rewriting engine calls changes the exported graph. To avoid damaging the program you passed in, export first stages its own copy of the graph structure, signatures, state containers, constants, and metadata. If rewriting or lowering fails, your original program is left alone. Tensor and TensorRT engine payloads are treated as read-only and shared with that staged copy, so a multi-gigabyte engine is not duplicated. Custom transform passes must not modify these shared payload objects. Metadata that describes symbolic shapes is a special case. It points back to the live shape environment that the exported program is guarded on, so it is shared rather than copied. Copying it would both fail (the shape environment holds fake tensors that cannot be duplicated) and detach the copy from the symbols the graph depends on. This is what keeps dynamic-shape models working. When a lifted TensorRT engine is replaced, the old placeholder, its graph-signature entry, and its constant are removed. Several calls that share one engine reuse a single materialized `uint8` payload buffer. ## Compatibility - Existing one-step `save(..., output_format="executorch")` behavior is preserved, including its rejection of non-list `partitioners` and `compile_specs`. - `"executorch"` is now included in the `output_format` type hint. It was already accepted at run time, so this only fixes the annotation. - `.pte` serialization and external `.ptd` persistence are preserved. - TensorRT remains the first partitioner. - Shared dynamic dimensions declared with `Input(shared_dims=...)` remain shared. - Zero-engine methods are allowed. A later partitioner may claim them, or portable operators may remain undelegated. ## Testing Tested on a CUDA GPU with TensorRT and ExecuTorch installed, building Torch-TensorRT from source so the C++ runtime was active. Verified by compiling a real model with TensorRT and then exporting it: - static-shape `export()` produces a loadable `.pte` - dynamic-shape `export()` produces a loadable `.pte`, matching the output of the previous one-step save path - the source `GraphModule` still holds its original engine nodes afterward - engine payloads are shared, not re-serialized - `save(..., output_format="executorch")` still writes a `.pte` - `save()` still raises `TypeError` for a non-list `partitioners` or `compile_specs` Also ran the ExecuTorch test directory (76 passed). Two GPU tests in `test_cuda_partitioner_composition.py` fail, and they fail the same way with this change reverted, so they are not caused by it. Static checks: formatting, import ordering, lint, and type checking on the changed files. --- .../runtime_performance/saving_models.rst | 31 + py/torch_tensorrt/_compile.py | 261 +----- py/torch_tensorrt/executorch/__init__.py | 3 + py/torch_tensorrt/executorch/_export_utils.py | 312 +++++++ py/torch_tensorrt/executorch/export.py | 350 +++++++ tests/py/dynamo/executorch/test_api.py | 11 + .../test_cuda_partitioner_composition.py | 111 ++- tests/py/dynamo/executorch/test_edge_cases.py | 80 +- tests/py/dynamo/executorch/test_export.py | 881 ++++++++++++++++++ 9 files changed, 1761 insertions(+), 279 deletions(-) create mode 100644 py/torch_tensorrt/executorch/_export_utils.py create mode 100644 py/torch_tensorrt/executorch/export.py create mode 100644 tests/py/dynamo/executorch/test_export.py diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 103cf19025..d32351d293 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -234,6 +234,7 @@ backend. It requires the ``executorch`` package (``pip install import torch import torch_tensorrt + import torch_tensorrt.executorch model = MyModel().eval().cuda() inputs = [torch.randn((1, 3, 224, 224)).cuda()] @@ -243,6 +244,36 @@ backend. It requires the ``executorch`` package (``pip install retrace=False, arg_inputs=inputs, ) +``save`` writes both the ``.pte`` and any external ``.ptd`` tensor-data files. +Advanced ExecuTorch users can stop at the standard Edge program boundary to +inspect the delegated graph, add metadata, or control final memory planning: + +.. code-block:: python + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=inputs, + retrace=False, + partitioners=extra_partitioners, + transform_passes=passes, + compile_config=edge_config, + constant_methods={"get_vocab_size": 256}, + ) + + # Inspect edge.exported_program() or apply additional Edge transforms here. + program = edge.to_executorch(config=backend_config) + with open("trt.pte", "wb") as output: + program.write_to_file(output) + program.write_tensor_data_to_file(".") + +``executorch.export`` accepts a TensorRT-compiled ``GraphModule``, an +engine-bearing ``ExportedProgram``, or a mapping of independently exported +methods. It always applies the TensorRT partitioner first, followed by caller +``partitioners`` in order, and returns ExecuTorch's native +``EdgeProgramManager``. Method mappings preserve independent entry points but +do not imply shared mutable state. Use ``save`` for the one-shot path where +Torch-TensorRT manages Edge lowering, finalization, and persistence. + **Coalesced TensorRT + CUDA .pte** To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 32ee801886..cf17c700a3 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -657,7 +657,7 @@ def save( *, extra_files: Optional[dict[str, str]] = None, output_format: Literal[ - "exported_program", "torchscript", "aot_inductor" + "exported_program", "torchscript", "aot_inductor", "executorch" ] = "exported_program", inputs: Optional[Sequence[torch.Tensor | Input]] = None, arg_inputs: Optional[Sequence[torch.Tensor | Input]] = None, @@ -1147,197 +1147,6 @@ def _extract_tensor(obj: Any) -> Any: ) -def _get_engine_info_from_state(engine_obj: Any) -> List[Any]: - """Normalize TensorRT engine state into the serialized engine-info list.""" - state = engine_obj.__getstate__() - engine_info = state[0] if isinstance(state, tuple) else state - return list(engine_info) - - -def _validate_executorch_engine_info( - engine_info: Sequence[Any], *, node_name: str = "" -) -> None: - """Reject engine configurations unsupported by the ExecuTorch export path.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - REQUIRES_OUTPUT_ALLOCATOR_IDX, - ) - - if ( - len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX - and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1" - ): - node_suffix = f" for node '{node_name}'" if node_name else "" - raise RuntimeError( - "ExecuTorch export does not support TensorRT engines that require " - "an output allocator (data-dependent output shapes)" - f"{node_suffix}." - ) - - -def _count_executorch_engine_nodes(exp_program: Any) -> int: - """Count TRT execute-engine nodes in an ExportedProgram.""" - count = 0 - for node in exp_program.graph_module.graph.nodes: - if node.op != "call_function": - continue - if node.target is torch.ops.tensorrt.execute_engine.default: - count += 1 - continue - target = node.target - if ( - hasattr(target, "_schema") - and target._schema.name == "tensorrt::no_op_placeholder_for_execute_engine" - ): - count += 1 - return count - - -def _replace_execute_engine_for_executorch(exp_program: Any) -> Any: - """Replace execute_engine nodes with no_op_placeholder_for_execute_engine. - - ExecuTorch's to_edge_transform_and_lower runs ExportPass subclasses that - dispatch through the C++ schema validator. The validator rejects the - ScriptObject engine arg (it arrives as a CustomObjArgument placeholder - rather than a real FakeScriptObject). Converting each execute_engine node - to no_op_placeholder_for_execute_engine (which carries all engine info as - plain strings) avoids the ScriptObject entirely so the passes succeed. - - The TRT engine bytes are stored as a ``torch.uint8`` buffer on the graph - module and referenced from the no_op call via a ``get_attr`` FX node. This - keeps the engine out of the FX-emitted Python source: CPython's tokenizer - cannot parse string literals larger than ~2 GB, so an inline base64 string - breaks ``gm.recompile()`` for any engine whose payload exceeds that limit. - """ - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - ENGINE_IDX, - SERIALIZATION_LEN, - ) - - gm = exp_program.graph_module - execute_engine_op = torch.ops.tensorrt.execute_engine.default - no_op = torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default - - nodes_to_replace = [ - n - for n in gm.graph.nodes - if n.op == "call_function" and n.target is execute_engine_op - ] - if not nodes_to_replace: - return exp_program - - for engine_idx_in_graph, node in enumerate(nodes_to_replace): - inputs_arg = node.args[0] - engine_node = node.args[1] - - # Retrieve the engine ScriptObject from the graph module or constants. - if engine_node.op == "get_attr": - engine_obj = getattr(gm, engine_node.target, None) - if engine_obj is None: - raise RuntimeError( - f"execute_engine node '{node.name}': get_attr target " - f"'{engine_node.target}' not found on graph module" - ) - elif engine_node.op == "placeholder": - from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj - - engine_obj = _resolve_lifted_custom_obj(exp_program, engine_node) - if engine_obj is None: - raise RuntimeError( - f"execute_engine node '{node.name}': placeholder engine " - f"'{engine_node.name}' did not resolve to a lifted " - f"custom-object constant (available: " - f"{sorted(getattr(exp_program, 'constants', {}) or {})})" - ) - else: - raise RuntimeError( - f"execute_engine node '{node.name}': unexpected engine arg op " - f"'{engine_node.op}'" - ) - - engine_info = _get_engine_info_from_state(engine_obj) - _validate_executorch_engine_info(engine_info, node_name=node.name) - # Ensure the engine bytes slot is a base64 string (no_op takes str args). - engine_bytes = engine_info[ENGINE_IDX] - if isinstance(engine_bytes, str): - # `_get_engine_info_from_state` returns the engine as a - # base64-encoded `str` when the engine arrived through the - # serialized TRT runtime round-trip path. Decode back to raw - # bytes so it can land in a uint8 buffer below. - import base64 - - engine_bytes = base64.b64decode(engine_bytes) - elif not isinstance(engine_bytes, (bytes, bytearray)): - engine_bytes = bytes(engine_bytes) - # Store engine payload as a uint8 buffer + get_attr ref. FX emits a - # name reference instead of an inline literal, sidestepping the - # tokenizer's >2 GB string-literal limit. - engine_tensor = torch.frombuffer(bytearray(engine_bytes), dtype=torch.uint8) - # Use FX's unique-attr-name helper so re-export passes (which may - # invoke this rewriter multiple times on the same `gm`) don't - # silently overwrite earlier engine buffers. - from torch.fx.experimental.const_fold import ( - get_unique_attr_name_in_module, - ) - - buffer_name = get_unique_attr_name_in_module(gm, "_trt_engine_0") - gm.register_buffer(buffer_name, engine_tensor, persistent=True) - exp_program.state_dict[buffer_name] = engine_tensor - - str_args = [ - str(x) if x is not None else "" for x in engine_info[:SERIALIZATION_LEN] - ] - # Build a FakeTensor mirror so downstream FX passes (FakeTensorProp, - # ExecuTorch lowering, export-serde) that read `node.meta["val"]` - # on the `get_attr` reference don't `KeyError`. Must reuse the - # graph's existing FakeTensorMode — creating a fresh one would - # fail downstream with "fake mode from input 0 doesn't match - # mode from input 1" the moment any pass mixes the two. - from torch._guards import detect_fake_mode - - fake_mode = detect_fake_mode( - [n.meta["val"] for n in gm.graph.nodes if "val" in n.meta] - ) - fake_engine = ( - fake_mode.from_tensor(engine_tensor) - if fake_mode is not None - else engine_tensor - ) - with gm.graph.inserting_before(node): - engine_attr_node = gm.graph.get_attr(buffer_name) - engine_attr_node.meta["val"] = fake_engine - no_op_args = ( - inputs_arg, - *str_args[:ENGINE_IDX], - engine_attr_node, - *str_args[ENGINE_IDX + 1 :], - ) - no_op_node = gm.graph.call_function(no_op, no_op_args) - no_op_node.meta["val"] = node.meta.get("val") - - node.replace_all_uses_with(no_op_node) - gm.graph.erase_node(node) - - # Erase the engine get_attr node if it is now unused. - if engine_node.op == "get_attr" and not engine_node.users: - gm.graph.erase_node(engine_node) - # Also drop the now-orphan attribute from the module so the - # original engine bytes aren't double-serialized into state_dict - # alongside the new uint8 buffer. Use FX's dotted-path helper - # so nested-target attrs (e.g. `submod.engine`) are deleted - # correctly — plain `delattr(gm, "a.b")` only works on - # top-level names. - from torch.fx.graph_module import _del_attr, _has_attr - - if _has_attr(gm, engine_node.target): - _del_attr(gm, engine_node.target) - - gm.graph.eliminate_dead_code() - gm.graph.lint() - gm.recompile() - - return exp_program - - def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None: """Write an ExecuTorch program's external named tensor data (``.ptd``) next to the ``.pte``. @@ -1362,69 +1171,35 @@ def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None: - """Save an ExportedProgram (with TensorRT execute_engine nodes) as an ExecuTorch .pte file. - - Partitions the graph by torch.ops.tensorrt.no_op_placeholder_for_execute_engine - (execute_engine is pre-converted to avoid schema type errors in edge passes), - serializes each engine to the same blob format as the TRT runtime (vector of - strings), and embeds it in the .pte. Requires the ``executorch`` package and - torch_tensorrt_runtime. See https://pytorch.org/executorch/stable/getting-started-setup.html - """ + """Save an engine-bearing ExportedProgram as an ExecuTorch program.""" if not ENABLED_FEATURES.torch_tensorrt_runtime: raise RuntimeError( "output_format='executorch' requires the Torch-TensorRT runtime " "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." ) try: - from executorch.exir import to_edge_transform_and_lower + from torch_tensorrt.executorch import export except ImportError: raise ImportError( "ExecuTorch is not installed. Install with: pip install " "\"torch_tensorrt[executorch]\" to use output_format='executorch'." ) import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 - from torch_tensorrt.executorch import ( - TensorRTPartitioner, - get_edge_compile_config, - ) - - extra_partitioners = kwargs.get("partitioners") or [] - if not isinstance(extra_partitioners, (list, tuple)): - raise TypeError( - "partitioners must be a list or tuple when using " - "output_format='executorch'" - ) - # Forward any caller-provided compile_specs to TensorRTPartitioner so users - # can override the default target_device ("cuda:0") by passing e.g. - # `compile_specs=[CompileSpec("target_device", b"cuda:1")]` to save(). - # When omitted, TensorRTPartitioner auto-appends the cuda:0 default. - executorch_compile_specs = kwargs.get("compile_specs") or [] - if not isinstance(executorch_compile_specs, (list, tuple)): - raise TypeError( - "compile_specs must be a list or tuple when using " - "output_format='executorch'" - ) - partitioners = [ - TensorRTPartitioner(compile_specs=list(executorch_compile_specs)) - ] + list(extra_partitioners) - - engine_count = _count_executorch_engine_nodes(exp_program) - if engine_count > 1: - logger.warning( - "%d TRT engines detected. Multi-engine .pte exports can incur extra " - "delegate boundary overhead.", - engine_count, - ) - # Replace execute_engine nodes (ScriptObject arg) with no_op_placeholder_for_execute_engine - # (string args only) before to_edge_transform_and_lower runs ExportPass subclasses that - # would fail trying to cast the CustomObjArgument placeholder to a real C++ Engine object. - exp_program = _replace_execute_engine_for_executorch(exp_program) + # save() only ever lowers a single method, so keep rejecting the per-method + # mappings that torch_tensorrt.executorch.export() accepts. + for option in ("partitioners", "compile_specs"): + value = kwargs.get(option) + if value is not None and not isinstance(value, (list, tuple)): + raise TypeError( + f"{option} must be a list or tuple when using " + "output_format='executorch'" + ) - edge_program = to_edge_transform_and_lower( + edge_program = export( exp_program, - partitioner=partitioners, - compile_config=get_edge_compile_config(), + partitioners=kwargs.get("partitioners"), + compile_specs=kwargs.get("compile_specs"), ) executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config")) with open(file_path, "wb") as f: @@ -1443,14 +1218,10 @@ def _normalize_engine_constants_to_python(exp_program: "ExportedProgram") -> Non import base64 from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ENGINE_IDX - from torch_tensorrt.dynamo.runtime._TRTEngine import ( - EngineSerializer, - TRTEngine, - ) + from torch_tensorrt.dynamo.runtime._TRTEngine import EngineSerializer, TRTEngine for fqn, constant in list(exp_program.constants.items()): if isinstance(constant, (torch._C.ScriptObject, TRTEngine)): - state = constant.__getstate__() if len(state) == 2 and ( state[1] == "TRTEngine" diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index 123eee846d..44ac5824b9 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -25,9 +25,11 @@ def __getattr__(name: str) -> NoReturn: "get_edge_compile_config", "TensorRTPartitioner", "TensorRTBackend", + "export", ] else: from torch_tensorrt.executorch.backend import TensorRTBackend + from torch_tensorrt.executorch.export import export from torch_tensorrt.executorch.partitioner import TensorRTPartitioner def get_edge_compile_config() -> "EdgeCompileConfig": @@ -40,4 +42,5 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "get_edge_compile_config", "TensorRTPartitioner", "TensorRTBackend", + "export", ] diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py new file mode 100644 index 0000000000..ebacd4bbce --- /dev/null +++ b/py/torch_tensorrt/executorch/_export_utils.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import copy +from typing import Any, Sequence + +import torch +from torch._subclasses.fake_tensor import is_fake +from torch.export.graph_signature import InputKind + + +def _is_graph_bound_metadata(value: Any) -> bool: + """True if a metadata value is tied to the exported program's shape environment. + + Fake tensors and symbolic sizes reach back to a live ``ShapeEnv`` that owns + tensors deepcopy refuses to touch, and cloning them would also detach the copy + from the symbols the graph is guarded on. + """ + for leaf in torch.utils._pytree.tree_leaves(value): + if isinstance(leaf, (torch.SymInt, torch.SymFloat, torch.SymBool)): + return True + if isinstance(leaf, torch.Tensor) and is_fake(leaf): + return True + return False + + +def get_engine_info_from_state(engine_obj: Any) -> list[Any]: + """Normalize TensorRT engine state into the serialized engine-info list.""" + state = engine_obj.__getstate__() + engine_info = state[0] if isinstance(state, tuple) else state + return list(engine_info) + + +def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> None: + """Reject engine configurations unsupported by ExecuTorch export.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + REQUIRES_OUTPUT_ALLOCATOR_IDX, + ) + + if ( + len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX + and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1" + ): + node_suffix = f" for node '{node_name}'" if node_name else "" + raise RuntimeError( + "ExecuTorch export does not support TensorRT engines that require " + "an output allocator (data-dependent output shapes)" + f"{node_suffix}." + ) + + +def _schema_name(node: Any) -> str: + target = node.target + return target._schema.name if hasattr(target, "_schema") else "" + + +def _resolve_engine_info(exported_program: Any, node: Any) -> list[Any]: + if node.target is not torch.ops.tensorrt.execute_engine.default: + return list(node.args[1:]) + + graph_module = exported_program.graph_module + engine_node = node.args[1] + if engine_node.op == "get_attr": + engine_obj = getattr(graph_module, engine_node.target, None) + if engine_obj is None: + raise RuntimeError( + f"execute_engine node '{node.name}': get_attr target " + f"'{engine_node.target}' not found on graph module" + ) + elif engine_node.op == "placeholder": + from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj + + engine_obj = _resolve_lifted_custom_obj(exported_program, engine_node) + if engine_obj is None: + raise RuntimeError( + f"execute_engine node '{node.name}': placeholder engine " + f"'{engine_node.name}' did not resolve to a lifted " + f"custom-object constant (available: " + f"{sorted(getattr(exported_program, 'constants', {}) or {})})" + ) + else: + raise RuntimeError( + f"execute_engine node '{node.name}': unexpected engine arg op " + f"'{engine_node.op}'" + ) + return get_engine_info_from_state(engine_obj) + + +def validate_engine_program(exported_program: Any) -> int: + """Validate all TensorRT engine nodes before any input program is mutated.""" + count = 0 + for node in exported_program.graph_module.graph.nodes: + if node.op != "call_function": + continue + if ( + node.target is not torch.ops.tensorrt.execute_engine.default + and _schema_name(node) != "tensorrt::no_op_placeholder_for_execute_engine" + ): + continue + validate_engine_info( + _resolve_engine_info(exported_program, node), node_name=node.name + ) + count += 1 + return count + + +def _payload_sharing_memo(exported_program: Any) -> dict[int, Any]: + memo: dict[int, Any] = {} + for value in ( + *exported_program.state_dict.values(), + *exported_program.constants.values(), + ): + memo[id(value)] = value + for module in exported_program.graph_module.modules(): + for value in ( + *module.parameters(recurse=False), + *module.buffers(recurse=False), + ): + memo[id(value)] = value + for value in module.__dict__.values(): + if isinstance(value, (torch.Tensor, torch.ScriptObject)): + memo[id(value)] = value + for module in exported_program.graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in module.graph.nodes: + for value in node.meta.values(): + for leaf in torch.utils._pytree.tree_leaves(value): + if isinstance(leaf, (torch.Tensor, torch.ScriptObject)): + memo[id(leaf)] = leaf + return memo + + +def _stage_graph_module( + graph_module: torch.fx.GraphModule, + payload_memo: dict[int, Any], +) -> torch.fx.GraphModule: + staged = copy.deepcopy(graph_module, payload_memo) + # FX gives each staged node a fresh meta dict but shares the values inside it, + # so nested containers need their own copy to keep the source unaffected. + for name, source_module in graph_module.named_modules(): + if not isinstance(source_module, torch.fx.GraphModule): + continue + staged_module = staged if not name else staged.get_submodule(name) + if not isinstance(staged_module, torch.fx.GraphModule): + continue + staged_nodes = {node.name: node for node in staged_module.graph.nodes} + source_nodes = {node.name: node for node in source_module.graph.nodes} + if staged_nodes.keys() != source_nodes.keys(): + raise RuntimeError( + f"Staged GraphModule {name or ''!r} changed node identities." + ) + for node_name, source_node in source_nodes.items(): + staged_nodes[node_name].meta = { + key: ( + value + if _is_graph_bound_metadata(value) + else copy.deepcopy(value, payload_memo) + ) + for key, value in source_node.meta.items() + } + return staged + + +def stage_exported_program(exported_program: Any) -> Any: + """Copy program structure while sharing immutable tensor and engine payloads.""" + graph_module = _stage_graph_module( + exported_program.graph_module, + _payload_sharing_memo(exported_program), + ) + return exported_program._update( + graph_module, + copy.deepcopy(exported_program.graph_signature), + state_dict=dict(exported_program.state_dict), + constants=dict(exported_program.constants), + ) + + +def _remove_lifted_engine_placeholder( + exported_program: Any, engine_node: torch.fx.Node +) -> None: + """Remove an unused lifted engine from the graph, signature, and constants.""" + if engine_node.op != "placeholder" or engine_node.users: + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} must be unused before cleanup." + ) + + matching_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CUSTOM_OBJ + and getattr(spec.arg, "name", None) == engine_node.name + ] + if len(matching_specs) != 1: + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} must have exactly one " + f"CUSTOM_OBJ input spec; found {len(matching_specs)}." + ) + + engine_spec = matching_specs[0] + engine_target = engine_spec.target + if not isinstance(engine_target, str): + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} has no constant target." + ) + + exported_program._graph_signature.input_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec is not engine_spec + ] + engine_node.graph.erase_node(engine_node) + + if not any( + spec.target == engine_target + for spec in exported_program.graph_signature.input_specs + ): + exported_program.constants.pop(engine_target, None) + + +def replace_execute_engine(exported_program: Any) -> Any: + """Replace execute_engine nodes with ExecuTorch-safe placeholder calls.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + ENGINE_IDX, + SERIALIZATION_LEN, + ) + + graph_module = exported_program.graph_module + execute_engine_op = torch.ops.tensorrt.execute_engine.default + no_op = torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + nodes_to_replace = [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is execute_engine_op + ] + if not nodes_to_replace: + return exported_program + + materialized_engines: dict[torch.fx.Node, tuple[torch.fx.Node, list[str]]] = {} + for node in nodes_to_replace: + inputs_arg = node.args[0] + engine_node = node.args[1] + materialized = materialized_engines.get(engine_node) + if materialized is None: + engine_info = _resolve_engine_info(exported_program, node) + engine_bytes = engine_info[ENGINE_IDX] + if isinstance(engine_bytes, str): + import base64 + + engine_bytes = base64.b64decode(engine_bytes) + elif not isinstance(engine_bytes, (bytes, bytearray)): + engine_bytes = bytes(engine_bytes) + engine_tensor = torch.frombuffer(bytearray(engine_bytes), dtype=torch.uint8) + + from torch.fx.experimental.const_fold import get_unique_attr_name_in_module + + buffer_name = get_unique_attr_name_in_module(graph_module, "_trt_engine_0") + graph_module.register_buffer(buffer_name, engine_tensor, persistent=True) + exported_program.state_dict[buffer_name] = engine_tensor + # The engine slot is replaced by engine_attr_node below, so skip it: + # str() on multi-gigabyte engine bytes would build a throwaway string + # roughly four times their size. + str_args = [ + ("" if index == ENGINE_IDX else str(value) if value is not None else "") + for index, value in enumerate(engine_info[:SERIALIZATION_LEN]) + ] + + from torch._guards import detect_fake_mode + + fake_mode = detect_fake_mode( + [ + graph_node.meta["val"] + for graph_node in graph_module.graph.nodes + if "val" in graph_node.meta + ] + ) + fake_engine = ( + fake_mode.from_tensor(engine_tensor) + if fake_mode is not None + else engine_tensor + ) + with graph_module.graph.inserting_before(node): + engine_attr_node = graph_module.graph.get_attr(buffer_name) + engine_attr_node.meta["val"] = fake_engine + materialized = (engine_attr_node, str_args) + materialized_engines[engine_node] = materialized + engine_attr_node, str_args = materialized + + with graph_module.graph.inserting_before(node): + no_op_args = ( + inputs_arg, + *str_args[:ENGINE_IDX], + engine_attr_node, + *str_args[ENGINE_IDX + 1 :], + ) + no_op_node = graph_module.graph.call_function(no_op, no_op_args) + no_op_node.meta["val"] = node.meta.get("val") + + node.replace_all_uses_with(no_op_node) + graph_module.graph.erase_node(node) + if engine_node.op == "get_attr" and not engine_node.users: + graph_module.graph.erase_node(engine_node) + from torch.fx.graph_module import _del_attr, _has_attr + + if _has_attr(graph_module, engine_node.target): + _del_attr(graph_module, engine_node.target) + elif engine_node.op == "placeholder" and not engine_node.users: + _remove_lifted_engine_placeholder(exported_program, engine_node) + + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + return exported_program diff --git a/py/torch_tensorrt/executorch/export.py b/py/torch_tensorrt/executorch/export.py new file mode 100644 index 0000000000..66e0ae5785 --- /dev/null +++ b/py/torch_tensorrt/executorch/export.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +import logging +import platform +from collections.abc import Mapping, Sequence +from inspect import signature +from typing import TYPE_CHECKING, Any + +import torch +from torch.export import ExportedProgram +from torch_tensorrt._Input import Input + +if TYPE_CHECKING: + from executorch.exir import EdgeCompileConfig, EdgeProgramManager + from executorch.exir.backend.compile_spec_schema import CompileSpec + from executorch.exir.backend.partitioner import Partitioner + +logger = logging.getLogger(__name__) + + +def _contains_dynamic_input(value: Any) -> bool: + if isinstance(value, Input): + return bool(value.shape_mode == Input._ShapeMode.DYNAMIC) + if isinstance(value, Mapping): + return any(_contains_dynamic_input(item) for item in value.values()) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return any(_contains_dynamic_input(item) for item in value) + return False + + +def _all_input_specs(value: Any) -> bool: + if isinstance(value, Input): + return True + if isinstance(value, Mapping): + return all(_all_input_specs(item) for item in value.values()) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return all(_all_input_specs(item) for item in value) + return False + + +def _dynamic_shapes_from_input(value: Any, dim_registry: dict[str, Any]) -> Any: + from torch_tensorrt.dynamo._tracer import get_dynamic_shapes + + if isinstance(value, Input): + return get_dynamic_shapes(value, dim_registry) + if isinstance(value, Mapping): + return { + key: _dynamic_shapes_from_input(item, dim_registry) + for key, item in value.items() + } + if isinstance(value, tuple): + return tuple(_dynamic_shapes_from_input(item, dim_registry) for item in value) + if isinstance(value, list): + return [_dynamic_shapes_from_input(item, dim_registry) for item in value] + raise TypeError(f"Unsupported input structure leaf: {type(value).__name__}") + + +def _to_torch_inputs(value: Any, device: torch.device) -> Any: + from torch_tensorrt.dynamo.utils import get_torch_tensor + + if isinstance(value, Input): + return get_torch_tensor(value, device) + if isinstance(value, torch.Tensor): + return value.to(device) + if isinstance(value, Mapping): + return {key: _to_torch_inputs(item, device) for key, item in value.items()} + if isinstance(value, tuple): + return tuple(_to_torch_inputs(item, device) for item in value) + if isinstance(value, list): + return [_to_torch_inputs(item, device) for item in value] + return value + + +def _prepare_graph_module( + module: torch.fx.GraphModule, + *, + arg_inputs: Sequence[Any] | None, + kwarg_inputs: Mapping[str, Any] | None, + dynamic_shapes: Any | None, + retrace: bool, +) -> ExportedProgram: + from torch_tensorrt.dynamo._defaults import default_device + from torch_tensorrt.dynamo._exporter import export as export_graph_module + from torch_tensorrt.dynamo._tracer import build_dim_registry + from torch_tensorrt.dynamo.utils import to_torch_device + + arguments = tuple(arg_inputs or ()) + keyword_arguments = dict(kwarg_inputs or {}) + all_specs = ( + bool(arguments or keyword_arguments) + and _all_input_specs(arguments) + and _all_input_specs(keyword_arguments) + ) + if dynamic_shapes is None and all_specs: + bound_arguments = ( + signature(module.forward).bind(*arguments, **keyword_arguments).arguments + ) + dim_registry = build_dim_registry(arguments, keyword_arguments) + dynamic_shapes = { + name: _dynamic_shapes_from_input(value, dim_registry) + for name, value in bound_arguments.items() + } + elif dynamic_shapes is None and ( + _contains_dynamic_input(arguments) or _contains_dynamic_input(keyword_arguments) + ): + raise ValueError( + "Mixed Tensor and dynamic Input arguments require explicit dynamic_shapes." + ) + + device = to_torch_device(default_device()) + torch_arguments = tuple(_to_torch_inputs(arguments, device)) + torch_keyword_arguments = _to_torch_inputs(keyword_arguments, device) + if retrace and not torch_arguments and not torch_keyword_arguments: + raise ValueError("retrace=True requires example inputs.") + + return export_graph_module( + module, + arg_inputs=torch_arguments, + kwarg_inputs=torch_keyword_arguments, + dynamic_shapes=dynamic_shapes, + use_legacy_exporter=not retrace, + ) + + +def _prepare_programs( + source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], + *, + arg_inputs: Sequence[Any] | None, + kwarg_inputs: Mapping[str, Any] | None, + dynamic_shapes: Any | None, + retrace: bool | None, +) -> tuple[ExportedProgram | dict[str, ExportedProgram], tuple[str, ...]]: + source_options_used = ( + arg_inputs is not None + or kwarg_inputs is not None + or dynamic_shapes is not None + or retrace is True + ) + if isinstance(source, ExportedProgram): + if source_options_used: + raise ValueError( + "Inputs, dynamic_shapes, and retrace=True are invalid for an " + "already-exported program." + ) + return source, ("forward",) + + if isinstance(source, torch.fx.GraphModule): + program = _prepare_graph_module( + source, + arg_inputs=arg_inputs, + kwarg_inputs=kwarg_inputs, + dynamic_shapes=dynamic_shapes, + retrace=False if retrace is None else retrace, + ) + return program, ("forward",) + + if isinstance(source, Mapping): + if not source: + raise ValueError("Method mapping must not be empty.") + if source_options_used: + raise ValueError( + "Method mappings accept only pre-exported programs; per-method " + "input and retrace options are not supported." + ) + programs: dict[str, ExportedProgram] = {} + program_names_by_id: dict[int, str] = {} + for name, program in source.items(): + if not isinstance(name, str) or not name: + raise ValueError("Every method name must be a non-empty string.") + if not isinstance(program, ExportedProgram): + raise TypeError( + "Method mappings must contain only ExportedProgram values." + ) + previous_name = program_names_by_id.get(id(program)) + if previous_name is not None: + raise ValueError( + "Method mapping contains the same ExportedProgram object for " + f"{previous_name!r} and {name!r}. Each method requires an " + "independent program." + ) + program_names_by_id[id(program)] = name + programs[name] = program + return programs, tuple(programs) + + raise TypeError( + "source must be a TensorRT-compiled GraphModule, an engine-bearing " + "ExportedProgram, or a mapping of method names to ExportedPrograms. " + "Compile nn.Module inputs with torch_tensorrt.compile() first." + ) + + +def _per_method_values( + value: Sequence[Any] | Mapping[str, Sequence[Any] | None] | None, + method_names: tuple[str, ...], + option_name: str, +) -> dict[str, list[Any]]: + if isinstance(value, Mapping): + unknown = set(value) - set(method_names) + if unknown: + raise ValueError( + f"{option_name} contains unknown methods: {sorted(unknown)}" + ) + normalized: dict[str, list[Any]] = {} + for name in method_names: + method_value = value.get(name) + if method_value is not None and ( + not isinstance(method_value, Sequence) + or isinstance(method_value, (str, bytes)) + ): + raise TypeError(f"{option_name}[{name!r}] must be a sequence or None.") + normalized[name] = list(method_value or ()) + return normalized + if value is not None and ( + not isinstance(value, Sequence) or isinstance(value, (str, bytes)) + ): + raise TypeError( + f"{option_name} must be a sequence, or a mapping of method name to " + "sequence." + ) + shared = list(value or ()) + return {name: list(shared) for name in method_names} + + +def export( + source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], + *, + inputs: Sequence[Any] | None = None, + arg_inputs: Sequence[Any] | None = None, + kwarg_inputs: Mapping[str, Any] | None = None, + dynamic_shapes: Any | None = None, + retrace: bool | None = None, + transform_passes: Any | None = None, + partitioners: ( + Sequence["Partitioner"] | Mapping[str, Sequence["Partitioner"] | None] | None + ) = None, + compile_specs: ( + Sequence["CompileSpec"] | Mapping[str, Sequence["CompileSpec"] | None] | None + ) = None, + compile_config: "EdgeCompileConfig | None" = None, + constant_methods: Mapping[str, Any] | None = None, + generate_etrecord: bool = False, +) -> "EdgeProgramManager": + """Prepare TensorRT-compiled programs for composable ExecuTorch lowering. + + TensorRT claims engine nodes first. Additional partitioners run afterward in + caller-provided order. The returned EdgeProgramManager is the standard + ExecuTorch inspection and customization boundary; call ``to_executorch()`` + on it when ready to perform final memory planning and serialization. + + Export stages independent graph and metadata containers while sharing tensor + and engine payload storage, avoiding copies of potentially multi-gigabyte + TensorRT engines. Transform passes must treat shared payload values as + immutable. Method mappings preserve independent entry points but do not imply + shared mutable state between them. + """ + from torch_tensorrt._features import ENABLED_FEATURES + + if platform.system() != "Linux": + raise ValueError( + f"The executorch format is only supported on Linux, {platform.system()} " + "is not a supported platform for this format" + ) + if not ENABLED_FEATURES.torch_tensorrt_runtime: + raise RuntimeError( + "ExecuTorch export requires the Torch-TensorRT runtime " + "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." + ) + if inputs is not None and arg_inputs is not None: + raise ValueError("inputs and arg_inputs are mutually exclusive.") + arguments = inputs if inputs is not None else arg_inputs + + import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 + from executorch.exir import to_edge_transform_and_lower + from torch_tensorrt.executorch import TensorRTPartitioner, get_edge_compile_config + from torch_tensorrt.executorch._export_utils import ( + replace_execute_engine, + stage_exported_program, + validate_engine_program, + ) + + programs, method_names = _prepare_programs( + source, + arg_inputs=arguments, + kwarg_inputs=kwarg_inputs, + dynamic_shapes=dynamic_shapes, + retrace=retrace, + ) + program_map = {"forward": programs} if not isinstance(programs, dict) else programs + extra_partitioners = _per_method_values(partitioners, method_names, "partitioners") + method_compile_specs = _per_method_values( + compile_specs, method_names, "compile_specs" + ) + + if constant_methods is not None: + collisions = set(constant_methods) & set(method_names) + if collisions: + raise ValueError( + f"constant_methods collide with executable methods: {sorted(collisions)}" + ) + + if isinstance(transform_passes, Mapping): + unknown = set(transform_passes) - set(method_names) + if unknown: + raise ValueError( + f"transform_passes contains unknown methods: {sorted(unknown)}" + ) + + engine_counts = { + name: validate_engine_program(program) for name, program in program_map.items() + } + # Zero-engine methods are allowed: later partitioners may claim their ops, + # or portable operators may remain undelegated. + staged_programs = { + name: stage_exported_program(program) for name, program in program_map.items() + } + rewritten: dict[str, ExportedProgram] = {} + method_partitioners: dict[str, list[Partitioner]] = {} + for name, program in staged_programs.items(): + if engine_counts[name] > 1: + logger.warning( + "%s contains %d TRT engines. Multi-engine .pte exports can incur " + "extra delegate boundary overhead.", + name, + engine_counts[name], + ) + rewritten[name] = replace_execute_engine(program) + method_partitioners[name] = [ + TensorRTPartitioner(compile_specs=method_compile_specs[name]), + *extra_partitioners[name], + ] + + edge_programs: ExportedProgram | dict[str, ExportedProgram] + partitioner_pipeline: list[Partitioner] | dict[str, list[Partitioner]] + if isinstance(programs, dict): + edge_programs = rewritten + partitioner_pipeline = method_partitioners + else: + edge_programs = rewritten["forward"] + partitioner_pipeline = method_partitioners["forward"] + + return to_edge_transform_and_lower( + edge_programs, + transform_passes=transform_passes, + partitioner=partitioner_pipeline, + constant_methods=dict(constant_methods) if constant_methods else None, + compile_config=( + compile_config if compile_config is not None else get_edge_compile_config() + ), + generate_etrecord=generate_etrecord, + ) diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 5d88822716..3763f7951e 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -13,7 +13,12 @@ @pytest.mark.unit def test_lazy_import_error_when_executorch_missing(monkeypatch): + import torch_tensorrt + original_module = sys.modules.pop("torch_tensorrt.executorch", None) + original_attribute = getattr(torch_tensorrt, "executorch", None) + if hasattr(torch_tensorrt, "executorch"): + delattr(torch_tensorrt, "executorch") original_find_spec = importlib.util.find_spec def fake_find_spec(name, package=None): @@ -30,6 +35,10 @@ def fake_find_spec(name, package=None): sys.modules.pop("torch_tensorrt.executorch", None) if original_module is not None: sys.modules["torch_tensorrt.executorch"] = original_module + if original_attribute is not None: + torch_tensorrt.executorch = original_attribute + elif hasattr(torch_tensorrt, "executorch"): + delattr(torch_tensorrt, "executorch") @pytest.mark.unit @@ -59,6 +68,8 @@ def test_public_api_symbols_present(): assert "get_edge_compile_config" in module.__all__ assert "TensorRTPartitioner" in module.__all__ assert "TensorRTBackend" in module.__all__ + assert "export" in module.__all__ + assert "to_executorch" not in module.__all__ _REPO_ROOT = Path(__file__).resolve().parents[4] diff --git a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py index 64233f45cd..455bbc16b2 100644 --- a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py +++ b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py @@ -73,11 +73,15 @@ def _cuda_partitioner(): return CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) -def _delegate_ids(pte_path): - """Backend ids of every delegate in the serialized program, in order.""" +def _deserialize_program(pte_path): from executorch.exir._serialize._program import deserialize_pte_binary - program = deserialize_pte_binary(pte_path.read_bytes()).program + return deserialize_pte_binary(pte_path.read_bytes()).program + + +def _delegate_ids(pte_path): + """Backend ids of every delegate in the serialized program, in order.""" + program = _deserialize_program(pte_path) return [ delegate.id for plan in program.execution_plan for delegate in plan.delegates ] @@ -176,7 +180,11 @@ def forward(self, x): partitioners=[_cuda_partitioner()], ) + assert out.exists() delegate_ids = _delegate_ids(out) + assert ( + "TensorRTBackend" in delegate_ids + ), f"tanh/cos were not delegated to TensorRT; delegates={delegate_ids}" assert ( "CudaBackend" in delegate_ids ), f"the pinned mm did not route to the CUDA backend; delegates={delegate_ids}" @@ -189,6 +197,103 @@ def forward(self, x): ), f"external .ptd data file is empty: {ptd_files}" +def test_multimethod_trt_export_preserves_methods(tmp_path): + """Independent TRT programs remain separate delegated methods after lowering.""" + import torch_tensorrt + + class Prefill(torch.nn.Module): + def forward(self, x): + return torch.cos(x) + + class Decode(torch.nn.Module): + def forward(self, x): + return torch.tanh(x) + + inputs = (torch.randn(16, 16, device="cuda"),) + programs = {} + source_snapshots = {} + for name, model in ( + ("prefill", Prefill()), + ("decode", Decode()), + ): + exported = torch.export.export(model.eval().to("cuda"), inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + program = torch_tensorrt.dynamo._exporter.export( + trt_gm, + arg_inputs=inputs, + use_legacy_exporter=False, + ) + programs[name] = program + source_snapshots[name] = ( + program.graph_module.code, + tuple(program.state_dict), + tuple(program.constants), + tuple(node.target for node in program.graph.nodes), + ) + + edge = torch_tensorrt.executorch.export(programs) + out = tmp_path / "multimethod_trt.pte" + with out.open("wb") as output: + edge.to_executorch().write_to_file(output) + + execution_plans = _deserialize_program(out).execution_plan + assert {plan.name for plan in execution_plans} == {"prefill", "decode"} + assert all( + any(delegate.id == "TensorRTBackend" for delegate in plan.delegates) + for plan in execution_plans + ) + for name, program in programs.items(): + assert source_snapshots[name] == ( + program.graph_module.code, + tuple(program.state_dict), + tuple(program.constants), + tuple(node.target for node in program.graph.nodes), + ) + assert not any(key.startswith("_trt_engine_") for key in program.state_dict) + + +def test_graph_module_export_preserves_source(tmp_path): + import torch_tensorrt + + class Model(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.tanh(x)) + + inputs = (torch.randn(16, 16, device="cuda"),) + exported = torch.export.export(Model().eval().to("cuda"), inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + source_code = trt_gm.code + source_state = tuple(trt_gm.state_dict()) + source_attrs = tuple(name for name, _ in trt_gm.named_modules()) + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=inputs, + retrace=False, + ) + out = tmp_path / "graph_module_trt.pte" + with out.open("wb") as output: + edge.to_executorch().write_to_file(output) + + assert "TensorRTBackend" in _delegate_ids(out) + assert trt_gm.code == source_code + assert tuple(trt_gm.state_dict()) == source_state + assert tuple(name for name, _ in trt_gm.named_modules()) == source_attrs + assert not any( + name.startswith("_trt_engine_") for name, _ in trt_gm.named_buffers() + ) + + def test_trt_only_writes_no_ptd(tmp_path): """A TRT-only program (no CudaPartitioner) carries no external data, so no .ptd is written -- exercises the real write_to_file / _tensor_data path.""" diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index 935b67604c..74b4a42c7c 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -3,21 +3,15 @@ from unittest.mock import MagicMock import pytest -import torch from torch_tensorrt._compile import ( - _count_executorch_engine_nodes, - _validate_executorch_engine_info, + _save_as_executorch, _write_external_tensor_data, ) from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( REQUIRES_OUTPUT_ALLOCATOR_IDX, SERIALIZATION_LEN, ) - - -class _SchemaTarget: - def __init__(self, name): - self._schema = SimpleNamespace(name=name) +from torch_tensorrt.executorch._export_utils import validate_engine_info @pytest.mark.unit @@ -26,36 +20,64 @@ def test_validate_executorch_engine_info_rejects_output_allocator(): engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "1" with pytest.raises(RuntimeError, match="output allocator"): - _validate_executorch_engine_info(engine_info, node_name="trt") + validate_engine_info(engine_info, node_name="trt") @pytest.mark.unit -def test_count_executorch_engine_nodes_handles_execute_and_placeholder(): - execute_node = SimpleNamespace( - op="call_function", - target=torch.ops.tensorrt.execute_engine.default, - ) - placeholder_node = SimpleNamespace( - op="call_function", - target=_SchemaTarget("tensorrt::no_op_placeholder_for_execute_engine"), +def test_save_as_executorch_uses_public_lowering_and_persists_data( + monkeypatch, tmp_path +): + import torch_tensorrt.executorch as executorch_api + + program = SimpleNamespace( + _tensor_data={"forward": b"weights"}, + write_to_file=MagicMock(), + write_tensor_data_to_file=MagicMock(), ) - other_node = SimpleNamespace( - op="call_function", - target=_SchemaTarget("aten::add"), + edge = SimpleNamespace(to_executorch=MagicMock(return_value=program)) + export = MagicMock(return_value=edge) + monkeypatch.setattr(executorch_api, "export", export) + + pte = tmp_path / "model.pte" + source = object() + partitioners = [object()] + compile_specs = [object()] + backend_config = object() + _save_as_executorch( + source, + str(pte), + partitioners=partitioners, + compile_specs=compile_specs, + backend_config=backend_config, ) - exp_program = SimpleNamespace( - graph_module=SimpleNamespace( - graph=SimpleNamespace(nodes=[execute_node, placeholder_node, other_node]) - ) + + export.assert_called_once_with( + source, + partitioners=partitioners, + compile_specs=compile_specs, ) + edge.to_executorch.assert_called_once_with(config=backend_config) + program.write_to_file.assert_called_once() + program.write_tensor_data_to_file.assert_called_once_with(str(tmp_path)) - assert _count_executorch_engine_nodes(exp_program) == 2 + +@pytest.mark.unit +@pytest.mark.parametrize("option", ["partitioners", "compile_specs"]) +def test_save_as_executorch_rejects_per_method_mapping(monkeypatch, tmp_path, option): + import torch_tensorrt.executorch as executorch_api + + export = MagicMock() + monkeypatch.setattr(executorch_api, "export", export) + + with pytest.raises(TypeError, match="must be a list or tuple"): + _save_as_executorch( + object(), str(tmp_path / "model.pte"), **{option: {"forward": []}} + ) + export.assert_not_called() @pytest.mark.unit def test_write_external_tensor_data_writes_when_present(tmp_path): - # A program with external named data (e.g. a CudaPartitioner delegate's - # weights) must have its .ptd written into the .pte's directory. prog = SimpleNamespace( _tensor_data={"forward": b"weights"}, write_tensor_data_to_file=MagicMock(), @@ -69,7 +91,6 @@ def test_write_external_tensor_data_writes_when_present(tmp_path): @pytest.mark.unit def test_write_external_tensor_data_noop_when_empty(tmp_path): - # TRT-only programs have empty _tensor_data (falsy) -> no .ptd written. prog = SimpleNamespace( _tensor_data={}, write_tensor_data_to_file=MagicMock(), @@ -80,9 +101,6 @@ def test_write_external_tensor_data_noop_when_empty(tmp_path): @pytest.mark.unit def test_write_external_tensor_data_fails_loud_without_attr(tmp_path): - # _tensor_data always exists on a real ExecutorchProgram; it is accessed - # directly (no getattr default) so a future rename fails loudly instead of - # silently skipping the .ptd write and reintroducing the null-weights crash. prog = SimpleNamespace(write_tensor_data_to_file=MagicMock()) with pytest.raises(AttributeError): _write_external_tensor_data(prog, str(tmp_path / "model.pte")) diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py new file mode 100644 index 0000000000..d1eae88b27 --- /dev/null +++ b/tests/py/dynamo/executorch/test_export.py @@ -0,0 +1,881 @@ +import importlib +from types import SimpleNamespace +from typing import NamedTuple +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("executorch.exir") + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 +from torch.export.graph_signature import ( # noqa: E402 + CustomObjArgument, + ExportGraphSignature, + InputKind, + InputSpec, + TensorArgument, +) +from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( # noqa: E402 + ENGINE_IDX, + REQUIRES_OUTPUT_ALLOCATOR_IDX, + SERIALIZATION_LEN, +) + + +class FakeExportedProgram: + pass + + +class FakeTensorRTPartitioner: + def __init__(self, compile_specs): + self.compile_specs = compile_specs + + +class _EngineState: + def __init__(self, payload=b"engine-bytes", requires_output_allocator=False): + self.info = [""] * SERIALIZATION_LEN + self.info[ENGINE_IDX] = payload + self.info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = str(int(requires_output_allocator)) + + def __getstate__(self): + return (self.info,) + + +def _engine_program(*, lifted, execute_count=1, requires_output_allocator=False): + graph = torch.fx.Graph() + engine = _EngineState(requires_output_allocator=requires_output_allocator) + if lifted: + engine_node = graph.placeholder("obj_engine") + root = torch.nn.Module() + engine_target = "engine_fqn" + input_specs = [ + InputSpec( + kind=InputKind.CUSTOM_OBJ, + arg=CustomObjArgument(name=engine_node.name, class_fqn=""), + target=engine_target, + ) + ] + constants = {engine_target: engine} + else: + root = torch.nn.Module() + root.engine = engine + engine_node = graph.get_attr("engine") + input_specs = [] + constants = {} + + input_node = graph.placeholder("x") + input_specs.append( + InputSpec( + kind=InputKind.USER_INPUT, + arg=TensorArgument(name=input_node.name), + target=None, + ) + ) + results = [ + graph.call_function( + torch.ops.tensorrt.execute_engine.default, + ([input_node], engine_node), + ) + for _ in range(execute_count) + ] + graph.output(results[0] if len(results) == 1 else tuple(results)) + graph_module = torch.fx.GraphModule(root, graph) + signature = ExportGraphSignature(input_specs=input_specs, output_specs=[]) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + state_dict={}, + constants=constants, + ) + return program, engine_node, input_node, engine + + +def _patch_lowering(monkeypatch, engine_counts=None): + import executorch.exir + import torch_tensorrt._features as features + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + monkeypatch.setattr( + features, + "ENABLED_FEATURES", + features.ENABLED_FEATURES._replace(torch_tensorrt_runtime=True), + ) + export_module = importlib.import_module("torch_tensorrt.executorch.export") + engine_counts = engine_counts or {} + lower = MagicMock(return_value=object()) + monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + monkeypatch.setattr(export_module, "ExportedProgram", FakeExportedProgram) + monkeypatch.setattr( + export_utils, + "validate_engine_program", + lambda program: engine_counts.get(program, 1), + ) + monkeypatch.setattr(export_utils, "stage_exported_program", lambda program: program) + monkeypatch.setattr( + export_utils, + "replace_execute_engine", + lambda program: ("rewritten", program), + ) + return export_module, lower + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +@pytest.mark.parametrize("lifted", [False, True]) +def test_validate_and_replace_execute_engine(lifted): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, input_node, _ = _engine_program(lifted=lifted) + + assert export_utils.validate_engine_program(program) == 1 + rewritten = export_utils.replace_execute_engine(program) + + assert rewritten is program + nodes = list(program.graph_module.graph.nodes) + assert not any( + node.target is torch.ops.tensorrt.execute_engine.default for node in nodes + ) + no_op_nodes = [ + node + for node in nodes + if node.target + is torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + ] + assert len(no_op_nodes) == 1 + buffer_node = no_op_nodes[0].args[1 + ENGINE_IDX] + engine_buffer = getattr(program.graph_module, buffer_node.target) + assert engine_buffer.dtype == torch.uint8 + assert engine_buffer.device.type == "cpu" + assert bytes(engine_buffer.tolist()) == b"engine-bytes" + assert program.state_dict[buffer_node.target] is engine_buffer + assert engine_node not in nodes + assert input_node in nodes + + if lifted: + assert program.graph_signature.inputs_to_lifted_custom_objs == {} + assert program.constants == {} + assert [spec.arg.name for spec in program.graph_signature.input_specs] == [ + input_node.name + ] + else: + assert not hasattr(program.graph_module, "engine") + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +@pytest.mark.parametrize("lifted", [False, True]) +def test_replace_execute_engine_cleans_shared_engine_after_last_use(lifted): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, input_node, _ = _engine_program( + lifted=lifted, execute_count=2 + ) + + assert export_utils.validate_engine_program(program) == 2 + export_utils.replace_execute_engine(program) + + nodes = list(program.graph_module.graph.nodes) + no_op_nodes = [ + node + for node in nodes + if node.target + is torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + ] + assert len(no_op_nodes) == 2 + buffer_nodes = {node.args[1 + ENGINE_IDX] for node in no_op_nodes} + assert len(buffer_nodes) == 1 + assert set(program.state_dict) == {next(iter(buffer_nodes)).target} + assert engine_node not in nodes + assert input_node in nodes + if lifted: + assert program.graph_signature.inputs_to_lifted_custom_objs == {} + assert program.constants == {} + else: + assert not hasattr(program.graph_module, "engine") + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +def test_validate_engine_program_rejects_output_allocator_without_mutation(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, _, engine = _engine_program( + lifted=True, requires_output_allocator=True + ) + original_nodes = list(program.graph_module.graph.nodes) + original_specs = list(program.graph_signature.input_specs) + + with pytest.raises(RuntimeError, match="output allocator"): + export_utils.validate_engine_program(program) + + assert list(program.graph_module.graph.nodes) == original_nodes + assert program.graph_signature.input_specs == original_specs + assert program.constants == {"engine_fqn": engine} + assert engine_node.users + assert program.state_dict == {} + + +@pytest.mark.unit +def test_export_returns_edge_and_forwards_all_options(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + extra_a = object() + extra_b = object() + compile_spec = object() + transform_pass = object() + edge_config = object() + constant_methods = {"get_vocab_size": 256} + partitioners = [extra_a, extra_b] + compile_specs = [compile_spec] + + result = export_module.export( + program, + transform_passes=[transform_pass], + partitioners=partitioners, + compile_specs=compile_specs, + compile_config=edge_config, + constant_methods=constant_methods, + generate_etrecord=True, + ) + + assert result is lower.return_value + assert lower.call_args.args == (("rewritten", program),) + assert lower.call_args.kwargs["transform_passes"] == [transform_pass] + assert lower.call_args.kwargs["compile_config"] is edge_config + assert lower.call_args.kwargs["constant_methods"] == constant_methods + assert lower.call_args.kwargs["generate_etrecord"] is True + lowered_partitioners = lower.call_args.kwargs["partitioner"] + assert isinstance(lowered_partitioners[0], FakeTensorRTPartitioner) + assert lowered_partitioners[0].compile_specs == [compile_spec] + assert lowered_partitioners[1:] == [extra_a, extra_b] + assert partitioners == [extra_a, extra_b] + assert compile_specs == [compile_spec] + + +@pytest.mark.unit +def test_export_preserves_independent_method_mapping(monkeypatch): + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + prefill_extra = object() + decode_extra = object() + prefill_spec = object() + decode_spec = object() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [prefill_extra], "decode": [decode_extra]}, + compile_specs={"prefill": [prefill_spec], "decode": [decode_spec]}, + ) + + assert lower.call_args.args == ( + {"prefill": ("rewritten", prefill), "decode": ("rewritten", decode)}, + ) + pipelines = lower.call_args.kwargs["partitioner"] + assert pipelines["prefill"][0].compile_specs == [prefill_spec] + assert pipelines["prefill"][1:] == [prefill_extra] + assert pipelines["decode"][0].compile_specs == [decode_spec] + assert pipelines["decode"][1:] == [decode_extra] + + +@pytest.mark.unit +def test_stage_exported_program_isolates_structure_and_shares_payloads(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2)) + + def forward(self, x): + return x + self.weight + + program = torch.export.export(Model(), (torch.ones(2),)) + + class Metadata(NamedTuple): + shape: tuple[int, ...] + dtype: torch.dtype + + first_node = next(iter(program.graph.nodes)) + first_node.meta["nested"] = {"items": ["source"]} + first_node.meta["tensor_meta"] = Metadata((2,), torch.float32) + program.graph_module.meta["nested"] = {"items": ["source"]} + staged = export_utils.stage_exported_program(program) + + assert staged is not program + assert staged.graph_module is not program.graph_module + assert staged.graph is not program.graph + assert staged.state_dict is not program.state_dict + assert staged.constants is not program.constants + assert staged.state_dict["weight"] is program.state_dict["weight"] + assert isinstance(next(iter(staged.graph.nodes)).meta["tensor_meta"], Metadata) + + staged.state_dict["staged_only"] = torch.zeros(1) + staged.graph_module.meta["staged_only"] = True + staged.graph_module.meta["nested"]["items"].append("staged") + next(iter(staged.graph.nodes)).meta["nested"]["items"].append("staged") + assert "staged_only" not in program.state_dict + assert "staged_only" not in program.graph_module.meta + assert program.graph_module.meta["nested"]["items"] == ["source"] + assert first_node.meta["nested"]["items"] == ["source"] + + +@pytest.mark.unit +def test_stage_exported_program_supports_dynamic_shapes(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(4, 4) + + def forward(self, x): + return self.lin(x) * x.shape[0] + + batch = torch.export.Dim("batch", min=1, max=32) + program = torch.export.export( + Model(), (torch.randn(2, 4),), dynamic_shapes={"x": {0: batch}} + ) + assert program.range_constraints + + staged = export_utils.stage_exported_program(program) + + # Symbolic metadata stays bound to the original ShapeEnv, so it must be shared + # rather than copied for the staged program to remain consistent. + source_nodes = {node.name: node for node in program.graph.nodes} + symbolic_shared = 0 + for node in staged.graph.nodes: + source_node = source_nodes[node.name] + for key, source_value in source_node.meta.items(): + if not export_utils._is_graph_bound_metadata(source_value): + continue + assert node.meta[key] is source_value + symbolic_shared += 1 + assert symbolic_shared + + assert staged.graph_module is not program.graph_module + assert staged.state_dict["lin.weight"] is program.state_dict["lin.weight"] + + +@pytest.mark.unit +def test_stage_exported_program_clones_nested_graph_modules(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Leaf(torch.nn.Module): + def forward(self, x): + return x + 1 + + root = torch.fx.Graph() + x = root.placeholder("x") + call = root.call_module("nested", (x,)) + root.output(call) + nested = torch.fx.symbolic_trace(Leaf()) + graph_module = torch.fx.GraphModule({"nested": nested}, root) + nested_node = next(iter(graph_module.nested.graph.nodes)) + nested_payload = torch.ones(1) + nested_node.meta["nested"] = { + "items": ["source"], + "payload": nested_payload, + } + payload_memo = export_utils._payload_sharing_memo( + SimpleNamespace(graph_module=graph_module, state_dict={}, constants={}) + ) + staged = export_utils._stage_graph_module(graph_module, payload_memo) + + assert staged.nested is not graph_module.nested + staged.nested.meta["staged_only"] = True + next(iter(staged.nested.graph.nodes)).meta["nested"]["items"].append("staged") + staged_nested_meta = next(iter(staged.nested.graph.nodes)).meta["nested"] + assert staged_nested_meta["payload"] is nested_payload + assert "staged_only" not in graph_module.nested.meta + assert nested_node.meta["nested"]["items"] == ["source"] + + +@pytest.mark.unit +def test_export_prepares_compiled_graph_module(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + prepared = FakeExportedProgram() + prepare = MagicMock(return_value=prepared) + monkeypatch.setattr(export_module, "_prepare_graph_module", prepare) + inputs = [torch.ones(1)] + + export_module.export(graph_module, arg_inputs=inputs, retrace=False) + + prepare.assert_called_once_with( + graph_module, + arg_inputs=inputs, + kwarg_inputs=None, + dynamic_shapes=None, + retrace=False, + ) + assert lower.call_args.args == (("rewritten", prepared),) + + +@pytest.mark.unit +@pytest.mark.parametrize("input_option", ["inputs", "arg_inputs"]) +def test_export_forwards_input_alias(monkeypatch, input_option): + export_module, _ = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + prepared = FakeExportedProgram() + prepare = MagicMock(return_value=prepared) + monkeypatch.setattr(export_module, "_prepare_graph_module", prepare) + example_inputs = [torch.ones(1)] + + export_module.export(graph_module, **{input_option: example_inputs}, retrace=False) + + prepare.assert_called_once_with( + graph_module, + arg_inputs=example_inputs, + kwarg_inputs=None, + dynamic_shapes=None, + retrace=False, + ) + + +@pytest.mark.unit +def test_export_rejects_inputs_and_arg_inputs(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + + with pytest.raises(ValueError, match="mutually exclusive"): + export_module.export( + graph_module, + inputs=[torch.ones(1)], + arg_inputs=[torch.ones(1)], + ) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_rejects_non_linux_platform(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + monkeypatch.setattr(export_module.platform, "system", lambda: "Windows") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + + with pytest.raises(ValueError, match="only supported on Linux"): + export_module.export(graph_module) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_prepare_graph_module_preserves_tensor_keyword_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + keyword_tensor = torch.ones(2, device="cuda") + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + monkeypatch.setattr( + "torch_tensorrt.dynamo._defaults.default_device", + lambda: torch.device("cuda"), + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(), + kwarg_inputs={"mask": keyword_tensor}, + dynamic_shapes=None, + retrace=True, + ) + is exported + ) + assert ( + export_graph_module.call_args.kwargs["kwarg_inputs"]["mask"] is keyword_tensor + ) + + +@pytest.mark.unit +def test_export_allows_zero_engine_program(monkeypatch): + program = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch, {program: 0}) + + assert export_module.export(program) is lower.return_value + lowered_partitioners = lower.call_args.kwargs["partitioner"] + assert len(lowered_partitioners) == 1 + assert isinstance(lowered_partitioners[0], FakeTensorRTPartitioner) + + +@pytest.mark.unit +def test_export_rejects_duplicate_program_identity_before_validation(monkeypatch): + program = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + validate = MagicMock(return_value=1) + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.validate_engine_program", validate + ) + + with pytest.raises(ValueError, match="same ExportedProgram object"): + export_module.export({"prefill": program, "decode": program}) + + validate.assert_not_called() + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_validates_all_methods_before_rewriting(monkeypatch): + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + validate = MagicMock(side_effect=[1, RuntimeError("decode is invalid")]) + replace = MagicMock() + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.validate_engine_program", validate + ) + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.replace_execute_engine", replace + ) + + with pytest.raises(RuntimeError, match="decode is invalid"): + export_module.export({"prefill": prefill, "decode": decode}) + + replace.assert_not_called() + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_rewrite_failure_leaves_sources_unchanged(monkeypatch): + import executorch.exir + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + class Model(torch.nn.Module): + def forward(self, x): + return x + 1 + + prefill = torch.export.export(Model(), (torch.ones(1),)) + decode = torch.export.export(Model(), (torch.ones(1),)) + original_codes = { + "prefill": prefill.graph_module.code, + "decode": decode.graph_module.code, + } + original_state_keys = { + "prefill": set(prefill.state_dict), + "decode": set(decode.state_dict), + } + + monkeypatch.setattr(export_utils, "validate_engine_program", lambda program: 1) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + lower = MagicMock() + monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) + call_count = 0 + + def fail_second_rewrite(program): + nonlocal call_count + call_count += 1 + program.state_dict["staged_only"] = torch.zeros(1) + program.graph_module.meta["staged_only"] = True + if call_count == 2: + raise RuntimeError("rewrite failed") + return program + + monkeypatch.setattr(export_utils, "replace_execute_engine", fail_second_rewrite) + + with pytest.raises(RuntimeError, match="rewrite failed"): + executorch_api.export({"prefill": prefill, "decode": decode}) + + assert prefill.graph_module.code == original_codes["prefill"] + assert decode.graph_module.code == original_codes["decode"] + assert set(prefill.state_dict) == original_state_keys["prefill"] + assert set(decode.state_dict) == original_state_keys["decode"] + assert "staged_only" not in prefill.graph_module.meta + assert "staged_only" not in decode.graph_module.meta + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_lowering_failure_leaves_source_unchanged(monkeypatch): + import executorch.exir + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + class Model(torch.nn.Module): + def forward(self, x): + return x + 1 + + source = torch.export.export(Model(), (torch.ones(1),)) + original_code = source.graph_module.code + original_state_keys = set(source.state_dict) + + monkeypatch.setattr(export_utils, "validate_engine_program", lambda program: 1) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + + def mutate_staged_program(program): + program.state_dict["staged_only"] = torch.zeros(1) + program.graph_module.meta["staged_only"] = True + return program + + monkeypatch.setattr(export_utils, "replace_execute_engine", mutate_staged_program) + monkeypatch.setattr( + executorch.exir, + "to_edge_transform_and_lower", + MagicMock(side_effect=RuntimeError("lowering failed")), + ) + + with pytest.raises(RuntimeError, match="lowering failed"): + executorch_api.export(source) + + assert source.graph_module.code == original_code + assert set(source.state_dict) == original_state_keys + assert "staged_only" not in source.graph_module.meta + + +@pytest.mark.unit +def test_export_accepts_false_retrace_for_exported_program(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + export_module.export(program, retrace=False) + + assert lower.call_args.args == (("rewritten", program),) + + +@pytest.mark.unit +def test_export_rejects_true_retrace_for_exported_program(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(ValueError, match="already-exported program"): + export_module.export(program, retrace=True) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_normalizes_none_per_method_values(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": None}, + compile_specs={"decode": None}, + ) + + pipelines = lower.call_args.kwargs["partitioner"] + assert len(pipelines["prefill"]) == 1 + assert len(pipelines["decode"]) == 1 + assert pipelines["prefill"][0].compile_specs == [] + assert pipelines["decode"][0].compile_specs == [] + + +@pytest.mark.unit +@pytest.mark.parametrize("container_type", [list, tuple]) +def test_prepare_graph_module_infers_nested_dynamic_shapes(monkeypatch, container_type): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class NestedModule(torch.nn.Module): + def forward(self, nested): + return nested[0] + + graph_module = torch.fx.symbolic_trace(NestedModule()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + dynamic_input = torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 2), + max_shape=(4, 2), + name="nested", + ) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(container_type([dynamic_input]),), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + nested_shapes = export_graph_module.call_args.kwargs["dynamic_shapes"]["nested"] + assert isinstance(nested_shapes, container_type) + assert len(nested_shapes) == 1 + assert set(nested_shapes[0]) == {0} + + +@pytest.mark.unit +def test_prepare_graph_module_preserves_shared_dynamic_dimensions(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class SharedBatchModule(torch.nn.Module): + def forward(self, left, right): + return left + right + + graph_module = torch.fx.symbolic_trace(SharedBatchModule()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + inputs = tuple( + torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 2), + max_shape=(4, 2), + name=name, + shared_dims={0: "batch"}, + ) + for name in ("left", "right") + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=inputs, + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + dynamic_shapes = export_graph_module.call_args.kwargs["dynamic_shapes"] + assert dynamic_shapes["left"][0] is dynamic_shapes["right"][0] + + +@pytest.mark.unit +def test_prepare_graph_module_requires_shapes_for_mixed_dynamic_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class MixedModule(torch.nn.Module): + def forward(self, dynamic, static): + return dynamic + static + + graph_module = torch.fx.symbolic_trace(MixedModule()) + dynamic_input = torch_tensorrt.Input( + min_shape=(1,), opt_shape=(2,), max_shape=(4,), name="dynamic" + ) + static_input = torch.ones(2) + + with pytest.raises(ValueError, match="require explicit dynamic_shapes"): + export_module._prepare_graph_module( + graph_module, + arg_inputs=(dynamic_input, static_input), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + explicit_shapes = ({0: torch.export.Dim("batch", min=1, max=4)}, None) + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(dynamic_input, static_input), + kwarg_inputs={}, + dynamic_shapes=explicit_shapes, + retrace=False, + ) + is exported + ) + assert export_graph_module.call_args.kwargs["dynamic_shapes"] is explicit_shapes + + +@pytest.mark.unit +def test_prepare_graph_module_does_not_infer_shapes_without_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + infer_args = MagicMock() + infer_kwargs = MagicMock() + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + monkeypatch.setattr( + "torch_tensorrt.dynamo._tracer.get_dynamic_shapes_args", infer_args + ) + monkeypatch.setattr( + "torch_tensorrt.dynamo._tracer.get_dynamic_shapes_kwargs", infer_kwargs + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + infer_args.assert_not_called() + infer_kwargs.assert_not_called() + assert export_graph_module.call_args.kwargs["dynamic_shapes"] is None + + +@pytest.mark.unit +def test_export_rejects_ambiguous_sources_and_options(monkeypatch): + export_module, _ = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(ValueError, match="already-exported program"): + export_module.export(program, arg_inputs=[torch.ones(1)]) + with pytest.raises(TypeError, match="Compile nn.Module inputs"): + export_module.export(torch.nn.Linear(1, 1)) + with pytest.raises(ValueError, match="unknown methods"): + export_module.export({"forward": program}, partitioners={"missing": [object()]}) + with pytest.raises(ValueError, match="collide"): + export_module.export({"forward": program}, constant_methods={"forward": 1}) + + +@pytest.mark.unit +@pytest.mark.parametrize("option_name", ["partitioners", "compile_specs"]) +def test_export_rejects_string_option_sequences(monkeypatch, option_name): + export_module, _ = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(TypeError, match=option_name): + export_module.export(program, **{option_name: "invalid"}) + + with pytest.raises(TypeError, match=option_name): + export_module.export( + {"forward": program}, **{option_name: {"forward": "invalid"}} + ) + + +@pytest.mark.unit +def test_prepare_graph_module_rejects_argument_mismatch(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class OneInput(torch.nn.Module): + def forward(self, x): + return x + + graph_module = torch.fx.symbolic_trace(OneInput()) + with pytest.raises(TypeError): + export_module._prepare_graph_module( + graph_module, + arg_inputs=( + torch_tensorrt.Input((1,), name="x"), + torch_tensorrt.Input((1,), name="extra"), + ), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + + +@pytest.mark.unit +def test_export_warns_for_multiple_engines(monkeypatch, caplog): + program = FakeExportedProgram() + export_module, _ = _patch_lowering(monkeypatch, {program: 2}) + + export_module.export(program) + assert "contains 2 TRT engines" in caplog.text