def METHOD_NAME(vthread_name):
    dtype = "int64"
    n = 100
    m = 4
    nthread = 2
    def get_vthread(name):
        tx = te.thread_axis(name)
        ty = te.thread_axis(name)
        ib = tvm.tir.ir_builder.create()
        with ib.for_range(0, n) as i:
            ib.scope_attr(tx, "virtual_thread", nthread)
            ib.scope_attr(ty, "virtual_thread", nthread)
            A = ib.allocate("float32", m, name="A", scope="shared")
            B = ib.allocate("float32", m, name="B", scope="shared")
            C = ib.allocate("float32", m, name="C", scope="shared")
            abuffer = A.asobject()
            bbuffer = B.asobject()
            cbuffer = C.asobject()
            A[tx] = tx + 1.0
            B[ty] = ty + 1.0
            ib.emit(
                tvm.tir.call_extern(
                    "int32",
                    "Run",
                    abuffer.access_ptr("r"),
                    bbuffer.access_ptr("r"),
                    cbuffer.access_ptr("rw"),
                )
            )
        return ib.get()
    if vthread_name == "vthread":
        A_expected_alloc = m * nthread
    elif vthread_name == "cthread":
        A_expected_alloc = m * nthread * nthread
    C_expected_alloc = m * nthread * nthread
    stmt = tvm.tir.transform.InjectVirtualThread()(
        tvm.IRModule.from_expr(
            tvm.tir.PrimFunc([], get_vthread(vthread_name)).with_attr("global_symbol", "main")
        )
    )["main"]
    assert list(stmt.body.body.extents) == [A_expected_alloc]
    assert list(stmt.body.body.body.body.extents) == [C_expected_alloc]