feat: support configurable constant-fold exclusions

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,
     }
 

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
 

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",
 }
 

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",
+]

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

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,

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

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