diff --git a/auto_tuning.py b/auto_tuning.py index 1e3b459..2c0a17a 100644 --- a/auto_tuning.py +++ b/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. """ diff --git a/evaluate.py b/evaluate.py index 634311d..828ef2e 100644 --- a/evaluate.py +++ b/evaluate.py @@ -2,8 +2,8 @@ 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 @@ -11,11 +11,11 @@ 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,27 +54,30 @@ 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 @@ -82,9 +85,9 @@ def evaluate_naive_func_execution_time(): 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: diff --git a/gemm_relu_add.py b/gemm_relu_add.py index 17be930..f62b3bb 100644 --- a/gemm_relu_add.py +++ b/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,9 +203,9 @@ 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: @@ -213,13 +213,13 @@ def 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. ... diff --git a/mlsys_hw2.ipynb b/mlsys_hw2.ipynb index 4d5eddb..6a6974b 100644 --- a/mlsys_hw2.ipynb +++ b/mlsys_hw2.ipynb @@ -183,26 +183,26 @@ "metadata": {}, "source": [ "```python\n", - "from tvm import tir\n", - "from tvm.script import tir as T\n", + "from tvm import s_tir\n", + "from tvm.script import tirx as T\n", "\n", "M = 2048\n", "N = 2048\n", "K = 2048\n", "\n", "\n", - "@T.prim_func\n", + "@T.prim_func(s_tir=True)\n", "def gemm_relu_add(\n", " A: T.Buffer((M, K), \"float32\"),\n", " B: T.Buffer((K, N), \"float32\"),\n", " C: T.Buffer((M, N), \"float32\"),\n", " D: T.Buffer((M, N), \"float32\"),\n", ") -> None:\n", - " matmul = T.alloc_buffer((M, N), \"float32\", scope=\"global\")\n", - " relu = T.alloc_buffer((M, N), \"float32\", scope=\"global\")\n", + " matmul = T.sblock_alloc_buffer((M, N), \"float32\", scope=\"global\")\n", + " relu = T.sblock_alloc_buffer((M, N), \"float32\", scope=\"global\")\n", " # Compute GeMM\n", " for i, j, k in T.grid(M, N, K):\n", - " with T.block(\"gemm\"):\n", + " with T.sblock(\"gemm\"):\n", " vi = T.axis.spatial(M, i)\n", " vj = T.axis.spatial(N, j)\n", " vk = T.axis.reduce(K, k)\n", @@ -211,19 +211,19 @@ " matmul[vi, vj] += A[vi, vk] * B[vk, vj]\n", " # Compute ReLU\n", " for i, j in T.grid(M, N):\n", - " with T.block(\"relu\"):\n", + " with T.sblock(\"relu\"):\n", " vi = T.axis.spatial(M, i)\n", " vj = T.axis.spatial(N, j)\n", " relu[vi, vj] = T.max(matmul[vi, vj], T.float32(0))\n", " # Compute add\n", " for i, j in T.grid(M, N):\n", - " with T.block(\"add\"):\n", + " with T.sblock(\"add\"):\n", " vi = T.axis.spatial(M, i)\n", " vj = T.axis.spatial(N, j)\n", " D[vi, vj] = relu[vi, vj] + C[vi, vj]\n", "\n", "\n", - "sch = tir.Schedule(gemm_relu_add)\n", + "sch = s_tir.Schedule(gemm_relu_add)\n", "```" ] }, @@ -250,9 +250,9 @@ " * The second one defines the ReLU computation, which element-wisely takes the maximum\n", " of `matmul` and zero.\n", " * The third one describes the element-wise addition of the ReLU result and the input tensor `C`.\n", - "* After defining the function, we create a `tir.Schedule` instance with regard to this function.\n", - "`tir.Schedule` is the core tool in this assignment you will use to optimize this function.\n", - "`tir.Schedule` provides the a set of function transformations (which we call \"schedule primitives\")\n", + "* After defining the function, we create a `s_tir.Schedule` instance with regard to this function.\n", + "`s_tir.Schedule` is the core tool in this assignment you will use to optimize this function.\n", + "`s_tir.Schedule` provides the a set of function transformations (which we call \"schedule primitives\")\n", "that help accelerate the function execution on hardware.\n", "We will introduce these schedule primitives later." ] @@ -391,19 +391,19 @@ "which contain a basic example on how to use the primitive.\n", "We also provide instructions in your TODO area in `gemm_relu_add.py`.\n", "\n", - "- `split` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.split)\n", + "- `split` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.split)\n", " - It splits one loop into multiple loops that collectively iterate the iteration space of the original loop.\n", " We use `split` for tiling.\n", - "- `reorder` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.reorder)\n", + "- `reorder` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.reorder)\n", " - It reorders the specified contiguous loops into a new order.\n", " We apply `reorder` for loops after splitting for tiling.\n", - "- `bind` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.bind)\n", + "- `bind` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.bind)\n", " - It binds a loop to `blockIdx.x/y/z` or `threadIdx.x/y/z` on GPU,\n", " so that we can specify the area to compute of a thread block or a single thread.\n", - "- `cache_read` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.cache_read)\n", + "- `cache_read` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.cache_read)\n", " - It generates a cache stage for the specified region in the specified memory scope.\n", " We use `cache_read` to generate the read stages from global memory to shared memory.\n", - "- `compute_at` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.compute_at)\n", + "- `compute_at` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.compute_at)\n", " - It moves a computation (e.g., the shared memory read stage) to the location under the specified loop.\n", " We use `compute_at` to move the shared memory read stages to the right location.\n", "\n", @@ -539,7 +539,7 @@ "we can write a common schedule function and apply it to the shared memory read stages\n", "of both `A` and `B`.\n", "The new schedule primitive you will use in this task is\n", - "`fuse` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.fuse),\n", + "`fuse` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.fuse),\n", "which, on the contrary to `split`, fuses multiple loops into a single loop." ] }, @@ -574,9 +574,9 @@ "source": [ "You need to implement the write cache in the `write_cache` function in `gemm_relu_add.py`.\n", "In this task, you will use the schedule primitive\n", - "`cache_write` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.cache_write)\n", + "`cache_write` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.cache_write)\n", "and\n", - "`reverse_compute_at` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.reverse_compute_at)\n", + "`reverse_compute_at` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.reverse_compute_at)\n", "to put the write cache stage at the right location.\n", "Similar to `cache_read` which generates a read cache stage for a computation,\n", "`cache_write` generates a write cache stage for a computation at the specified scope.\n", @@ -618,8 +618,8 @@ "metadata": {}, "source": [ "You need to implement the epilogue fusion in the `epilogue_fusion` function in `gemm_relu_add.py`,\n", - "where you will use `get_block` to retrieve the computation of addition and ReLU,\n", - "and use `reverse_compute_inline` [(docs)](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.reverse_compute_inline)\n", + "where you will use `get_sblock` to retrieve the computation of addition and ReLU,\n", + "and use `reverse_compute_inline` [(docs)](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.reverse_compute_inline)\n", "to inline the addition and ReLU into the write-back stage of GeMM." ] }, @@ -843,8 +843,8 @@ "offer better performance.\n", "\n", "If you want to define a search space that goes beyond a single schedule,\n", - "the primitive [`Schedule.sample_perfect_tile`](https://tvm.apache.org/docs/reference/api/python/tir/schedule.html#tvm.tir.schedule.Schedule.sample_perfect_tile)\n", - "in `tir.Schedule` helps you to sample the tile sizes from the specified loop.\n", + "the primitive [`Schedule.sample_perfect_tile`](https://tvm.apache.org/docs/reference/api/python/s_tir/schedule.html#tvm.s_tir.schedule.Schedule.sample_perfect_tile)\n", + "in `s_tir.Schedule` helps you to sample the tile sizes from the specified loop.\n", "For example, for a given schedule `sch`, a given loop `i` with loop length `length`,\n", "```python\n", "tile0, tile1, tile2, tile3 = sch.sample_perfect_tile(i, n=4)\n", @@ -923,9 +923,9 @@ "```\n", "in the end, and a function `apply_trace` printed out:\n", "```python\n", - "# from tvm import tir\n", - "def apply_trace(sch: tir.Schedule) -> None:\n", - " b0 = sch.get_block(name=\"gemm\", func_name=\"main\")\n", + "# from tvm import s_tir\n", + "def apply_trace(sch: s_tir.Schedule) -> None:\n", + " b0 = sch.get_sblock(name=\"gemm\", func_name=\"main\")\n", " # ...\n", "```\n", "\n", @@ -1070,4 +1070,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/schedule_example.ipynb b/schedule_example.ipynb index fe5cde7..692e00c 100644 --- a/schedule_example.ipynb +++ b/schedule_example.ipynb @@ -39,24 +39,25 @@ "outputs": [], "source": [ "import tvm\n", - "from tvm import tir\n", - "from tvm.script import tir as T\n", + "from tvm import s_tir\n", + "from tvm.script import tirx as T\n", + "from tvm.runtime import cuda, empty, tensor\n", "\n", "\n", - "@T.prim_func\n", + "@T.prim_func(s_tir=True)\n", "def sum_broadcast(\n", " A: T.Buffer((256,), \"float32\"),\n", " B: T.Buffer((256,), \"float32\"),\n", ") -> None:\n", " for i, k in T.grid(256, 256):\n", - " with T.block(\"sum_broadcast\"):\n", + " with T.sblock(\"sum_broadcast\"):\n", " vi = T.axis.spatial(256, i)\n", " vk = T.axis.reduce(256, k)\n", " with T.init():\n", " B[vi] = T.float32(0)\n", " B[vi] += A[vk]\n", "\n", - "sch = tir.Schedule(sum_broadcast)" + "sch = s_tir.Schedule(sum_broadcast)" ] }, { @@ -99,7 +100,7 @@ "outputs": [], "source": [ "# Fetch the computation block of \"sum_broadcast\".\n", - "block = sch.get_block(\"sum_broadcast\")\n", + "block = sch.get_sblock(\"sum_broadcast\")\n", "# Fetch the i, j loops of the computation.\n", "i, j = sch.get_loops(block)\n", "# Split i into two loops.\n", @@ -124,23 +125,24 @@ "data": { "text/html": [ "
# from tvm.script import ir as I\n",
-                            "# from tvm.script import tir as T\n",
+                            "# from tvm.script import tirx as T\n",
+                            "# from tvm.tirx.layout import Axis\n",
                             "\n",
-                            "@I.ir_module\n",
-                            "class Module:\n",
-                            "    @T.prim_func\n",
-                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
-                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
-                            "        # with T.block("root"):\n",
-                            "        for i_0, i_1, k in T.grid(2, 128, 256):\n",
-                            "            with T.block("sum_broadcast"):\n",
-                            "                vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
-                            "                vk = T.axis.reduce(256, k)\n",
-                            "                T.reads(A[vk])\n",
-                            "                T.writes(B[vi])\n",
-                            "                with T.init():\n",
-                            "                    B[vi] = T.float32(0)\n",
-                            "                B[vi] = B[vi] + A[vk]\n",
+                            "@I.ir_module\n",
+                            "class Module:\n",
+                            "    @T.prim_func(s_tir=True)\n",
+                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
+                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
+                            "        # with T.sblock("root"):\n",
+                            "        for i_0, i_1, k in T.grid(2, 128, 256):\n",
+                            "            with T.sblock("sum_broadcast"):\n",
+                            "                vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
+                            "                vk = T.axis.reduce(256, k)\n",
+                            "                T.reads(A[vk])\n",
+                            "                T.writes(B[vi])\n",
+                            "                with T.init():\n",
+                            "                    B[vi] = T.float32(0.0)\n",
+                            "                B[vi] = B[vi] + A[vk]\n",
                             "
\n" ], "text/plain": [ @@ -153,11 +155,11 @@ { "data": { "text/html": [ - "
# from tvm import tir\n",
-                            "def apply_trace(sch: tir.Schedule) -> None:\n",
-                            "  b0 = sch.get_block(name="sum_broadcast", func_name="main")\n",
-                            "  l1, l2 = sch.get_loops(block=b0)\n",
-                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True)\n",
+                            "
# from tvm import s_tir\n",
+                            "def apply_trace(sch: s_tir.Schedule) -> None:\n",
+                            "  b0 = sch.get_sblock(name="sum_broadcast", func_name="main")\n",
+                            "  l1, l2 = sch.get_loops(block=b0)\n",
+                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True, disable_predication=False)\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -188,25 +190,26 @@
                     "data": {
                         "text/html": [
                             "
# from tvm.script import ir as I\n",
-                            "# from tvm.script import tir as T\n",
+                            "# from tvm.script import tirx as T\n",
+                            "# from tvm.tirx.layout import Axis\n",
                             "\n",
-                            "@I.ir_module\n",
-                            "class Module:\n",
-                            "    @T.prim_func\n",
-                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
-                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
-                            "        # with T.block("root"):\n",
-                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
-                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
+                            "@I.ir_module\n",
+                            "class Module:\n",
+                            "    @T.prim_func(s_tir=True)\n",
+                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
+                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
+                            "        # with T.sblock("root"):\n",
+                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
+                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
                             "                for k in range(256):\n",
-                            "                    with T.block("sum_broadcast"):\n",
-                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
-                            "                        vk = T.axis.reduce(256, k)\n",
-                            "                        T.reads(A[vk])\n",
-                            "                        T.writes(B[vi])\n",
-                            "                        with T.init():\n",
-                            "                            B[vi] = T.float32(0)\n",
-                            "                        B[vi] = B[vi] + A[vk]\n",
+                            "                    with T.sblock("sum_broadcast"):\n",
+                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
+                            "                        vk = T.axis.reduce(256, k)\n",
+                            "                        T.reads(A[vk])\n",
+                            "                        T.writes(B[vi])\n",
+                            "                        with T.init():\n",
+                            "                            B[vi] = T.float32(0.0)\n",
+                            "                        B[vi] = B[vi] + A[vk]\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -219,13 +222,13 @@
                 {
                     "data": {
                         "text/html": [
-                            "
# from tvm import tir\n",
-                            "def apply_trace(sch: tir.Schedule) -> None:\n",
-                            "  b0 = sch.get_block(name="sum_broadcast", func_name="main")\n",
-                            "  l1, l2 = sch.get_loops(block=b0)\n",
-                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True)\n",
-                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
-                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
+                            "
# from tvm import s_tir\n",
+                            "def apply_trace(sch: s_tir.Schedule) -> None:\n",
+                            "  b0 = sch.get_sblock(name="sum_broadcast", func_name="main")\n",
+                            "  l1, l2 = sch.get_loops(block=b0)\n",
+                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True, disable_predication=False)\n",
+                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
+                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -262,6 +265,42 @@
                         "Test passed.\n",
                         "CUDA source code:\n",
                         "\n",
+                        "#ifdef __CUDACC_RTC__\n",
+                        "  #include \n",
+                        "  using cuda::std::uint8_t;\n",
+                        "  using cuda::std::uint16_t;\n",
+                        "  using cuda::std::uint32_t;\n",
+                        "  using cuda::std::uint64_t;\n",
+                        "  using cuda::std::int8_t;\n",
+                        "  using cuda::std::int16_t;\n",
+                        "  using cuda::std::int32_t;\n",
+                        "  using cuda::std::int64_t;\n",
+                        "\n",
+                        "  #include \n",
+                        "  namespace std {\n",
+                        "    using cuda::std::is_same;\n",
+                        "    using cuda::std::is_same_v;\n",
+                        "    using cuda::std::is_integral;\n",
+                        "    using cuda::std::is_signed;\n",
+                        "    using cuda::std::is_unsigned;\n",
+                        "    using cuda::std::is_floating_point;\n",
+                        "    using cuda::std::enable_if;\n",
+                        "    using cuda::std::conditional;\n",
+                        "  }\n",
+                        "\n",
+                        "  // NVRTC uses asm/volatile instead of __asm__/__volatile__ (gcc extension).\n",
+                        "  #ifndef __asm__\n",
+                        "  #define __asm__ asm\n",
+                        "  #endif\n",
+                        "  #ifndef __volatile__\n",
+                        "  #define __volatile__ volatile\n",
+                        "  #endif\n",
+                        "#else\n",
+                        "  #include \n",
+                        "  #include \n",
+                        "  #include \n",
+                        "#endif\n",
+                        "\n",
                         "#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \\\n",
                         "     (__CUDACC_VER_MAJOR__ > 11))\n",
                         "#define TVM_ENABLE_L2_PREFETCH 1\n",
@@ -279,16 +318,15 @@
                         "  #define uint unsigned int\n",
                         "  #define uchar unsigned char\n",
                         "  #define ushort unsigned short\n",
-                        "  #define int64_t long long\n",
-                        "  #define uint64_t unsigned long long\n",
                         "#endif\n",
-                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A, float* __restrict__ B);\n",
-                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A, float* __restrict__ B) {\n",
+                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr);\n",
+                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {\n",
                         "  for (int k = 0; k < 256; ++k) {\n",
+                        "    int cse_v1 = ((((int)blockIdx.x) * 128) + ((int)threadIdx.x));\n",
                         "    if (k == 0) {\n",
-                        "      B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = 0.000000e+00f;\n",
+                        "      B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = 0x0p+0f/*0.000000e+00*/;\n",
                         "    }\n",
-                        "    B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = (B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] + A[k]);\n",
+                        "    B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = (B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] + A_ptr[k]);\n",
                         "  }\n",
                         "}\n",
                         "\n",
@@ -297,7 +335,7 @@
                 }
             ],
             "source": [
-                "def build_and_test(sch: tir.Schedule) -> None:\n",
+                "def build_and_test(sch: s_tir.Schedule) -> None:\n",
                 "    import numpy as np\n",
                 "\n",
                 "    # Build the scheduled function.\n",
@@ -308,9 +346,9 @@
                 "    b_np = np.broadcast_to(a_np.sum(keepdims=True), shape=(256,))\n",
                 "\n",
                 "    # Run the function we scheduled and built.\n",
-                "    device = tvm.cuda()\n",
-                "    a_tvm = tvm.nd.array(a_np, device=device)\n",
-                "    b_tvm = tvm.nd.empty((256,), \"float32\", device=device)\n",
+                "    device = cuda()\n",
+                "    a_tvm = tensor(a_np, device=device)\n",
+                "    b_tvm = empty((256,), \"float32\", device=device)\n",
                 "    f(a_tvm, b_tvm)\n",
                 "\n",
                 "    # Validate the result correctness.\n",
@@ -318,7 +356,7 @@
                 "    print(\"Test passed.\")\n",
                 "\n",
                 "    # Print out the CUDA source code of the function we scheduled.\n",
-                "    print(f\"CUDA source code:\\n{f.imported_modules[0].get_source()}\")\n",
+                "    print(f\"CUDA source code:\\n{f.imports[0].inspect_source()}\")\n",
                 "\n",
                 "\n",
                 "# Run building and testing.\n",
@@ -353,32 +391,33 @@
                     "data": {
                         "text/html": [
                             "
# from tvm.script import ir as I\n",
-                            "# from tvm.script import tir as T\n",
+                            "# from tvm.script import tirx as T\n",
+                            "# from tvm.tirx.layout import Axis\n",
                             "\n",
-                            "@I.ir_module\n",
-                            "class Module:\n",
-                            "    @T.prim_func\n",
-                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
-                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
-                            "        # with T.block("root"):\n",
-                            "        A_shared = T.alloc_buffer((256,), scope="shared")\n",
+                            "@I.ir_module\n",
+                            "class Module:\n",
+                            "    @T.prim_func(s_tir=True)\n",
+                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
+                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
+                            "        # with T.sblock("root"):\n",
+                            "        A_shared = T.sblock_alloc_buffer((256,), scope="shared")\n",
                             "        for ax0 in range(256):\n",
-                            "            with T.block("A_shared"):\n",
-                            "                v0 = T.axis.spatial(256, ax0)\n",
-                            "                T.reads(A[v0])\n",
-                            "                T.writes(A_shared[v0])\n",
-                            "                A_shared[v0] = A[v0]\n",
-                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
-                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
+                            "            with T.sblock("A_shared"):\n",
+                            "                v0 = T.axis.spatial(256, ax0)\n",
+                            "                T.reads(A[v0])\n",
+                            "                T.writes(A_shared[v0])\n",
+                            "                A_shared[v0] = A[v0]\n",
+                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
+                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
                             "                for k in range(256):\n",
-                            "                    with T.block("sum_broadcast"):\n",
-                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
-                            "                        vk = T.axis.reduce(256, k)\n",
-                            "                        T.reads(A_shared[vk])\n",
-                            "                        T.writes(B[vi])\n",
-                            "                        with T.init():\n",
-                            "                            B[vi] = T.float32(0)\n",
-                            "                        B[vi] = B[vi] + A_shared[vk]\n",
+                            "                    with T.sblock("sum_broadcast"):\n",
+                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
+                            "                        vk = T.axis.reduce(256, k)\n",
+                            "                        T.reads(A_shared[vk])\n",
+                            "                        T.writes(B[vi])\n",
+                            "                        with T.init():\n",
+                            "                            B[vi] = T.float32(0.0)\n",
+                            "                        B[vi] = B[vi] + A_shared[vk]\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -391,14 +430,14 @@
                 {
                     "data": {
                         "text/html": [
-                            "
# from tvm import tir\n",
-                            "def apply_trace(sch: tir.Schedule) -> None:\n",
-                            "  b0 = sch.get_block(name="sum_broadcast", func_name="main")\n",
-                            "  l1, l2 = sch.get_loops(block=b0)\n",
-                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True)\n",
-                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
-                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
-                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
+                            "
# from tvm import s_tir\n",
+                            "def apply_trace(sch: s_tir.Schedule) -> None:\n",
+                            "  b0 = sch.get_sblock(name="sum_broadcast", func_name="main")\n",
+                            "  l1, l2 = sch.get_loops(block=b0)\n",
+                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True, disable_predication=False)\n",
+                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
+                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
+                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -442,32 +481,33 @@
                     "data": {
                         "text/html": [
                             "
# from tvm.script import ir as I\n",
-                            "# from tvm.script import tir as T\n",
+                            "# from tvm.script import tirx as T\n",
+                            "# from tvm.tirx.layout import Axis\n",
                             "\n",
-                            "@I.ir_module\n",
-                            "class Module:\n",
-                            "    @T.prim_func\n",
-                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
-                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
-                            "        # with T.block("root"):\n",
-                            "        A_shared = T.alloc_buffer((256,), scope="shared")\n",
-                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
-                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
+                            "@I.ir_module\n",
+                            "class Module:\n",
+                            "    @T.prim_func(s_tir=True)\n",
+                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
+                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
+                            "        # with T.sblock("root"):\n",
+                            "        A_shared = T.sblock_alloc_buffer((256,), scope="shared")\n",
+                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
+                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
                             "                for ax0 in range(256):\n",
-                            "                    with T.block("A_shared"):\n",
-                            "                        v0 = T.axis.spatial(256, ax0)\n",
-                            "                        T.reads(A[v0])\n",
-                            "                        T.writes(A_shared[v0])\n",
-                            "                        A_shared[v0] = A[v0]\n",
+                            "                    with T.sblock("A_shared"):\n",
+                            "                        v0 = T.axis.spatial(256, ax0)\n",
+                            "                        T.reads(A[v0])\n",
+                            "                        T.writes(A_shared[v0])\n",
+                            "                        A_shared[v0] = A[v0]\n",
                             "                for k in range(256):\n",
-                            "                    with T.block("sum_broadcast"):\n",
-                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
-                            "                        vk = T.axis.reduce(256, k)\n",
-                            "                        T.reads(A_shared[vk])\n",
-                            "                        T.writes(B[vi])\n",
-                            "                        with T.init():\n",
-                            "                            B[vi] = T.float32(0)\n",
-                            "                        B[vi] = B[vi] + A_shared[vk]\n",
+                            "                    with T.sblock("sum_broadcast"):\n",
+                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
+                            "                        vk = T.axis.reduce(256, k)\n",
+                            "                        T.reads(A_shared[vk])\n",
+                            "                        T.writes(B[vi])\n",
+                            "                        with T.init():\n",
+                            "                            B[vi] = T.float32(0.0)\n",
+                            "                        B[vi] = B[vi] + A_shared[vk]\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -480,15 +520,15 @@
                 {
                     "data": {
                         "text/html": [
-                            "
# from tvm import tir\n",
-                            "def apply_trace(sch: tir.Schedule) -> None:\n",
-                            "  b0 = sch.get_block(name="sum_broadcast", func_name="main")\n",
-                            "  l1, l2 = sch.get_loops(block=b0)\n",
-                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True)\n",
-                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
-                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
-                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
-                            "  sch.compute_at(block=b5, loop=l4, preserve_unit_loops=False, index=-1)\n",
+                            "
# from tvm import s_tir\n",
+                            "def apply_trace(sch: s_tir.Schedule) -> None:\n",
+                            "  b0 = sch.get_sblock(name="sum_broadcast", func_name="main")\n",
+                            "  l1, l2 = sch.get_loops(block=b0)\n",
+                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True, disable_predication=False)\n",
+                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
+                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
+                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
+                            "  sch.compute_at(block=b5, loop=l4, preserve_unit_loops=False, index=-1)\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -529,33 +569,34 @@
                     "data": {
                         "text/html": [
                             "
# from tvm.script import ir as I\n",
-                            "# from tvm.script import tir as T\n",
+                            "# from tvm.script import tirx as T\n",
+                            "# from tvm.tirx.layout import Axis\n",
                             "\n",
-                            "@I.ir_module\n",
-                            "class Module:\n",
-                            "    @T.prim_func\n",
-                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
-                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
-                            "        # with T.block("root"):\n",
-                            "        A_shared = T.alloc_buffer((256,), scope="shared")\n",
-                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
-                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
+                            "@I.ir_module\n",
+                            "class Module:\n",
+                            "    @T.prim_func(s_tir=True)\n",
+                            "    def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):\n",
+                            "        T.func_attr({"global_symbol": "sum_broadcast"})\n",
+                            "        # with T.sblock("root"):\n",
+                            "        A_shared = T.sblock_alloc_buffer((256,), scope="shared")\n",
+                            "        for i_0 in T.thread_binding(2, thread="blockIdx.x"):\n",
+                            "            for i_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
                             "                for ax0_0 in range(2):\n",
-                            "                    for ax0_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
-                            "                        with T.block("A_shared"):\n",
-                            "                            v0 = T.axis.spatial(256, ax0_0 * 128 + ax0_1)\n",
-                            "                            T.reads(A[v0])\n",
-                            "                            T.writes(A_shared[v0])\n",
-                            "                            A_shared[v0] = A[v0]\n",
+                            "                    for ax0_1 in T.thread_binding(128, thread="threadIdx.x"):\n",
+                            "                        with T.sblock("A_shared"):\n",
+                            "                            v0 = T.axis.spatial(256, ax0_0 * 128 + ax0_1)\n",
+                            "                            T.reads(A[v0])\n",
+                            "                            T.writes(A_shared[v0])\n",
+                            "                            A_shared[v0] = A[v0]\n",
                             "                for k in range(256):\n",
-                            "                    with T.block("sum_broadcast"):\n",
-                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
-                            "                        vk = T.axis.reduce(256, k)\n",
-                            "                        T.reads(A_shared[vk])\n",
-                            "                        T.writes(B[vi])\n",
-                            "                        with T.init():\n",
-                            "                            B[vi] = T.float32(0)\n",
-                            "                        B[vi] = B[vi] + A_shared[vk]\n",
+                            "                    with T.sblock("sum_broadcast"):\n",
+                            "                        vi = T.axis.spatial(256, i_0 * 128 + i_1)\n",
+                            "                        vk = T.axis.reduce(256, k)\n",
+                            "                        T.reads(A_shared[vk])\n",
+                            "                        T.writes(B[vi])\n",
+                            "                        with T.init():\n",
+                            "                            B[vi] = T.float32(0.0)\n",
+                            "                        B[vi] = B[vi] + A_shared[vk]\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -568,18 +609,18 @@
                 {
                     "data": {
                         "text/html": [
-                            "
# from tvm import tir\n",
-                            "def apply_trace(sch: tir.Schedule) -> None:\n",
-                            "  b0 = sch.get_block(name="sum_broadcast", func_name="main")\n",
-                            "  l1, l2 = sch.get_loops(block=b0)\n",
-                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True)\n",
-                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
-                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
-                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
-                            "  sch.compute_at(block=b5, loop=l4, preserve_unit_loops=False, index=-1)\n",
-                            "  l6, l7, l8 = sch.get_loops(block=b5)\n",
-                            "  l9, l10 = sch.split(loop=l8, factors=[None, 128], preserve_unit_iters=True)\n",
-                            "  sch.bind(loop=l10, thread_axis="threadIdx.x")\n",
+                            "
# from tvm import s_tir\n",
+                            "def apply_trace(sch: s_tir.Schedule) -> None:\n",
+                            "  b0 = sch.get_sblock(name="sum_broadcast", func_name="main")\n",
+                            "  l1, l2 = sch.get_loops(block=b0)\n",
+                            "  l3, l4 = sch.split(loop=l1, factors=[2, 128], preserve_unit_iters=True, disable_predication=False)\n",
+                            "  sch.bind(loop=l3, thread_axis="blockIdx.x")\n",
+                            "  sch.bind(loop=l4, thread_axis="threadIdx.x")\n",
+                            "  b5 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")\n",
+                            "  sch.compute_at(block=b5, loop=l4, preserve_unit_loops=False, index=-1)\n",
+                            "  l6, l7, l8 = sch.get_loops(block=b5)\n",
+                            "  l9, l10 = sch.split(loop=l8, factors=[None, 128], preserve_unit_iters=True, disable_predication=False)\n",
+                            "  sch.bind(loop=l10, thread_axis="threadIdx.x")\n",
                             "
\n"
                         ],
                         "text/plain": [
@@ -619,6 +660,42 @@
                         "Test passed.\n",
                         "CUDA source code:\n",
                         "\n",
+                        "#ifdef __CUDACC_RTC__\n",
+                        "  #include \n",
+                        "  using cuda::std::uint8_t;\n",
+                        "  using cuda::std::uint16_t;\n",
+                        "  using cuda::std::uint32_t;\n",
+                        "  using cuda::std::uint64_t;\n",
+                        "  using cuda::std::int8_t;\n",
+                        "  using cuda::std::int16_t;\n",
+                        "  using cuda::std::int32_t;\n",
+                        "  using cuda::std::int64_t;\n",
+                        "\n",
+                        "  #include \n",
+                        "  namespace std {\n",
+                        "    using cuda::std::is_same;\n",
+                        "    using cuda::std::is_same_v;\n",
+                        "    using cuda::std::is_integral;\n",
+                        "    using cuda::std::is_signed;\n",
+                        "    using cuda::std::is_unsigned;\n",
+                        "    using cuda::std::is_floating_point;\n",
+                        "    using cuda::std::enable_if;\n",
+                        "    using cuda::std::conditional;\n",
+                        "  }\n",
+                        "\n",
+                        "  // NVRTC uses asm/volatile instead of __asm__/__volatile__ (gcc extension).\n",
+                        "  #ifndef __asm__\n",
+                        "  #define __asm__ asm\n",
+                        "  #endif\n",
+                        "  #ifndef __volatile__\n",
+                        "  #define __volatile__ volatile\n",
+                        "  #endif\n",
+                        "#else\n",
+                        "  #include \n",
+                        "  #include \n",
+                        "  #include \n",
+                        "#endif\n",
+                        "\n",
                         "#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \\\n",
                         "     (__CUDACC_VER_MAJOR__ > 11))\n",
                         "#define TVM_ENABLE_L2_PREFETCH 1\n",
@@ -636,21 +713,21 @@
                         "  #define uint unsigned int\n",
                         "  #define uchar unsigned char\n",
                         "  #define ushort unsigned short\n",
-                        "  #define int64_t long long\n",
-                        "  #define uint64_t unsigned long long\n",
                         "#endif\n",
-                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A, float* __restrict__ B);\n",
-                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A, float* __restrict__ B) {\n",
-                        "  __shared__ float A_shared[256];\n",
+                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr);\n",
+                        "extern \"C\" __global__ void __launch_bounds__(128) sum_broadcast_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {\n",
+                        "  __shared__ alignas(64) float A_ptr_shared[256];\n",
                         "  for (int ax0_0 = 0; ax0_0 < 2; ++ax0_0) {\n",
-                        "    A_shared[((ax0_0 * 128) + ((int)threadIdx.x))] = A[((ax0_0 * 128) + ((int)threadIdx.x))];\n",
+                        "    int cse_v1 = ((ax0_0 * 128) + ((int)threadIdx.x));\n",
+                        "    A_ptr_shared[((ax0_0 * 128) + ((int)threadIdx.x))] = A_ptr[((ax0_0 * 128) + ((int)threadIdx.x))];\n",
                         "  }\n",
                         "  __syncthreads();\n",
                         "  for (int k = 0; k < 256; ++k) {\n",
+                        "    int cse_v2 = ((((int)blockIdx.x) * 128) + ((int)threadIdx.x));\n",
                         "    if (k == 0) {\n",
-                        "      B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = 0.000000e+00f;\n",
+                        "      B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = 0x0p+0f/*0.000000e+00*/;\n",
                         "    }\n",
-                        "    B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = (B[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] + A_shared[k]);\n",
+                        "    B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] = (B_ptr[((((int)blockIdx.x) * 128) + ((int)threadIdx.x))] + A_ptr_shared[k]);\n",
                         "  }\n",
                         "}\n",
                         "\n",
diff --git a/trace_submission.py b/trace_submission.py
index ee1b42f..435f7f8 100644
--- a/trace_submission.py
+++ b/trace_submission.py
@@ -1,5 +1,5 @@
-from tvm import tir
+from tvm import s_tir
 
-def apply_trace(sch: tir.Schedule) -> None:
+def apply_trace(sch: s_tir.Schedule) -> None:
   # Paste your `apply_trace` function here...
   ...