# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E501, F401, F841

import sys

import numpy as np
import pytest

import tvm
import tvm.testing
from tvm import tirx
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


def opt_gemm_lower():
    @tvm.script.ir_module
    class Module:
        @T.prim_func(s_tir=True)
        def mmult(A: T.handle, B: T.handle, C: T.handle) -> None:
            # function attr dict
            T.func_attr({"tirx.noalias": True})
            A_1 = T.match_buffer(A, [16384], elem_offset=0, align=64, offset_factor=1)
            B_1 = T.match_buffer(B, [1024, 1024], elem_offset=0, align=64, offset_factor=1)
            C_1 = T.match_buffer(C, [16384], elem_offset=0, align=64, offset_factor=1)
            # body
            packedB = T.alloc_buffer((32768,))
            for x in T.parallel(0, 32):
                for y in T.serial(0, 1024):
                    packedB[T.ramp(((x * 32768) + (y * 32)), 1, 32)] = B_1[y, T.ramp(x * 32, 1, 32)]
            for x_outer in T.parallel(0, 32):
                C_global = T.alloc_buffer((1024,))
                for y_outer in T.serial(0, 32):
                    for x_c_init in T.serial(0, 32):
                        C_global[T.ramp((x_c_init * 32), 1, 32)] = T.broadcast(T.float32(0), 32)
                    for k_outer in T.serial(0, 256):
                        for x_c in T.serial(0, 32):
                            C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
                                T.ramp((x_c * 32), 1, 32)
                            ] + (
                                T.broadcast(
                                    A_1[(((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4))],
                                    32,
                                )
                                * packedB[T.ramp(((y_outer * 32768) + (k_outer * 128)), 1, 32)]
                            )
                            C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
                                T.ramp((x_c * 32), 1, 32)
                            ] + (
                                T.broadcast(
                                    A_1[
                                        ((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 1),
                                    ],
                                    32,
                                )
                                * packedB[
                                    T.ramp((((y_outer * 32768) + (k_outer * 128)) + 32), 1, 32)
                                ]
                            )
                            C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
                                T.ramp((x_c * 32), 1, 32)
                            ] + (
                                T.broadcast(
                                    A_1[
                                        ((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 2),
                                    ],
                                    32,
                                )
                                * packedB[
                                    T.ramp((((y_outer * 32768) + (k_outer * 128)) + 64), 1, 32)
                                ]
                            )
                            C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
                                T.ramp((x_c * 32), 1, 32)
                            ] + (
                                T.broadcast(
                                    A_1[
                                        ((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 3),
                                    ],
                                    32,
                                )
                                * packedB[
                                    T.ramp((((y_outer * 32768) + (k_outer * 128)) + 96), 1, 32)
                                ]
                            )
                    for x_inner in T.serial(0, 32):
                        for y_inner in T.serial(0, 32):
                            C_1[
                                (
                                    (((x_outer * 32768) + (x_inner * 1024)) + (y_outer * 32))
                                    + y_inner
                                )
                            ] = C_global[((x_inner * 32) + y_inner)]

    return Module


def launch_env_thread():
    @T.prim_func(s_tir=True)
    def main(inputs: T.Buffer((64, 2, 4), "float32")) -> None:
        bx = T.launch_thread("blockIdx.x", 64)
        for i, j in T.grid(2, 4):
            T.evaluate(inputs[bx, i, j])

    return main


def opt_conv_tensorcore_lower():
    @T.prim_func(s_tir=True)
    def func(
        A: T.Buffer((16, 14, 14, 16, 16, 16), "float16"),
        W: T.Buffer((3, 3, 16, 32, 16, 16), "float16"),
        Conv: T.Buffer((16, 14, 14, 32, 16, 16), "float32"),
    ) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
        # body
        A_1 = T.decl_buffer([12845056], dtype="float16", data=A.data)
        W_1 = T.decl_buffer([1179648], dtype="float16", data=W.data)
        Conv_1 = T.decl_buffer([25690112], data=Conv.data)
        bx = T.env_thread("blockIdx.x")
        by = T.env_thread("blockIdx.y")
        bz = T.env_thread("blockIdx.z")
        tx = T.env_thread("threadIdx.x")
        ty = T.env_thread("threadIdx.y")
        tz = T.env_thread("threadIdx.z")
        T.launch_thread(bz, 196)
        Conv_wmma_accumulator = T.alloc_buffer((2048,), scope="wmma.accumulator")
        Apad_shared = T.alloc_buffer((12288,), "float16", scope="shared")
        W_shared = T.alloc_buffer((12288,), "float16", scope="shared")
        Apad_shared_wmma_matrix_a = T.alloc_buffer((512,), "float16", scope="wmma.matrix_a")
        W_shared_wmma_matrix_b = T.alloc_buffer((1024,), "float16", scope="wmma.matrix_b")
        T.launch_thread(bx, 2)
        T.launch_thread(by, 4)
        T.launch_thread(ty, 4)
        T.launch_thread(tz, 2)
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 0, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 1, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 2, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 3, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 4, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 5, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 6, T.float32(0), dtype="handle"
            )
        )
        T.evaluate(
            T.tvm_fill_fragment(
                Conv_wmma_accumulator.data, 16, 16, 16, 7, T.float32(0), dtype="handle"
            )
        )
        for ic_outer in T.serial(0, 8):
            for kh in T.serial(0, 3):
                for ax2 in T.serial(0, 3):
                    with T.launch_thread(tx, 32):
                        Apad_shared[((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61440
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 32)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61408
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 64)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61376
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 96)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61344
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 128)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61312
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 160)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61280
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 192)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61248
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 224)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61216
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 256)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61184
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 288)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61152
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 320)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61120
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 352)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61088
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 384)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61056
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 416)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            (1 <= (T.floordiv(bz, 14) + kh))
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 61024
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    with T.launch_thread(tx, 32):
                        Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 448)] = (
                            T.if_then_else(
                                (
                                    (
                                        (
                                            1 <= (T.floordiv(bz, 14) + kh)
                                            and ((T.floordiv(bz, 14) + kh) < 15)
                                        )
                                        and (1 <= (ax2 + T.floormod(bz, 14)))
                                    )
                                    and ((ax2 + T.floormod(bz, 14)) < 15)
                                ),
                                A_1[
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            (
                                                                ((bx * 6422528) + (ty * 1605632))
                                                                + (tz * 802816)
                                                            )
                                                            + (kh * 57344)
                                                        )
                                                        + (bz * 4096)
                                                    )
                                                    + (ax2 * 4096)
                                                )
                                                + (ic_outer * 512)
                                            )
                                            + tx
                                        )
                                        - 60992
                                    ),
                                ],
                                T.float16(0),
                                dtype="float16",
                            )
                        )
                    T.launch_thread(tx, 32)
                    Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 480)] = (
                        T.if_then_else(
                            (
                                (
                                    (
                                        (1 <= (T.floordiv(bz, 14) + kh))
                                        and ((T.floordiv(bz, 14) + kh) < 15)
                                    )
                                    and (1 <= (ax2 + T.floormod(bz, 14)))
                                )
                                and ((ax2 + T.floormod(bz, 14)) < 15)
                            ),
                            A_1[
                                (
                                    (
                                        (
                                            (
                                                (
                                                    (
                                                        (
                                                            ((bx * 6422528) + (ty * 1605632))
                                                            + (tz * 802816)
                                                        )
                                                        + (kh * 57344)
                                                    )
                                                    + (bz * 4096)
                                                )
                                                + (ax2 * 4096)
                                            )
                                            + (ic_outer * 512)
                                        )
                                        + tx
                                    )
                                    - 60960
                                ),
                            ],
                            T.float16(0),
                            dtype="float16",
                        )
                    )
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp((((ty * 512) + (tz * 256)) + (tx * 8)), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                        + (ty * 512)
                                    )
                                    + (tz * 256)
                                )
                                + (tx * 8)
                            ),
                            1,
                            8,
                        )
                    ]
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 2048), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (
                                            (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                            + (ty * 512)
                                        )
                                        + (tz * 256)
                                    )
                                    + (tx * 8)
                                )
                                + 8192
                            ),
                            1,
                            8,
                        )
                    ]
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 4096), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (
                                            (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                            + (ty * 512)
                                        )
                                        + (tz * 256)
                                    )
                                    + (tx * 8)
                                )
                                + 131072
                            ),
                            1,
                            8,
                        )
                    ]
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 6144), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (
                                            (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                            + (ty * 512)
                                        )
                                        + (tz * 256)
                                    )
                                    + (tx * 8)
                                )
                                + 139264
                            ),
                            1,
                            8,
                        )
                    ]
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 8192), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (
                                            (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                            + (ty * 512)
                                        )
                                        + (tz * 256)
                                    )
                                    + (tx * 8)
                                )
                                + 262144
                            ),
                            1,
                            8,
                        )
                    ]
                with T.launch_thread(tx, 32):
                    W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 10240), 1, 8)] = W_1[
                        T.ramp(
                            (
                                (
                                    (
                                        (
                                            (((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
                                            + (ty * 512)
                                        )
                                        + (tz * 256)
                                    )
                                    + (tx * 8)
                                )
                                + 270336
                            ),
                            1,
                            8,
                        )
                    ]
                for ic_inner in T.serial(0, 2):
                    for kw in T.serial(0, 3):
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                Apad_shared_wmma_matrix_a.data,
                                16,
                                16,
                                16,
                                0,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    Apad_shared.data,
                                    (((ty * 3072) + (kw * 512)) + (ic_inner * 256)),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                Apad_shared_wmma_matrix_a.data,
                                16,
                                16,
                                16,
                                1,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    Apad_shared.data,
                                    ((((ty * 3072) + (kw * 512)) + (ic_inner * 256)) + 1536),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                W_shared_wmma_matrix_b.data,
                                16,
                                16,
                                16,
                                0,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    W_shared.data,
                                    (((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                W_shared_wmma_matrix_b.data,
                                16,
                                16,
                                16,
                                1,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    W_shared.data,
                                    ((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 256),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                W_shared_wmma_matrix_b.data,
                                16,
                                16,
                                16,
                                2,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    W_shared.data,
                                    ((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 512),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_load_matrix_sync(
                                W_shared_wmma_matrix_b.data,
                                16,
                                16,
                                16,
                                3,
                                T.tvm_access_ptr(
                                    T.type_annotation(dtype="float16"),
                                    W_shared.data,
                                    ((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 768),
                                    256,
                                    1,
                                    dtype="handle",
                                ),
                                16,
                                "row_major",
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                0,
                                Apad_shared_wmma_matrix_a.data,
                                0,
                                W_shared_wmma_matrix_b.data,
                                0,
                                Conv_wmma_accumulator.data,
                                0,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                1,
                                Apad_shared_wmma_matrix_a.data,
                                0,
                                W_shared_wmma_matrix_b.data,
                                1,
                                Conv_wmma_accumulator.data,
                                1,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                2,
                                Apad_shared_wmma_matrix_a.data,
                                0,
                                W_shared_wmma_matrix_b.data,
                                2,
                                Conv_wmma_accumulator.data,
                                2,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                3,
                                Apad_shared_wmma_matrix_a.data,
                                0,
                                W_shared_wmma_matrix_b.data,
                                3,
                                Conv_wmma_accumulator.data,
                                3,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                4,
                                Apad_shared_wmma_matrix_a.data,
                                1,
                                W_shared_wmma_matrix_b.data,
                                0,
                                Conv_wmma_accumulator.data,
                                4,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                5,
                                Apad_shared_wmma_matrix_a.data,
                                1,
                                W_shared_wmma_matrix_b.data,
                                1,
                                Conv_wmma_accumulator.data,
                                5,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                6,
                                Apad_shared_wmma_matrix_a.data,
                                1,
                                W_shared_wmma_matrix_b.data,
                                2,
                                Conv_wmma_accumulator.data,
                                6,
                                dtype="handle",
                            )
                        )
                        T.evaluate(
                            T.tvm_mma_sync(
                                Conv_wmma_accumulator.data,
                                7,
                                Apad_shared_wmma_matrix_a.data,
                                1,
                                W_shared_wmma_matrix_b.data,
                                3,
                                Conv_wmma_accumulator.data,
                                7,
                                dtype="handle",
                            )
                        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                0,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                        + (tz * 1024)
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                1,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 256
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                2,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 512
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                3,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 768
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                4,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 1605632
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                5,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 1605888
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                6,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 1606144
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )
        T.evaluate(
            T.tvm_store_matrix_sync(
                Conv_wmma_accumulator.data,
                16,
                16,
                16,
                7,
                T.tvm_access_ptr(
                    T.type_annotation(dtype="float32"),
                    Conv_1.data,
                    (
                        (
                            ((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
                            + (tz * 1024)
                        )
                        + 1606400
                    ),
                    256,
                    2,
                    dtype="handle",
                ),
                16,
                "row_major",
                dtype="handle",
            )
        )

    return func


def opt_conv_tensorcore_mod_host():
    @T.prim_func(s_tir=True)
    def opt_conv_tensorcore_mod_host(
        args: T.handle,
        arg_type_ids: T.Buffer((3,), "int32"),
        num_args: T.int32,
        out_ret_value: T.handle,
        out_ret_tcode: T.handle,
        resource_handle: T.handle,
    ) -> T.int32:
        # function attr dict
        T.func_attr(
            {
                "tirx.noalias": True,
                "global_symbol": "default_function",
                "tirx.is_entry_func": True,
                "calling_conv": 1,
            }
        )
        # body
        stack_tcode_data: T.let[T.handle("int32")] = T.tvm_stack_alloca(
            "arg_tcode", 10, dtype="handle"
        )
        stack_tcode = T.decl_buffer([9], "int32", data=stack_tcode_data)
        stack_value: T.let[T.handle] = T.tvm_stack_alloca("arg_value", 10, dtype="handle")
        assert num_args == 3, "default_function: num_args should be 3"
        arg0: T.let[T.handle] = T.tvm_struct_get(args, 0, 12, dtype="handle")
        arg0_code: T.let[T.int32] = arg_type_ids[0]
        arg1: T.let[T.handle] = T.tvm_struct_get(args, 1, 12, dtype="handle")
        arg1_code: T.let[T.int32] = arg_type_ids[1]
        arg2: T.let[T.handle] = T.tvm_struct_get(args, 2, 12, dtype="handle")
        arg2_code: T.let[T.int32] = arg_type_ids[2]

        A: T.let[T.handle] = T.tvm_struct_get(arg0, 0, 1, dtype="handle")
        T.attr(A, "storage_alignment", 128)
        arg0_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg0, 0, 2, dtype=T.handle("int64").ty
        )
        arg0_shape = T.decl_buffer([6], "int64", data=arg0_shape_data)
        arg0_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg0, 0, 3, dtype=T.handle("int64").ty
        )
        arg0_strides = T.decl_buffer([6], "int64", data=arg0_strides_data)

        dev_id: T.let[T.int32] = T.tvm_struct_get(arg0, 0, 9, dtype="int32")

        W: T.let[T.handle] = T.tvm_struct_get(arg1, 0, 1, dtype="handle")
        T.attr(W, "storage_alignment", 128)
        arg1_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg1, 0, 2, dtype=T.handle("int64").ty
        )
        arg1_shape = T.decl_buffer([6], "int64", data=arg1_shape_data)
        arg1_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg1, 0, 3, dtype=T.handle("int64").ty
        )
        arg1_strides = T.decl_buffer([6], "int64", data=arg1_strides_data)

        Conv: T.let[T.handle] = T.tvm_struct_get(arg2, 0, 1, dtype="handle")
        T.attr(Conv, "storage_alignment", 128)
        arg2_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg2, 0, 2, dtype=T.handle("int64").ty
        )
        arg2_shape = T.decl_buffer([6], "int64", data=arg2_shape_data)
        arg2_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
            arg2, 0, 3, dtype=T.handle("int64").ty
        )
        arg2_strides = T.decl_buffer([6], "int64", data=arg2_strides_data)

        assert (((arg0_code == 3) or (arg0_code == 13)) or (arg0_code == 7)) or (arg0_code == 4), (
            "default_function: Expect arg[0] to be pointer"
        )
        assert (((arg1_code == 3) or (arg1_code == 13)) or (arg1_code == 7)) or (arg1_code == 4), (
            "default_function: Expect arg[1] to be pointer"
        )
        assert (((arg2_code == 3) or (arg2_code == 13)) or (arg2_code == 7)) or (arg2_code == 4), (
            "default_function: Expect arg[2] to be pointer"
        )
        assert 6 == T.tvm_struct_get(arg0, 0, 4, dtype="int32"), "arg0.ndim is expected to equal 6"
        assert 6 == T.tvm_struct_get(arg0, 0, 4, dtype="int32"), "arg0.ndim is expected to equal 6"
        assert (
            (T.tvm_struct_get(arg0, 0, 5, dtype="uint8") == T.uint8(2))
            and (T.tvm_struct_get(arg0, 0, 6, dtype="uint8") == T.uint8(16))
        ) and (T.tvm_struct_get(arg0, 0, 7, dtype="uint16") == T.uint16(1)), (
            "arg0.dtype is expected to be float16"
        )
        assert 16 == T.cast(arg0_shape[0], "int32"), (
            "Argument arg0.shape[0] has an unsatisfied constraint"
        )
        assert 14 == T.cast(arg0_shape[1], "int32"), (
            "Argument arg0.shape[1] has an unsatisfied constraint"
        )
        assert 14 == T.cast(arg0_shape[2], "int32"), (
            "Argument arg0.shape[2] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg0_shape[3], "int32"), (
            "Argument arg0.shape[3] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg0_shape[4], "int32"), (
            "Argument arg0.shape[4] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg0_shape[5], "int32"), (
            "Argument arg0.shape[5] has an unsatisfied constraint"
        )
        if not (T.isnullptr(arg0_strides.data, dtype="bool")):
            assert (
                (
                    (
                        (
                            (1 == T.cast(arg0_strides[5], "int32"))
                            and (16 == T.cast(arg0_strides[4], "int32"))
                        )
                        and (256 == T.cast(arg0_strides[3], "int32"))
                    )
                    and (4096 == T.cast(arg0_strides[2], "int32"))
                )
                and (57344 == T.cast(arg0_strides[1], "int32"))
            ) and (802816 == T.cast(arg0_strides[0], "int32")), (
                "arg0.strides: expected to be compact array"
            )
            T.evaluate(0)
        assert T.uint64(0) == T.tvm_struct_get(arg0, 0, 8, dtype="uint64"), (
            "Argument arg0.byte_offset has an unsatisfied constraint"
        )
        assert 2 == T.tvm_struct_get(arg0, 0, 10, dtype="int32"), (
            "Argument arg0.device_type has an unsatisfied constraint"
        )
        assert 6 == T.tvm_struct_get(arg1, 0, 4, dtype="int32"), "arg1.ndim is expected to equal 6"
        assert 6 == T.tvm_struct_get(arg1, 0, 4, dtype="int32"), "arg1.ndim is expected to equal 6"
        assert (
            (T.tvm_struct_get(arg1, 0, 5, dtype="uint8") == T.uint8(2))
            and (T.tvm_struct_get(arg1, 0, 6, dtype="uint8") == T.uint8(16))
        ) and (T.tvm_struct_get(arg1, 0, 7, dtype="uint16") == T.uint16(1)), (
            "arg1.dtype is expected to be float16"
        )
        assert 3 == T.cast(arg1_shape[0], "int32"), (
            "Argument arg1.shape[0] has an unsatisfied constraint"
        )
        assert 3 == T.cast(arg1_shape[1], "int32"), (
            "Argument arg1.shape[1] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg1_shape[2], "int32"), (
            "Argument arg1.shape[2] has an unsatisfied constraint"
        )
        assert 32 == T.cast(arg1_shape[3], "int32"), (
            "Argument arg1.shape[3] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg1_shape[4], "int32"), (
            "Argument arg1.shape[4] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg1_shape[5], "int32"), (
            "Argument arg1.shape[5] has an unsatisfied constraint"
        )
        if not (T.isnullptr(arg1_strides.data, dtype="bool")):
            assert (
                (
                    (
                        (
                            (1 == T.cast(arg1_strides[5], "int32"))
                            and (16 == T.cast(arg1_strides[4], "int32"))
                        )
                        and (256 == T.cast(arg1_strides[3], "int32"))
                    )
                    and (8192 == T.cast(arg1_strides[2], "int32"))
                )
                and (131072 == T.cast(arg1_strides[1], "int32"))
            ) and (393216 == T.cast(arg1_strides[0], "int32")), (
                "arg1.strides: expected to be compact array"
            )
            T.evaluate(0)
        assert T.uint64(0) == T.tvm_struct_get(arg1, 0, 8, dtype="uint64"), (
            "Argument arg1.byte_offset has an unsatisfied constraint"
        )
        assert 2 == T.tvm_struct_get(arg1, 0, 10, dtype="int32"), (
            "Argument arg1.device_type has an unsatisfied constraint"
        )
        assert dev_id == T.tvm_struct_get(arg1, 0, 9, dtype="int32"), (
            "Argument arg1.device_id has an unsatisfied constraint"
        )
        assert 6 == T.tvm_struct_get(arg2, 0, 4, dtype="int32"), "arg2.ndim is expected to equal 6"
        assert 6 == T.tvm_struct_get(arg2, 0, 4, dtype="int32"), "arg2.ndim is expected to equal 6"
        assert (
            (T.tvm_struct_get(arg2, 0, 5, dtype="uint8") == T.uint8(2))
            and (T.tvm_struct_get(arg2, 0, 6, dtype="uint8") == T.uint8(32))
        ) and (T.tvm_struct_get(arg2, 0, 7, dtype="uint16") == T.uint16(1)), (
            "arg2.dtype is expected to be float32"
        )
        assert 16 == T.cast(arg2_shape[0], "int32"), (
            "Argument arg2.shape[0] has an unsatisfied constraint"
        )
        assert 14 == T.cast(arg2_shape[1], "int32"), (
            "Argument arg2.shape[1] has an unsatisfied constraint"
        )
        assert 14 == T.cast(arg2_shape[2], "int32"), (
            "Argument arg2.shape[2] has an unsatisfied constraint"
        )
        assert 32 == T.cast(arg2_shape[3], "int32"), (
            "Argument arg2.shape[3] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg2_shape[4], "int32"), (
            "Argument arg2.shape[4] has an unsatisfied constraint"
        )
        assert 16 == T.cast(arg2_shape[5], "int32"), (
            "Argument arg2.shape[5] has an unsatisfied constraint"
        )
        if not (T.isnullptr(arg2_strides.data, dtype="bool")):
            assert (
                (
                    (
                        (
                            (1 == T.cast(arg2_strides[5], "int32"))
                            and (16 == T.cast(arg2_strides[4], "int32"))
                        )
                        and (256 == T.cast(arg2_strides[3], "int32"))
                    )
                    and (8192 == T.cast(arg2_strides[2], "int32"))
                )
                and (114688 == T.cast(arg2_strides[1], "int32"))
            ) and (1605632 == T.cast(arg2_strides[0], "int32")), (
                "arg2.strides: expected to be compact array"
            )
            T.evaluate(0)
        assert T.uint64(0) == T.tvm_struct_get(arg2, 0, 8, dtype="uint64"), (
            "Argument arg2.byte_offset has an unsatisfied constraint"
        )
        assert 2 == T.tvm_struct_get(arg2, 0, 10, dtype="int32"), (
            "Argument arg2.device_type has an unsatisfied constraint"
        )
        assert dev_id == T.tvm_struct_get(arg2, 0, 9, dtype="int32"), (
            "Argument arg2.device_id has an unsatisfied constraint"
        )
        T.evaluate(T.tvm_struct_set(stack_value, 0, 12, T.cast(2, "int64"), dtype="int32"))
        stack_tcode[0] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 1, 12, T.cast(dev_id, "int64"), dtype="int32"))
        stack_tcode[1] = 0
        T.evaluate(
            T.tvm_call_packed_lowered(
                "__tvm_set_device", stack_value, stack_tcode.data, 0, 2, dtype="int32"
            )
        )
        T.attr(0, "compute_scope", "default_function_compute_")
        T.evaluate(T.tvm_struct_set(stack_value, 0, 12, A, dtype="int32"))
        stack_tcode[0] = 3
        T.evaluate(T.tvm_struct_set(stack_value, 1, 12, W, dtype="int32"))
        stack_tcode[1] = 3
        T.evaluate(T.tvm_struct_set(stack_value, 2, 12, Conv, dtype="int32"))
        stack_tcode[2] = 3
        T.evaluate(T.tvm_struct_set(stack_value, 3, 12, T.cast(196, "int64"), dtype="int32"))
        stack_tcode[3] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 4, 12, T.cast(2, "int64"), dtype="int32"))
        stack_tcode[4] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 5, 12, T.cast(4, "int64"), dtype="int32"))
        stack_tcode[5] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 6, 12, T.cast(4, "int64"), dtype="int32"))
        stack_tcode[6] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 7, 12, T.cast(2, "int64"), dtype="int32"))
        stack_tcode[7] = 0
        T.evaluate(T.tvm_struct_set(stack_value, 8, 12, T.cast(32, "int64"), dtype="int32"))
        stack_tcode[8] = 0
        T.evaluate(
            T.tvm_call_packed_lowered(
                "default_function_kernel0", stack_value, stack_tcode.data, 0, 9, dtype="int32"
            )
        )

    return opt_conv_tensorcore_mod_host


