[Bugfix][TVMScript] Preserve foldable constants while parsing

[Lunderberg]

Avoid folding constants while parsing TVMScript. This looks like a bug, which was introduced as part of the parsing changes implemented in #14200.

The operator overloads in python delegate to the C++ operator overloads, which perform constant folding. While useful in most circumstances, the script parsing should avoid performing any manipulations while parsing, as this can invalidate unit tests that rely on known TIR inputs. This wasn't caught by any of the existing unit tests in test_tvmscript_roundtrip.py, as the inputs to those unit tests are generated by the parsing, and already have their constants folded.

There are a few edge cases that would still be constant-folded while parsing the TVMScript, but these only impact subexpressions that do not contain TVM-specific constructs. For example, T.evaluate(1 + 2) would be parsed as T.evaluate(3), because the 1 + 2 subexpression is evaluated during parsing.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[junrushao]

@cyx-6 introduced a behavior in the printer where foldable expressions (e.g 1 + 2) are instead printed as it's complex form (e.g. T.Add(1, 2)). Therefore the dispatch was removed accordingly. Do you think it's a reasonable move?

[Lunderberg]

Thank you, and that makes sense to use T.Add(1,2) when printing integers, so that it can correctly round-trip from TIR -> TVMScript -> TIR.

That said, I think this dispatch still makes sense to keep, at least for PrimExpr argument types, so that it also handles hand-written TVMScript as expected. Writing TVMScript by hand has been very useful, both when debugging and developing new passes, as a human-readable mapping with TIR. (This has also been a huge benefit as compared to using the Python API directly, because 0 < x in Python gets translated as tir.GT(x, 0) due to a limitation in Python's operator overloading, whereas TVMScript correctly generates tir.LT(0, x).)

While the changes to printing could be replicated in the hand-written TVMScript, it's an unexpected change that needs to be remembered, especially when i + offset could have structurally different TIR depending on whether offset is defined within the TVMScript or as a dynamic parameter.

[cyx-6]

Hi @Lunderberg Here, we migrate from dispatch to the direct computation based on following concerns:

1. Methods like T.Add do not support the feature of auto-value-promotion, eg. T.int64 + 1.
2. For your concerns about 0 < x, current TVMScript parser does dispatch these logic ops, like <, >, !=. For those direct computation operators, Python support to overload the __radd__, as 1 + x will be parsed as T.Add(1, x). So we do not worry about the implicit swap in binary operators over non-PrimExpr op PrimExpr.
3. As for unittests, there are very few cases when dispatch is neccessary, like x + 0 generated by meta-schedule. Both dispatch and direct computation make no difference in most of our unittests set. Refer to [Fix][TVMScript]TVMScript BinOP printing refactor  #14200, there are only 5 unittests making difference.
4. And direct computation keeps consistent with the PrimExpr original logic as much as possible. And we do offer directly calling T.Add in TVMScript as backdoor, if needed in unittests.
5. Actually, the constants folding only occurs between non-PrimExpr, like 1 + 1, and some IntImm, FloatImm.

[junrushao]

The primary reason we are changing this is for consistency with python interface. In python, if we have

a = IntImm("int64", 1)
b = IntImm("int64", 2)
print(a + b)

then it const-folds to 3. Therefore, we kinda favor having the same behavior in TVMScript. What do you think?

[Lunderberg]

Thank you @cyx-6 , and those are definitely good reasons. I like the general simplifications, and that especially makes sense to re-use the auto-value promotion.

Actually, the constants folding only occurs between non-PrimExpr, like 1 + 1, and some IntImm, FloatImm.

As a slight correction here, you can also get constant folding between two PrimExpr, such as tvm.tir.Var('x','int32') + tvm.tir.IntImm('int32',0), but your general point holds.

The primary reason we are changing this is for consistency with python interface.

Hmm, good point. I had been seeing TVMScript as primarily a representation of the underlying TIR/Relax, rather than a scripting language in itself. As a representation, it should match the underlying TIR/Relax as much as possible and avoid constant folding, whereas a scripting language should produce the best output and perform constant folding.

I suppose my key question is: should TVMScript be viewed as a compact representation of TIR/Relax, or as a scripting language?

[tqchen]

i think one tradeoff here is that we do introduce some form of sugaring, like meta-programming etc.

But we also need to serve some purpose of roundtripping, in which case the printer should detect such cases and operate accordingly. So in that regard, because we already have T.Add which explicitly constructs and do not run constant fold, we do have a tool for roundtripping as well.

So we can view operator+ as a magic op that may constant fold, while having the printer to be careful and print out explicit T.Add if there is a desire to not to.

[Lunderberg]

Good point. I had been assuming that the meta-programming was restricted to T.meta_var and references to externally-defined values. Since it is a more general feature, having the consistency with the Python API makes sense.

I think I'm comfortable closing this PR, so long as #14320 can be merged in. The constant-folding makes sense, so long as its effects are purely local. My main confusion as a user was the interaction of constant-folding (this PR) and variable renaming for var1 = var2 (#14320). Between the two, the specific value of a dynamic parameter could result in variables being renamed across the PrimFunc, which was rather unexpected.

[Lunderberg]

Looks like #14320 has been merged (thank you @junrushao !), closing this PR.