Batch cheap FX cleanups and defer post_lowering recompile

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)
 

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

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

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

py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py

@@ -3,17 +3,55 @@
 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()
     gm.recompile()
     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:

py/torch_tensorrt/dynamo/lowering/passes/remove_assert_nodes.py

@@ -9,22 +9,28 @@
 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
         ):
             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

py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py

@@ -9,23 +9,34 @@
 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
             and len(node.users) == 0
             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

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