def vthread_func():
    @T.prim_func(s_tir=True)
    def vthread_func(a: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, [256], "float32")
        C = T.match_buffer(c, [256], "float32")

        i0 = T.env_thread("blockIdx.x")
        i1 = T.env_thread("threadIdx.x")
        i2 = T.env_thread("vthread")

        T.launch_thread(i0, 4)
        T.launch_thread(i1, 2)
        T.launch_thread(i2, 2)
        B = T.alloc_buffer((16,), scope="local")
        for j in range(16):
            B[j] = A[i0 * 64 + i1 * 32 + i2 * 16 + j] + T.float32(1)
        for j in range(16):
            C[i0 * 64 + i1 * 32 + i2 * 16 + j] = B[j] * T.float32(2)

    return vthread_func


def matmul():
    @T.prim_func(s_tir=True)
    def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, [128, 128])
        B = T.match_buffer(b, [128, 128])
        C = T.match_buffer(c, [128, 128])

        for i, j, k in T.grid(128, 128, 128):
            with T.sblock("update"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    C[vi, vj] = T.float32(0)
                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]

    return matmul


def matmul_original():
    @T.prim_func(s_tir=True)
    def matmul_original(a: T.handle, b: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, [128, 128])
        B = T.match_buffer(b, [128, 128])
        C = T.match_buffer(c, [128, 128])

        for i, j in T.grid(128, 128):
            with T.sblock("init"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.float32(0)

            for k in range(128):
                with T.sblock("update"):
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]

    return matmul_original


