You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/05/26 15:47:35 UTC
[GitHub] [tvm-rfcs] yelite commented on a diff in pull request #74: [RFC] TUNIP: TVMScript Unified Printer
yelite commented on code in PR #74:
URL: https://github.com/apache/tvm-rfcs/pull/74#discussion_r882819598
##########
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
+
+@T.prim_func
+def func_a(A: T.Buffer[(1,), "int32"]):
+ A[0] = 0
+
+@T.prim_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
+
+@T.prim_func
+def func_a(A: T.Buffer[(1,), "int32"]) -> None:
+ ^^^^
+ A[0] = 0
+
+@T.prim_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
+of translating a TVM IR to Doc tree. As an example, `tir.For` is translated to
+`ForDoc` without having to worry about the underlying language. Note that some
+complicated nodes in TVM IR, for example, `PrimFunc`, could be translated to
+multiple IR elements, including `FunctionDoc` and a few `StmtDoc`.
+
+During the translation from IR to Doc tree, it is possible that some statement
+influences the syntax of its children or vice verse, especially for syntactic
+sugars and declaring undefined variables in IR fragment printing. Therefore, a
+generic data structure `Frame` is introduced to allow retrieval and
+manipulation the relevant context information.
+
+**Stage 2. [Doc => target language]**. On the second stage, Doc tree is then
+honestly translated to the target language in text format. For example, when
+the target language is python, `ForDoc` is translated to python’s for loop
+syntax:
+
+```python
+for ... in ...:
+ ...
+```
+
+When the target language becomes python IRBuilder, `ForDoc` is translated to:
+
+```cpp
+with T.For(...):
+ ...
+```
+
+For generality, the Doc tree is designed to select minimal elements that exist
+in languages used in developing TVM. A full spec of the Doc could be found in
+the next section.
+
+## Distributed Registration
+
+As a major engineering challenge for TVMScript to scale to multiple IRs, the
+existing printing logic has to be engineered, maintained and re-engineered in a
+single file, which has brought significant confusion for developing multi-level
+IRs for TVM Unity.
+
+Inspired by the pass infrastructure, as well as the ReprPrinter in TVM, we
+propose to develop the infrastructure to enable distributed registration, and
+further allows printer for different levels of IR to be registered in separate
+translation units, and in the meantime keeps the capability to be mixed
+together at various level, for example, Relax uses TIR expression in its
+function bodies, and TIR calls back to Relax function.
+
+## Diagnostics and Reproducibility
+
+Existing error reporting mechanisms have not taken IR structure and
+reproducibility into consideration. Usually it reports a single line error
+message without providing necessary context of how the IR looks like during
+compilation. For example, when comparing whether two TIRs are structurally
+equivalent, the system may report:
+
+```cpp
+ValueError: StructuralEqual check failed, caused by lhs:
+{slow_memory_3_var: buffer(slow_memory_3_buffer_var, 0x501bf80), fast_memory_2_var: buffer(fast_memory_2_buffer_var, 0x501bd80), placeholder_3: buffer(placeholder_5, 0x50138a0), placeholder_2: buffer(placeholder_4, 0x5012b60), T_subtract: buffer(T_subtract_1, 0x5014390)}
+and rhs:
+{}
+```
+
+which lacks necessary information for users to understand where the mismatch
+is.
+
+As a recent effort, structural error reporting in TIR scheduling provides
+relevant and reproducible context, as demonstrated below:
+
+```cpp
+@tvm.script.ir_module
+class Module:
+ @tir.prim_func
+ def main(a: tir.handle, b: tir.handle) -> None:
+ A = tir.match_buffer(a, [128, 128, 128, 128], dtype="float32")
+ B = tir.match_buffer(b, [128, 128, 128, 128], dtype="float32")
+ # body
+ # with tir.block("root")
+ for i, j, k, l in tir.grid(128, 128, 128, 8):
+ tir.Block#0
+ with tir.block("B"):
+ ^^^^^^^^^^^^^^^^^^^^
+ vi, vj, vk = tir.axis.remap("SSS", [i, j, k])
+ vl = tir.axis.spatial(128, l * 16)
+ tir.reads([A[vi, vj, vk, vl]])
+ tir.writes([B[vi, vj, vk, vl]])
+ B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * tir.float32(2)
+
+Error: ...
+```
+
+However, the underlying mechanism supports only S-TIR and error reporting on
+`tir.ForNode` and `tir.BlockNode`, and is less extensible for generic cases.
+
+To generalize this UX across the TVM stack, during the first stage in
+translation, the following steps is additionally executed:
+
+- Each Doc node is optionally attached to a node in TVM IR
+- After the 1st stage is finished, collect all IR nodes that gets attached to
+ Doc into a map, whose key is IR node and value is a list of Doc nodes.
+- For each IR node that has diagnostic message, trace back through its parent
+ until it reaches to an IR node in the map collected in previous step. Then it
+ can produce a map from Doc node to diagnostic message.
+- In the 2nd stage, diagnostic message will be printed as doc is being printed
+ into target language
+
+# Reference-level explanation
+
+## Doc Spec
+
Review Comment:
Types will be represented by `ExprDoc`. For example, `Tensor((2, 3), "float32")` will be
```
CallDoc(
IdDoc("Tensor"),
[
TupleDoc([
LiteralDoc(2),
LiteralDoc(3)
]),
LiteralDoc("float32")
]
)
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org