This is an automated email from the ASF dual-hosted git repository.
expye pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e3638e772d [TensorIR][Doc] Docstring of `reorder_block_iter_var`
(#14504)
e3638e772d is described below
commit e3638e772d54e9360990066ca7f423f765fd2e5d
Author: Zihao Ye <[email protected]>
AuthorDate: Sat Apr 15 01:04:29 2023 -0700
[TensorIR][Doc] Docstring of `reorder_block_iter_var` (#14504)
* add docstring
* pylint
* trigger
---
python/tvm/tir/schedule/schedule.py | 51 +++++++++++++++++++++++++++++++++++++
1 file changed, 51 insertions(+)
diff --git a/python/tvm/tir/schedule/schedule.py
b/python/tvm/tir/schedule/schedule.py
index 11fa7ae12e..7c7af998be 100644
--- a/python/tvm/tir/schedule/schedule.py
+++ b/python/tvm/tir/schedule/schedule.py
@@ -868,6 +868,57 @@ class Schedule(Object):
The block to be transformed.
new_order : List[int]
The new block itervar order.
+
+ Examples
+ --------
+
+ Before reorder_block_iter_var, in TensorIR, the IR is:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def matmul(
+ A: T.Buffer((128, 128), "float32"),
+ B: T.Buffer((128, 128), "float32"),
+ C: T.Buffer((128, 128), "float32"),
+ ) -> None:
+ for i, j, k in T.grid(128, 128, 128):
+ with T.block("C"):
+ vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+ with T.init():
+ C[vi, vj] = 0.0
+ C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+ Create the schedule and do reorder_block_iter_var:
+
+ .. code-block:: python
+
+ sch = tir.Schedule(matmul)
+ C = sch.get_block("C")
+ sch.reorder_block_iter_var(C, [2, 1, 0])
+
+ After applying reorder_block_iter_var, the IR becomes:
+
+ .. code-block:: python
+
+ @T.prim_func
+ def matmul_after_reorder_block_iter_var(
+ A: T.Buffer((128, 128), "float32"),
+ B: T.Buffer((128, 128), "float32"),
+ C: T.Buffer((128, 128), "float32"),
+ ):
+ for i, j, k in T.grid(128, 128, 128):
+ with T.block("C"):
+ vk, vj, vi = T.axis.remap("RSS", [k, j, i])
+ T.reads(A[vi, vk], B[vj, vk])
+ T.writes(C[vi, vj])
+ with T.init():
+ C[vi, vj] = T.float32(0)
+ C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+ See Also
+ --------
+ reorder
"""
_ffi_api.ScheduleReorderBlockIterVar(self, block, new_order) # type:
ignore # pylint: disable=no-member