def element_wise():
    @T.prim_func(s_tir=True)
    def element_wise(a: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (128, 128), "float32")
        C = T.match_buffer(c, (128, 128), "float32")
        B = T.sblock_alloc_buffer((128, 128), "float32")

        for i, j in T.grid(128, 128):
            with T.sblock("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = A[vi, vj] * T.float32(2)
        for i, j in T.grid(128, 128):
            with T.sblock("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = B[vi, vj] + T.float32(1)

    return element_wise


def predicate():
    @T.prim_func(s_tir=True)
    def predicate(b: T.handle, c: T.handle) -> None:
        B = T.match_buffer(b, (16, 16), "float32")
        C = T.match_buffer(c, (16, 16), "float32")

        for i, jo, ji in T.grid(16, 4, 5):
            with T.sblock("update"):
                vi = T.axis.S(16, i)
                vj = T.axis.S(16, jo * 4 + ji)
                T.where(jo * 4 + ji < 16)
                C[vi, vj] = B[vi, vj] + T.float32(1)

    return predicate


def test_module_define():
    func1 = tvm.ir.IRModule({"matmul": matmul()})["matmul"]
    func2 = tvm.ir.IRModule({"element_wise": element_wise()})["element_wise"]
    func3 = tvm.ir.IRModule({"predicate": predicate()})["predicate"]
    mod1 = tvm.ir.IRModule({"func1": func1, "func2": func2, "func3": func3})
    mod2 = tvm.ir.IRModule({"func1": matmul(), "func2": element_wise(), "func3": predicate()})
    tvm.ir.assert_structural_equal(mod1, mod2)


def test_matmul_original():
    func = matmul_original()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
    assert isinstance(rt_func.body.block.body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.SeqStmt)
    assert isinstance(rt_func.body.block.body.body.body[0].block, tirx.stmt.SBlock)
    assert isinstance(rt_func.body.block.body.body.body[1], tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body.body[1].body.block, tirx.stmt.SBlock)


def test_element_wise():
    func = element_wise()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
    assert isinstance(rt_func.body.block.body, tirx.stmt.SeqStmt)
    assert isinstance(rt_func.body.block.body[0], tirx.stmt.For)
    assert isinstance(rt_func.body.block.body[0].body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body[0].body.body.block, tirx.stmt.SBlock)

    assert isinstance(rt_func.body.block.body[1], tirx.stmt.For)
    assert isinstance(rt_func.body.block.body[1].body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body[1].body.body.block, tirx.stmt.SBlock)


def test_predicate():
    func = predicate()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
    assert isinstance(rt_func.body.block.body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.For)
    assert isinstance(rt_func.body.block.body.body.body.body.block, tirx.stmt.SBlock)


def for_thread_binding():
    @T.prim_func(s_tir=True)
    def for_thread_binding(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16, 16), "float32")
        B = T.match_buffer(b, (16, 16), "float32")

        for i in T.thread_binding(0, 16, thread="threadIdx.x"):
            for j in T.thread_binding(
                0, 16, thread="threadIdx.y", annotations={"attr_key": "attr_value"}
            ):
                A[i, j] = B[i, j] + T.float32(1)

    return for_thread_binding


def test_for_thread_binding():
    func = for_thread_binding()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body, tirx.stmt.For)
    assert rt_func.body.kind == 4
    assert rt_func.body.thread_binding.thread_tag == "threadIdx.x"
    assert isinstance(rt_func.body.body, tirx.stmt.For)
    assert rt_func.body.body.kind == 4
    assert rt_func.body.body.thread_binding.thread_tag == "threadIdx.y"
    assert rt_func.body.body.annotations["attr_key"] == "attr_value"


def match_buffer_region():
    @T.prim_func(s_tir=True)
    def match_buffer_region(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16, 16, 16), "float32")
        B = T.match_buffer(b, (1), "float32")

        for i, j in T.grid(16, 4):
            with T.sblock():
                vi, vj = T.axis.remap("SS", [i, j])
                C = T.match_buffer(A[0:16, vi, vj * 4 : vj * 4 + 4], (16, 1, 4))
                for ii in range(4):
                    with T.sblock():
                        vii = T.axis.S(4, ii)
                        D = T.match_buffer(C[vii * 4 : vii * 4 + 4, 0, 0:4], (4, 1, 4))
                        for i, j in T.grid(4, 4):
                            B[0] += D[i, 0, j]

    return match_buffer_region


