From 7c180314c668742d9d6003ac8a5c03a16fb12d3d Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:38:41 -0700 Subject: [PATCH] feat: support configurable constant-fold exclusions --- py/torch_tensorrt/dynamo/_compiler.py | 18 ++++ py/torch_tensorrt/dynamo/_defaults.py | 1 + py/torch_tensorrt/dynamo/_settings.py | 35 ++++++++ .../constant_fold_exclusions/__init__.py | 13 +++ .../constant_fold_exclusions/_core.py | 61 +++++++++++++ .../lowering/passes/_aten_lowering_pass.py | 4 + .../lowering/passes/constant_folding.py | 7 +- .../passes/mark_constant_fold_exclusions.py | 40 +++++++++ .../lowering/test_constant_fold_exclusions.py | 85 +++++++++++++++++++ 9 files changed, 263 insertions(+), 1 deletion(-) create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py create mode 100644 py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py create mode 100644 py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py create mode 100644 tests/py/dynamo/lowering/test_constant_fold_exclusions.py diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..b6a1321342 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -117,6 +117,7 @@ def cross_compile_for_windows( cpu_memory_budget: Optional[int] = _defaults.CPU_MEMORY_BUDGET, dynamically_allocate_resources: bool = _defaults.DYNAMICALLY_ALLOCATE_RESOURCES, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS, **kwargs: Any, @@ -201,6 +202,10 @@ def cross_compile_for_windows( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -357,6 +362,7 @@ def cross_compile_for_windows( "cpu_memory_budget": cpu_memory_budget, "dynamically_allocate_resources": dynamically_allocate_resources, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, "fallback_data_dependent_ops": fallback_data_dependent_ops, } @@ -482,6 +488,7 @@ def compile( enable_resource_partitioning: bool = _defaults.ENABLE_RESOURCE_PARTITIONING, dynamically_allocate_resources: bool = _defaults.DYNAMICALLY_ALLOCATE_RESOURCES, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, fallback_data_dependent_ops: bool = _defaults.FALLBACK_DATA_DEPENDENT_OPS, **kwargs: Any, @@ -576,6 +583,10 @@ def compile( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, @@ -765,6 +776,7 @@ def compile( "cpu_memory_budget": cpu_memory_budget, "dynamically_allocate_resources": dynamically_allocate_resources, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, "fallback_data_dependent_ops": fallback_data_dependent_ops, } @@ -1715,6 +1727,7 @@ def convert_exported_program_to_serialized_trt_engine( offload_module_to_cpu: bool = _defaults.OFFLOAD_MODULE_TO_CPU, use_distributed_mode_trace: bool = _defaults.USE_DISTRIBUTED_MODE_TRACE, decompose_attention: bool = _defaults.DECOMPOSE_ATTENTION, + disabled_constant_fold_exclusions: Collection[str] = _defaults.DISABLED_CONSTANT_FOLD_EXCLUSIONS, attn_bias_is_causal: bool = _defaults.ATTN_BIAS_IS_CAUSAL, lift_mutable_buffers: bool = False, arg_input_binding_names: Any = None, @@ -1813,6 +1826,10 @@ def convert_exported_program_to_serialized_trt_engine( instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of + Torch-TensorRT rules that exclude matching FX nodes from constant + folding. Naming a rule here turns it off, so the nodes it would have + kept become foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. **kwargs: Any, Returns: @@ -1978,6 +1995,7 @@ def convert_exported_program_to_serialized_trt_engine( "offload_module_to_cpu": offload_module_to_cpu, "use_distributed_mode_trace": use_distributed_mode_trace, "decompose_attention": decompose_attention, + "disabled_constant_fold_exclusions": disabled_constant_fold_exclusions, "attn_bias_is_causal": attn_bias_is_causal, } diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 9c8a1f9f90..ce6295cf66 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -65,6 +65,7 @@ CPU_MEMORY_BUDGET = None DYNAMICALLY_ALLOCATE_RESOURCES = False DECOMPOSE_ATTENTION = False +DISABLED_CONSTANT_FOLD_EXCLUSIONS = frozenset[str]() ATTN_BIAS_IS_CAUSAL = True FALLBACK_DATA_DEPENDENT_OPS = False diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 7ff5f6bdff..243c9bb94b 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -18,6 +18,7 @@ CACHE_BUILT_ENGINES, CPU_MEMORY_BUDGET, DECOMPOSE_ATTENTION, + DISABLED_CONSTANT_FOLD_EXCLUSIONS, DISABLE_TF32, DLA_GLOBAL_DRAM_SIZE, DLA_LOCAL_DRAM_SIZE, @@ -58,6 +59,20 @@ ) +def _normalize_disabled_constant_fold_exclusions( + rule_ids: Collection[str], +) -> Set[str]: + # Deferred import: torch_tensorrt.dynamo.lowering.passes imports this module, + # so importing the rule registry at module scope would be circular. Reaching + # for the core directly is safe from any point of that cycle because it + # depends on nothing but torch. + from torch_tensorrt.dynamo.lowering.constant_fold_exclusions._core import ( + validate_disabled_constant_fold_exclusions, + ) + + return validate_disabled_constant_fold_exclusions(rule_ids) + + @dataclass class CompilationSettings: """Compilation settings for Torch-TensorRT Dynamo Paths @@ -119,6 +134,10 @@ class CompilationSettings: instead of using the attention converters. When combined with ``use_fp32_acc=True``, decomposed FP16 attention keeps its intermediate calculation in FP32 and casts only the final output back to FP16. + disabled_constant_fold_exclusions (Collection[str]): IDs of Torch-TensorRT + rules that exclude matching FX nodes from constant folding. Naming a + rule here turns it off, so the nodes it would have kept become + foldable again. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. """ @@ -178,9 +197,19 @@ class CompilationSettings: cpu_memory_budget: Optional[int] = CPU_MEMORY_BUDGET dynamically_allocate_resources: bool = DYNAMICALLY_ALLOCATE_RESOURCES decompose_attention: bool = DECOMPOSE_ATTENTION + disabled_constant_fold_exclusions: Collection[str] = ( + DISABLED_CONSTANT_FOLD_EXCLUSIONS + ) attn_bias_is_causal: bool = ATTN_BIAS_IS_CAUSAL fallback_data_dependent_ops: bool = FALLBACK_DATA_DEPENDENT_OPS + def __post_init__(self) -> None: + self.disabled_constant_fold_exclusions = ( + _normalize_disabled_constant_fold_exclusions( + self.disabled_constant_fold_exclusions + ) + ) + def __getstate__(self) -> dict[str, Any]: from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( ConverterRegistry, @@ -195,6 +224,11 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: state.pop("use_python_runtime", None) + state["disabled_constant_fold_exclusions"] = ( + _normalize_disabled_constant_fold_exclusions( + state.get("disabled_constant_fold_exclusions", ()) + ) + ) state.setdefault("fallback_data_dependent_ops", FALLBACK_DATA_DEPENDENT_OPS) self.__dict__.update(state) @@ -221,6 +255,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: "autocast_max_depth_of_reduction", "autocast_calibration_dataloader", "decompose_attention", + "disabled_constant_fold_exclusions", "attn_bias_is_causal", } diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py new file mode 100644 index 0000000000..c591d312ab --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py @@ -0,0 +1,13 @@ +from ._core import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + ConstantFoldExclusionRule, + register_constant_fold_exclusion_rule, + validate_disabled_constant_fold_exclusions, +) + +__all__ = [ + "CONSTANT_FOLD_EXCLUSION_META_KEY", + "ConstantFoldExclusionRule", + "register_constant_fold_exclusion_rule", + "validate_disabled_constant_fold_exclusions", +] diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py new file mode 100644 index 0000000000..7d38f385c6 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py @@ -0,0 +1,61 @@ +from typing import Callable, Collection, Iterable + +import torch + +CONSTANT_FOLD_EXCLUSION_META_KEY = "_torch_tensorrt_constant_fold_exclusions" + +ConstantFoldExclusionRule = Callable[[torch.fx.Node], Iterable[torch.fx.Node]] +_CONSTANT_FOLD_EXCLUSION_RULES: dict[str, ConstantFoldExclusionRule] = {} + + +def _mark_constant_fold_exclusion(nodes: Iterable[torch.fx.Node], rule_id: str) -> None: + """Record which rule wants ``nodes`` kept out of constant folding. + + The marks carry their rule ID rather than a bare flag so the exclusion pass + can revoke the ones belonging to disabled rules, whichever marking path + produced them. + """ + for node in nodes: + node.meta.setdefault(CONSTANT_FOLD_EXCLUSION_META_KEY, set()).add(rule_id) + + +def register_constant_fold_exclusion_rule( + rule_id: str, +) -> Callable[[ConstantFoldExclusionRule], ConstantFoldExclusionRule]: + """Register a named rule that selects FX nodes to exclude from folding.""" + if not isinstance(rule_id, str) or not rule_id: + raise ValueError("A constant-fold exclusion rule ID must be a non-empty string") + + def register(rule: ConstantFoldExclusionRule) -> ConstantFoldExclusionRule: + if rule_id in _CONSTANT_FOLD_EXCLUSION_RULES: + raise ValueError( + f"Constant-fold exclusion rule {rule_id!r} is already registered" + ) + + _CONSTANT_FOLD_EXCLUSION_RULES[rule_id] = rule + return rule + + return register + + +def validate_disabled_constant_fold_exclusions( + rule_ids: Collection[str], +) -> set[str]: + """Validate disabled rule IDs and return them as a set.""" + if isinstance(rule_ids, str): + raise TypeError( + "disabled_constant_fold_exclusions must be a collection of rule IDs, " + "not a single string" + ) + + disabled_rule_ids = set(rule_ids) + unknown_rule_ids = disabled_rule_ids - _CONSTANT_FOLD_EXCLUSION_RULES.keys() + if unknown_rule_ids: + available_rule_ids = ", ".join(sorted(_CONSTANT_FOLD_EXCLUSION_RULES)) + raise ValueError( + "Unknown constant-fold exclusion rule IDs: " + f"{sorted(unknown_rule_ids)}. Available rule IDs: " + f"[{available_rule_ids}]" + ) + + return disabled_rule_ids 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 4f41d3e6b7..d141feedd1 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -6,6 +6,9 @@ from torch_tensorrt._utils import is_tegra_platform from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.lowering.passes._FakeTensorUpdater import FakeTensorUpdater +from torch_tensorrt.dynamo.lowering.passes.mark_constant_fold_exclusions import ( + mark_constant_fold_exclusions, +) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( trace_intermediate_node_outputs, ) @@ -37,6 +40,7 @@ post_lowering_pass_list = [ replace_fused_rms_norm, remove_input_alias_fixing_clones, + mark_constant_fold_exclusions, constant_fold, repair_input_as_output, fuse_prims_broadcast, diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 60dec56f3b..a88aed4a9c 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -4,6 +4,9 @@ import torch from torch_tensorrt._utils import sanitized_torch_version from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, +) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( clean_up_graph_after_modifications, ) @@ -125,4 +128,6 @@ def is_impure(self, node: torch.fx.node.Node) -> bool: if node.target in self.quantization_ops: return True - return False + # The meta value holds the IDs of the rules that marked this node; it is + # empty once every rule that claimed it has been disabled. + return bool(node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY)) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py new file mode 100644 index 0000000000..a6856680e6 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py @@ -0,0 +1,40 @@ +from typing import Any, Optional + +import torch +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions._core import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + _CONSTANT_FOLD_EXCLUSION_RULES, + _mark_constant_fold_exclusion, + validate_disabled_constant_fold_exclusions, +) + + +def mark_constant_fold_exclusions( + gm: torch.fx.GraphModule, settings: Optional[Any] = None +) -> torch.fx.GraphModule: + """Apply the registered rules that exclude FX nodes from constant folding. + + This pass is the single authority on which rules are in effect. It runs + immediately before ``constant_fold`` and is the only marking path that sees + ``settings``: rules that mark nodes while a decomposition is traced run + during ``run_decompositions``, long before a settings object is reachable. + Those marks are therefore revoked here rather than suppressed where they are + made, so a caller only has to communicate the disabled rules once. + """ + disabled_rule_ids = validate_disabled_constant_fold_exclusions( + settings.disabled_constant_fold_exclusions if settings is not None else () + ) + + for node in gm.graph.nodes: + for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items(): + if rule_id in disabled_rule_ids: + continue + _mark_constant_fold_exclusion(rule(node), rule_id) + + if disabled_rule_ids: + for node in gm.graph.nodes: + marking_rule_ids = node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY) + if marking_rule_ids: + marking_rule_ids -= disabled_rule_ids + + return gm diff --git a/tests/py/dynamo/lowering/test_constant_fold_exclusions.py b/tests/py/dynamo/lowering/test_constant_fold_exclusions.py new file mode 100644 index 0000000000..b421003d83 --- /dev/null +++ b/tests/py/dynamo/lowering/test_constant_fold_exclusions.py @@ -0,0 +1,85 @@ +import torch +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo._settings import ( + CompilationSettings, + settings_are_compatible, +) +from torch_tensorrt.dynamo.lowering.constant_fold_exclusions import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + register_constant_fold_exclusion_rule, +) +from torch_tensorrt.dynamo.lowering.passes.mark_constant_fold_exclusions import ( + mark_constant_fold_exclusions, +) + +TEST_RULE_ID = "test_arbitrary_node" + + +def custom_target(): + return torch.ones(1) + + +@register_constant_fold_exclusion_rule(TEST_RULE_ID) +def custom_rule(node): + return (node,) if node.target is custom_target else () + + +class TestConstantFoldExclusionMechanics(TestCase): + def _custom_graph(self): + graph = torch.fx.Graph() + custom_node = graph.call_function(custom_target) + graph.output(custom_node) + return torch.fx.GraphModule({}, graph), custom_node + + def test_disabled_rules_default_to_empty(self): + self.assertEqual(CompilationSettings().disabled_constant_fold_exclusions, set()) + self.assertEqual( + CompilationSettings( + disabled_constant_fold_exclusions=[TEST_RULE_ID] + ).disabled_constant_fold_exclusions, + {TEST_RULE_ID}, + ) + with self.assertRaisesRegex(TypeError, "collection of rule IDs"): + CompilationSettings(disabled_constant_fold_exclusions=TEST_RULE_ID) + + def test_old_serialized_setting_defaults_to_no_disabled_rules(self): + state = CompilationSettings().__dict__.copy() + state.pop("disabled_constant_fold_exclusions") + restored = CompilationSettings.__new__(CompilationSettings) + restored.__setstate__(state) + self.assertEqual(restored.disabled_constant_fold_exclusions, set()) + + def test_setting_changes_engine_compatibility(self): + compatible, incompatible_settings = settings_are_compatible( + CompilationSettings(), + CompilationSettings(disabled_constant_fold_exclusions={TEST_RULE_ID}), + ) + self.assertFalse(compatible) + self.assertIn( + "disabled_constant_fold_exclusions", + incompatible_settings, + ) + + def test_registered_rule_can_mark_an_arbitrary_node(self): + gm, custom_node = self._custom_graph() + mark_constant_fold_exclusions(gm) + self.assertTrue(custom_node.meta[CONSTANT_FOLD_EXCLUSION_META_KEY]) + + def test_disabled_rule_does_not_mark_a_node(self): + gm, custom_node = self._custom_graph() + mark_constant_fold_exclusions( + gm, + CompilationSettings(disabled_constant_fold_exclusions={TEST_RULE_ID}), + ) + self.assertFalse(custom_node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY, False)) + + def test_unknown_disabled_rule_is_rejected(self): + with self.assertRaisesRegex( + ValueError, + "Unknown constant-fold exclusion rule IDs", + ): + CompilationSettings(disabled_constant_fold_exclusions={"unknown_rule"}) + + +if __name__ == "__main__": + run_tests()