Skip executing large aliased constant folds that 4531 will not install.

py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

@@ -1,8 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import inspect
 import logging
-from typing import Any, Set
+from operator import attrgetter
+from typing import Any, Callable, Optional, Set
 
 import torch
 from torch_tensorrt._utils import sanitized_torch_version
@@ -31,6 +33,85 @@
 # can be required for TRT legality (int64 indices never reach the converter).
 _MAX_CONSTANT_FOLD_BYTES = 1 << 20  # 1 MiB
 
+# True views that never allocate. ``aten.contiguous`` is ``is_view`` but can
+# materialize; keep executing those so 4531 can still install the copy.
+_ALIASING_VIEW_OP_NAMES = (
+    "permute.default",
+    "transpose.int",
+    "t.default",
+    "view.default",
+    "reshape.default",
+    "squeeze.default",
+    "squeeze.dim",
+    "squeeze.dims",
+    "unsqueeze.default",
+    "expand.default",
+    "slice.Tensor",
+    "select.int",
+    "detach.default",
+    "alias.default",
+    "as_strided.default",
+    "flatten.using_ints",
+    "unflatten.int",
+    "movedim.int",
+    "movedim.intlist",
+)
+
+
+def _resolve_aten_op(name: str) -> Optional[Any]:
+    op: Any = torch.ops.aten
+    for part in name.split("."):
+        op = getattr(op, part, None)
+        if op is None:
+            return None
+    return op
+
+
+_ALIASING_VIEW_OPS: Set[Any] = {
+    op for name in _ALIASING_VIEW_OP_NAMES if (op := _resolve_aten_op(name)) is not None
+}
+
+
+def _named_attr(gm: torch.fx.GraphModule, target: Any) -> Any:
+    if not isinstance(target, str):
+        return None
+    try:
+        return attrgetter(target)(gm)
+    except (AttributeError, ValueError):
+        return None
+
+
+def _source_attr_nbytes(gm: torch.fx.GraphModule, node: torch.fx.Node) -> int:
+    """Bytes of the get_attr tensor at the root of an aliasing-view chain, else 0."""
+    seen: Set[torch.fx.Node] = set()
+    cur: Optional[torch.fx.Node] = node
+    while cur is not None and cur not in seen:
+        seen.add(cur)
+        if cur.op == "get_attr":
+            tensor = _named_attr(gm, cur.target)
+            if isinstance(tensor, torch.Tensor):
+                return int(tensor.numel() * tensor.element_size())
+            return 0
+        if cur.op != "call_function" or cur.target not in _ALIASING_VIEW_OPS:
+            return 0
+        tensor_args = [arg for arg in cur.args if isinstance(arg, torch.fx.Node)]
+        if len(tensor_args) != 1:
+            return 0
+        cur = tensor_args[0]
+    return 0
+
+
+def skip_large_aliased_view_fold(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool:
+    """Inductor ``skip_folding_node_fn``: True → do not execute this fold.
+
+    Matches the 4531 install skip: large view/permute of a module tensor that
+    would only alias existing storage. Materializing ops (add, contiguous, …)
+    still run.
+    """
+    if node.op != "call_function" or node.target not in _ALIASING_VIEW_OPS:
+        return False
+    return _source_attr_nbytes(gm, node) > _MAX_CONSTANT_FOLD_BYTES
+
 
 def _tensor_reuses_module_storage(
     gm: torch.fx.GraphModule, constant: torch.Tensor
@@ -60,7 +141,11 @@ def constant_fold(
 
     Modifies the graph in-place and replaces node with constants
     """
-    cf = _TorchTensorRTConstantFolder(gm, skip_constructors=False)
+    cf = _TorchTensorRTConstantFolder(
+        gm,
+        skip_constructors=False,
+        skip_folding_node_fn=lambda node: skip_large_aliased_view_fold(gm, node),
+    )
     cf.run()
 
     # The constants are created on CPU to save GPU memory for TensorRT compilation.
@@ -154,7 +239,14 @@ def replace_node_with_constant(
 # https://github.com/pytorch/pytorch/blob/4b881b0da390c1290bb12850ef9daad6f6eb2cb6/torch/_inductor/constant_folding.py#L53-L63
 class _TorchTensorRTConstantFolder(ConstantFolder):  # type: ignore[misc]
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
+        skip_fn = kwargs.get("skip_folding_node_fn")
+        init_params = inspect.signature(ConstantFolder.__init__).parameters
+        if "skip_folding_node_fn" not in init_params:
+            kwargs.pop("skip_folding_node_fn", None)
+            super().__init__(*args, **kwargs)
+            self.skip_folding_node_fn = skip_fn
+        else:
+            super().__init__(*args, **kwargs)
         # Set of known quantization ops to be excluded from constant folding.
         # Currently, we exclude all quantization ops coming from modelopt library.
         self.quantization_ops: Set[torch._ops.OpOverload] = set()
@@ -171,6 +263,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         except Exception as e:
             pass
 
+    def run_node(self, node: torch.fx.node.Node) -> Any:
+        # Inductor only consults skip_folding_node_fn when lifted_constant_names
+        # is set. Our cf.run() path has none, so honor the callback here and
+        # return unknown without executing the op.
+        skip_fn: Optional[Callable[[torch.fx.Node], bool]] = getattr(
+            self, "skip_folding_node_fn", None
+        )
+        if skip_fn is not None and node.op == "call_function" and skip_fn(node):
+            logger.debug(
+                "Skipping constant-fold execute for aliased view %s", node.name
+            )
+            return self.unknown_value
+        return super().run_node(node)
+
     # TODO: Update this function when quantization is added
     def is_impure(self, node: torch.fx.node.Node) -> bool:
 

tests/py/dynamo/lowering/test_skip_aliased_fold_execute.py

@@ -0,0 +1,149 @@
+import unittest
+
+import torch
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.lowering.passes.constant_folding import (
+    _MAX_CONSTANT_FOLD_BYTES,
+    _TorchTensorRTConstantFolder,
+    constant_fold,
+    skip_large_aliased_view_fold,
+)
+
+
+def _exported_gm(model: torch.nn.Module, example: torch.Tensor) -> torch.fx.GraphModule:
+    return torch.export.export(model, (example,)).module()
+
+
+def _permute_nodes(gm: torch.fx.GraphModule) -> list[torch.fx.Node]:
+    return [
+        node for node in gm.graph.nodes if node.target is torch.ops.aten.permute.default
+    ]
+
+
+def _call_targets(gm: torch.fx.GraphModule) -> set[object]:
+    return {node.target for node in gm.graph.nodes if node.op == "call_function"}
+
+
+class _WeightPermute(torch.nn.Module):
+    def __init__(self, rows: int) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(rows, rows))
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return value + self.weight.permute(1, 0)
+
+
+class _WeightAdd(torch.nn.Module):
+    def __init__(self, rows: int = 32) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(rows, rows))
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return value + (self.weight + self.weight)
+
+
+class _WeightContiguous(torch.nn.Module):
+    def __init__(self, rows: int) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(rows, rows))
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        # Tensor.contiguous() is DCE'd by export when the param is already
+        # contiguous; the ATen op stays in the graph.
+        return value + torch.ops.aten.contiguous.default(self.weight)
+
+
+class _CountingFolder(_TorchTensorRTConstantFolder):
+    def __init__(self, *args: object, **kwargs: object) -> None:
+        super().__init__(*args, **kwargs)
+        self.executed_targets: list[object] = []
+
+    def run_node(self, node: torch.fx.Node) -> object:
+        skip_fn = getattr(self, "skip_folding_node_fn", None)
+        skipping = skip_fn is not None and node.op == "call_function" and skip_fn(node)
+        if node.op == "call_function" and not skipping:
+            self.executed_targets.append(node.target)
+        return super().run_node(node)
+
+
+class TestSkipAliasedFoldExecute(unittest.TestCase):
+    def test_skip_fn_true_for_large_permute(self) -> None:
+        rows = 1024
+        self.assertGreater(rows * rows * 4, _MAX_CONSTANT_FOLD_BYTES)
+        gm = _exported_gm(_WeightPermute(rows), torch.ones(rows, rows))
+        permute_nodes = _permute_nodes(gm)
+        self.assertEqual(len(permute_nodes), 1)
+        self.assertTrue(skip_large_aliased_view_fold(gm, permute_nodes[0]))
+
+    def test_skip_fn_false_for_small_permute(self) -> None:
+        gm = _exported_gm(_WeightPermute(8), torch.ones(8, 8))
+        permute_nodes = _permute_nodes(gm)
+        self.assertEqual(len(permute_nodes), 1)
+        self.assertFalse(skip_large_aliased_view_fold(gm, permute_nodes[0]))
+
+    def test_skip_fn_false_for_large_contiguous(self) -> None:
+        rows = 1024
+        gm = _exported_gm(_WeightContiguous(rows), torch.ones(rows, rows))
+        contiguous_nodes = [
+            node
+            for node in gm.graph.nodes
+            if node.target is torch.ops.aten.contiguous.default
+        ]
+        self.assertEqual(len(contiguous_nodes), 1)
+        self.assertFalse(skip_large_aliased_view_fold(gm, contiguous_nodes[0]))
+
+    def test_large_permute_is_not_executed(self) -> None:
+        gm = _exported_gm(_WeightPermute(1024), torch.ones(1024, 1024))
+        permute_node = _permute_nodes(gm)[0]
+        folder = _CountingFolder(
+            gm,
+            skip_constructors=False,
+            skip_folding_node_fn=lambda node: skip_large_aliased_view_fold(gm, node),
+        )
+        folder.run()
+        self.assertNotIn(torch.ops.aten.permute.default, folder.executed_targets)
+        self.assertNotIn(permute_node, folder.node_replacements)
+
+    def test_small_permute_is_executed(self) -> None:
+        gm = _exported_gm(_WeightPermute(8), torch.ones(8, 8))
+        permute_node = _permute_nodes(gm)[0]
+        folder = _CountingFolder(
+            gm,
+            skip_constructors=False,
+            skip_folding_node_fn=lambda node: skip_large_aliased_view_fold(gm, node),
+        )
+        folder.run()
+        self.assertIn(torch.ops.aten.permute.default, folder.executed_targets)
+        self.assertIn(permute_node, folder.node_replacements)
+
+    def test_constant_fold_keeps_large_permute_in_graph(self) -> None:
+        gm = _exported_gm(_WeightPermute(1024), torch.ones(1024, 1024))
+        folded = constant_fold(gm, CompilationSettings())
+        self.assertIn(torch.ops.aten.permute.default, _call_targets(folded))
+
+    def test_constant_fold_still_folds_small_permute(self) -> None:
+        model = _WeightPermute(8)
+        example = torch.ones(8, 8)
+        gm = _exported_gm(model, example)
+        folded = constant_fold(gm, CompilationSettings())
+        self.assertNotIn(torch.ops.aten.permute.default, _call_targets(folded))
+        torch.testing.assert_close(folded(example), model(example))
+
+    def test_materialized_weight_add_still_folds(self) -> None:
+        model = _WeightAdd()
+        example = torch.ones(32, 32)
+        gm = _exported_gm(model, example)
+        adds_before = sum(
+            1 for node in gm.graph.nodes if node.target is torch.ops.aten.add.Tensor
+        )
+        folded = constant_fold(gm, CompilationSettings())
+        adds_after = sum(
+            1 for node in folded.graph.nodes if node.target is torch.ops.aten.add.Tensor
+        )
+        self.assertEqual(adds_before, 2)
+        self.assertEqual(adds_after, 1)
+        torch.testing.assert_close(folded(example), model(example))
+
+
+if __name__ == "__main__":
+    unittest.main()