def test_match_buffer_region():
    func = match_buffer_region()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body, tirx.stmt.SBlockRealize)
    root = rt_func.body.block

    assert isinstance(root.body, tirx.stmt.For)
    assert isinstance(root.body.body, tirx.stmt.For)
    assert isinstance(root.body.body.body, tirx.stmt.SBlockRealize)
    outer_block = root.body.body.body.block
    assert len(outer_block.match_buffers) == 1
    buffer_C = outer_block.match_buffers[0].buffer
    tvm.ir.assert_structural_equal(buffer_C.shape, [T.int32(16), T.int32(1), T.int32(4)])

    assert isinstance(outer_block.body, tirx.stmt.For)
    assert isinstance(outer_block.body.body, tirx.stmt.SBlockRealize)
    inner_block = outer_block.body.body.block
    assert len(inner_block.match_buffers) == 1
    buffer_D = inner_block.match_buffers[0].buffer
    tvm.ir.assert_structural_equal(buffer_D.shape, [T.int32(4), T.int32(1), T.int32(4)])


def block_elements():
    @T.prim_func(s_tir=True)
    def block_elements(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16, 16), "float32")
        B = T.match_buffer(b, (1, 1), "float32")

        with T.sblock("update"):
            vi = T.axis.S(1, 0)
            T.where(True)
            T.reads(A[0:16, 0:16])
            T.writes(B[0, 0])
            T.sblock_attr({"attr_key": "attr_value"})
            C = T.sblock_alloc_buffer((4, 4), dtype="float32")
            D = T.match_buffer(A[0:4, 0], (4, 1))
            with T.init():
                B[0, 0] = T.float32(0)
            B[0, 0] = A[0, 0] + B[0, 0] + C[1, 1] + D[2, 0]

    return block_elements


def test_block_elements():
    func = block_elements()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
    assert isinstance(rt_func.body.block.body, tirx.stmt.SBlockRealize)
    assert isinstance(rt_func.body.block.body.block, tirx.stmt.SBlock)
    block = rt_func.body.block.body.block
    assert isinstance(block.body, tirx.stmt.BufferStore)
    assert isinstance(block.init, tirx.stmt.BufferStore)
    assert len(block.annotations) == 1
    assert block.annotations["attr_key"] == "attr_value"


def opaque_block():
    @T.prim_func(s_tir=True)
    def opaque_block(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16, 16), "float32")
        B = T.match_buffer(b, (16, 16), "float32")

        for i in range(16):
            for j in range(16):
                with T.sblock():
                    T.reads([])
                    T.writes(A[i, j])
                    A[i, j] = T.float32(0)
            with T.sblock():
                T.reads([A[i, 0:16]])
                T.writes([B[i, 0:16]])
                for j in range(16):
                    B[i, j] = A[i, j]

    return opaque_block


def test_opaque_block():
    func = opaque_block()
    rt_func = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, rt_func)

    root_block = rt_func.body.block
    assert isinstance(root_block, tirx.stmt.SBlock)
    assert isinstance(root_block.body, tirx.stmt.For)
    assert isinstance(root_block.body.body[0], tirx.stmt.For)
    assert isinstance(root_block.body.body[0].body, tirx.stmt.SBlockRealize)
    assert isinstance(root_block.body.body[0].body.block, tirx.stmt.SBlock)
    assert len(root_block.body.body[0].body.block.iter_vars) == 0
    assert isinstance(root_block.body.body[1], tirx.stmt.SBlockRealize)
    assert isinstance(root_block.body.body[1].block, tirx.stmt.SBlock)
    assert len(root_block.body.body[1].block.iter_vars) == 0


def rank0():
    @T.prim_func(s_tir=True)
    def rank0(a: T.handle) -> None:
        A = T.match_buffer(a, (), "float32")
        B = T.sblock_alloc_buffer((), "float32")
        A[()] = 2
        B[()] = A[()]

    return rank0


def rank0_block():
    @T.prim_func(s_tir=True)
    def rank0_block(a: T.handle) -> None:
        A = T.match_buffer(a, (), "float32")
        B = T.sblock_alloc_buffer((), "float32")
        B[()] = A[()]

        with T.sblock("update"):
            T.reads([A[()]])
            T.writes([B[()]])
            for i in range(1):
                B[()] = A[()]

    return rank0_block


def select():
    @T.prim_func(s_tir=True)
    def select(a: T.handle) -> None:
        A = T.match_buffer(a, (), "float32")
        A[()] = T.Select(True, 1, 2)

    return select


def minmax():
    @T.prim_func(s_tir=True)
    def minmax(a: T.handle) -> None:
        A = T.match_buffer(a, (), "float32")
        A[()] = T.min(1, 2)
        A[()] = T.max(1, 2)

    return minmax


def abs():
    @T.prim_func(s_tir=True)
    def abs(a: T.handle) -> None:
        A = T.match_buffer(a, (128, 128), "float32")

        for i, j in T.grid(128, 128):
            with T.sblock("A"):
                vi, vj = T.axis.remap("SS", [i, j])
                A[vi, vj] = T.abs(A[vi, vj])

    return abs


def constant_folding():
    @T.prim_func(s_tir=True)
    def constant_folding(a: T.handle) -> None:
        A = T.match_buffer(a, (), "float32")
        A[()] = T.min(2.2, 5.2)
        A[()] = T.max(T.float32(2.2), T.float32(T.float32(5.2)))
        A[()] = T.min(2.2, 5.0)

    return constant_folding


def simplify_bracket():
    # uninitialized variables
    @T.prim_func(check_well_formed=False, s_tir=True)
    def simplify_bracket() -> None:
        a = T.int32()
        b = T.int32()
        c = T.int32()
        d = T.int32()
        T.evaluate(a + b * (c + d))

    return simplify_bracket


def var_with_same_name():
    @T.prim_func(s_tir=True)
    def var_with_same_name(a: T.handle) -> None:
        A = T.match_buffer(a, (16, 16), "float32")
        for i, j in T.grid(16, 16):
            with T.sblock():
                vi, vj = T.axis.remap("SS", [i, j])
                A[vi, vj] = 0
        for i, j in T.grid(16, 16):
            with T.sblock():
                vi, vj = T.axis.remap("SS", [i, j])
                A[vi, vj] = 0

    return var_with_same_name


def test_same_name_var():
    func = var_with_same_name()
    out_str = func.script()
    rt_func = tvm.script.from_source(out_str)
    tvm.ir.assert_structural_equal(func, rt_func)
    assert out_str.count("for i, j in T.grid(16, 16)") == 2
    assert out_str.find("i_") == -1
    assert out_str.find("i_") == -1


def while_loop():
    @T.prim_func(s_tir=True)
    def while_loop(a: T.handle, b: T.handle) -> None:
        A = T.match_buffer(a, (16,), "float32")
        B = T.match_buffer(b, (16,), "float32")
        i = T.sblock_alloc_buffer((), "int32", scope="local")
        for ii in range(16):
            with T.sblock():
                vi = T.axis.S(16, ii)
                B[vi] = 0
            while i[()] < 10:
                for j in range(16):
                    B[j] += A[j]

    return while_loop


# fmt: off
def primfunc_with_allocate_annotations():
    @T.prim_func(s_tir=True)
    def primfunc_with_allocate_annotations(placeholder_28: T.handle, T_cast_6: T.handle) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "tvmgen_default_fused_nn_max_pool2d_cast", "tirx.noalias": True})
        placeholder_29 = T.match_buffer(placeholder_28, [802816], dtype="uint8", elem_offset=0, align=64, offset_factor=1)
        T_cast_7 = T.match_buffer(T_cast_6, [200704], dtype="int16", elem_offset=0, align=64, offset_factor=1)
        # body
        tensor_2 = T.alloc_buffer((200704,), "uint8", annotations={"attr1_key": "attr1_value"})
        for ax0_ax1_fused_4 in T.serial(0, 56):
            for ax2_4 in T.serial(0, 56):
                for ax3_init in T.serial(0, 64):
                    tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_init)] = T.uint8(0)
                for rv0_rv1_fused_1, ax3_2 in T.grid(9, 64):
                    tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_2)] = T.max(tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_2)], T.if_then_else(((((ax0_ax1_fused_4*2) + T.floordiv(rv0_rv1_fused_1, 3)) < 112) and (((ax2_4*2) + T.floormod(rv0_rv1_fused_1, 3)) < 112)), placeholder_29[(((((ax0_ax1_fused_4*14336) + (T.floordiv(rv0_rv1_fused_1, 3)*7168)) + (ax2_4*128)) + (T.floormod(rv0_rv1_fused_1, 3)*64)) + ax3_2)], T.uint8(0), dtype="uint8"))
        for ax0_ax1_fused_5 in T.serial(0, 56):
            for ax2_5, ax3_3 in T.grid(56, 64):
                T_cast_7[(((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3)] = T.cast(tensor_2[(((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3)], "int16")

    return primfunc_with_allocate_annotations
# fmt: on


# fmt: off
def comm_reducer_single_reduce_group():
    @T.prim_func(s_tir=True)
    def comm_reducer_single_reduce_group(a: T.handle, b: T.handle) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        threadIdx_x = T.env_thread("threadIdx.x")
        A = T.match_buffer(a, [16384], dtype="float32")
        for i in T.serial(0, 128):
            T.launch_thread(threadIdx_x, 128)
            reduce_temp0 = T.alloc_buffer((1,), scope="local")
            with T.attr(T.comm_reducer(lambda x, y: x + y, [T.float32(0)]), "reduce_scope", T.int32(0)):
                T.evaluate(T.tvm_thread_allreduce(T.uint32(1), A[i * 128 + threadIdx_x], True, reduce_temp0.data, threadIdx_x, dtype="handle"))

    return comm_reducer_single_reduce_group


def comm_reducer_multiple_reduce_groups():
    @T.prim_func(s_tir=True)
    def comm_reducer_multiple_reduce_groups(a: T.handle, b: T.handle) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        threadIdx_x = T.env_thread("threadIdx.x")
        A = T.match_buffer(a, [16384], dtype="float32")
        for i in T.serial(0, 128):
            T.launch_thread(threadIdx_x, 128)
            reduce_temp0 = T.alloc_buffer((1,), scope="local")
            with T.attr(T.comm_reducer(lambda x0, x1, y0, y1: (T.Select((x1 >= y1), x0, y0), T.Select((x1 >= y1), x1, y1)), [T.int32(-1), T.min_value("float32")]), "reduce_scope", T.int32(0)):
                T.evaluate(T.tvm_thread_allreduce(T.uint32(1), A[i * 128 + threadIdx_x], True, reduce_temp0.data, threadIdx_x, dtype="handle"))

    return comm_reducer_multiple_reduce_groups


def multiple_commreducer():
    # normal_reduce_temp0 is treated as uninitialized value
    @T.prim_func(check_well_formed=False, s_tir=True)
    def multiple_commreducer() -> None:
        normal_reduce_temp0 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
        normal_reduce_temp1 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
        reduce_temp0 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
        reduce_temp1 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
        for ax0_1 in T.thread_binding(0, 32, thread="threadIdx.x"):
            with T.sblock("T_softmax_maxelem_cross_thread_reduction"):
                T.attr(T.comm_reducer(lambda x, y: T.max(x, y), [T.min_value("float32")]), "reduce_scope", T.int32(0))
                T.evaluate(T.tvm_thread_allreduce(T.uint32(1), normal_reduce_temp0[0], True, reduce_temp0.data, ax0_1, dtype="handle"))
        for ax0_1 in T.thread_binding(0, 32, thread="threadIdx.x"):
            with T.sblock("T_softmax_expsum_cross_thread_reduction"):
                T.attr(T.comm_reducer(lambda x, y: x + y, [T.float32(0)]), "reduce_scope", T.int32(0))
                T.evaluate(T.tvm_thread_allreduce(T.uint32(1), normal_reduce_temp1[0], True, reduce_temp1.data, ax0_1, dtype="handle"))

    return multiple_commreducer
# fmt: on


def func_div_mod():
    # not well-formed: free variables
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func_div_mod():
        a = T.int32()
        b = T.int32()
        T.evaluate(a // b)
        T.evaluate(a % b)
        T.evaluate(T.truncmod(a, b))

    return func_div_mod


def test_div_mod():
    func = func_div_mod()
    rt_func = tvm.script.from_source(func.script(), check_well_formed=False)
    tvm.ir.assert_structural_equal(func, rt_func, True)

    assert isinstance(func.body[0].value, tvm.tirx.FloorDiv)
    assert isinstance(func.body[1].value, tvm.tirx.FloorMod)
    assert isinstance(func.body[2].value, tvm.tirx.Mod)


def loop_extent_dependent():
    @T.prim_func(s_tir=True)
    def loop_extent_dependent(a: T.handle) -> None:
        A = T.match_buffer(a, [], dtype="int32")
        for i in T.serial(0, 128):
            for j in T.serial(0, i):
                A[()] = A[()] + j

    return loop_extent_dependent


def nontrivial_range_axis():
    @T.prim_func(s_tir=True)
    def nontrivial_range_axis(a: T.handle) -> None:
        A = T.match_buffer(a, (10), "float32")
        for i in range(10):
            with T.sblock("block"):
                vi = T.axis.spatial((1, 11), i + 1)
                A[vi - 1] = A[vi - 1] + 1.0

    return nontrivial_range_axis


def func_with_target_spec_by_config():
    @T.prim_func(s_tir=True)
    def func_with_target_spec_by_config() -> None:
        T.func_attr(
            {
                "kTarget": T.target(
                    {
                        "max_num_threads": 1024,
                        "arch": "sm_70",
                        "thread_warp_size": 32,
                        "kind": "cuda",
                        "tag": "",
                        "keys": ["cuda", "gpu"],
                        "host": T.target({"kind": "llvm", "tag": "", "keys": ["cpu"]}),
                    }
                )
            }
        )
        T.evaluate(0)

    return func_with_target_spec_by_config


def func_with_target_spec_by_str():
    @T.prim_func(s_tir=True)
    def func_with_target_spec_by_str() -> None:
        T.func_attr({"kTarget": T.target("nvidia/nvidia-a100")})
        T.evaluate(0)

    return func_with_target_spec_by_str


def func_with_target_and_host_spec_by_str():
    @T.prim_func(s_tir=True)
    def func():
        T.func_attr({"target": T.target("nvidia/nvidia-a100", host="llvm")})
        T.evaluate(0)

    return func


def func_root_attr():
    @T.prim_func(s_tir=True)
    def func_root_attr():
        with T.sblock("root"):
            T.sblock_attr({"a": "0"})
            T.evaluate(0)

    return func_root_attr


def func_trivial_root_block():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer(1, "int32")):
        with T.sblock("root"):
            A[0] = 0

    return func


def func_nested_root_block():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer(1, "int32")):
        with T.sblock("root"):
            with T.sblock("block"):
                A[0] = 0

    return func


def func_T_ptr_let_statement():
    @T.prim_func(s_tir=True)
    def func_T_ptr_let_statement(
        args: T.handle, arg_type_ids_handle: T.handle("int32"), num_args: T.int32
    ) -> None:
        # The T.Ptr declaration in the parameter list should parse
        # correctly, and should be usable as the data pointer in a buffer.
        arg_type_ids = T.decl_buffer([2], dtype="int32", data=arg_type_ids_handle)

        arg0: T.let[T.handle] = T.tvm_struct_get(args, 0, 12, dtype="handle")
        arg1: T.let[T.handle] = T.tvm_struct_get(args, 1, 12, dtype="handle")

        # The ABI field is an opaque pointer.  Retag it explicitly before
        # binding it to the buffer's exact element pointer type.
        A_data: T.let[T.handle("float32")] = T.reinterpret(
            T.handle("float32").ty,
            T.tvm_struct_get(arg0, 0, 1, dtype="handle"),
        )

        # The buffer declaration has a data pointer defined earlier in
        # this function.  It should only be defined after the data pointer
        # has been defined, and should not be hoisted into the header of
        # the function as other buffer_decl statements can be.
        A = T.decl_buffer([1024], dtype="float32", data=A_data)
        B_data: T.let[T.handle("float32")] = T.reinterpret(
            T.handle("float32").ty,
            T.tvm_struct_get(arg1, 0, 1, dtype="handle"),
        )
        B = T.decl_buffer([1024], dtype="float32", data=B_data)

        B[0] = A[0]

    return func_T_ptr_let_statement


def func_T_ptr_allocate():
    @T.prim_func(s_tir=True)
    def func_T_ptr_allocate() -> None:
        A = T.alloc_buffer((1024,))
        A[0] = 0.0

    return func_T_ptr_allocate


def llvm_intrin_call():
    @T.prim_func(s_tir=True)
    def ctpop(A: T.Buffer((16,), "uint8"), B: T.Buffer((16,), "uint8")) -> None:
        for i in range(0, 16):
            with T.sblock("A"):
                vi = T.axis.remap(
                    "S",
                    [
                        i,
                    ],
                )
                B[vi] = T.call_llvm_pure_intrin(
                    T.llvm_lookup_intrinsic_id("llvm.ctpop.i8"),
                    A[vi],
                    dtype="uint8",
                )

    return ctpop


def parse_bufferslice_as_range_bound():
    # apparently the use of i in the "outer" block when it is defined outside of a block is wrong
    @T.prim_func(check_well_formed=False, s_tir=True)
    def segment_sum(
        A_ptr: T.handle, B_ptr: T.handle, indptr_ptr: T.handle, n: T.int32, m: T.int32
    ) -> None:
        A = T.match_buffer(A_ptr, [m], dtype="float32")
        B = T.match_buffer(B_ptr, [n], dtype="float32")
        indptr = T.match_buffer(indptr_ptr, [n + 1], dtype="int32")
        for i in T.serial(n):
            with T.sblock("outer"):
                vi = T.axis.spatial(n, i)
                T.reads(indptr[i : i + 2], B[vi], A[indptr[i] : indptr[i + 1]])
                T.writes(B[vi])
                for j in T.serial(indptr[i], indptr[i + 1]):
                    with T.sblock("inner"):
                        vj = T.axis.reduce(m, j)
                        T.reads(B[vi], A[vj])
                        T.writes(B[vi])
                        with T.init():
                            B[vi] = T.float32(0)
                        B[vi] = B[vi] + A[vj]

    return segment_sum


def int64_support():
    @T.prim_func(s_tir=True)
    def elementwise_shape_int64(a: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (T.int64(128), T.int64(128)), dtype="float32")
        B = T.sblock_alloc_buffer((T.int64(128), T.int64(128)), dtype="float32")
        C = T.match_buffer(c, (T.int64(128), T.int64(128)), dtype="float32")
        for i, j in T.grid(128, 128):
            with T.sblock("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = A[vi, vj] * 2.0
        for i, j in T.grid(T.int64(128), T.int64(128)):
            with T.sblock("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = B[vi, vj] + 1.0

    return elementwise_shape_int64


def string_annotation_escaping():
    @T.prim_func(s_tir=True)
    def string_annotation_of_special_chars():
        T.func_attr(
            {
                "key1": '"\'hello\t\r"',
                "key2": """
            %1 = add i32 %0, %0
            %2 = add i32 %0, %1
            %3 = add i32 %1, %2
            """,
            }
        )
        T.evaluate(0)

    return string_annotation_of_special_chars


def pointer_type():
    @T.prim_func(s_tir=True)
    def func_with_ptr_type_annotations(x: T.handle("int32"), y: T.handle("int32", "shared")):
        xx = T.alloc_buffer((16,), "int32")
        yy = T.alloc_buffer((16,), "int32", scope="shared")
        a: T.let[T.handle("int32")] = T.address_of(xx[0], dtype="handle")
        b: T.let[T.handle("int32", "shared")] = T.address_of(yy[0], dtype="handle")
        T.evaluate(T.call_extern("copy", a, b, dtype=""))

    return func_with_ptr_type_annotations


def buffer_axis_separator():
    @T.prim_func(s_tir=True)
    def element_wise(a: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (128, 128), "float32", axis_separators=[1])
        C = T.match_buffer(c, (128, 128), "float32")
        B = T.sblock_alloc_buffer((128, 128), "float32", axis_separators=[1])

        for i, j in T.grid(128, 128):
            with T.sblock("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = A[vi, vj] * T.float32(2)
        for i, j in T.grid(128, 128):
            with T.sblock("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = B[vi, vj] + T.float32(1)

    return element_wise


def buffer_ramp_access_as_slice_index():
    @T.prim_func(s_tir=True)
    def buffer_ramp_access(a: T.handle, b: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (128,), "float32")
        B = T.match_buffer(b, (128,), "float32")
        C = T.match_buffer(c, (128,), "float32")
        for i in range(128):
            A[i : i + 1 : 1] = i
        for i in range(4):
            B[i * 32 : i * 32 + 32] = A[i * 32 : i * 32 + 32 : 1] + T.broadcast(1.0, 32)
        for i in range(4):
            C[i : i + 128 : 4] = B[i : i + 128 : 4] + T.broadcast(1.0, 32)

    return buffer_ramp_access


def ramp_int64():
    @T.prim_func(s_tir=True)
    def func() -> None:
        T.evaluate(T.Ramp(T.int64(0), 1, 3))

    return func


def scalable_vectors():
    @T.prim_func(s_tir=True)
    def func(a: T.handle):
        A = T.match_buffer(a, (200,), "float32")
        A[T.Ramp(11, 2, 4 * tirx.vscale())] = T.Broadcast(125, 4 * tirx.vscale())

    return func


def predicated_buffer_load_store():
    @T.prim_func(s_tir=True)
    def func(a: T.handle, b: T.handle):
        A = T.match_buffer(a, (4,), "float32")
        B = T.match_buffer(b, (8,), "float32")
        for i_0 in range(4):
            load_a = T.meta_var(
                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4))
            )
            B.vstore([T.Ramp(0, 2, 4)], load_a, predicate=T.Broadcast(T.bool(True), 4))

    return func


def let_expression():
    @T.prim_func(s_tir=True)
    def func():
        x = T.int32()
        T.evaluate(T.Let(x + 1, where={x: 1}))

    return func


def test_void_ptr_vs_handle():
    """An untyped handle is the canonical void-pointer type."""

    # Generates PointerType(PrimType::Void())
    @T.prim_func(s_tir=True)
    def void_ptr(out_ret_value: T.handle("void")):
        T.evaluate(out_ret_value)

    # Generates PointerType::VoidPointerTy()
    @T.prim_func(s_tir=True)
    def handle(out_ret_value: T.handle):
        T.evaluate(out_ret_value)

    tvm.ir.assert_structural_equal(void_ptr.params[0].ty, handle.params[0].ty)
    script = void_ptr.script()
    assert "out_ret_value: T.handle" in script
    assert 'T.handle("void")' not in script
    tvm.ir.assert_structural_equal(void_ptr, tvm.script.from_source(script))

    @T.prim_func(s_tir=True)
    def scoped_void_ptr(out_ret_value: T.handle("void", "shared")):
        T.evaluate(out_ret_value)

    scoped_script = scoped_void_ptr.script()
    assert 'out_ret_value: T.handle(storage_scope="shared")' in scoped_script
    assert 'T.handle("void"' not in scoped_script
    tvm.ir.assert_structural_equal(scoped_void_ptr, tvm.script.from_source(scoped_script))


def void_ptr():
    @T.prim_func(s_tir=True)
    def func(out_ret_value: T.handle("void")):
        T.evaluate(out_ret_value)

    return func


def decl_buffer():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")) -> None:
        A_flattened = T.decl_buffer(data=A.data, shape=(256,), dtype="float32")
        B_flattened = T.decl_buffer(data=B.data, shape=(256,), dtype="float32")
        C_alias = T.decl_buffer(data=A_flattened.data, shape=(256,), dtype="float32")
        for i in range(256):
            B_flattened[i] = A_flattened[i] + C_alias[i] + T.float32(1.0)

    return func


def allocate_and_decl_buffer():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")) -> None:
        D = T.alloc_buffer((16,))
        for i in range(4):
            C = T.alloc_buffer((4,))
            for j in range(4):
                C[j] = A[i * 4 + j] + T.float32(1.0)
            for j in range(4):
                D[j] = C[j]
            for j in range(4):
                B[i * 4 + j] = D[j]

    return func


def alloc_buffer_example():
    @T.prim_func(s_tir=True)
    def func(a: T.handle, c: T.handle):
        A = T.match_buffer(a, (128,), "float32")
        C = T.match_buffer(c, (128,), "float32")
        B = T.alloc_buffer((128,), "float32")
        for i in range(128):
            B[i] = A[i] * T.float32(2)
        for i in range(128):
            C[i] = B[i] + T.float32(1)

    return func


def float_infinity():
    @T.prim_func(s_tir=True)
    def func(
        placeholder: T.Buffer((1, 512, 768), "float32"), T_isinf: T.Buffer((1, 512, 768), "bool")
    ) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        # body
        # with T.sblock("root")
        for i0, i1, i2 in T.grid(1, 512, 768):
            with T.sblock("T_isinf"):
                ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                T.reads(placeholder[ax0, ax1, ax2])
                T.writes(T_isinf[ax0, ax1, ax2])
                T_isinf[ax0, ax1, ax2] = T.fabs(
                    placeholder[ax0, ax1, ax2], dtype="float32"
                ) == T.float32("inf") and not (T.isnan(placeholder[ax0, ax1, ax2], dtype="bool"))

    return func


def minimal_i32_literal():
    @T.prim_func(s_tir=True)
    def func() -> None:
        T.evaluate(T.int32(-2147483648))
        T.evaluate(-T.int64(2147483648))

    return func


def boolean_argument():
    @T.prim_func(s_tir=True)
    def func(a: T.boolean) -> None:
        T.evaluate(a)

    return func


def bool_argument():
    @T.prim_func(s_tir=True)
    def func(a: T.bool) -> None:
        T.evaluate(a)

    return func


def bool_variable_annotation():
    @T.prim_func(s_tir=True)
    def func() -> None:
        a: T.let[T.bool] = T.call_extern("dummy", dtype="bool")
        T.evaluate(0)

    return func


def return_none():
    @T.prim_func(s_tir=True)
    def func():
        T.evaluate(0)

    return func


def bool_primitive():
    @T.prim_func(s_tir=True)
    def func() -> None:
        T.evaluate(T.bool(True))

    return func


def bool_cast():
    # uninitialized var
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func() -> None:
        a = T.bool()
        T.evaluate(T.bool(T.int32(0)))
        T.evaluate(a == T.bool(False))

    return func


def implicit_evaluate():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer(1, "int32")):
        T.evaluate(T.assume(A[0] == 5))
        A[0] = 10

    return func


def if_true_else():
    @T.prim_func(s_tir=True)
    def func() -> None:
        if True:
            T.evaluate(0)
        else:
            T.evaluate(1)

    return func


def elif_chain_without_else():
    @T.prim_func(s_tir=True)
    def func(i: T.int32) -> None:
        if i == 0:
            T.evaluate(0)
        elif i == 1:
            T.evaluate(1)
        elif i == 2:
            T.evaluate(2)

    return func


def elif_chain_with_else():
    @T.prim_func(s_tir=True)
    def func(i: T.int32) -> None:
        if i == 0:
            T.evaluate(0)
        elif i == 1:
            T.evaluate(1)
        elif i == 2:
            T.evaluate(2)
        else:
            T.evaluate(3)

    return func


def nested_boolean_expressions():
    expressions = {
        "and_lhs_and": lambda i, j, k: tirx.all(tirx.all(i, j), k),
        "and_rhs_and": lambda i, j, k: tirx.all(i, tirx.all(j, k)),
        "and_lhs_or": lambda i, j, k: tirx.all(tirx.any(i, j), k),
        "and_rhs_or": lambda i, j, k: tirx.all(i, tirx.any(j, k)),
        "or_lhs_and": lambda i, j, k: tirx.any(tirx.all(i, j), k),
        "or_rhs_and": lambda i, j, k: tirx.any(i, tirx.all(j, k)),
        "or_lhs_or": lambda i, j, k: tirx.any(tirx.any(i, j), k),
        "or_rhs_or": lambda i, j, k: tirx.any(i, tirx.any(j, k)),
        "and_of_ors": lambda i, j, k: tirx.all(
            tirx.any(i, j), tirx.any(j, k), tirx.any(i, k), i, j, k
        ),
        "or_of_ands": lambda i, j, k: tirx.any(
            tirx.all(i, j), tirx.all(j, k), tirx.all(i, k), i, j, k
        ),
    }

    def make_ir_generator(name, expression):
        def inner():
            @T.prim_func(s_tir=True)
            def func(A: T.Buffer(1, "bool"), i: T.bool, j: T.bool, k: T.bool):
                A[0] = expression(i, j, k)

            return func

        inner.__name__ = f"nested_boolean_expr_{name}"
        return inner

    for name, expression in expressions.items():
        generator = make_ir_generator(name, expression)

        yield generator


def multi_env_threads():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer(128, "float32"), C: T.Buffer(128, "float32")):
        B = T.sblock_alloc_buffer([128], dtype="float32")
        for i in T.thread_binding(128, thread="threadIdx.x"):
            B[i] = A[i] + 1.0
        for i in T.thread_binding(128, thread="threadIdx.x"):
            C[i] = B[i] + 2.0

    mod = tvm.s_tir.transform.LowerOpaqueBlock()(
        tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
    )
    return mod["main"]


def intrinsic_pow():
    @T.prim_func(s_tir=True)
    def func():
        T.pow(T.float32(1), T.float32(1))

    return func


def bind_var():
    @T.prim_func(s_tir=True)
    def func():
        x = T.bind(0)
        y = T.bind(0)
        T.evaluate(0)
        T.evaluate(0)

    return func


def string_stride():
    @T.prim_func(s_tir=True)
    def main(a: T.handle, b: T.handle):
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        n = T.int32()
        A = T.match_buffer(a, (n,), strides=("A_s0",), buffer_type="auto")
        B = T.match_buffer(b, (n,), strides=("B_s0",), buffer_type="auto")
        blockIdx_x = T.launch_thread("blockIdx.x", (n + 63) // 64)
        threadIdx_x = T.launch_thread("threadIdx.x", 64)
        if T.likely(blockIdx_x * 64 + threadIdx_x < n):
            B2 = T.decl_buffer((B.strides[0] * n,), data=B.data)
            A2 = T.decl_buffer((A.strides[0] * n,), data=A.data)
            B2[(blockIdx_x * 64 + threadIdx_x) * B.strides[0]] = A2[
                (blockIdx_x * 64 + threadIdx_x) * A.strides[0]
            ] * T.float32(2)

    return main


def string_stride_int64():
    @T.prim_func(s_tir=True)
    def main(a: T.handle, b: T.handle):
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        n = T.int64()
        A_s0 = T.int64()
        B_s0 = T.int64()
        A = T.match_buffer(a, (n,), strides=(A_s0,), buffer_type="auto")
        B = T.match_buffer(b, (n,), strides=(B_s0,), buffer_type="auto")
        for i in range(n):
            B[i] = A[i]

    return main


def merge_shape_var_def():
    # uninitialized vars
    @T.prim_func(check_well_formed=False, s_tir=True)
    def main(A: T.handle, B: T.handle):
        # fmt: off
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        m, n = T.int32(), T.int32()
        A_1 = T.match_buffer(A, (m, n), strides=("A_1_s0", "A_1_s1"), buffer_type="auto")
        B_1 = T.match_buffer(B, (m, n), strides=("B_1_s0", "B_1_s1"), buffer_type="auto")
        for i_outer, j_outer, i_inner in T.grid((m + 9) // 10, (n + 4) // 5, 10):
            if T.likely(i_outer * 10 + i_inner < m):
                for j_inner in range(5):
                    if T.likely(j_outer * 5 + j_inner < n):
                        cse_v2: T.let[T.int32] = j_outer * 5 + j_inner
                        cse_v1: T.let[T.int32] = i_outer * 10 + i_inner
                        B_2 = T.decl_buffer(
                            (B_1.strides[0] * m,),
                            data=B_1.data,
                            strides=("B_2_s0",),
                            buffer_type="auto",
                        )
                        A_2 = T.decl_buffer(
                            (A_1.strides[0] * m,),
                            data=A_1.data,
                            strides=("A_2_s0",),
                            buffer_type="auto",
                        )
                        B_2[cse_v1 * B_1.strides[0] + cse_v2 * B_1.strides[1]] = A_2[
                            cse_v1 * A_1.strides[0] + cse_v2 * A_1.strides[1]
                        ]
        # fmt: on

    return main


def if_then_else_var():
    @T.prim_func(s_tir=True)
    def main(n: T.int32):
        if n == 0:
            x = 5
            T.evaluate(x)
        else:
            x = 10
            T.evaluate(x)

    return main


def tvm_shfl_builtins():
    @T.prim_func(s_tir=True)
    def func(
        A: T.handle("float32"),
        B: T.handle("float32"),
        C: T.handle("float32"),
    ):
        blockIdx_x = T.launch_thread("blockIdx.x", 1)
        threadIdx_x = T.launch_thread("threadIdx.x", 32)
        A_warp = T.alloc_buffer((1,), scope="local")
        B_warp = T.alloc_buffer((1,), scope="local")
        red_buf0 = T.alloc_buffer((1,), scope="local")
        A_warp_1 = T.decl_buffer((32,), data=A_warp.data, scope="local")
        A_1 = T.decl_buffer((32,), data=A)  # A is a handle param
        A_warp_1[0] = A_1[threadIdx_x]
        B_warp_1 = T.decl_buffer((32,), data=B_warp.data, scope="local")
        T.tvm_storage_sync("warp")
        B_warp_1[0] = T.tvm_warp_shuffle(
            T.tvm_warp_activemask(), A_warp_1[0], threadIdx_x % 4 * 8 + threadIdx_x // 4, 32, 32
        ) + T.float32(1)
        red_buf0_1 = T.decl_buffer((1,), data=red_buf0.data, scope="local")
        with T.attr(
            T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
            "reduce_scope",
            T.int32(0),
        ):
            mask = T.alloc_buffer((1,), "uint32", scope="local")
            t0 = T.alloc_buffer((1,), scope="local")
            red_buf0_1[0] = A_warp_1[0]
            mask_1 = T.decl_buffer((1,), "uint32", data=mask.data, scope="local")
            mask_1[0] = T.tvm_warp_activemask()
            t0_1 = T.decl_buffer((1,), data=t0.data, scope="local")
            t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 16, 32, 32)
            red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
            t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 8, 32, 32)
            red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
            t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 4, 32, 32)
            red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
            t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 2, 32, 32)
            red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
            t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 1, 32, 32)
            red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
            red_buf0_1[0] = T.tvm_warp_shuffle(mask_1[0], red_buf0_1[0], 0, 32, 32)
            # NOTE(Zihao): test tvm_warp_shuffle_up
            red_buf0_1[0] = T.tvm_warp_shuffle_up(mask_1[0], red_buf0_1[0], 0, 32, 32)
        if threadIdx_x == 0:
            C_1 = T.decl_buffer((1,), data=C)
            C_1[0] = red_buf0_1[0]
        B_1 = T.decl_buffer((32,), data=B)
        B_1[threadIdx_x] = B_warp_1[0]

    return func


def make_packed_api_result():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer(64, "float32")):
        T.func_attr({"global_symbol": "main", "target": T.target("cuda")})
        bx = T.launch_thread("blockIdx.x", 64)
        T.evaluate(A[bx])

    mod = tvm.IRModule.from_expr(func)
    return tvm.tirx.transform.MakePackedAPI()(mod)


def tvm_struct_set_generated_in_cpp():
    """Ensure same dtype for tvm_struct_set in Python/C++

    The TVMStructSet method in C++, used internally by
    LowerTVMBuiltin, and the Python method `T.tvm_struct_set`, used
    when parsing TVMScript should use the same dtype "int32".
    """

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def tir_packed_call(A: T.Buffer(16)):
            T.attr(0, "device_id", 0)
            T.attr(0, "device_type", 0)
            T.evaluate(
                T.tvm_call_cpacked(
                    "tvm_test_cpacked",
                    T.tvm_stack_make_array(
                        A.data,
                        T.tvm_stack_make_shape(16, dtype="handle"),
                        T.reinterpret(T.uint64(0), dtype="handle"),
                        T.uint32(1),
                        T.Cast("float32", 0),
                        0,
                        dtype="handle",
                    ),
                    dtype="int32",
                )
            )

    return tvm.tirx.transform.LowerTVMBuiltin()(Module)


