fix: migrate TVM API from tir to s_tir/tirx for apache-tvm 0.25

auto_tuning.py

@@ -1,25 +1,25 @@
 import tempfile
 
 import tvm
-from tvm import meta_schedule as ms
-from tvm import tir
-from tvm.meta_schedule.space_generator import ScheduleFn
+from tvm import s_tir
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.meta_schedule.space_generator import ScheduleFn
 
 from evaluate import test_numerical_correctness
 from gemm_relu_add import gemm_relu_add
 
 
-def auto_tuning_schedule(sch: tir.Schedule) -> tir.Schedule:
+def auto_tuning_schedule(sch: s_tir.Schedule) -> s_tir.Schedule:
     """The function that defines the schedule space for automatic tuning.
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         An empty schedule of the GeMM + ReLU + add workload.
 
     Returns
     -------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The updated schedule of the GeMM + ReLU + add workload.
     """
 

evaluate.py

@@ -2,20 +2,20 @@
 
 import numpy as np
 import tvm
-import tvm.testing
-from tvm import tir
+from tvm import s_tir
+from tvm.runtime import cuda, empty, tensor
 
 from gemm_relu_add import K, M, N, manual_schedule
 from trace_submission import apply_trace
 
 np.random.seed(0)
 
 
-def build_sch(sch: tir.Schedule) -> tvm.runtime.Module:
+def build_sch(sch: s_tir.Schedule) -> tvm.runtime.Module:
     return tvm.build(sch.mod, target="cuda")
 
 
-def test_numerical_correctness(sch: tir.Schedule, num_rounds: int = 5):
+def test_numerical_correctness(sch: s_tir.Schedule, num_rounds: int = 5):
     f = build_sch(sch)
 
     for i in range(num_rounds):
@@ -25,25 +25,25 @@ def test_numerical_correctness(sch: tir.Schedule, num_rounds: int = 5):
 
         D_std = np.maximum(A_np @ B_np, 0) + C_np
 
-        device = tvm.cuda()
-        A_tvm = tvm.nd.array(A_np, device)
-        B_tvm = tvm.nd.array(B_np, device)
-        C_tvm = tvm.nd.array(C_np, device)
-        D_tvm = tvm.nd.array(np.zeros((M, N), dtype="float32"), device)
+        device = cuda()
+        A_tvm = tensor(A_np, device=device)
+        B_tvm = tensor(B_np, device=device)
+        C_tvm = tensor(C_np, device=device)
+        D_tvm = tensor(np.zeros((M, N), dtype="float32"), device=device)
         f(A_tvm, B_tvm, C_tvm, D_tvm)
-        tvm.testing.assert_allclose(D_tvm.numpy(), D_std, rtol=1e-4, atol=1e-4)
+        np.testing.assert_allclose(D_tvm.numpy(), D_std, rtol=1e-4, atol=1e-4)
         print(f"Passing test round {i}...")
-    print(f"Passed all tests.")
+    print("Passed all tests.")
 
 
-def evaluate_execution_time(sch: tir.Schedule):
+def evaluate_execution_time(sch: s_tir.Schedule):
     f = build_sch(sch)
 
