Batch cheap FX cleanups and defer post_lowering recompile

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:

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/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)
 
 

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: