sunggg commented on code in PR #74: URL: https://github.com/apache/tvm-rfcs/pull/74#discussion_r883071372
########## rfcs/0074-tvmscript-unified-printer.md: ########## @@ -0,0 +1,475 @@ +- Feature Name: TUNIP: TVMScript Unified Printer +- Start Date: 05/25/2022 +- RFC PR: [apache/tvm-rfcs#74](https://github.com/apache/tvm-rfcs/pull/74) +- GitHub Issue: [apache/tvm#0000](https://github.com/apache/tvm/issues/0000) +- Co-Authors: Lite Ye ([**@yelite**](https://github.com/yelite)), Yong Wu + ([**@yongwww**](https://github.com/yongwww)), Yuchen Jin + ([**@YuchenJin**](https://github.com/YuchenJin)) + +# Summary +[summary]: #summary + +This RFC proposes to modularize and infrastructuralize the existing TVMScript +printer, to develop unified printing mechanism across TVM stack, where TIR, +Relax and any future vendor-specific IR are all treated equally as dialects and +could be printed together without potential conflict in engineering. + +# Motivation +[motivation]: #motivation + +TVMScript, as a roundtrippable python-based text format, is +the central piece of TVM performance productivity. As the frontend of TVM, it +enables end users to directly construct the TVM IR, either TIR or Relax, in a +pragmatic approach. From Relax to MetaSchedule and TIR, TVMScript enables +inspectability and reproducibility at any level of compilation and +optimization. Furthermore, based on TVMScript, developers are empowered to +intercept, manipulate and customize the compiler behavior in a principled way. + +While TVMScript is gaining traction and buy-in from the open source community, +the TVMScript printer suffers from multiple profound design issues: +- Not supporting IR fragment printing requires users to jump in-between + TVMScript syntax and TIRText syntax +- The lack of modularity leads to practical inability to scale up to and + maintain multiple IRs without engineering conflicts +- Enhancing co-existence of multi-level IRs often leads to re-engineering of + existing features. + +**Goal.** This RFC introduces Tvmscript UNIfied Printer (TUNIP), a systematic +redesign to +address those engineering, usability and scalability issues above. The goal of +this re-design includes: + +**Goal 1 [Unified Representation].** Become the unified roundtrippable +representation of TIR and Relax, allowing systematic mixing of IRs or IR +fragments (Relax + TIR) in the same IRModule in the target language (for +example, python, C++). + +Currently TVMScript priner is designed specifically for TIR, and printing +multiple dialects together was not a design goal at that time. Therefore, +supporting Relax requires ad-hoc hack around the system (for +instance, [relax#149](https://github.com/tlc-pack/relax/pull/149) added support +of printing `T.cast` and `T.max` in an ad-hoc way, without reusing the printing +code for TIR). The unified printer in this RFC addresses this issue by having a +unified approach for printing IR tree to TVMScript. Engineers will be able to +implement a fully-fledged printer for Relax, TIR and any potential IR in the +future with minimal effort. + +The folder structure that we want to pursue is: +```bash +include/tvm/script/printer/ +└── ... # Public headers for the core infra +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) + └── ... +``` + +**Goal 2 [Third-Party IRs in Multi-Stage Compilation].** Modularize and +infrastructuralize the printer to support more future IRs or third-party IRs at +any level with maintainability, for example, IRs at lower-level than TIR, or +Relax VM executable. + +The current TVMScript printer is tightly coupled with TIR by being a subclass +of TIR-specific functors +([link](https://github.com/apache/tvm/blob/main/src/printer/tvmscript_printer.cc#L129)). +This design isn’t scalable when we want to support more IRs. More importantly, +it’s impossible for the current approach to support third-party IR bteing +registered in a dynamic library. + +**Goal 3 [Reproducibility and Error Reporting].** Expand reproducibility and +flexible rendering of diagnostic messages during any level of IR +transformation. + +For example, the following snippet runs and produces an error. + +```py +import tvm + [email protected]_func +def func_a(A: T.Buffer[(1,), "int32"]): + A[0] = 0 + [email protected]_func +def func_b(A: T.Buffer[(8,), "int32"]): + A[0] = 0 + +tvm.ir.assert_structural_equal(func_a, func_b) +``` + +The current error message indicates what the difference was, but not where it +occurred. This can sometimes be inferred from a stack trace, but becomes +increasingly difficult with larger IR graphs. + +``` +ValueError: StructuralEqual check failed, caused by lhs: +1 +and rhs: +8 +``` + +TUNIP should enable individual utilities and IR passes to have error messages +directing the user to exact locations in the IR representation. + +``` +ValueError: StructuralEqual check failed, first delta highlighted below + [email protected]_func +def func_a(A: T.Buffer[(1,), "int32"]) -> None: + ^^^^ + A[0] = 0 + [email protected]_func +def func_b(A: T.Buffer[(8,), "int32"]) -> None: + ^^^^ + A[0] = 0 +``` + + +# Guide-level explanation +[guide-level-explanation]: #guide-level-explanation + +This section introduces the design philosophy of the printer, and demonstrates +the proposed user-facing APIs where users means IR developers. + +## Two-Stage Translation + +Traditionally in TVM stack, printing is a single-stage process. The printer +assumes certain syntax of the target language, and therefore, so far there are +3 different printers all for TIR: ReprPrinter, TIRTextPrinter, +TVMScriptPrinter. + +We extend the idea of the existing Doc class at +[src/printer/doc.h#L67](https://github.com/apache/tvm/blob/main/src/printer/doc.h#L67) +to allow better consistency and scalability. An IR, which could +be TIR, Relax or any other ones developed by third-party vendors, is first +translated to an intermediate Doc node tree, and then the Doc tree is mapped to a target +language, for example, Python, C++ IRBuilder API, or Rust. + +**Stage 1 [TVM IR => Doc]**. On the first stage, the printer needs to take care Review Comment: +1 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