def ir_module_with_attrs():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_attrs({"attr": 10})

        @T.prim_func(s_tir=True)
        def tir_func(A: T.Buffer(16, "int32"), B: T.Buffer(16, "int32")):
            for i in range(16):
                B[i] = A[i]

    return Module


def nested_seqstmt():
    """Nested SeqStmt should be normalized to flat SeqStmt

    Nested SeqStmt are representable in the TIR structures, but are
    flattened when converted to TVMScript.  Previously, this could
    cause failures to round-trip through TVMScript, including
    erroneous use of TVMScript's concise-scoping rules.  This was
    resolved by normalizing nested SeqStmt in TIR, such that the use
    of `tirx.SeqStmt` below results in a single flat `tirx.SeqStmt`
    containing the three `tirx.Evaluate` calls.
    """
    func = tvm.tirx.PrimFunc(
        params=[],
        body=tvm.tirx.SeqStmt(
            [
                tvm.tirx.SeqStmt([tvm.tirx.Evaluate(0), tvm.tirx.Evaluate(1)]),
                tvm.tirx.Evaluate(2),
            ]
        ),
    )

    return func


def subroutine_call():
    """A GlobalVar may reference other functions in the module"""

    @I.ir_module(s_tir=True)
    class mod:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer(16, "float32")):
            mod.subroutine(A.data, T.int32(16))

        @T.prim_func(s_tir=True)
        def subroutine(A_data: T.handle("float32"), n: T.int32):
            T.evaluate(0)

    return mod


def subroutine_call_returning_int():
    """An internal function call may return non-void"""

    @I.ir_module(s_tir=True)
    class mod:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer(2, "float32")):
            mod.subroutine(A[0]) + mod.subroutine(A[1])

        @T.prim_func(s_tir=True)
        def subroutine(x: T.float32) -> T.float32:
            T.ret(x * x)

    return mod


def undefined_data_ptr_in_decl_buffer():
    """The T.decl_buffer syntax should not introduce an Allocate

    While T.decl_buffer can be used to represent an
    Allocate/DeclBuffer pair, performing a round-trip through
    TVMScript should not introduce an Allocate node.
    """

    # uninitialized var
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func():
        data_ptr = T.handle("float32")
        buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr)
        T.evaluate(buf[0])

    return func


def undefined_shape_in_decl_buffer():
    # uninitialized var
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func():
        size = T.int32()
        buf = T.decl_buffer(shape=[size], dtype="float32")
        T.evaluate(buf[0])

    return func


def undefined_stride_in_decl_buffer():
    # uninitialized var
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func():
        stride = T.int32()
        data_ptr = T.handle("float32")
        buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr, strides=[stride])
        T.evaluate(buf[0])

    return func


def undefined_elem_offset_in_decl_buffer():
    # uninitialized var
    @T.prim_func(check_well_formed=False, s_tir=True)
    def func():
        elem_offset = T.int32()
        data_ptr = T.handle("float32")
        buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr, elem_offset=elem_offset)
        T.evaluate(buf[0])

    return func


def subroutine_call_without_arguments():
    @I.ir_module(s_tir=True)
    class mod:
        @T.prim_func(s_tir=True)
        def main():
            # Should be equivalent to the bare "mod.subroutine()", but
            # that relies on `GlobalVar.__call__` returning the
            # correct IR type.
            tirx.call_tir(mod.subroutine)

        @T.prim_func(s_tir=True)
        def subroutine():
            T.evaluate(0)

    return mod


def return_zero():
    @T.prim_func(s_tir=True)
    def func() -> T.int32:
        T.ret(0)

    return func


def return_zero_private():
    @T.prim_func(private=True, s_tir=True)
    def func() -> T.int32:
        T.ret(0)

    return func


def return_zero_private_with_attr():
    @T.prim_func(private=True, s_tir=True)
    def func() -> T.int32:
        T.func_attr({"greeting": "hello"})
        T.ret(0)

    return func


def func_attr_with_list():
    @T.prim_func(s_tir=True)
    def func(
        A: T.Buffer((128, 128), "float32"),
        B: T.Buffer((128, 128), "float32"),
        D: T.Buffer((128, 128), "float32"),
    ) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True, "layout_free_buffers": [1]})
        C = T.sblock_alloc_buffer([128, 128], dtype="float32")
        for i0, i1, i2 in T.grid(128, 128, 128):
            with T.sblock("C"):
                x, y, k = T.axis.remap("SSR", [i0, i1, i2])
                with T.init():
                    C[x, y] = T.float32(0)
                C[x, y] = C[x, y] + A[x, k] * B[y, k]
        for i0, i1 in T.grid(128, 128):
            with T.sblock("D"):
                T.sblock_attr({"layout_free_placeholders": [C]})
                x, y = T.axis.remap("SS", [i0, i1])
                D[x, y] = C[x, y] + T.float32(1)

    return func


def func_with_loop_jumps():
    @T.prim_func(s_tir=True)
    def func(In: T.Buffer((1,), "int32"), Out: T.Buffer((2,), "int32")):
        Out[0] = 0
        Out[1] = 0
        for i in range(1000):
            if i % 13 == 0:
                Out[1] = Out[1] + 1
                continue
            Out[0] = Out[0] + 1
            if Out[0] >= In[0]:
                break

    return func


def func_with_loop_steps():
    @T.prim_func(s_tir=True)
    def func(
        A: T.Buffer((1024,)), B: T.Buffer((1024,)), C: T.Buffer((1024,)), tid: T.int32, v: T.int32
    ):
        for i in T.serial(tid, 1024, step=2):
            C[i] = A[i] + B[i]
        for i in T.unroll(tid, 1024, step=3):
            C[i] = A[i] + B[i]
        for i in T.vectorized(tid, 1024, step=4):
            C[i] = A[i] + B[i]
        for i in T.parallel(tid, 1024, step=5):
            C[i] = A[i] + B[i]
        for i in range(tid, 1024, 6):
            C[i] = A[i] + B[i]

    return func


def op_of_literal():
    op_list = [
        (T.exp, 0),
        (T.exp2, 0),
        (T.exp10, 0),
        (T.erf, 0.0),
        (T.tanh, 0.0),
        (T.sigmoid, 0.0),
        (T.log, 0.0),
        (T.log2, 0.0),
        (T.log1p, 0.0),
        (T.tan, 0.0),
        (T.cos, 0.0),
        (T.acos, 0.0),
        (T.acosh, 0.0),
        (T.sin, 0.0),
        (T.sinh, 0.0),
        (T.asin, 0.0),
        (T.asinh, 0.0),
        (T.atan, 0.0),
        (T.atanh, 0.0),
        (T.atan2, (1.0, 0.0)),
        (T.sqrt, 0.0),
        (T.rsqrt, 1.0),
        (T.nextafter, (0.0, 1.0)),
        (T.hypot, (1.0, 1.0)),
        (T.copysign, (1.0, 1.0)),
        (T.popcount, 0),
        (T.fmod, (1.0, 1.0)),
    ]

    def make_ir_generator(op, arg):
        def inner():
            call_expr = op(*arg) if isinstance(arg, tuple) else op(arg)

            @T.prim_func(s_tir=True)
            def func():
                T.evaluate(call_expr)

            return func

        inner.__name__ = f"{op.__name__}_of_literal"
        return inner

    for op, arg in op_list:
        yield make_ir_generator(op, arg)


def relax_extern_func():
    @R.function
    def func(A: R.Tensor([10, 20], "float32")):
        func = R.ExternFunc("dummy_func")

        B: R.Tensor([10, 20], "float32") = R.call_dps_packed(
            func, [A], out_ty=R.Tensor([10, 20], "float32")
        )

        C: R.Tensor(ndim=2, dtype="float32") = R.call_dps_packed(
            func, [B], out_ty=R.Tensor([10, 20], "float32")
        )

        return C

    return func


def relax_match_cast_ty_proxy():
    """TypeProxy subclasses may be used as expressions

    This is a regression test.  The TVMScript parser allows Type
    to be specified using a default-constructible class
    (e.g. `R.Tensor` or `R.Shape`) rather than an instance of that
    class (e.g. `R.Tensor()` or `R.Shape()`).  In previous
    implementations, this was only handled when the `Type` was
    used in an annotation context.  However, a `Type` may also
    appear as an argument, which is passed to `R.match_cast`.  Use of
    a default-constructible class must be handled in this context as
    well.
    """

    def make_ir_generator(proxy_subclass):
        def inner():
            @R.function
            def func(A: R.Any):
                B = R.match_cast(A, proxy_subclass)
                return B

            return func

        inner.__name__ = subclass.__name__
        return inner

    # Not all subclasses of TypeProxy are default-constructible.
    # This list is a subset of `TypeProxy.__subclasses__()`,
    # excluding `PrimProxy` and `DTensorProxy`.
    subclasses = [
        tvm.script.parser.relax.entry.AnyProxy,
        tvm.script.parser.relax.entry.TensorProxy,
        tvm.script.parser.relax.entry.CallableProxy,
        tvm.script.parser.relax.entry.TupleProxy,
        tvm.script.parser.relax.entry.ShapeProxy,
    ]

    for subclass in subclasses:
        yield make_ir_generator(subclass)


def relax_symbolic_var():
    """Relax tensors may use symbolic variables."""
    N = tvm.tirx.Var("N", "int64")

    @R.function
    def func(A: R.Tensor([N], "float16")):
        B: R.Tensor([N], "float16") = A
        return B

    return func


def relax_float_symbolic_var():
    """Relax scalar variables may use any dtype."""

    @R.function
    def func(value: R.Prim("float16")):
        return value

    return func


ir_generator = tvm.testing.parameter(
    launch_env_thread,
    opt_gemm_lower,
    opt_conv_tensorcore_lower,
    opt_conv_tensorcore_mod_host,
    vthread_func,
    matmul,
    rank0,
    rank0_block,
    select,
    minmax,
    abs,
    constant_folding,
    simplify_bracket,
    while_loop,
    primfunc_with_allocate_annotations,
    comm_reducer_single_reduce_group,
    comm_reducer_multiple_reduce_groups,
    multiple_commreducer,
    loop_extent_dependent,
    nontrivial_range_axis,
    func_with_target_spec_by_config,
    func_with_target_spec_by_str,
    func_with_target_and_host_spec_by_str,
    func_root_attr,
    func_trivial_root_block,
    func_nested_root_block,
    func_T_ptr_let_statement,
    func_T_ptr_allocate,
    llvm_intrin_call,
    parse_bufferslice_as_range_bound,
    int64_support,
    string_annotation_escaping,
    pointer_type,
    buffer_axis_separator,
    buffer_ramp_access_as_slice_index,
    ramp_int64,
    scalable_vectors,
    predicated_buffer_load_store,
    let_expression,
    void_ptr,
    decl_buffer,
    allocate_and_decl_buffer,
    alloc_buffer_example,
    float_infinity,
    minimal_i32_literal,
    boolean_argument,
    bool_argument,
    bool_variable_annotation,
    bool_primitive,
    bool_cast,
    return_none,
    implicit_evaluate,
    if_true_else,
    elif_chain_without_else,
    elif_chain_with_else,
    *nested_boolean_expressions(),
    multi_env_threads,
    intrinsic_pow,
    bind_var,
    string_stride,
    string_stride_int64,
    merge_shape_var_def,
    if_then_else_var,
    tvm_shfl_builtins,
    make_packed_api_result,
    tvm_struct_set_generated_in_cpp,
    ir_module_with_attrs,
    nested_seqstmt,
    subroutine_call,
    subroutine_call_returning_int,
    undefined_data_ptr_in_decl_buffer,
    undefined_shape_in_decl_buffer,
    undefined_stride_in_decl_buffer,
    undefined_elem_offset_in_decl_buffer,
    subroutine_call_without_arguments,
    return_zero,
    return_zero_private,
    return_zero_private_with_attr,
    func_attr_with_list,
    func_with_loop_jumps,
    func_with_loop_steps,
    *op_of_literal(),
    *relax_match_cast_ty_proxy(),
    relax_symbolic_var,
    relax_float_symbolic_var,
)

relax_ir_generator = tvm.testing.parameter(
    relax_extern_func,
)

show_all_relax_ty = tvm.testing.parameter(
    by_dict={
        "show_all_ty": True,
        "hide_inferable_ty": False,
    }
)


_NOT_ROUNDTRIP_STABLE: set[str] = set()


def test_roundtrip(ir_generator):
    if getattr(ir_generator, "__name__", "") in _NOT_ROUNDTRIP_STABLE:
        import pytest

        pytest.skip(f"{ir_generator.__name__}: not round-trip stable here")
    original = ir_generator()
    after_roundtrip = tvm.script.from_source(
        original.script(show_meta=True), check_well_formed=False
    )
    tvm.ir.assert_structural_equal(original, after_roundtrip, True)


def test_relax_roundtrip(relax_ir_generator, show_all_relax_ty):
    original = relax_ir_generator()
    after_roundtrip = tvm.script.from_source(
        original.script(
            show_meta=True,
            show_all_ty=show_all_relax_ty,
        )
    )
    tvm.ir.assert_structural_equal(original, after_roundtrip, True)


def test_return_none_no_trailing_type():
    func = return_none()
    script = func.script()
    assert "-> None" not in script


def test_address_of_buffer():
    @T.prim_func(s_tir=True)
    def func(a: T.handle):
        A = T.match_buffer(a, (128, 128), "float32")
        T.evaluate(T.address_of(A))

    assert "T.address_of(A[0, 0])" in func.script()


def test_assert_stmt_roundtrip_runtime_error():
    """RuntimeError assert roundtrips through print->parse."""

    @T.prim_func(s_tir=True)
    def func(x: T.int32):
        assert x > 0, ("RuntimeError", ["x must be positive"])

    script = func.script(show_meta=True)
    roundtrip = tvm.script.from_source(script, check_well_formed=False)
    tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)


def test_assert_stmt_roundtrip_value_error():
    """ValueError assert roundtrips through print->parse."""

    @T.prim_func(s_tir=True)
    def func(x: T.int32):
        assert x > 0, ("ValueError", ["Shape mismatch"])

    script = func.script(show_meta=True)
    roundtrip = tvm.script.from_source(script, check_well_formed=False)
    tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)


def test_assert_stmt_roundtrip_type_error():
    """TypeError assert roundtrips through print->parse."""

    @T.prim_func(s_tir=True)
    def func(x: T.int32):
        assert x > 0, ("TypeError", ["Expected Tensor but got int"])

    script = func.script(show_meta=True)
    roundtrip = tvm.script.from_source(script, check_well_formed=False)
    tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)


def test_assert_stmt_roundtrip_multi_parts():
    """Multi-part message assert roundtrips with structural equality."""

    @T.prim_func(s_tir=True)
    def func(x: T.int32):
        assert x > 0, ("TypeError", ["Expected ", "Tensor", " but got ", "int"])

    script = func.script(show_meta=True)
    roundtrip = tvm.script.from_source(script, check_well_formed=False)
    tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)


if __name__ == "__main__":
    tvm.testing.main()