[tvm-rfcs] branch main updated: [RFC][TVMScript] New namespace for tvmscript (#36)

[ju...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-rfcs.git


The following commit(s) were added to refs/heads/main by this push:
     new 4695608  [RFC][TVMScript] New namespace for tvmscript (#36)
4695608 is described below

commit 46956080aef49836b8466ccda212475c444c2e31
Author: Siyuan Feng <Hz...@sjtu.edu.cn>
AuthorDate: Thu Sep 30 00:37:59 2021 +0800

    [RFC][TVMScript] New namespace for tvmscript (#36)
    
    * init script rfc
    
    * update pr number
---
 rfcs/0036-tvm-script-namespace.md | 89 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/rfcs/0036-tvm-script-namespace.md b/rfcs/0036-tvm-script-namespace.md
new file mode 100644
index 0000000..c002582
--- /dev/null
+++ b/rfcs/0036-tvm-script-namespace.md
@@ -0,0 +1,89 @@
+- Feature Name: TVM Script Namespace
+- Start Date: 2021-09-23
+- RFC PR: [apache/tvm-rfcs#0036](https://github.com/apache/tvm-rfcs/pull/36)
+- GitHub Issue: [apache/tvm#0000](https://github.com/apache/tvm/issues/0000)
+
+# Summary
+
+[summary]: #summary
+
+This is a RFC for the new namespace and user interface for TVM script.
+
+- Use tvm.script as the `root` namespace for all TVM script related stuff
+- Use `tvm.script.tir` for `TIR`, and idiomatically import it as `T`, like Keras is usually imported as K
+- To be consistent with the names of their resulting types, use
+  - `tvm.script.ir_module` for `IRModule`
+  - `T.prim_func` for `tir.PrimFunc`
+
+# Motivation
+
+[motivation]: #motivation
+
+TVMScript is a fancy tools to write TIR in Python syntax. However, it's not fully compatible with native Python tools (e.g., pylint) because it is not a runnable program in Python.
+
+We have following pain points:
+
+- No Python auto-completion support
+- Usually conflicts with pylint
+- APIs scatter in namespaces like tvm.script, tvm.tir, tvm.script.ty
+- Somewhat non-trivial to understand at first glance what the decorator generates
+
+With the new proposal, we are able to provide type stubs that provides users with TVM scripts that work well with linting and auto-completion.
+
+# Guide-level explanation
+
+[guide-level-explanation]: #guide-level-explanation
+
+Let's see an example for the same TIR program but in different syntax.
+
+## Existing Syntax
+
+```Python
+from tvm import tir
+# ^ here tir is not related to the script but we need to import it. Otherwise there are many lint errors.
+from tvm.script import ty
+# ^ ty is a sub-module for only type system, which makes no sense to import it global
+
+@tvm.script.tir
+class Module:
+  def func(a: ty.handle, b: ty.handle) -> None:
+    A = tir.match_buffer(a, [128, 128], dtype="float32")  # pylint error: tir does not have attr match_buffer
+    B = tir.match_buffer(b, [128], dtype="float32")
+    with tir.block([128, tir.reduce_axis(0, 128)], "C") as [i, k]:
+        B[i] += A[i, k]
+```
+
+## New Syntax
+
+```Python
+from tvm.script import tir as T
+# ^ there is a broadly accepted precedence in doing this in the python community:
+#   from keras import backend as K
+
+@tvm.script.ir_module                                   # it generates an IRModule
+class Module:
+  @T.prim_func                                          # it generates a PrimFunc explicitly
+  def func(a: T.handle, b: T.handle):                   # return type is not necessary for PrimFunc
+    A = T.match_buffer(a, [128, 128], dtype="float32")  # no pylint errors for match_buffer
+    B = T.match_buffer(b, [128], dtype="float32")
+    with T.block([128, T.reduce_axis(0, 128)], "C") as [i, k]:
+        B[i] += A[i, k]
+```
+
+# Reference-level explanation
+
+[reference-level-explanation]: #reference-level-explanation
+
+Pylint and Python auto-completion will look the symbols (function definitions) in the module (e.g. `tvm.script.tir`). Just write (or generate) a concrete function in the module, then can enable auto-completion and pass Pylint checks.
+
+# Drawbacks
+
+[drawbacks]: #drawbacks
+
+Here are some existing works based on current TVM Script syntax. It need a huge refactor to migrate it to the new one.
+
+# Future possibilities
+
+[future-possibilities]: #future-possibilities
+
+Possible to write relay functions in the same IRModule if possible.