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