From f0695ce139ba9d5a03b87aebe020690066d28182 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Wed, 26 Aug 2026 15:03:17 -0700 Subject: [PATCH] Skip executing large aliased constant folds that 4531 will not install. Inductor's skip_folding_node_fn never fires on our cf.run() path; honor it so large weight permutes/views are not computed, then discarded. --- .../lowering/passes/constant_folding.py | 112 ++++++++++++- .../test_skip_aliased_fold_execute.py | 149 ++++++++++++++++++ 2 files changed, 258 insertions(+), 3 deletions(-) create mode 100644 tests/py/dynamo/lowering/test_skip_aliased_fold_execute.py diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 7edc85f4ec..4f12a03b5c 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/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: diff --git a/tests/py/dynamo/lowering/test_skip_aliased_fold_execute.py b/tests/py/dynamo/lowering/test_skip_aliased_fold_execute.py new file mode 100644 index 0000000000..efae337d3d --- /dev/null +++ b/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()