Complete revamp of float/promotion sympy handling

[ezyang]

Stack from ghstack (oldest at bottom):

• Introduce int_oo #127693
• -> Complete revamp of float/promotion sympy handling #126905

At a high level, the idea behind this PR is:

• Make it clearer what the promotion and int/float rules for various Sympy operations are. Operators that previously were polymorphic over int/float are now split into separate operators for clarity. We never do mixed int/float addition/multiplication etc in sympy, instead, we always promote to the appropriate operator. (However, equality is currently not done correctly.)
• Enforce strict typing on ValueRanges: if you have a ValueRange for a float, the lower and upper MUST be floats, and so forth for integers.

The story begins in torch/utils/_sympy/functions.py. Here, I make some changes to how we represent certain operations in sympy expressions:

• FloorDiv now only supports integer inputs; to do float floor division, do a truediv and then a trunc. Additionally, we remove the divide out addition by gcd optimization, because sympy gcd is over fields and is willing to generate rationals (but rationals are bad for ValueRange strict typing).
• ModularIndexing, LShift, RShift now assert they are given integer inputs.
• Mod only supports integer inputs; eventually we will support FloatMod (left for later work, when we build out Sympy support for floating operations). Unfortunately, I couldn't assert integer inputs here, because of a bad interaction with sympy's inequality solver that is used by the offline solver
• TrueDiv is split into FloatTrueDiv and IntTrueDiv. This allows for us to eventually generate accurate code for Python semantics IntTrueDiv, which is written in a special way to preserve precision when the inputs are >= 2**53 beyond what first coercing the integer to floats and then doing true division.
• Trunc is split to TruncToFloat and TruncToInt.
• Round is updated to return a float, not an int, making it consistent with the round op handler in Inductor. To get Python-style conversion to int, we call TruncToInt on the result.
• RoundDecimal updated to consistently only ever return a float
• Add ToFloat for explicit coercion to float (required so we can enforce strict ValueRanges typing)

In torch/init.py, we modify SymInt and SymFloat to appropriately call into new bindings that route to these refined sympy operations. Also, we modify torch.sym_min and torch.sym_max to have promotion semantics (if one argument is a float, the return result is always a float), making them inconsistent with builtins.min/max, but possible to do type analysis without runtime information.

We also need to introduce some new op handlers in torch/_inductor/ops_handler.py:

• to_int for truncation to int64, directly corresponding to TruncToInt; this can be implemented by trunc and dtype, but with a dedicated handler it is more convenient for roundtripping in Sympy
• int_truediv for Python-style integer true division, which has higher precision than casting to floats and then running truediv

These changes have consequences. First, we need to make some administrative changes:

• Actually wire up these Sympy functions from SymInt/SymFloat in torch/fx/experimental/sym_node.py, including the new promotion rules (promote2)
• Add support for new Sympy functions in torch/utils/_sympy/interp.py, torch/utils/_sympy/reference.py
  • In particular, in torch.utils._sympy.reference, we have a strong preference to NOT do nontrivial compute, instead, everything in ops handler should map to a singular sympy function
  • TODO: I chose to roundtrip mod back to our Mod function, but I think I'm going to have to deal with the C/Python inconsistency this to fix tests here
• Add printer support for the Sympy functions in torch/_inductor/codegen/common.py, torch/_inductor/codegen/cpp_utils.py, torch/_inductor/codegen/triton.py. int_truediv and mixed precision equality is currently not implemented soundly, so we will lose precision in codegen for large values. TODO: The additions here are not exhaustive yet
• Update ValueRanges logic to use new sympy functions in torch/utils/_sympy/value_ranges.py. In general, we prefer to use the new Sympy function rather than try to roll things by hand, which is what was done previously for many VR analysis functions.

In torch/fx/experimental/symbolic_shapes.py we need to make some symbolic reasoning adjustments:

• Avoid generation of rational subexpressions by removing simplification of x // y into floor(x / y). This simplification then triggers an addition simplification rule (x + y) / c --> x / c + y / c which is bad because x / c is a rational number now
• _assert_bound_is_rational is no more, we no longer generate rational bounds
• Don't intersect non-int value ranges with the int_range
• Support more sympy Functions for guard SYMPY_INTERP
• Assert the type of value range is consistent with the variable type

The new asserts uncovered necessary bug fixes:

• torch/_inductor/codegen/cpp.py, torch/_inductor/select_algorithm.py, torch/_inductor/sizevars.py - Ensure Wild/Symbol manually allocated in Inductor is marked is_integer so it's accepted to build expressions
• torch/_inductor/utils.py - make sure you actually pass in sympy.Expr to these functions
• torch/_inductor/ir.py - make_contiguous_strides_for takes int/SymInt, not sympy.Expr!
• torch/export/dynamic_shapes.py - don't use infinity to represent int ranges, instead use sys.maxsize - 1

Because of the removal of some symbolic reasoning that produced rationals, some of our symbolic reasoning has gotten worse and we are unable to simplify some guards. Check the TODO at test/test_proxy_tensor.py

Reland notes. This requires this internal fbcode diff https://www.internalfb.com/phabricator/paste/view/P1403322587 but I cannot prepare the diff codev due to https://fb.workplace.com/groups/osssupport/posts/26343544518600814/

It also requires this Executorch PR pytorch/executorch#3911 but the ET PR can be landed prior to this landing.

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @gchanan @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @albanD @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126905

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rebase your PRs: Unstable CUDA signal in CI caused by cudnn 9 update

✅ You can merge normally! (1 Unrelated Failure)

As of commit 3861dde with merge base f681e36 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal, unstable) (gh) (#126993)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ezyang]

Judging from early testing results, the removal of FloorDiv(x, y) --> floor(x / y) needs to be compensated for, need to find some alternate reasoning to solve for it.

[ezyang]

This one is odd, I got a completely incorrect lowering with operator.truediv was sent to lowerings. This is clearly correct, but it's probably also worth figuring out why we have incorrect lowerings in the lowering dict.

[lezcano]

tracking issue

[ezyang]

#127789

[ezyang]

I'm not sure this actually does anything lol

[lezcano]

?
Don't we have tests that test these things?

[ezyang]

I changed it from to truediv to floordiv because there was a test that was failing at the time, but then I changed some more stuff so I don't remember if this actually matters now. floordiv is nice though because it produces an int, so the typing is more accurate this way.

[ezyang]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[ezyang]

These are the new changes from the internal revert. It turned out that the old logic (that discarded solutions here which were not integral) impeded replacements entirely when you had a nontrivial value range on the symbol being substituted. So I reallow rational intermediate values and round them at the very end (since fractional solutions cannot be real entries in the value range). This also works because I deleted all of the is integer asserts in our custom functions.

More discussion at #128053

[ezyang]

This PR is BC breaking in the following way: if you had an exported program which had floating point size compute (e.g., doing true division on a shape), we may now export torch.sym_float operator whereas previously we relied on implicit promotion in truediv operator. This change in export format is why I needed to adjust Executorch repository.

[shazqadeer]

FloorDiv now only supports integer inputs; to do float floor division, do a truediv and then a trunc.

The above sentence is from the PR summary.

Wondering if the new treatment of float floor division is accurate from the point of view of Python semantics?

[ezyang]

Wondering if the new treatment of float floor division is accurate from the point of view of Python semantics?

In fact I did check CPython sources for how floordiv on floats was implemented and it is done precisely this way.