From 62eae199c222d220f4b2772855120c3be857d0e6 Mon Sep 17 00:00:00 2001 From: Zihao Date: Tue, 4 Apr 2023 19:04:11 -0700 Subject: [PATCH 1/3] add docstring --- 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 8113097003fa..4a6df83bef72 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -840,6 +840,57 @@ def reorder_block_iter_var(self, block: BlockRV, new_order: List[int]) -> None: 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 From 15d58642d022ea65d77572bef83fb172f2ab7c4d Mon Sep 17 00:00:00 2001 From: Zihao Date: Wed, 5 Apr 2023 01:17:45 -0700 Subject: [PATCH 2/3] pylint --- python/tvm/tir/schedule/schedule.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 4a6df83bef72..cb62bbd406f5 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -840,7 +840,7 @@ def reorder_block_iter_var(self, block: BlockRV, new_order: List[int]) -> None: The block to be transformed. new_order : List[int] The new block itervar order. - + Examples -------- @@ -860,7 +860,7 @@ def matmul( 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 @@ -868,7 +868,7 @@ def matmul( 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 From ab87ec4d3d339cab286586287ce3336bb5bba7e0 Mon Sep 17 00:00:00 2001 From: Zihao Date: Fri, 14 Apr 2023 20:56:39 -0700 Subject: [PATCH 3/3] trigger