Problem

Inductor's FX backend receives sympy expressions for Triton launch grids, and passes these to a tracer to generate equivalent FX IR. However, the tracer does not support all possible sympy expressions. In particular, it can't handle ops like floor and Pow which would be found in an expression like floor(x / y). Instead, it expects FloorDiv(x, y), which has the advantage that all intermediate values are integers, unlike x / y.

Inductor's Python backend uses a trick where ceil(x / y) is computed in Python as -(x // -y), which is faster when evaluating Python launch grids at runtime. However, this trick generates more complex sympy expressions, so the FX backend introduced a "python_slow" mode using a more familiar form of ceil division. However, this mode is slower to evaluate, which increased production CPU usage. (Internal reviewers see T237853632.)

Solution

To get the best of both worlds, this PR removes "python_slow" mode, and generalizes the replace_floor_div function to handle the more complex expressions resulting from the "python" grid mode. The new algorithm is conceptually similar to the existing one, except instead of analyzing only the first argument to a sympy.Mul op, it checks all factors, so it can handle expressions containing both Rational and Pow ops, among other cases. It also uses Mul.make_args to handle the case when the argument to floor is not a Mul. Finally, it uses expr.is_positive to check the sign of symbolic exponents.

This new algorithm is guaranteed to convert all floor ops to an equivalent expression using FloorDiv. (To see this, consider that floor(x) == FloorDiv(x, 1).) Note it may not remove all Pow ops, with a counterexample being floor(x / (2 + z ** y)), but it covers everything we've seen in practice for symbolic launch grids. In particular, it covers the typical case where Pow is a factor of the argument to floor, and the exponent is -1. Is this situation, we move the Pow to the denominator of FloorDiv and the exponent becomes 1, eliminating the Pow op.

Test plan

This PR adds an end-to-end test for static padding with dynamic outer dimensions, which creates a difficult sympy expression that the existing algorithm would not be able to handle.

This PR also adds some unit tests for the replace_floor_div function. It can be difficult to construct end-to-end tests that expose all the trickiest expressions, as those tests have to pass through a number of other systems handling dynamic shapes. Therefore, it's easier to expose the edge cases with these new unit tests. The tests check that we can replace all floor ops in the input expression with FloorDiv, then they expand FloorDiv back to floor and check equality with the original expression.

Note this PR also requires some MTIA changes to pass internal tests. Those will be stacked onto the imported diff.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben