Repository: apache/tvm

Issue:
[Bugfix] Disable SingleEnvThreadVerifier
During TensorIR scheduling, the `IterVar`s that represent environment threads may duplicate, i.e. it is legal to have two env threads with the same name tag, which may fail the `SingleEnvThreadVerifier` check during schedule creation. This PR disables this check in this case. In the future, it may be worthwhile to bring it back against post-scheduling TIR.

It's related to [this commit](https://github.com/apache/tvm/commit/eb15d04c3bff76062e26d5647fb8af0323de1bed). CC: @jinhongyii @Lunderberg 


Discussion:
Oh, interesting.  I initially added the check due to encountering [this error](https://github.com/apache/tvm/blob/main/src/tir/transforms/lower_warp_memory.cc#L209) in `tir.transform.LowerWarpMemory`, concluding that the requirement for a single TIR variable per thread index was a TIR limitation and wanting to raise the error as early as possible.

Can we add a test case that includes multiple `T.env_thread` assignments within the same kernel, verifying that it can be lowered/built without error?  That way, it would be clear that this behavior should be valid (so long as the multiple bindings all have the same extent, that is).  I'm thinking something like the following:

```python
@T.prim_func
def func(A: T.Buffer(16, "float32"), B: T.Buffer(16, "float32")):
    tx_A = T.launch_thread("threadIdx.x", 16)
    tx_B = T.launch_thread("threadIdx.x", 16)

    B[tx_B] = A[tx_A]


tvm.build(func, target="cuda")
```
Hey thanks for getting back to me @Lunderberg!

The fundamental reason is that this pass is invoked in TIR schedule creation to validate if TIR is well-formed, and at that stage, env variables are allowed to duplicate. Below is an example of a cumulative sum operator we use when enabling [MegaBlock](https://arxiv.org/abs/2211.15841) style MoE Mixtral serving. (See also PR: https://github.com/mlc-ai/mlc-llm/pull/1529). It's the first step in computing `indptr` and `indices` in CSR/COO sparse format.

```python
@T.prim_func(private=True)
def main(var_A: T.handle, var_T_add: T.handle, seq_len: T.int64):
    T.func_attr({"tir.noalias": T.bool(True)})
    A = T.match_buffer(var_A, (seq_len * T.int64(8),), "int32")
    T_add = T.match_buffer(var_T_add, (seq_len * T.int64(8),), "int32")
    # with T.block("root"):
    T_expand_dims = T.alloc_buffer((T.int64(1), seq_len * T.int64(8)), "int32")
    output_buf = T.alloc_buffer((T.int64(1), seq_len * T.int64(8)), "int32", align=8)
    T_squeeze = T.alloc_buffer((seq_len * T.int64(8),), "int32")
    for ax0, ax1 in T.grid(T.int64(1), seq_len * T.int64(8)):
        with T.block("T_expand_dims"):
            v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(A[v_ax1])
            T.writes(T_expand_dims[v_ax0, v_ax1])
            T_expand_dims[v_ax0, v_ax1] = A[v_ax1]
    with T.block("exclusive_scan"):
        T.reads()
        T.writes()
        if seq_len * T.int64(8) == T.int64(0):
            blockIdx_x = T.launch_thread("blockIdx.x", T.int64(1))
            if blockIdx_x < T.int64(1):
                T.evaluate(0)
        else:
            with T.launch_thread("threadIdx.x", T.int64(1024)) as threadIdx_x:
                blockIdx_x = T.launch_thread("blockIdx.x", T.max(T.int64(1), (seq_len * T.int64(8) + T.int64(1023)) // T.int64(1024)))
                blockIdx_y = T.launch_thread("blockIdx.y", T.int64(1))
                if blockIdx_x * T.int64(1024) + threadIdx_x < seq_len * T.int64(8):
                    output_buf[(blockIdx_y * (seq_len * T.int64(8)) + (blockIdx_x * T.int64(1024) + threadIdx_x)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + (blockIdx_x * T.int64(1024) + threadIdx_x)) % (seq_len * T.int64(8))] = T_expand_dims[(blockIdx_y * (seq_len * T.int64(8)) + (blockIdx_x * T.int64(1024) + threadIdx_x)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + (blockIdx_x * T.int64(1024) + threadIdx_x)) % (seq_len * T.int64(8))]
            for i in range(T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8)))))):
                threadIdx_x = T.launch_thread("threadIdx.x", 1024)
                blockIdx_x = T.launch_thread("blockIdx.x", T.max(1, T.Cast("int32", (seq_len * T.int64(8) + (T.int64(1024) * T.shift_left(T.int64(2), i) - T.int64(1))) // (T.int64(1024) * T.shift_left(T.int64(2), i)))))
                blockIdx_y = T.launch_thread("blockIdx.y", T.int64(1))
                start = T.allocate([T.int64(1)], "int64", "local")
                middle = T.allocate([T.int64(1)], "int64", "local")
                end = T.allocate([T.int64(1)], "int64", "local")
                start_1 = T.Buffer((1,), "int64", data=start, scope="local")
                start_1[T.int64(0)] = T.shift_left(T.int64(2), i) * T.Cast("int64", blockIdx_x * 1024 + threadIdx_x)
                if start_1[T.int64(0)] < seq_len * T.int64(8):
                    middle_1 = T.Buffer((1,), "int64", data=middle, scope="local")
                    middle_1[T.int64(0)] = start_1[T.int64(0)] + T.shift_left(T.int64(2), i) // T.int64(2)
                    end_1 = T.Buffer((1,), "int64", data=end, scope="local")
                    end_1[T.int64(0)] = T.min(start_1[T.int64(0)] + T.shift_left(T.int64(2), i), seq_len * T.int64(8))
                    if middle_1[T.int64(0)] < seq_len * T.int64(8):
                        output_buf[(blockIdx_y * (seq_len * T.int64(8)) + end_1[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + end_1[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))] = output_buf[(blockIdx_y * (seq_len * T.int64(8)) + end_1[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + end_1[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))] + output_buf[(blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))]
            with T.launch_thread("blockIdx.x", T.int64(1)) as blockIdx_x:
                if blockIdx_x < T.int64(1):
                    output_buf[((blockIdx_x + T.int64(1)) * (seq_len * T.int64(8)) - T.int64(1)) // (seq_len * T.int64(8)), ((blockIdx_x + T.int64(1)) * (seq_len * T.int64(8)) - T.int64(1)) % (seq_len * T.int64(8))] = 0
            for j in range(T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8)))))):
                threadIdx_x = T.launch_thread("threadIdx.x", 1024)
                blockIdx_x = T.launch_thread("blockIdx.x", T.max(1, T.Cast("int32", (seq_len * T.int64(8) + (T.int64(1024) * T.shift_left(T.int64(2), T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8))))) - j - T.int64(1)) - T.int64(1))) // (T.int64(1024) * T.shift_left(T.int64(2), T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8))))) - j - T.int64(1))))))
                blockIdx_y = T.launch_thread("blockIdx.y", T.int64(1))
                start = T.allocate([T.int64(1)], "int64", "local")
                middle = T.allocate([T.int64(1)], "int64", "local")
                end = T.allocate([T.int64(1)], "int64", "local")
                end_1 = T.allocate([T.int64(1)], "int32", "local")
                start_1 = T.Buffer((1,), "int64", data=start, scope="local")
                start_1[T.int64(0)] = T.shift_left(T.int64(2), T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8))))) - j - T.int64(1)) * T.Cast("int64", blockIdx_x * 1024 + threadIdx_x)
                if start_1[T.int64(0)] < seq_len * T.int64(8):
                    middle_1 = T.Buffer((1,), "int64", data=middle, scope="local")
                    middle_1[T.int64(0)] = start_1[T.int64(0)] + T.shift_left(T.int64(2), T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8))))) - j - T.int64(1)) // T.int64(2)
                    end_2 = T.Buffer((1,), "int64", data=end, scope="local")
                    end_2[T.int64(0)] = T.min(start_1[T.int64(0)] + T.shift_left(T.int64(2), T.Cast("int64", T.ceil(T.log2(T.Cast("float64", seq_len * T.int64(8))))) - j - T.int64(1)), seq_len * T.int64(8))
                    if middle_1[T.int64(0)] < seq_len * T.int64(8):
                        end_3 = T.Buffer((1,), "int32", data=end_1, scope="local")
                        end_3[T.int64(0)] = output_buf[(blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))]
                        output_buf[(blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + middle_1[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))] = output_buf[(blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))]
                        output_buf[(blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))] = output_buf[(blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) // (seq_len * T.int64(8)), (blockIdx_y * (seq_len * T.int64(8)) + end_2[T.int64(0)] - T.int64(1)) % (seq_len * T.int64(8))] + end_3[T.int64(0)]
    for ax0 in range(seq_len * T.int64(8)):
        with T.block("T_squeeze"):
            v_ax0 = T.axis.spatial(seq_len * T.int64(8), ax0)
            T.reads(output_buf[T.int64(0), v_ax0])
            T.writes(T_squeeze[v_ax0])
            T_squeeze[v_ax0] = output_buf[T.int64(0), v_ax0]
    for ax0 in range(seq_len * T.int64(8)):
        with T.block("T_add"):
            v_ax0 = T.axis.spatial(seq_len * T.int64(8), ax0)
            T.reads(A[v_ax0], T_squeeze[v_ax0])
            T.writes(T_add[v_ax0])
            T_add[v_ax0] = A[v_ax0] + T_squeeze[v_ax0]
```

As you could see from the IR, `threadIdx.x` nests twice in this IR, and this is actually valid for codegen :((
Ah, I see.  It looks like this originates from using `sch.bind` multiple times for the same thread index.  With that usage, I'm less sure whether this should be correct behavior.  When binding a thread index that isn't currently bound, the total number of loop iterations remains the same.  When binding a thread index in a context where it is already bound, it constrains the two loop indices to always be the same.

```python
# Before
for i in range(16):
    for j in range(16):
        B[i,j] = A[i,j]

# First thread binding, still have 16*16 = 256 total iterations.  B is
# set at all indices.
for i in T.thread_binding('threadIdx.x', 16):
    for j in range(16):
        B[i,j] = A[i,j]

# Second thread binding, now only have 16 iterations.  B is only set
# for indices where i==j, along the diagonal.
for i in T.thread_binding('threadIdx.x', 16):
    for j in T.thread_binding('threadIdx.x', 16):
        B[i,j] = A[i,j]
```
On CUDA, when fetching data from global memory to shared, there's something we use called cooperative fetching, where the data is divided evenly and all threads cooperatively fetch their own portions. To implement this logic in TVM, we have to bind `threadIdx.x` for an extra time:
- Originally, it binds an outer loop to `threadIdx.x` as part of tiling process
- It adds `sch.cache_read` + `sch.compute_at` that moves data copy under a loop in TIR
- It partitions the data copy loop by the length of `threadIdx.x` and then do loop binding

It's the default schedule strategy used in MetaSchedule's multi-level tiling and here's an example of this usecase: https://github.com/apache/tvm/blob/main/tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py#L71
Thank you on the details for cooperative fetching, and I agree that this check should be removed to make sure the use case is preserved, at least for the time being.  I think so long as we have a unit test showing the functionality, this PR can be merged.

My worry wasn't with the overall sequence that produces cooperative fetching in the scheduled module, more that the individual step of producing a nested thread binding doesn't feel sound to me.  The scheduling primitives shouldn't be able to impact the output values, only how the way in which the outputs are computed.  Since there are definitely cases in which producing a nested thread binding changes the result, I was hoping there would be an alternative way to express cooperative fetching without requiring the nested thread binding.
> The scheduling primitives shouldn't be able to impact the output values, only how the way in which the outputs are computed. Since there are definitely cases in which producing a nested thread binding changes the result, I was hoping there would be an alternative way to express cooperative fetching without requiring the nested thread binding.

@Lunderberg This is a good point! Actually when designing TensorIR, @spectrometerHBH has put lots of thoughts in it specifically for cooperative fetching, and he has a proof that nesting thread bindings won't change the semantics of a TensorIR during scheduling stage if it satisfies "compact dataflow" condition. Therefore, I wouldn't worry too much in this particular case.

In fact, in a simplified non-nested case below, the TIR describes two separate kernels, but when creating a TIR schedule, it errors out because `blockIdx.x` are different in different blocks. This is actually a valid case because the split-host-device pass later will split it into two kernels.

```python
@T.prim_func
def two_kernels(var_A: T.handle, var_B: T.handle, seq_len: T.int32):
    T.func_attr({"tir.noalias": T.bool(True)})
    A = T.match_buffer(var_A, (1, seq_len * 8), "int32")
    B = T.match_buffer(var_B, (1, seq_len * 8), "int32", align=8)
    with T.block("exclusive_scan"):
        T.reads()
        T.writes()
        s8: T.int32 = seq_len * 8
        if s8 == 0:
            blockIdx_x = T.launch_thread("blockIdx.x", 1)
        else:
            with T.launch_thread("threadIdx.x", 1024) as threadIdx_x:
                blockIdx_x = T.launch_thread("blockIdx.x", T.ceildiv(s8, 1024))
                i: T.int32 = blockIdx_x * 1024 + threadIdx_x
                if i < s8:
                    B[i // s8, i % s8] = A[i // s8, i % s8]
```

I added this as a testcase
There's regression on unity CI https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-unity/detail/unity/868/pipeline, we will need this fix
To quickly unblock the Unity CI, @vinx13 could you cherry-pick this commit to unity?

Resolving patch:
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -116,6 +116,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
             tir_call = self.builder_.call_te(
                 te_func,
                 call.args[0],
+                k=call.attrs.k,
                 axis=call.attrs.axis,
                 ret_type=call.attrs.ret_type,
                 is_ascend=not call.attrs.largest,


Test patch:
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py b/tests/python/relax/test_backend_dispatch_sort_scan.py
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -38,84 +38,27 @@ class Before:
         def foo(x: R.Tensor((2, 3), "float32", "llvm")):
             with R.dataflow():
                 lv0 = R.cumsum(x, axis=1, dtype="float64", exclusive=False)
-                gv = R.cumprod(lv0, axis=1, dtype="float64", exclusive=False)
+                lv1 = R.cumprod(lv0, axis=1, dtype="float64", exclusive=False)
+                gv = lv1
                 R.output(gv)
             return gv
 
-    @I.ir_module
-    class Expected:
-        I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
+    mod = DispatchSortScan()(Before)
 
-        @T.prim_func(private=True)
-        def cumsum(var_A: T.handle, out_buf: T.Buffer((T.int64(2), T.int64(3)), "float64")):
-            T.func_attr({"tir.noalias": T.bool(True)})
-            A = T.match_buffer(var_A, (T.int64(2), T.int64(3)), offset_factor=1)
-            with T.block("cumsum_generic"):
-                for fused in T.parallel(T.int64(2)):
-                    out_buf[
-                        fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3)
-                    ] = T.Cast(
-                        "float64",
-                        A[fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3)],
-                    )
-                    for _k in range(T.int64(2)):
-                        out_buf[
-                            (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3),
-                            (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3),
-                        ] = out_buf[
-                            (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) // T.int64(3),
-                            (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) % T.int64(3),
-                        ] + T.Cast(
-                            "float64",
-                            A[
-                                (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3),
-                                (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3),
-                            ],
-                        )
-
-        @T.prim_func(private=True)
-        def cumprod(var_A: T.handle, out_buf: T.Buffer((T.int64(2), T.int64(3)), "float64")):
-            T.func_attr({"tir.noalias": T.bool(True)})
-            A = T.match_buffer(var_A, (T.int64(2), T.int64(3)), "float64", offset_factor=1)
-            with T.block("cumprod_generic"):
-                T.reads(A[T.int64(0) : T.int64(2), T.int64(0) : T.int64(3)])
-                T.writes(out_buf[T.int64(0) : T.int64(2), T.int64(0) : T.int64(3)])
-                for fused in T.parallel(T.int64(2)):
-                    out_buf[fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3)] = A[
-                        fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3)
-                    ]
-                    for _k in range(T.int64(2)):
-                        out_buf[
-                            (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3),
-                            (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3),
-                        ] = (
-                            out_buf[
-                                (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) // T.int64(3),
-                                (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) % T.int64(3),
-                            ]
-                            * A[
-                                (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3),
-                                (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3),
-                            ]
-                        )
+    vdevices = [I.vdevice("llvm", 0)]
+    x = relax.Var("x", R.Tensor((2, 3), "float32", vdevices[0]))
+    bb = relax.BlockBuilder()
 
-        @R.function
-        def foo(
-            x: R.Tensor((2, 3), dtype="float32", vdevice="llvm")
-        ) -> R.Tensor((2, 3), dtype="float64", vdevice="llvm"):
-            cls = Expected
-            with R.dataflow():
-                lv0 = R.call_tir(cls.cumsum, (x,), out_sinfo=R.Tensor((2, 3), "float64", "llvm"))
-                gv = R.call_tir(
-                    cls.cumprod,
-                    (lv0,),
-                    out_sinfo=R.Tensor((2, 3), dtype="float64", vdevice="llvm"),
-                )
-                R.output(gv)
-            return gv
+    with bb.function("foo", (x,), {"global_symbol": "foo"}):
+        with bb.dataflow():
+            lv0 = bb.emit_te(topi.cumsum, x, axis=1, dtype="float64", exclusive=False)
+            out = bb.emit_te(topi.cumprod, lv0, axis=1, dtype="float64", exclusive=False)
+            out = bb.emit_output(out)
+        bb.emit_func_output(out)
+    expected_mod = bb.finalize()
+    expected_mod.update_global_info("vdevice", vdevices)
 
-    mod = DispatchSortScan()(Before)
-    assert_structural_equal(mod, Expected)
+    assert_structural_equal(mod, expected_mod)
 
 
 def test_dispatch_scanop_cuda():
@@ -172,60 +115,26 @@ class Before:
         def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
             m = T.int64()
             with R.dataflow():
-                gv = R.sort(x, axis=1, descending=False)
+                lv = R.sort(x, axis=1, descending=False)
+                gv = lv
                 R.output(gv)
             return gv
 
-    @I.ir_module
-    class Expected:
-        I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
-
-        @T.prim_func(private=True)
-        def sort(var_A: T.handle, var_sort_cpu: T.handle):
-            T.func_attr({"tir.noalias": T.bool(True)})
-            m = T.int64()
-            data_buf = T.match_buffer(var_A, (m, T.int64(3)), align=8)
-            out_buf = T.match_buffer(var_sort_cpu, (m, T.int64(3)), align=8)
-            with T.block("sort_cpu"):
-                T.reads()
-                T.writes()
-                T.call_packed(
-                    "tvm.contrib.sort.sort",
-                    T.tvm_stack_make_array(
-                        data_buf.data,
-                        T.tvm_stack_make_shape(m, T.int64(3)),
-                        0,
-                        2,
-                        T.float32(0),
-                        T.int64(0),
-                    ),
-                    T.tvm_stack_make_array(
-                        out_buf.data,
-                        T.tvm_stack_make_shape(m, T.int64(3)),
-                        0,
-                        2,
-                        T.float32(0),
-                        T.int64(0),
-                    ),
-                    1,
-                    T.bool(True),
-                )
+    vdevices = [I.vdevice("llvm", 0)]
+    m = tir.Var("m", "int64")
+    x = relax.Var("x", R.Tensor((m, 3), "float32", vdevices[0]))
+    bb = relax.BlockBuilder()
 
-        @R.function
-        def foo(
-            x: R.Tensor(("m", 3), dtype="float32", vdevice="llvm")
-        ) -> R.Tensor(("m", 3), dtype="float32", vdevice="llvm"):
-            m = T.int64()
-            cls = Expected
-            with R.dataflow():
-                gv = R.call_tir(
-                    cls.sort, (x,), out_sinfo=R.Tensor((m, 3), dtype="float32", vdevice="llvm")
-                )
-                R.output(gv)
-            return gv
+    with bb.function("foo", (x,), {"global_symbol": "foo"}):
+        with bb.dataflow():
+            out = bb.emit_te(topi.sort, x, axis=1, is_ascend=True)
+            out = bb.emit_output(out)
+        bb.emit_func_output(out)
+    expected_mod = bb.finalize()
+    expected_mod.update_global_info("vdevice", vdevices)
 
     mod = DispatchSortScan()(Before)
-    assert_structural_equal(mod, Expected)
+    assert_structural_equal(mod, expected_mod)
 
 
 def test_dispatch_sort_cuda():
@@ -295,55 +204,26 @@ class Before:
         def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
             m = T.int64()
             with R.dataflow():
-                gv = R.argsort(x, axis=1, descending=False)
+                lv = R.argsort(x, axis=1, descending=False, dtype="int32")
+                gv = lv
                 R.output(gv)
             return gv
 
-    @I.ir_module
-    class Expected:
-        I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
-
-        @T.prim_func(private=True)
-        def argsort(var_A: T.handle, var_argsort_cpu: T.handle):
-            T.func_attr({"tir.noalias": T.bool(True)})
-            m = T.int64()
-            data_buf = T.match_buffer(var_A, (m, T.int64(3)), align=8)
-            out_buf = T.match_buffer(var_argsort_cpu, (m, T.int64(3)), "int32", align=8)
-            with T.block("argsort_cpu"):
-                T.reads(data_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
-                T.writes(out_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
-                T.call_packed(
-                    "tvm.contrib.sort.argsort",
-                    T.tvm_stack_make_array(
-                        data_buf.data,
-                        T.tvm_stack_make_shape(m, T.int64(3)),
-                        0,
-                        2,
-                        T.float32(0),
-                        T.int64(0),
-                    ),
-                    T.tvm_stack_make_array(
-                        out_buf.data, T.tvm_stack_make_shape(m, T.int64(3)), 0, 2, 0, T.int64(0)
-                    ),
-                    1,
-                    T.bool(True),
-                )
+    vdevices = [I.vdevice("llvm", 0)]
+    m = tir.Var("m", "int64")
+    x = relax.Var("x", R.Tensor((m, 3), "float32", vdevices[0]))
+    bb = relax.BlockBuilder()
 
-        @R.function
-        def foo(
-            x: R.Tensor(("m", 3), dtype="float32", vdevice="llvm")
-        ) -> R.Tensor(("m", 3), dtype="int32", vdevice="llvm"):
-            m = T.int64()
-            cls = Expected
-            with R.dataflow():
-                gv = R.call_tir(
-                    cls.argsort, (x,), out_sinfo=R.Tensor((m, 3), dtype="int32", vdevice="llvm")
-                )
-                R.output(gv)
-            return gv
+    with bb.function("foo", (x,), {"global_symbol": "foo"}):
+        with bb.dataflow():
+            out = bb.emit_te(topi.argsort, x, axis=1, is_ascend=True, dtype="int32")
+            out = bb.emit_output(out)
+        bb.emit_func_output(out)
+    expected_mod = bb.finalize()
+    expected_mod.update_global_info("vdevice", vdevices)
 
     mod = DispatchSortScan()(Before)
-    assert_structural_equal(mod, Expected)
+    assert_structural_equal(mod, expected_mod)
 
 
 def test_dispatch_argsort_cuda():
@@ -410,71 +290,26 @@ class Before:
         def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
             m = T.int64()
             with R.dataflow():
-                gv = R.topk(x, k=2, axis=1, largest=True)
+                lv = R.topk(x, k=2, axis=1, largest=True)
+                gv = lv
                 R.output(gv)
             return gv
 
-    @I.ir_module
-    class Expected:
-        I.module_global_infos({"vdevice": [I.vdevice("llvm", 0)]})
-
-        @T.prim_func(private=True)
-        def topk(var_A: T.handle, var_topk_cpu_v0: T.handle, var_topk_cpu_v1: T.handle):
-            T.func_attr({"tir.noalias": T.bool(True)})
-            m = T.int64()
-            data_buf = T.match_buffer(var_A, (m, T.int64(3)), align=8)
-            value_buf = T.match_buffer(var_topk_cpu_v0, (m, T.int64(1)), align=8)
-            indices_buf = T.match_buffer(var_topk_cpu_v1, (m, T.int64(1)), "int32", align=8)
-            with T.block("topk_cpu"):
-                T.reads(data_buf[T.int64(0) : m, T.int64(0) : T.int64(3)])
-                T.writes(
-                    value_buf[T.int64(0) : m, T.int64(0)], indices_buf[T.int64(0) : m, T.int64(0)]
-                )
-                T.call_packed(
-                    "tvm.contrib.sort.topk",
-                    T.tvm_stack_make_array(
-                        data_buf.data,
-                        T.tvm_stack_make_shape(m, T.int64(3)),
-                        0,
-                        2,
-                        T.float32(0),
-                        T.int64(0),
-                    ),
-                    T.tvm_stack_make_array(
-                        value_buf.data, T.tvm_stack_make_shape(m, 1), 0, 2, T.float32(0), T.int64(0)
-                    ),
-                    T.tvm_stack_make_array(
-                        indices_buf.data, T.tvm_stack_make_shape(m, 1), 0, 2, 0, T.int64(0)
-                    ),
-                    1,
-                    1,
-                    "both",
-                    T.bool(False),
-                )
+    vdevices = [I.vdevice("llvm", 0)]
+    m = tir.Var("m", "int64")
+    x = relax.Var("x", R.Tensor((m, 3), "float32", vdevices[0]))
+    bb = relax.BlockBuilder()
 
-        @R.function
-        def foo(
-            x: R.Tensor(("m", 3), dtype="float32", vdevice="llvm")
-        ) -> R.Tuple(
-            R.Tensor(("m", 1), dtype="float32", vdevice="llvm"),
-            R.Tensor(("m", 1), dtype="int32", vdevice="llvm"),
-        ):
-            m = T.int64()
-            cls = Expected
-            with R.dataflow():
-                gv = R.call_tir(
-                    cls.topk,
-                    (x,),
-                    out_sinfo=[
-                        R.Tensor((m, 1), dtype="float32", vdevice="llvm"),
-                        R.Tensor((m, 1), dtype="int32", vdevice="llvm"),
-                    ],
-                )
-                R.output(gv)
-            return gv
+    with bb.function("foo", (x,), {"global_symbol": "foo"}):
+        with bb.dataflow():
+            out = bb.emit_te(topi.topk, x, k=2, axis=1, is_ascend=False, dtype="int32")
+            out = bb.emit_output(out)
+        bb.emit_func_output(out)
+    expected_mod = bb.finalize()
+    expected_mod.update_global_info("vdevice", vdevices)
 
     mod = DispatchSortScan()(Before)
-    assert_structural_equal(mod, Expected)
+    assert_structural_equal(mod, expected_mod)
 
 
 def test_dispatch_topk_cuda():
@@ -494,12 +329,11 @@ def foo(x: R.Tensor((2, 3), "float32", "cuda")):
 
     vdevices = [I.vdevice("cuda", 0)]
     x = relax.Var("x", R.Tensor((2, 3), "float32", vdevices[0]))
-    y = relax.Var("y", R.Tensor((2, 3), "float32"))
     bb = relax.BlockBuilder()
     with target:
         with bb.function("foo", (x,), {"global_symbol": "foo"}):
             with bb.dataflow():
-                out = bb.emit_te(topi.cuda.topk, x, axis=1, is_ascend=False, dtype="int32")
+                out = bb.emit_te(topi.cuda.topk, x, k=2, axis=1, is_ascend=False, dtype="int32")
                 out = bb.emit_output(out)
             bb.emit_func_output(out)
     expected_mod = bb.finalize()