[RFC] TUNIP: TVMScript Unified Printer

[yelite]

Rendered

[sunggg]

Thank you for the great RFC!
Just a few questions.

[YuchenJin]

LGTM! The unified representation across multiple dialects and better error reporting will greatly help the development of Relax and future IRs.

Thank you @yelite and @junrushao1994 for the great design and RFC!

[areusch]

very sorry for the delay here.

@junrushao1994 i guess i'm still a little fuzzy on the rationale for why we want to have two ways to parse TVMScript. I know this is a bit bigger than just this RFC, but with this RFC we're further pushing the cart down this path so I'd like us to be explicit about why.

Relay has a single roundtrippable serialization format, as do most languages. I think we benefit from this in that we only have one set of tests to maintain.

With TIR, my understanding is that we used Python as an AST parser since language constructs in TIR somewhat closely matched to Python. I'm not sure I see why we should want to have two different ways to do this--shouldn't we just have a single way to serialize TIR and parse it that always works? It seems like we could still expose this parser to various frontends via PackedFunc.

[junrushao]

Relay has a single roundtrippable serialization format, as do most languages. I think we benefit from this in that we only have one set of tests to maintain.

To clarify, Relay has two roundtrippable serialization formats: text format and json. people use text format for readability, and in most usecases, go for json reliable serialization.

why we want to have two ways to parse TVMScript

First of all, the RFC focuses on the printer, so would be great if we talk about the parser in the incoming RFC, but to clarify, we are not having two ways to parse TVMScript in the same language. Instead, for metaprogramming capability, one may need to interleave host language-based IRBuilder with the parser.

language constructs in TIR somewhat closely matched to Python.

To clarify, TIR is a DSL independent of any existing language. It's somewhat close to HalideIR but has been evolving for years with lots of new features, for example, blocks. Python syntax is an add-on for usability but TIR design is in no ways matched to python.

Anyways, I'm happy to talk more on the upcoming metaprogramming RFC.

[Hzfengsy]

LGTM. Thanks @yelite !!!

[areusch]

@junrushao1994

people use text format for readability, and in most usecases, go for json reliable serialization.

what do you mean by "reliable" here? if they're truly roundtrippable, aren't they both reliable? just trying to understand :)

Instead, for metaprogramming capability, one may need to interleave host language-based IRBuilder with the parser.

Is the goal then that you'd be able to serialize an IRModule which contains fragments that need to be repeated via meta-programming? then, if you did this using a host language's IRBuilder, you could quickly write some meta-programming steps in that host language to expand the program.

To clarify, TIR is a DSL independent of any existing language. ... Python syntax is an add-on for usability but TIR design is in no ways matched to python.

I agree with your statement, but I thought the repr form of TIR wasn't so easy to work with and thus the motivation for TVMScript. While I agree that TIR is technically separate, if TVMScript is a core way in which TIR is used, I'd argue we should treat them conceptually as joined (e.g. TVMScript as the recommended roundtrip text format for TIR). What are your thoughts there?

[junrushao]

Thanks @areusch for your response!

if TVMScript is a core way in which TIR is used, I'd argue we should treat them conceptually as joined (e.g. TVMScript as the recommended roundtrip text format for TIR). What are your thoughts there?

Let's phrase it this way: TVMScript serves as a frontend of easily constructing and manipulating TIR, Relax and any third-party IRs. It offers readability, type checking, metaprogramming, etc. In terms of serialization, my standing is that reflection-based JSON serialization is still the best way to go as it's the least error-prone if we don't need readability and manipulability.

what do you mean by "reliable" here? if they're truly roundtrippable, aren't they both reliable? just trying to understand :)

In practice we did notice that there could be some ad-hoc issues (which are definitely fixable, not a big deal, so I'm not saying it's not good) that probably be overlooked when implementing the parser. For example, @mbs-octoml fixed a dtype issue very recently (apache/tvm#11224). Instead, reflection-based JSON serialization mechanism in TVM is definitely less bug-prone because it doesn't need to handle edge cases - the downside is that it's almost unreadable though.

Is the goal then that you'd be able to serialize an IRModule which contains fragments that need to be repeated via meta-programming? then, if you did this using a host language's IRBuilder, you could quickly write some meta-programming steps in that host language to expand the program.

Yeah with metaprogramming we are able to write bunch of code that require host-language control. The next RFC is all about metaprogramming so stay tuned :-)

[junrushao]

To summarize offline discussion with @areusch:

Q: Is this going to unify current fragmented printing formats?
A: Yes for TIR. After following the standard deprecation procedure (see Section "Upgrade Plan"), TVMScript will be the only frontend (i.e. user-facing printer/parsing) for TIR, while the original non-roundtrippable text format will eventually be deprecated. For serialization/deserialization, reflection-based JSON format is preserved without change, which is basically like protobuf for reliably data sharing.

Q: Is JSON serialization going to be deprecated?
A: No. It is a separate issue. TVM's JSON serialization is conceptually similar to protobuf, which is designed to be reliably persist data, while not necessarily human-readable, which is different than our goal in this RFC (frontend).

Q: Is this going to introduce new functionalities that make TVM more fragmented?
A: No. It's a refactoring of existing TIR TVMScript printer. With future deprecations done, it actually reduces the fragmentation on TIR side.

Q: Why is this called "unified" printer?
A: It provides a unified infrastructure that allows IR developers to extend to other IRs, for example, Relax, without interfering with each other. More concretely speaking, as pointed out in the folder structure in the RFC text:

src/script/printer/
├── core # Core infra, which is IR-agnostic
│   ├── ir_docsifier.cc
│   └── ...
├── tir # TIR dialect 
│   ├── expr.cc
│   ├── stmt.cc
│   └── ...
└── relax # Hypothetical Relax dialect (not part of our RFC)
    └── ...

TIR developers only need to maintain the tir folder, and relax developers only need to maintain relax folder. The IRModule that contains both TIR and other dialects is allowed to be printed with this infrastructure.

[areusch]

thanks for summarizing @junrushao1994 . to provide some additional perspective:

• I'd really like us to get to a place where we have 1 reliable and readable format to serialize TIR. i don't think we should build all of this infrastructure just to provide a "best-effort" serialization--anything expressable in TIR should roundtrip in any serialization format we provide. there's no reason we have to sacrifice reliability for readability, either--it's just a matter of putting effort into the unit tests.
• right now i believe we have 3 TIR formats: repr(), TVMScript, and JSON. This RFC looks to provide infra that allows for generation of more formats e.g. so that Python doesn't have to be the frontend for TIR, or so we can leverage other APIs or frontend languages to write TIR. i'm okay with adopting this design to separate indentation from serialization, but if we do follow through and introduce an IRBuilder serialization, that would count as a fourth way to serialize TIR--doing that is only going to make it harder to achieve the previous goal because it splits our human and unit-test attention across multiple formats. i don't see a good reason for doing that laid out in the RFC.
• this PR does provide considerations for Relax, although the language spec hasn't been RFC'd or introduced yet. i don't want to block the PR over that, but it's plain that the eventual goal is to leverage this infrastructure for Relax when it lands. then that gets us into a scenario where we have one parser/serializer for Relay and another for everything else. that's also more complex than it needs to be. admittedly, Relax and TIR benefit more from meta-programming than Relay, so it makes sense that the old Relay parser/serializer isn't a great fit here. but if we're going to discuss Relax here, it would be great to see this design considering how it could replace all of the printers in the codebase, not just overhaul the TIR one. anyhow, this discussion can be left for a future time when the Relax merge plan is clearer.

[junrushao]

Hey @areusch thanks for elaborating your points and these are all definitely great questions to me!

right now i believe we have 3 TIR formats: repr(), TVMScript, and JSON. This RFC looks to provide infra that allows for generation of more formats e.g. so that Python doesn't have to be the frontend for TIR, or so we can leverage other APIs or frontend languages to write TIR

Oh I think I get your concern here!!

No, we are not building more and more and more frontends to cause fragmentation. Here is the thing - we only need two forms for serialization:

• JSON for protobuf-like reliable serialization;
• TVMScript python syntax.
  And we defragment the system by retiring repr() for TVMScript. Therefore, as a python user, the only interface you touch is TVMScript python syntax after the RFC landed.

Imagine you are a rust user who doesn't want to use python, our proposal makes it possible to develop a frontend in pure rust. And of course, it's not going to be our priority, i'm just stating the possibility.

but if we do follow through and introduce an IRBuilder serialization, that would count as a fourth way to serialize TIR

Ah no, IRBuilder is not serialization, sorry for the confusion. Let's imagine a common usecases in company's prod env where python is strictly prohibited (e.g. on-device training), but our developers wrote most of the operators in python, which is the reality. How do we quickly migrate those operators to C++? Our RFC provides an answer to that - just write a codegen from Doc to C++ IRBuilder :-) Of course again, the RFC aims to open possibilities instead of implementing those functionalities.

this PR does provide considerations for Relax, although the language spec hasn't been RFC'd or introduced yet. i don't want to block the PR over that, but it's plain that the eventual goal is to leverage this infrastructure for Relax when it lands

Yeah it's definitely fair to say the PR considers Relax a lot. On the other hand, Relax doesn't need this RFC to use TVMScript. Take the current codebase as an example, it's not using the design of this proposal and it's totally fine given it's assuming there are only 2 IRs. Instead, I'm trying to convey is that we build an open infra for any vendors to integrate with their own IRs, which can be represented and compiled in the same IRModule with Relax, TIR, etc.

[areusch]

@junrushao1994 ah thanks for clarifying! a few follow-ups then:

Imagine you are a rust user who doesn't want to use python, our proposal makes it possible to develop a frontend in pure rust. And of course, it's not going to be our priority, i'm just stating the possibility.

Does this mean you can import the JSON-based serialization using only libtvm.so (e.g. rust could import JSON)? If so, then, in the hypothetical prod scenario:

Let's imagine a common usecases in company's prod env where python is strictly prohibited (e.g. on-device training), but our developers wrote most of the operators in python, which is the reality. How do we quickly migrate those operators to C++?

, what's the motivation for someone to use IRBuilder instead of just serializing the TVMScript to JSON via parse/print and then ingesting the JSON in their prod env using the JSON parser?

it's not using the design of this proposal and it's totally fine given it's assuming there are only 2 IRs. Instead, I'm trying to convey is that we build an open infra for any vendors to integrate with their own IRs, which can be represented and compiled in the same IRModule with Relax, TIR, etc.

That makes sense, thanks for clarifying :)

[junrushao]

@areusch Thanks for following up!

what's the motivation for someone to use IRBuilder instead of just serializing the TVMScript to JSON via parse/print

JSON is mostly for serializing an IR after it's constructed (which users cannot manipulate), and the TVMScript format is for users to construct the IR themselves. For example, to implement a ReLU operator, TVMScript users could write a python/rust code, while it's pretty infeasible for them to write JSON with bare hand.

[junrushao]

Hey I'm happy to discuss more, and let's keep this RFC open until the end of this week

[areusch]

@junrushao1994 ok, just want to clarify one question here:

TVMScript users could write a python/rust code, while it's pretty infeasible for them to write JSON with bare hand.
what functionality does JSON provide that TVMScript doesn't do right now as it's designed? Is there a reason to have two formats to begin with? imo it is not worth having two formats here unless there's a clear need for structural analysis of TIR outside of TVM.

i have heard some folks in the community interested in converting TIR to JSON (I think @SebastianBoblestETAS and @MichaelJKlaiber if memory serves), but I also think that the use cases they're interested in could be accomplished if we had a user-friendly, robust Python interface to the TIR.

[junrushao]

Thanks @areusch for pointing me to the thread! Definitely happy to read the discussion, and glad to see that @vinx13 unblocks the SaveJSON method for NDArrays :-) completely agreed that TVMScript could be the usecase which provides readability if this is the case that @SebastianBoblestETAS and @MichaelJKlaiber wanted.

what functionality does JSON provide that TVMScript doesn't do right now as it's designed? Is there a reason to have two formats to begin with? imo it is not worth having two formats here unless there's a clear need for structural analysis of TIR outside of TVM.

Yeah JSON doesn't provide any extra functionality, but it's pretty stable and reliable protobuf-like serialization, which could be faster to parse/dump if we only want to save the IRs or send them around between processes without the need of readability and manipulatability.

For general users, our suggestion should be like, just use TVMScript to write TIR and forget about JSON. For advanced developers, for example, we want to store TIR in tuning database, cache TIR somewhere for later usage, JSON is definitely the "to-go" approach.

[yelite]

what functionality does JSON provide that TVMScript doesn't do right now as it's designed? Is there a reason to have two formats to begin with? imo it is not worth having two formats here unless there's a clear need for structural analysis of TIR outside of TVM.

A good analogy is the ProtoBuf text format (unfortunately no public spec outside google) vs. wire format (binary format). People write proto in text format when they need to create proto manually, while the binary format is used in RPC payload for its portability and efficiency.

Back to our own use case, TVMScript is the choice if user want to construct IR graph manually, but it's not portable because the TVMScript parser has a strong dependency on Python, especially after the Metaprogramming RFC (#79). On contrast, the JSON format will be portable because it only depends on the IR node definition in C++.

[junrushao]

I'm merging this RFC as it seems that our discussion has reached consensus, but feel free to follow-up any time!