feat: support configurable constant-fold exclusions

[fs-eire]

Description

This PR introduces a new lowering pass mark_constant_fold_exclusions just before the constant_fold pass. The new pass is used to mark specific nodes so that they will not be constant folded in the next pass.

The feature uses a rule based, configurable approach. There are pre-defined rules, and by default all rules are applied. User can use a disabled_constant_fold_exclusions setting to disable one or more rules by specifying the rule name.

Currently, this PR should not affect the behavior because the rule list is empty. I will work on add a new rule related to attention in the next PR, which depends on the current PR.

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[zewenli98]

Is it a must for users to register rules?
From the perspective of users, is there any principle for them to say in which scenarios they need to register rules? Can we predefine some rules?

[fs-eire]

In this PR, there are no rule yet. #4452 will introduce the first rule.

Usually, users do not need to register rules. I expect that in most cases, the rules are pre-registered.

[micwill755]

Approved

[shoumikhin]

Heads-up on a packaging gap from this PR — the new subpackage is not in setup.py, so it does not make it into the built wheel.

setup.py enumerates packages explicitly rather than using find_packages(), so a new subpackage has to be added in two places. constant_fold_exclusions is currently in neither (grep -c constant_fold_exclusions setup.py on main returns 0), while its sibling lowering.passes is in both.

That matters because the import happens at import torch_tensorrt time:

py/torch_tensorrt/dynamo/lowering/__init__.py:7            from .passes import *
  passes/_aten_lowering_pass.py:9                          from ...passes.mark_constant_fold_exclusions import ...
    passes/mark_constant_fold_exclusions.py:4              from torch_tensorrt.dynamo.lowering.constant_fold_exclusions._core import ...

So it works from a source checkout, where the directory is present on disk, and fails only in an installed wheel:

ModuleNotFoundError: No module named 'torch_tensorrt.dynamo.lowering.constant_fold_exclusions'

which is why it would not have shown up in PR CI.

Fix is one line in each list, alongside the existing lowering.passes entries:

 # dynamo_packages, after line 607
     "torch_tensorrt.dynamo.lowering.passes",
+    "torch_tensorrt.dynamo.lowering.constant_fold_exclusions",

 # dynamo_package_dir, after line 645
     "torch_tensorrt.dynamo.lowering.passes": "py/torch_tensorrt/dynamo/lowering/passes",
+    "torch_tensorrt.dynamo.lowering.constant_fold_exclusions": "py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions",

I checked every package on disk against both lists and this is the only one missing, so it looks like a one-off rather than a broader pattern.

Leaving this to you rather than sending a drive-by patch, since you mentioned a follow-up PR that builds on this one, and you would know whether the intent is for this subpackage to ship or for the import to be restructured instead. Happy to put up the two-line PR if that is easier — just say so.

[shoumikhin]

Opened #4547 with the two-line fix, since this is now failing 331 jobs on main (I sampled standard, rtx, python-only, executorch-runtime-build and build-docs, and they all report this same ModuleNotFoundError).

Happy to close it if you would rather restructure the import so the subpackage stays internal, which is the other valid shape here. Just did not want to leave main red while that gets decided.