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