diff --git a/rfcs/0036-tvm-script-namespace.md b/rfcs/0036-tvm-script-namespace.md new file mode 100644 index 00000000..c0025826 --- /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.