You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2021/10/09 20:29:16 UTC
[tvm-rfcs] branch main updated: [RFC] Update script block syntax
(#41)
This is an automated email from the ASF dual-hosted git repository.
tqchen 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 400988b [RFC] Update script block syntax (#41)
400988b is described below
commit 400988bc0cfcc8bb0da7b2f6f7cdc389f7dffc2c
Author: Siyuan Feng <Hz...@sjtu.edu.cn>
AuthorDate: Sun Oct 10 04:29:08 2021 +0800
[RFC] Update script block syntax (#41)
* init block syntax rfc
* fix typo
Co-authored-by: Junru Shao <ju...@gmail.com>
Co-authored-by: Zihao Ye <zi...@gmail.com>
Co-authored-by: Junru Shao <ju...@gmail.com>
Co-authored-by: Zihao Ye <zi...@gmail.com>
---
rfcs/0041-script-block-syntax.md | 142 +++++++++++++++++++++++++++++++++++++++
1 file changed, 142 insertions(+)
diff --git a/rfcs/0041-script-block-syntax.md b/rfcs/0041-script-block-syntax.md
new file mode 100644
index 0000000..1a6274f
--- /dev/null
+++ b/rfcs/0041-script-block-syntax.md
@@ -0,0 +1,142 @@
+- Feature Name: Update TVM Script block syntax
+- Start Date: 2021-10-06
+- RFC PR: [apache/tvm-rfcs#0041](https://github.com/apache/tvm-rfcs/pull/41)
+
+# Summary
+
+[summary]: #summary
+
+This is a RFC for the new syntax for blocks in TVM Script:
+
+- Disable auto-complete nesting loops
+- Use `T.axis.S` and `T.axis.R` for block var defining and value binding.
+- Use `T.axis.remap` for trivial bindings.
+
+# Motivation
+
+[motivation]: #motivation
+
+Block is the core data structure in TensorIR, meanwhile, TVMScript is one of the major input to TensorIR. Current block syntax in TVMScript does a good job but still can be better.
+
+We have following pain points:
+
+## Lines can be very long if a block has many block var
+
+```Python
+# An example block for conv2d on NHWCnc (packed layout for TensorCore)
+with tir.block([2, 14, 14, 4, tir.reduce_axis(0, 2), tir.reduce_axis(0, 3),
+ tir.reduce_axis(0, 3), 16, 16, tir.reduce_axis(0, 16)], "Conv") as \
+ [n, h, w, o, ic, kh, kw, nn, oo, ii]:
+ with tir.init():
+ C[n, h, w, o, nn, oo] = tir.float32(0)
+ C[n, h, w, o, nn, oo] = C[n, h, w, o, nn, oo] \
+ + tir.cast(Apad[n, h + kh, w + kw, ic, nn, ii], "float32") \
+ * tir.cast(W[kh, kw, ic, o, ii, oo], "float32")
+```
+
+## Unreasonable loop completion
+
+In order to make TVMScript easy to write, we enable auto-completion to blocks. Currently, we have two loop completion rules:
+
+- Auto map trivial values: if the number of block vars is equal to the number of nested loops, bind them.
+
+ ```Python
+ for i, j in T.grid(16, 16):
+ with T.block([16, 16]) as [vi, vj]:
+ # T.bind(i, vi) <- auto-completion
+ # T.bind(j, vj) <- auto-completion
+ ...
+ ```
+
+- Auto generate nested loops: generate loop nesting and bind them if there is no loop out of block.
+
+ ```Python
+ # for i, j in T.grid(16, 16): <- auto-completion
+ with T.block([16, 16]) as [vi, vj]:
+ # T.bind(i, vi) <- auto-completion
+ # T.bind(j, vj) <- auto-completion
+ ...
+ ```
+
+Both rules are too *SMART*, which may confuse the users.
+
+# Guide-level explanation
+
+[guide-level-explanation]: #guide-level-explanation
+
+Based on those two pain points, we design a new block syntax for TensorIR, which no longer has too *SMART* completion and too long lines but also easy to write.
+
+## Complete Form
+
+```Python
+for i, j, k in T.grid(512, 512, 512):
+ with T.block("name"):
+ vi = T.axis.spatial((0, 512), i)
+ # (0, 512) for the block var iter_dom, can be write as 512 if starts from 0
+ vj = T.axis.spatial(512, j)
+ # vj = T.axis.S(512, j) <- we can use `S` for spatial.
+ vk = T.axis.reduce(512, k)
+ # vk = T.axis.R(512, k) <- we can use `R` for reduce.
+ T.reads(...) # <- access region still can be detected.
+ T.writes(...)
+ ...
+```
+
+## A sugar for trivial bindings
+
+```Python
+for i, j, k in T.grid(512, 512, 512):
+ with T.block("name"):
+ # SSR means [spatial, spatial, reduce] for three vars
+ # Only trivial bindings are allowed here since we need to detect iter_dom from the loops
+ vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+ ...
+```
+
+## A Mixed format
+
+```Python
+for io, ii, j, k in T.grid(16, 32, 512, 512):
+ with T.block("name"):
+ vi = T.axis.S(512, io * 32 + ii)
+ vj, vk = T.axis.remap("SR", [j, k])
+ ...
+```
+
+# Reference-level explanation
+
+[reference-level-explanation]: #reference-level-explanation
+
+It's almost an user interface change, so might not have many technical explanations. Only one thing notable: the block var is a ordered list rather than a list. See an example:
+
+```Python
+for i, jo, ji, k in T.grid(512, 32, 16, 512):
+ with T.block("A"):
+ vi = T.axis.S(512, i)
+ vj = T.axis.S(512, jo * 32 + ji)
+ vk = T.axis.R(512, k)
+ ...
+
+for i, jo, ji, k in T.grid(512, 32, 16, 512):
+ with T.block("B"):
+ vi, vk = T.axis.remap("SR", [i, k])
+ vj = T.axis.S(512, jo * 32 + ji)
+ ...
+```
+
+`block A` (block vars:`[vi, vj, vk]`) is different from `block B` (block vars:`[vi, vk, vj]`)
+
+
+# Drawbacks
+
+[drawbacks]: #drawbacks
+
+- Here are some existing works based on current TVM Script syntax. It need some refactor to migrate it to the new one.
+
+- Some early developers get used to the old format, may bring some extra effort to move to the new one.
+
+# Future possibilities
+
+[future-possibilities]: #future-possibilities
+
+Iter domain may be detected from any PrimExpr which is affine.