From d5041eb1842c0b922e56ca04d2de384e7a2bae19 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 19 Aug 2026 14:47:12 -0700 Subject: [PATCH 1/3] Skip adjacency partitioning when the graph is fully supported. When require_full_compilation already has full converter coverage, wrap the graph as one TRT block instead of repeating AccNodesFinder, fusion, and adjacency split. --- py/torch_tensorrt/dynamo/_compiler.py | 4 +- .../partitioning/_adjacency_partitioner.py | 73 ++++++++++++++++--- 2 files changed, 64 insertions(+), 13 deletions(-) diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..dc82d93c8f 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -5,7 +5,6 @@ import os import platform import warnings - from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union import sympy @@ -1263,6 +1262,9 @@ def preserve_module_specs( torch_executed_ops=settings.torch_executed_ops, require_full_compilation=settings.require_full_compilation, skip_fusion=(num_supported_ops == total_ops), + assume_full_support=( + settings.require_full_compilation and num_supported_ops == total_ops + ), ) except torch.fx.passes.splitter_base.FxNetSplitterInternalError: diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 098e9b2685..979d2c476f 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -138,6 +138,7 @@ def __init__( require_full_compilation: bool = REQUIRE_FULL_COMPILATION, return_tuple: bool = False, skip_fusion: bool = False, + assume_full_support: bool = False, ): """ Preprocesses graph before splitting: @@ -157,20 +158,30 @@ def __init__( skip_fusion=skip_fusion, ) self.operator_support = operator_support - - # Get all accelerated nodes based on operator support conditions - self.acc_nodes = FxNetAccNodesFinder( - self.module, self.operator_support, self.settings.allow_non_tensor - )() - - if self.settings.skip_fusion: - self.fusions = {} + self.assume_full_support = assume_full_support + + if self.assume_full_support: + # The caller already walked the graph and verified converter support. + # Avoid repeating support discovery and dependency construction. + self.acc_nodes = { + node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS + } + self.fusions: Dict[torch.fx.Node, NodeSet] = {} + self.deps: Dict[torch.fx.Node, NodeSet] = {} else: - self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))() + # Get all accelerated nodes based on operator support conditions + self.acc_nodes = FxNetAccNodesFinder( + self.module, self.operator_support, self.settings.allow_non_tensor + )() - # Modify deps to add more deps for fused nodes - self.deps = self.find_deps() - self.update_deps_for_fusions() + if self.settings.skip_fusion: + self.fusions = {} + else: + self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))() + + # Modify deps to add more deps for fused nodes + self.deps = self.find_deps() + self.update_deps_for_fusions() self.non_acc_submodule_name = "_run_on_gpu_" self._node_submodule_map: Dict[str, str] = {} @@ -223,6 +234,40 @@ def partition_graph(self) -> torch.fx.GraphModule: Returns a GraphModule with submodules for each segment """ + unsupported = getattr(self.operator_support, "unsupported_operators", None) + # The explicit assumption comes from the compiler's earlier support walk. + # Otherwise, an empty dict means AccNodesFinder found no unsupported ops. + fully_supported = self.assume_full_support or ( + isinstance(unsupported, dict) and len(unsupported) == 0 + ) + + # Fast path: user demanded a single TRT engine and every op is convertible. + # Skip adjacency splitting; emit one ACC block and tag/split as usual so the + # rest of compile still sees `_run_on_acc_*`. + if self.require_full_compilation and fully_supported: + if self.settings.min_acc_module_size != MIN_BLOCK_SIZE: + logger.warning( + "Detected both require_full_compilation and min_block_size compilation " + "arguments were specified. Disregarding min_block_size argument for " + "fully supported model." + ) + nodes = [ + node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS + ] + if not nodes: + raise AssertionError( + "require_full_compilation=True was specified, but no accelerated " + "operators were found in the graph" + ) + logger.info( + "require_full_compilation + full operator support: " + "wrapping graph as a single TRT submodule (skipping adjacency split)" + ) + subgraphs = [Subgraph(is_acc=True, nodes=nodes)] + self.num_trt_accelerated_subgraphs = 1 + self.tag(subgraphs) + return self.split(remove_tag=True) + # Delegate nodes based on operator coverage subgraphs = self.put_nodes_into_subgraphs() @@ -280,6 +325,7 @@ def partition( torch_executed_ops: Collection[Target] = set(), require_full_compilation: bool = REQUIRE_FULL_COMPILATION, skip_fusion: bool = False, + assume_full_support: bool = False, ) -> Tuple[torch.fx.GraphModule, OpSupportTester]: """Partition an FX GraphModule with aten ops into TRT engines Partitioning is based on converter operator support @@ -290,6 +336,8 @@ def partition( torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage require_full_compilation: Require that all computational operators be run in TRT skip_fusion: Skip fusions found by FxNetAccFusionsFinder + assume_full_support: Skip repeated support/dependency discovery because + the caller already verified that every computational op is supported Returns: torch.fx.GraphModule, OpSupportTester """ @@ -306,6 +354,7 @@ def partition( min_block_size=min_block_size, require_full_compilation=require_full_compilation, skip_fusion=skip_fusion, + assume_full_support=assume_full_support, ) partitioned_graph = partitioner.partition_graph() From 486fc5b25a1a8369c232a15c562697c6a78c70cf Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 19 Aug 2026 14:47:34 -0700 Subject: [PATCH 2/3] Skip installing large aliased constant-folded weights. Avoid cpu().contiguous() copies of folded tensors that already share storage with module weights under offload_module_to_cpu, while still installing materialized folds. --- .../lowering/passes/constant_folding.py | 57 +++++++++++++++++-- 1 file changed, 51 insertions(+), 6 deletions(-) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 60dec56f3b..8f91ee9606 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -18,6 +18,30 @@ logger = logging.getLogger(__name__) +# Skip installing large folded tensors that already share storage with an +# existing module attribute (typically weight permutes/views). On Flux NVFP4 +# those dominate replace cost via cpu().contiguous() under offload_module_to_cpu. +# Do NOT skip materialized folds (e.g. VLA position embedding outputs) — those +# can be required for TRT legality (int64 indices never reach the converter). +_MAX_CONSTANT_FOLD_BYTES = 1 << 20 # 1 MiB + + +def _tensor_reuses_module_storage( + gm: torch.fx.GraphModule, constant: torch.Tensor +) -> bool: + try: + ptr = constant.untyped_storage().data_ptr() + except Exception: + return False + for tensor in gm.state_dict().values(): + if isinstance(tensor, torch.Tensor): + try: + if tensor.untyped_storage().data_ptr() == ptr: + return True + except Exception: + continue + return False + @torch.utils._python_dispatch._disable_current_modes() # type: ignore def constant_fold( @@ -35,17 +59,38 @@ def constant_fold( # The constants are created on CPU to save GPU memory for TensorRT compilation. # For TRT INetwork construction the constants are moved to CPU in get_attr call. + skipped_alias = 0 for node, constant in cf.node_replacements.items(): + if isinstance(constant, torch.Tensor): + nbytes = int(constant.numel() * constant.element_size()) + if nbytes > _MAX_CONSTANT_FOLD_BYTES and _tensor_reuses_module_storage( + gm, constant + ): + skipped_alias += 1 + logger.debug( + "Skipping constant-fold install for aliased %s (%d bytes > %d)", + node.name, + nbytes, + _MAX_CONSTANT_FOLD_BYTES, + ) + continue + # Register folded values as plain tensors (buffers), matching Inductor. if settings.offload_module_to_cpu: replace_node_with_constant( gm, node, - torch.nn.Parameter(constant.cpu().contiguous(), requires_grad=False), + constant.cpu().contiguous(), ) else: - replace_node_with_constant( - gm, node, torch.nn.Parameter(constant, requires_grad=False) - ) + replace_node_with_constant(gm, node, constant) + + if skipped_alias: + logger.info( + "Skipped installing %d large aliased folded constant(s) (>%d bytes); " + "leaving original view/permute ops in the graph", + skipped_alias, + _MAX_CONSTANT_FOLD_BYTES, + ) erased_params = [] for node in gm.graph.nodes: @@ -71,7 +116,7 @@ def replace_node_with_constant( """Adapted from: https://github.com/pytorch/pytorch/blob/bcf35c6ae62bb6560befa3550e37a8283944e5f4/torch/_inductor/constant_folding.py#L17-L43 - Modified to register parameters, instead of buffers for frozen constants + Registers frozen constants as buffers (same as Inductor), not Parameters. """ g = gm.graph @@ -95,7 +140,7 @@ def replace_node_with_constant( g.erase_node(node) # Needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning - gm.register_parameter(qualname, constant) + gm.register_buffer(qualname, constant) setattr(gm, qualname, constant) From b287f3aeaf98cb08da8f8eb67dd013be0dc42760 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 19 Aug 2026 14:47:35 -0700 Subject: [PATCH 3/3] Batch cheap FX cleanups and defer graph recompile to one post_lowering flush. Run non-conflicting node repairs together and skip repeated DCE/lint/recompile so post_lowering pays for cleanup once. --- .../lowering/passes/_aten_lowering_pass.py | 21 +++++--- .../passes/batch_cheap_fx_cleanups.py | 52 +++++++++++++++++++ .../passes/eliminate_sym_min_int64_max.py | 30 +++++++---- .../passes/normalize_negative_slice_stop.py | 29 +++++++---- .../dynamo/lowering/passes/pass_utils.py | 48 +++++++++++++++-- .../lowering/passes/remove_assert_nodes.py | 22 +++++--- .../passes/remove_num_users_is_0_nodes.py | 25 ++++++--- .../lowering/passes/replace_fused_rms_norm.py | 5 +- 8 files changed, 181 insertions(+), 51 deletions(-) create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/batch_cheap_fx_cleanups.py diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index f12eef79a7..71ca9a57cd 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -7,23 +7,23 @@ from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.lowering.passes._FakeTensorUpdater import FakeTensorUpdater from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + flush_deferred_graph_cleanup, + set_defer_graph_cleanup, trace_intermediate_node_outputs, ) from .annotate_fp8_sdpa import annotate_fp8_sdpa +from .batch_cheap_fx_cleanups import batch_cheap_fx_cleanups from .complex_graph_rewrite import complex_graph_detection from .constant_folding import constant_fold from .decompose_dynamic_slice_scatter import decompose_dynamic_slice_scatter -from .eliminate_sym_min_int64_max import eliminate_sym_min_int64_max from .force_causal_efficient_attention import force_causal_efficient_attention from .fuse_pad_into_convolution import fuse_pad_into_convolution from .fuse_prims_broadcast import fuse_prims_broadcast -from .normalize_negative_slice_stop import normalize_negative_slice_stop from .pass_manager import DynamoPassManager from .remove_assert_nodes import remove_assert_nodes from .remove_detach import remove_detach from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones -from .remove_num_users_is_0_nodes import remove_num_users_is_0_nodes from .repair_input_as_output import repair_input_as_output from .replace_fused_rms_norm import replace_fused_rms_norm from .replace_max_pool_with_indices import replace_max_pool_with_indices @@ -43,12 +43,10 @@ fuse_prims_broadcast, replace_max_pool_with_indices, fuse_pad_into_convolution, - remove_assert_nodes, - remove_num_users_is_0_nodes, + # One cleanup cycle for non-conflicting node repairs. + batch_cheap_fx_cleanups, complex_graph_detection, force_causal_efficient_attention, - eliminate_sym_min_int64_max, - normalize_negative_slice_stop, annotate_fp8_sdpa, decompose_dynamic_slice_scatter, ] @@ -146,7 +144,14 @@ def post_lowering( ) fake_mode = torch._export.utils._detect_fake_mode_from_gm(gm) fake_tensor_updater = FakeTensorUpdater(gm) - gm = ATEN_POST_LOWERING_PASSES(gm, settings) + # Batch DCE/lint/recompile across passes: each pass may still call + # clean_up_graph_after_modifications, but only the final flush pays for it. + set_defer_graph_cleanup(True) + try: + gm = ATEN_POST_LOWERING_PASSES(gm, settings) + gm = flush_deferred_graph_cleanup(gm) + finally: + set_defer_graph_cleanup(False) if fake_mode is not None: fake_tensor_updater.incremental_update(fake_mode) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/batch_cheap_fx_cleanups.py b/py/torch_tensorrt/dynamo/lowering/passes/batch_cheap_fx_cleanups.py new file mode 100644 index 0000000000..c6794ed19c --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/batch_cheap_fx_cleanups.py @@ -0,0 +1,52 @@ +"""Batch cheap, non-conflicting FX cleanups into one graph repair cycle. + +Several post-lowering passes only delete/rewrite a few node kinds, then each +calls ``clean_up_graph_after_modifications`` (DCE + lint + recompile). On large +graphs that recompile dominates and is paid repeatedly. + +Naren's guidance: use one iteration / one cleanup for repairs that do not +conflict. This pass runs those mutations back-to-back and cleans up once. +""" + +from __future__ import annotations + +import logging + +import torch +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.eliminate_sym_min_int64_max import ( + apply_eliminate_sym_min_int64_max, +) +from torch_tensorrt.dynamo.lowering.passes.normalize_negative_slice_stop import ( + apply_normalize_negative_slice_stop, +) +from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + clean_up_graph_after_modifications, +) +from torch_tensorrt.dynamo.lowering.passes.remove_assert_nodes import ( + apply_remove_assert_nodes, +) +from torch_tensorrt.dynamo.lowering.passes.remove_num_users_is_0_nodes import ( + apply_remove_num_users_is_0_nodes, +) + +logger = logging.getLogger(__name__) + + +def batch_cheap_fx_cleanups( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Apply cheap FX cleanups with a single trailing graph cleanup.""" + del settings # unused; kept for DynamoPassManager signature + modified = False + modified |= apply_remove_assert_nodes(gm) + # Dead-user removal after assert erasure so newly orphaned nodes go away. + modified |= apply_remove_num_users_is_0_nodes(gm) + modified |= apply_eliminate_sym_min_int64_max(gm) + modified |= apply_normalize_negative_slice_stop(gm) + + if modified: + gm = clean_up_graph_after_modifications(gm) + logger.debug("Graph after batch_cheap_fx_cleanups:\n%s", gm.graph) + + return gm diff --git a/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py b/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py index d1d40bda2e..31706de24e 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py @@ -13,17 +13,10 @@ def _is_int64_max(x: object) -> bool: return isinstance(x, int) and x in (sys.maxsize, _INT64_MAX) -def eliminate_sym_min_int64_max( - gm: GraphModule, settings: object = None -) -> GraphModule: - """Remove no-op sym_min nodes where one operand is INT64_MAX. - - torch.export may emit sym_min(sym, INT64_MAX) for an effectively unbounded - symbolic value. That expression is equivalent to sym, and leaving it in the - graph can produce runtime calls to torch.sym_min with Tensor inputs. - """ +def apply_eliminate_sym_min_int64_max(gm: GraphModule) -> bool: + """Remove no-op sym_min(INT64_MAX) nodes in-place. Returns if changed.""" if _SYM_MIN is None: - return gm + return False modified = False for node in list(gm.graph.nodes): @@ -46,4 +39,19 @@ def eliminate_sym_min_int64_max( gm.graph.erase_node(node) modified = True - return clean_up_graph_after_modifications(gm) if modified else gm + return modified + + +def eliminate_sym_min_int64_max( + gm: GraphModule, settings: object = None +) -> GraphModule: + """Remove no-op sym_min nodes where one operand is INT64_MAX. + + torch.export may emit sym_min(sym, INT64_MAX) for an effectively unbounded + symbolic value. That expression is equivalent to sym, and leaving it in the + graph can produce runtime calls to torch.sym_min with Tensor inputs. + """ + del settings + if apply_eliminate_sym_min_int64_max(gm): + return clean_up_graph_after_modifications(gm) + return gm diff --git a/py/torch_tensorrt/dynamo/lowering/passes/normalize_negative_slice_stop.py b/py/torch_tensorrt/dynamo/lowering/passes/normalize_negative_slice_stop.py index b94d88d46b..08d9a2013c 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/normalize_negative_slice_stop.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/normalize_negative_slice_stop.py @@ -1,10 +1,9 @@ import operator from typing import Optional, cast -from torch_tensorrt.dynamo.lowering._SubgraphBuilder import SubgraphBuilder - import torch from torch.fx import GraphModule, Node +from torch_tensorrt.dynamo.lowering._SubgraphBuilder import SubgraphBuilder from .pass_utils import clean_up_graph_after_modifications @@ -31,14 +30,8 @@ def _rank(x: Node) -> Optional[int]: return None -def normalize_negative_slice_stop( - gm: GraphModule, settings: object = None -) -> GraphModule: - """Normalize negative symbolic slice bounds to positive dim-relative bounds. - - Python slicing accepts negative bounds such as x[-n:] or x[:-n]. TensorRT - shape expressions need the equivalent positive bound, dim_size - n. - """ +def apply_normalize_negative_slice_stop(gm: GraphModule) -> bool: + """Normalize negative symbolic slice bounds in-place. Returns if changed.""" modified = False for node in list(gm.graph.nodes): @@ -83,4 +76,18 @@ def normalize_negative_slice_stop( node.args = tuple(args) modified = True - return clean_up_graph_after_modifications(gm) if modified else gm + return modified + + +def normalize_negative_slice_stop( + gm: GraphModule, settings: object = None +) -> GraphModule: + """Normalize negative symbolic slice bounds to positive dim-relative bounds. + + Python slicing accepts negative bounds such as x[-n:] or x[:-n]. TensorRT + shape expressions need the equivalent positive bound, dim_size - n. + """ + del settings + if apply_normalize_negative_slice_stop(gm): + return clean_up_graph_after_modifications(gm) + return gm diff --git a/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py b/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py index 2a51bbac6f..69adcb8221 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py @@ -3,10 +3,33 @@ import torch from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES +# When True, clean_up_graph_after_modifications only marks the graph dirty and +# skips DCE/lint/recompile. post_lowering enables this around the pass list so +# multiple cheap FX repairs share one cleanup cycle (Naren: "one iteration"). +_defer_graph_cleanup: bool = False +_graph_cleanup_pending: bool = False -def clean_up_graph_after_modifications( + +def set_defer_graph_cleanup(enabled: bool) -> None: + """Enable/disable deferred FX cleanup for the current post_lowering run.""" + global _defer_graph_cleanup, _graph_cleanup_pending + _defer_graph_cleanup = enabled + if not enabled: + _graph_cleanup_pending = False + + +def flush_deferred_graph_cleanup( gm: torch.fx.GraphModule, ) -> torch.fx.GraphModule: + """Run a real cleanup if any deferred clean_up call happened.""" + global _graph_cleanup_pending + if _graph_cleanup_pending: + _graph_cleanup_pending = False + return _run_graph_cleanup(gm) + return gm + + +def _run_graph_cleanup(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: """Runs dead-code elimination, linting, and recompilation for graph, in-place""" gm.graph.eliminate_dead_code() gm.graph.lint() @@ -14,6 +37,21 @@ def clean_up_graph_after_modifications( return gm +def clean_up_graph_after_modifications( + gm: torch.fx.GraphModule, +) -> torch.fx.GraphModule: + """Runs dead-code elimination, linting, and recompilation for graph, in-place. + + If deferred cleanup is enabled (see ``set_defer_graph_cleanup``), records that + a cleanup is needed and returns immediately so callers can batch mutations. + """ + global _graph_cleanup_pending + if _defer_graph_cleanup: + _graph_cleanup_pending = True + return gm + return _run_graph_cleanup(gm) + + def get_tensor_placeholders( gm: torch.fx.GraphModule, ) -> List[torch.fx.Node]: @@ -32,16 +70,16 @@ def get_tensor_placeholders( return placeholders -def find_complex_nodes(gm: torch.fx.GraphModule): - complex_nodes = [] - complexNodes = {} +def find_complex_nodes(gm: torch.fx.GraphModule) -> List[torch.fx.Node]: + complex_nodes: List[torch.fx.Node] = [] + complexNodes: Dict[str, bool] = {} for node in gm.graph.nodes: if is_node_complex(node, complexNodes): complex_nodes.append(node) return complex_nodes -def is_node_complex(node: torch.fx.Node, complexNodes): +def is_node_complex(node: Any, complexNodes: Dict[str, bool]) -> bool: if not isinstance(node, torch.fx.Node): return False if node.name in complexNodes: diff --git a/py/torch_tensorrt/dynamo/lowering/passes/remove_assert_nodes.py b/py/torch_tensorrt/dynamo/lowering/passes/remove_assert_nodes.py index 890391e280..09c97b93e3 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/remove_assert_nodes.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/remove_assert_nodes.py @@ -9,12 +9,10 @@ logger = logging.getLogger(__name__) -def remove_assert_nodes( - gm: torch.fx.GraphModule, settings: CompilationSettings -) -> torch.fx.GraphModule: - """Remove assert_scalar ops in the graph""" +def apply_remove_assert_nodes(gm: torch.fx.GraphModule) -> bool: + """Remove assert ops in-place. Returns True if the graph changed.""" count = 0 - for node in gm.graph.nodes: + for node in list(gm.graph.nodes): if ( node.target == torch.ops.aten._assert_scalar.default or node.target == torch.ops.aten._assert_tensor_metadata.default @@ -22,9 +20,17 @@ def remove_assert_nodes( gm.graph.erase_node(node) count += 1 - if count > 0: - gm = clean_up_graph_after_modifications(gm) + if count: + logger.debug("Removed %d assert nodes", count) + return count > 0 - logger.debug(f"Removed {count} assert_scalar nodes:\n{gm.graph}") +def remove_assert_nodes( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Remove assert_scalar ops in the graph""" + del settings + if apply_remove_assert_nodes(gm): + gm = clean_up_graph_after_modifications(gm) + logger.debug("Graph after remove_assert_nodes:\n%s", gm.graph) return gm diff --git a/py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py b/py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py index a9b7c48ec2..f9de8937b3 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py @@ -9,13 +9,14 @@ logger = logging.getLogger(__name__) -def remove_num_users_is_0_nodes( - gm: torch.fx.GraphModule, settings: CompilationSettings -) -> torch.fx.GraphModule: - """Remove ops that [num_users=0] in the graph""" +def apply_remove_num_users_is_0_nodes(gm: torch.fx.GraphModule) -> bool: + """Remove unused ops in-place. Returns True if the graph changed.""" nodes = list(gm.graph.nodes) - output_node = nodes[-1] + if not nodes: + return False + output_node = nodes[-1] + erased = 0 for node in nodes[::-1]: if ( node != output_node @@ -23,9 +24,19 @@ def remove_num_users_is_0_nodes( and len(node.all_input_nodes) > 0 ): gm.graph.erase_node(node) + erased += 1 - gm = clean_up_graph_after_modifications(gm) + if erased: + logger.debug("Removed %d num_users=0 nodes", erased) + return erased > 0 - logger.debug(f"Removed ops that [num_users=0] nodes:\n{gm.graph}") +def remove_num_users_is_0_nodes( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Remove ops that [num_users=0] in the graph""" + del settings + if apply_remove_num_users_is_0_nodes(gm): + gm = clean_up_graph_after_modifications(gm) + logger.debug("Graph after remove_num_users_is_0_nodes:\n%s", gm.graph) return gm diff --git a/py/torch_tensorrt/dynamo/lowering/passes/replace_fused_rms_norm.py b/py/torch_tensorrt/dynamo/lowering/passes/replace_fused_rms_norm.py index b1bb768a5f..1553e6ae92 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/replace_fused_rms_norm.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/replace_fused_rms_norm.py @@ -22,7 +22,10 @@ def replace_fused_rms_norm( logger.debug(f"Replaced {count} fused rms norm nodes:\n{gm.graph}") - gm = clean_up_graph_after_modifications(gm) + # Avoid a full DCE/lint/recompile when nothing changed — on large graphs + # that cleanup alone is tens of ms. + if count > 0: + gm = clean_up_graph_after_modifications(gm) return gm