-    device = tvm.cuda()
-    A_tvm = tvm.nd.empty((M, K), "float32", device)
-    B_tvm = tvm.nd.empty((K, N), "float32", device)
-    C_tvm = tvm.nd.empty((M, N), "float32", device)
-    D_tvm = tvm.nd.empty((M, N), "float32", device)
+    device = cuda()
+    A_tvm = empty((M, K), "float32", device)
+    B_tvm = empty((K, N), "float32", device)
+    C_tvm = empty((M, N), "float32", device)
+    D_tvm = empty((M, N), "float32", device)
 
     t = f.time_evaluator(f.entry_name, device, number=3, repeat=10, min_repeat_ms=100)(
         A_tvm, B_tvm, C_tvm, D_tvm
@@ -54,37 +54,40 @@ def evaluate_execution_time(sch: tir.Schedule):
 def evaluate_naive_func_execution_time():
     from gemm_relu_add import gemm_relu_add
 
-    sch = tir.Schedule(gemm_relu_add)
-    i, j, _ = sch.get_loops("gemm")
+    sch = s_tir.Schedule(gemm_relu_add)
+    gemm_block = sch.get_sblock("gemm")
+    i, j, _ = sch.get_loops(gemm_block)
     io, ii = sch.split(i, [None, 32])
     jo, ji = sch.split(j, [None, 32])
     sch.bind(io, "blockIdx.x")
     sch.bind(jo, "blockIdx.y")
     sch.bind(ii, "threadIdx.x")
     sch.bind(ji, "threadIdx.y")
-    sch.reverse_compute_at("relu", ji)
-    sch.reverse_compute_inline("add")
-    sch.set_scope("gemm", 0, "local")
+    relu_block = sch.get_sblock("relu")
+    sch.reverse_compute_at(relu_block, ji)
+    add_block = sch.get_sblock("add")
+    sch.reverse_compute_inline(add_block)
+    sch.set_scope(gemm_block, 0, "local")
     # Uncomment the line below to check the naive function.
     # sch.show()
 
     f = build_sch(sch)
 
-    device = tvm.cuda()
-    A_tvm = tvm.nd.empty((M, K), "float32", device)
-    B_tvm = tvm.nd.empty((K, N), "float32", device)
-    C_tvm = tvm.nd.empty((M, N), "float32", device)
-    D_tvm = tvm.nd.empty((M, N), "float32", device)
+    device = cuda()
+    A_tvm = empty((M, K), "float32", device)
+    B_tvm = empty((K, N), "float32", device)
+    C_tvm = empty((M, N), "float32", device)
+    D_tvm = empty((M, N), "float32", device)
 
     t = f.time_evaluator(f.entry_name, device, number=3, repeat=10, min_repeat_ms=100)(
         A_tvm, B_tvm, C_tvm, D_tvm
     ).mean
     print("Naive function execution time: %.2f ms" % (t * 1e3))
 
 
-def show_cuda(sch: tir.Schedule):
+def show_cuda(sch: s_tir.Schedule):
     f = build_sch(sch)
-    print(f.imported_modules[0].get_source())
+    print(f.imports[0].inspect_source())
 
 
 if __name__ == "__main__":
@@ -105,7 +108,7 @@ def show_cuda(sch: tir.Schedule):
     if parsed.evaluate_tuned:
         from gemm_relu_add import gemm_relu_add
 
-        sch = tir.Schedule(gemm_relu_add)
+        sch = s_tir.Schedule(gemm_relu_add)
         apply_trace(sch)
         evaluate_execution_time(sch)
     if parsed.evaluate_naive:

gemm_relu_add.py

@@ -1,26 +1,26 @@
 from typing import Tuple
 
-from tvm import tir
-from tvm.script import tir as T
-from tvm.tir.schedule import BlockRV
+from tvm import s_tir
+from tvm.s_tir.schedule import SBlockRV
+from tvm.script import tirx as T
 
 M = 2048
 N = 2048
 K = 2048
 
 
-@T.prim_func
+@T.prim_func(s_tir=True)
 def gemm_relu_add(
     A: T.Buffer((M, K), "float32"),
     B: T.Buffer((K, N), "float32"),
     C: T.Buffer((M, N), "float32"),
     D: T.Buffer((M, N), "float32"),
 ) -> None:
-    matmul = T.alloc_buffer((M, N), "float32", scope="global")
-    relu = T.alloc_buffer((M, N), "float32", scope="global")
+    matmul = T.sblock_alloc_buffer((M, N), "float32", scope="global")
+    relu = T.sblock_alloc_buffer((M, N), "float32", scope="global")
     # Compute GeMM
     for i, j, k in T.grid(M, N, K):
-        with T.block("gemm"):
+        with T.sblock("gemm"):
             vi = T.axis.spatial(M, i)
             vj = T.axis.spatial(N, j)
             vk = T.axis.reduce(K, k)
@@ -29,25 +29,25 @@ def gemm_relu_add(
             matmul[vi, vj] += A[vi, vk] * B[vk, vj]
     # Compute ReLU
     for i, j in T.grid(M, N):
-        with T.block("relu"):
+        with T.sblock("relu"):
             vi = T.axis.spatial(M, i)
             vj = T.axis.spatial(N, j)
             relu[vi, vj] = T.max(matmul[vi, vj], T.float32(0))
     # Compute add
     for i, j in T.grid(M, N):
-        with T.block("add"):
+        with T.sblock("add"):
             vi = T.axis.spatial(M, i)
             vj = T.axis.spatial(N, j)
             D[vi, vj] = relu[vi, vj] + C[vi, vj]
 
 
-def manual_schedule() -> tir.Schedule:
+def manual_schedule() -> s_tir.Schedule:
     """The function that manually schedules and optimizes
     the GeMM + ReLU + add workload.
 
     Returns
     -------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule of the GeMM + ReLU + add workload.
 
     Note
@@ -56,7 +56,7 @@ def manual_schedule() -> tir.Schedule:
     scheduling so far at any time.
     """
     # Create a schedule from the workload.
-    sch = tir.Schedule(gemm_relu_add)
+    sch = s_tir.Schedule(gemm_relu_add)
     # Define the shared memory tile sizes and register tile sizes.
     tile_x, tile_y, tile_k = 64, 64, 64
     thread_tile_x, thread_tile_y, thread_tile_k = 4, 4, 1
@@ -78,13 +78,13 @@ def manual_schedule() -> tir.Schedule:
 
 
 def shared_memory_tiling(
-    sch: tir.Schedule, tile_x: int, tile_y: int, tile_k: int
-) -> Tuple[BlockRV, BlockRV]:
+    sch: s_tir.Schedule, tile_x: int, tile_y: int, tile_k: int
+) -> Tuple[SBlockRV, SBlockRV]:
     """The implementation of shared memory tiling.
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule instance.
 
     tile_x : int
@@ -98,11 +98,11 @@ def shared_memory_tiling(
 
     Returns
     -------
-    A_shared : tir.schedule.BlockRV
+    A_shared : SBlockRV
         The generated shared memory read stage of `A`.
         It is returned for the cooperative fetching in later tasks.
 
-    B_shared : tir.schedule.BlockRV
+    B_shared : SBlockRV
         The generated shared memory read stage of `B`.
         It is returned for the cooperative fetching in later tasks.
 
@@ -112,7 +112,7 @@ def shared_memory_tiling(
     scheduling so far at any time.
     - We do not return `sch`, because it is in-place updated during scheduling.
     """
-    block_gemm = sch.get_block("gemm")
+    block_gemm = sch.get_sblock("gemm")
     # Fetch the loops outside the "gemm" block.
     i, j, k = sch.get_loops(block_gemm)
 
@@ -157,7 +157,7 @@ def shared_memory_tiling(
 
 
 def register_tiling(
-    sch: tir.Schedule,
+    sch: s_tir.Schedule,
     thread_tile_x: int,
     thread_tile_y: int,
     thread_tile_k: int,
@@ -166,7 +166,7 @@ def register_tiling(
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule instance.
 
     thread_tile_x : int
@@ -184,7 +184,7 @@ def register_tiling(
     scheduling so far at any time.
     - We do not return `sch`, because it is in-place updated during scheduling.
     """
-    block_gemm = sch.get_block("gemm")
+    block_gemm = sch.get_sblock("gemm")
     # Fetch the last three loops of the "gemm" block,
     # which are exactly the `i_inner`, `j_inner` and `k_inner` you get
     # in the last task.
@@ -203,23 +203,23 @@ def register_tiling(
 
 
 def cooperative_fetching(
-    sch: tir.Schedule,
-    A_shared: BlockRV,
-    B_shared: BlockRV,
+    sch: s_tir.Schedule,
+    A_shared: SBlockRV,
+    B_shared: SBlockRV,
     thread_extent_x: int,
     thread_extent_y: int,
 ) -> None:
     """The implementation of cooperative fetching.
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule instance.
 
-    A_shared : tir.schedule.BlockRV
+    A_shared : SBlockRV
         The shared memory read stage of `A` generated by shared memory tiling.
 
-    B_shared : tir.schedule.BlockRV
+    B_shared : SBlockRV
         The shared memory read stage of `B` generated by shared memory tiling.
 
     thread_extent_x : int
@@ -237,7 +237,7 @@ def cooperative_fetching(
     - We do not return `sch`, because it is in-place updated during scheduling.
     """
 
-    def _cooperative_fetching_impl(block: BlockRV):
+    def _cooperative_fetching_impl(block: SBlockRV):
         # TODO: Fetch the loops of the read stage with `get_loops`.
         # Think about what loops and how many we want to fetch here?
         ...
@@ -254,12 +254,12 @@ def _cooperative_fetching_impl(block: BlockRV):
     _cooperative_fetching_impl(B_shared)
 
 
-def write_cache(sch: tir.Schedule) -> None:
+def write_cache(sch: s_tir.Schedule) -> None:
     """The implementation of write cache.
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule instance.
 
     Note
@@ -268,7 +268,7 @@ def write_cache(sch: tir.Schedule) -> None:
     scheduling so far at any time.
     - We do not return `sch`, because it is in-place updated during scheduling.
     """
-    block_gemm = sch.get_block("gemm")
+    block_gemm = sch.get_sblock("gemm")
     # TODO: Use `sch.get_loops` to find out the location of inserting write cache.
     loop_index = ...
     write_cache_loc = sch.get_loops(block_gemm)[loop_index]
@@ -279,12 +279,12 @@ def write_cache(sch: tir.Schedule) -> None:
     ...
 
 
-def epilogue_fusion(sch: tir.Schedule) -> None:
+def epilogue_fusion(sch: s_tir.Schedule) -> None:
     """The implementation of epilogue_fusion.
 
     Parameters
     ----------
-    sch : tir.Schedule
+    sch : s_tir.Schedule
         The schedule instance.
 
     Note
@@ -293,7 +293,7 @@ def epilogue_fusion(sch: tir.Schedule) -> None:
     scheduling so far at any time.
     - We do not return `sch`, because it is in-place updated during scheduling.
     """
-    # TODO: Use `get_block` to retrieve the addition computation and ReLU computation.
+    # TODO: Use `get_sblock` to retrieve the addition computation and ReLU computation.
     ...
     # TODO: Use `reverse_compute_inline` to fuse addition into ReLU, and fuse ReLU into GeMM.
     ...