# Licensed to the Apache Software Foundation (ASF) under one
# ruff: noqa: E501, E741, F401, F841

# 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.

import numpy as np
import pytest

import tvm
import tvm.support.nvcc
import tvm.testing
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.support.nvcc import have_bf16, have_fp16, have_int8
from tvm.testing import env


@pytest.fixture(autouse=True, params=["nvcc", "nvrtc"])
def setup_cuda_compile_mode(request):
    mode = request.param
    if mode == "nvrtc":
        try:
            from cuda.bindings import nvrtc
        except ImportError:
            pytest.skip("cuda-python not available, skipping nvrtc tests")

    orig_func = tvm.support.nvcc.tvm_callback_cuda_compile

    def compile_mode_wrapper(code):
        if mode == "nvcc":
            return tvm.support.nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc")
        elif mode == "nvrtc":
            return tvm.support.nvcc.compile_cuda(code, target_format="cubin", compiler="nvrtc")
        else:
            raise ValueError(f"Unknown mode: {mode}")

    tvm.register_global_func("tvm_callback_cuda_compile", compile_mode_wrapper, override=True)
    # yield back to the original function so that each test runs twice
    yield
    tvm.register_global_func("tvm_callback_cuda_compile", orig_func, override=True)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_add():
    num_thread = 8

    def check_cuda(dtype, n, lanes):
        if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
            print("Skip because gpu does not have fp16 support")
            return
        if dtype == "int8" and not have_int8(tvm.cuda(0).compute_version):
            print("skip because gpu does not support int8")
            return
        vec_dtype = f"{dtype}x{lanes}"
        one = tvm.tirx.const(1, vec_dtype)
        num_blocks = (n + num_thread - 1) // num_thread

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
                            T.where(i_0 * num_thread + i_1 < n)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = A[v_i] + one

        fun = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np.random.uniform(size=(n, lanes)))
            c = tvm.runtime.empty((n,), vec_dtype, dev)
            fun(a, c)
            tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda("float32", 64, 2)
    check_cuda("float32", 64, 3)
    check_cuda("float32", 64, 4)
    check_cuda("int8", 64, 2)
    check_cuda("int8", 64, 3)
    check_cuda("int8", 64, 4)
    check_cuda("uint8", 64, 2)
    check_cuda("uint8", 64, 3)
    check_cuda("uint8", 64, 4)
    check_cuda("float16", 64, 2)
    check_cuda("float16", 64, 4)
    check_cuda("float16", 64, 6)
    check_cuda("float16", 64, 8)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_bf16_vectorize_add():
    if not have_bf16(tvm.cuda(0).compute_version):
        print("skip because gpu does not support bf16")
        return
    num_thread = 8

    def np_float2np_bf16(arr):
        """Convert a numpy array of float to a numpy array
        of bf16 in uint16"""
        orig = arr.view("<u4")
        bias = np.bitwise_and(np.right_shift(orig, 16), 1) + 0x7FFF
        return np.right_shift(orig + bias, 16).astype("uint16")

    def np_bf162np_float(arr):
        """Convert a numpy array of bf16 (uint16) to a numpy array
        of float"""
        u32 = np.left_shift(arr.astype("uint32"), 16)
        return u32.view("<f4")

    def check_cuda(n, lanes):
        vec_dtype = f"bfloat16x{lanes}"
        num_blocks = n // num_thread
        one = tvm.tirx.Broadcast(tvm.tirx.const(1, "bfloat16"), lanes)

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = A[v_i] + one

        with tvm.transform.PassContext(
            disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
        ):
            fun = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            np_a = np.random.uniform(size=(n, lanes)).astype("float32")
            np_a = np_bf162np_float(np_float2np_bf16(np_a))
            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_float2np_bf16(np_a))
            c = tvm.runtime.empty((n,), vec_dtype, dev)
            fun(a, c)
            c = tvm.runtime.empty((n, lanes), "uint16", dev).copyfrom(c)
            tvm.testing.assert_allclose(c.numpy(), np_float2np_bf16(np_a + 1))

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda(64, 2)
    check_cuda(64, 4)
    check_cuda(64, 6)
    check_cuda(64, 8)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_multiply_add():
    num_thread = 8

    def check_cuda(dtype, n, lanes):
        if dtype == "int8" and not have_int8(tvm.cuda(0).compute_version):
            print("skip because gpu does not support int8")
            return
        vec_dtype = f"{dtype}x{lanes}"
        num_blocks = n // num_thread

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(
                A: T.Buffer((n,), vec_dtype),
                B: T.Buffer((n,), vec_dtype),
                C: T.Buffer((n,), "int32"),
                D: T.Buffer((n,), "int32"),
            ):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
                        with T.sblock("D"):
                            v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
                            T.reads(A[v_i], B[v_i], C[v_i])
                            T.writes(D[v_i])
                            D[v_i] = T.call_pure_extern("int32", "__dp4a", A[v_i], B[v_i], C[v_i])

        fun = tvm.compile(Module, target="cuda")

        np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
        np_b = np.random.randint(low=-128, high=127, size=(n, lanes))
        np_c = np.random.randint(low=0, high=127, size=(n,))
        np_d = [sum(x * y) + z for x, y, z in zip(np_a, np_b, np_c)]

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_a)
            b = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_b)
            c = tvm.runtime.empty((n,), "int32", dev).copyfrom(np_c)
            d = tvm.runtime.empty((n,), "int32", dev)
            fun(a, b, c, d)
            tvm.testing.assert_allclose(d.numpy(), np_d)

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda("int8", 64, 4)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_load():
    num_thread = 8

    def check_cuda(dtype, n, lanes):
        vec_dtype = f"{dtype}x{lanes}"
        num_blocks = n // num_thread

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(num_thread, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = A[v_i]

        fun = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
            a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_a)
            b = tvm.runtime.empty((n,), vec_dtype, dev)
            fun(a, b)
            tvm.testing.assert_allclose(a.numpy(), b.numpy())

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda("int8", 64, 2)
    check_cuda("int8", 64, 3)
    check_cuda("int8", 64, 4)
    check_cuda("int8", 64, 8)
    check_cuda("int8", 64, 16)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_make_int8():
    def check_cuda(n, value, lanes):
        dtype = "int8"
        const_value = tvm.tirx.const(value, dtype=dtype)

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n, lanes), dtype)):
                T.func_attr({"tirx.noalias": True})
                for i in T.thread_binding(n, thread="blockIdx.x"):
                    for j in T.vectorized(lanes):
                        with T.sblock("A"):
                            v_i, v_j = T.axis.remap("SS", [i, j])
                            T.reads()
                            T.writes(A[v_i, v_j])
                            A[v_i, v_j] = const_value

        fun = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            np_a = np.full((n, lanes), value, dtype=dtype)
            a = tvm.runtime.empty(np_a.shape, dtype, dev)
            fun(a)
            np.testing.assert_equal(a.numpy(), np_a)

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda(64, np.uint8(0xAB).view(np.int8), 4)
    check_cuda(64, 0, 4)
    check_cuda(64, -3, 4)
    check_cuda(64, np.uint8(0xAB).view(np.int8), 3)
    check_cuda(64, 0, 3)
    check_cuda(64, -3, 3)
    check_cuda(64, np.uint8(0xAB).view(np.int8), 2)
    check_cuda(64, 0, 2)
    check_cuda(64, -3, 2)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_inf_nan():
    def check_inf_nan(n, value, dtype):
        inf_value = tvm.tirx.const(value, dtype=dtype)

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), dtype), C: T.Buffer((n,), dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(8, thread="threadIdx.x"):
                        with T.sblock("C"):
                            v_i = T.axis.spatial(n, i_0 * 8 + i_1)
                            T.where(i_0 * 8 + i_1 < n)
                            T.reads()
                            T.writes(C[v_i])
                            C[v_i] = inf_value

        fun = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.device("cuda", 0)
            a = tvm.runtime.empty((n,), dtype, dev)
            c = tvm.runtime.empty((n,), dtype, dev)
            fun(a, c)

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_inf_nan(1, -float("inf"), "float32")
    check_inf_nan(1, -float("inf"), "float64")
    check_inf_nan(1, float("inf"), "float32")
    check_inf_nan(1, float("inf"), "float64")
    check_inf_nan(1, float("nan"), "float32")
    check_inf_nan(1, float("nan"), "float64")


@pytest.mark.parametrize(
    "target",
    [
        pytest.param("cuda", marks=pytest.mark.gpu),
        pytest.param("rocm", marks=pytest.mark.gpu),
    ],
)
def test_crossthread_reduction1(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    def sched(nthd):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(var_A: T.handle, var_B: T.handle):
                T.func_attr({"tirx.noalias": True})
                n, m = T.int32(), T.int32()
                A = T.match_buffer(var_A, (n, m))
                B = T.match_buffer(var_B, (n,))
                for i in T.thread_binding(n, thread="blockIdx.x"):
                    for m_0 in T.thread_binding(nthd, thread="threadIdx.x"):
                        for m_1 in range((m + nthd - 1) // nthd):
                            with T.sblock("B"):
                                v_i = T.axis.spatial(n, i)
                                v_m = T.axis.reduce(m, m_0 * ((m + nthd - 1) // nthd) + m_1)
                                T.where(m_0 * ((m + nthd - 1) // nthd) + m_1 < m)
                                T.reads(A[v_i, v_m])
                                T.writes(B[v_i])
                                with T.init():
                                    B[v_i] = T.float32(0.0)
                                B[v_i] = B[v_i] + A[v_i, v_m]

        fun = tvm.compile(Module, target="cuda")
        return fun

    def verify(nthd):
        func = sched(nthd)
        nn = 3
        # checks three typical cases
        vals = [nthd - 1, nthd, nthd + 1]

        def run_and_check():
            dev = tvm.device(target)
            for kk in vals:
                size = (nn, kk)
                a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
                b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
                func(a, b)
                tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=1), rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)

    verify(16)
    verify(32)
    verify(64)


@pytest.mark.parametrize(
    "target",
    [
        pytest.param("cuda", marks=pytest.mark.gpu),
        pytest.param("rocm", marks=pytest.mark.gpu),
    ],
)
def test_crossthread_reduction2(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    def sched(nthdx, nthdy):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(var_A: T.handle, var_B: T.handle):
                T.func_attr({"tirx.noalias": True})
                n, k0, k1 = T.int32(), T.int32(), T.int32()
                A = T.match_buffer(var_A, (n, k0, k1))
                B = T.match_buffer(var_B, (n,))
                for i in T.thread_binding(n, thread="blockIdx.x"):
                    for k0_0 in T.thread_binding(nthdx, thread="threadIdx.x"):
                        for k0_1 in range((k0 + nthdx - 1) // nthdx):
                            for k1_0 in T.thread_binding(nthdy, thread="threadIdx.y"):
                                for k1_1 in range((k1 + nthdy - 1) // nthdy):
                                    with T.sblock("B"):
                                        v_i = T.axis.spatial(n, i)
                                        v_k0 = T.axis.reduce(
                                            k0, k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1
                                        )
                                        v_k1 = T.axis.reduce(
                                            k1, k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1
                                        )
                                        T.where(
                                            k0_0 * ((k0 + nthdx - 1) // nthdx) + k0_1 < k0
                                            and k1_0 * ((k1 + nthdy - 1) // nthdy) + k1_1 < k1
                                        )
                                        T.reads(A[v_i, v_k0, v_k1])
                                        T.writes(B[v_i])
                                        with T.init():
                                            B[v_i] = T.float32(0.0)
                                        B[v_i] = B[v_i] + A[v_i, v_k0, v_k1]

        func = tvm.compile(Module, target="cuda")
        return func

    def verify(nthdx, nthdy):
        func = sched(nthdx, nthdy)
        nn = 3
        # checks three typical cases
        vx = [nthdx - 1, nthdx, nthdx + 1]
        vy = [nthdy - 1, nthdy, nthdy + 1]

        def run_and_check():
            dev = tvm.device(target)
            for kk0, kk1 in [(x, y) for x in vx for y in vy]:
                size = (nn, kk0, kk1)
                a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
                b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
                func(a, b)
                tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=(1, 2)), rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)

    verify(16, 16)
    verify(32, 32)
    verify(16, 32)
    verify(32, 16)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_reduction_binding():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((96, 32), "float32"), B: T.Buffer((96,), "float32")):
            T.func_attr({"tirx.noalias": True})
            for k in range(32):
                for m_0 in T.thread_binding(3, thread="blockIdx.x"):
                    for m_1 in range(32):
                        with T.sblock("B"):
                            v_m = T.axis.spatial(96, m_0 * 32 + m_1)
                            v_k = T.axis.reduce(32, k)
                            T.reads(A[v_m, v_k])
                            T.writes(B[v_m])
                            with T.init():
                                B[v_m] = T.float32(0.0)
                            B[v_m] = B[v_m] + A[v_m, v_k]

    func = tvm.compile(Module, target="cuda")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_const_float_to_half():
    # This import is required to use nvcc to perform code gen;
    # otherwise it is found that the code gen is done by nvrtc.

    half_const = tvm.tirx.const(0.5, dtype="float16")

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(a: T.Buffer((2, 3, 4), "float16"), C: T.Buffer((2, 3, 4), "bool")):
            T.func_attr({"tirx.noalias": True})
            for i_j_k_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
                for i_j_k_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
                    with T.sblock("C"):
                        v_i = T.axis.spatial(2, (i_j_k_fused_0 * 64 + i_j_k_fused_1) // 12)
                        v_j = T.axis.spatial(3, (i_j_k_fused_0 * 64 + i_j_k_fused_1) % 12 // 4)
                        v_k = T.axis.spatial(4, (i_j_k_fused_0 * 64 + i_j_k_fused_1) % 4)
                        T.where(i_j_k_fused_0 * 64 + i_j_k_fused_1 < 24)
                        T.reads(a[v_i, v_j, v_k])
                        T.writes(C[v_i, v_j, v_k])
                        C[v_i, v_j, v_k] = half_const < a[v_i, v_j, v_k]

    func = tvm.compile(Module, target="cuda")

    shape = (2, 3, 4)
    a_np = np.random.uniform(size=shape).astype("float16")
    c_np = np.zeros(shape=shape, dtype="bool")

    def run_and_check():
        dev = tvm.cuda(0)
        a = tvm.runtime.tensor(a_np, dev)
        c = tvm.runtime.tensor(c_np, dev)
        func(a, c)
        np.testing.assert_equal(c.numpy(), a_np > 0.5)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_floordiv_with_vectorization():
    with tvm.target.Target("cuda"):
        # B[i] = A[floordiv(i, k)]
        n = 256
        k = 37

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
                    for i_1_0 in T.thread_binding(64, thread="threadIdx.x"):
                        for i_1_1 in T.vectorized(4):
                            with T.sblock("B"):
                                v_i = T.axis.spatial(256, i_0 * 256 + i_1_0 * 4 + i_1_1)
                                T.reads(A[v_i // 37])
                                T.writes(B[v_i])
                                B[v_i] = A[v_i // 37]

        func = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            a_np = np.random.uniform(size=(n,)).astype("float32")
            b_np = np.array([a_np[i // k] for i in range(0, n)])
            a_nd = tvm.runtime.tensor(a_np, dev)
            b_nd = tvm.runtime.tensor(np.zeros(b_np.shape, dtype=b_np.dtype), dev)
            func(a_nd, b_nd)
            tvm.testing.assert_allclose(b_nd.numpy(), b_np, rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_floormod_with_vectorization():
    with tvm.target.Target("cuda"):
        # B[i] = A[floormod(i, k)]
        n = 256
        k = 37

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((256,), "float32"), B: T.Buffer((256,), "float32")):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
                    for i_1_0 in T.thread_binding(64, thread="threadIdx.x"):
                        for i_1_1 in T.vectorized(4):
                            with T.sblock("B"):
                                v_i = T.axis.spatial(256, i_0 * 256 + i_1_0 * 4 + i_1_1)
                                T.reads(A[v_i % 37])
                                T.writes(B[v_i])
                                B[v_i] = A[v_i % 37]

        func = tvm.compile(Module, target="cuda")

        def run_and_check():
            dev = tvm.cuda(0)
            a_np = np.random.uniform(size=(n,)).astype("float32")
            b_np = np.array([a_np[i % k] for i in range(0, n)])
            a_nd = tvm.runtime.tensor(a_np, dev)
            b_nd = tvm.runtime.tensor(np.zeros(b_np.shape, dtype=b_np.dtype), dev)
            func(a_nd, b_nd)
            tvm.testing.assert_allclose(b_nd.numpy(), b_np, rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vectorized_casts():
    def check(t0, t1, factor):
        if (t0 == "float16" or t1 == "float16") and not have_fp16(tvm.cuda(0).compute_version):
            print("Skip because gpu does not have fp16 support")
            return

        n = 128
        num_thread = n // factor

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), t0), B: T.Buffer((n,), t1), C: T.Buffer((n,), t0)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_thread, thread="threadIdx.x"):
                    for i_1 in T.vectorized(factor):
                        with T.sblock("C"):
                            v_i = T.axis.spatial(n, i_0 * factor + i_1)
                            T.reads(A[v_i], B[v_i])
                            T.writes(C[v_i])
                            C[v_i] = A[v_i] + T.Cast(t0, B[v_i])

        func = tvm.compile(Module, target="cuda")

        # correctness
        def run_and_check():
            dev = tvm.cuda(0)
            low, high = (0, 20) if t0.startswith("u") or t1.startswith("u") else (-10, 10)
            a_np = np.random.randint(low, high, size=n).astype(t0)
            b_np = np.random.randint(low, high, size=n).astype(t1)
            c_np = (a_np + b_np).astype(t0)
            a_nd = tvm.runtime.tensor(a_np, dev)
            b_nd = tvm.runtime.tensor(b_np, dev)
            c_nd = tvm.runtime.tensor(np.zeros(c_np.shape, dtype=c_np.dtype), dev)
            func(a_nd, b_nd, c_nd)
            tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)

    def skip(t0, t1):
        if t0 == t1:
            return True
        # CUDA does support cast between {u}int8 and fp16.
        skip_set = {"float16", "uint8", "int8"}
        if t0 in skip_set and t1 in skip_set:
            return True
        return False

    types_4 = [
        "float16",
        "float32",
        "int8",
        "uint8",
        "int16",
        "uint16",
        "int32",
        "uint32",
        "float64",
        "int64",
        "uint64",
    ]
    types_8 = ["float16", "float32", "int8", "uint8", "int16", "uint16", "int32", "uint32"]
    for t0, t1 in [(x, y) for x in types_4 for y in types_4 if not skip(x, y)]:
        check(t0, t1, 4)
    for t0, t1 in [(x, y) for x in types_8 for y in types_8 if not skip(x, y)]:
        check(t0, t1, 8)
    check("int8", "uint8", 16)
    check("uint8", "int8", 16)


def sched(compute_fn, dtype, n=128):
    """Create a vectorized CUDA module with the given compute function.

    The schedule structure is: split [1, None] -> split [32, None] -> split [None, 4]
    then vectorize innermost, bind blockIdx.x and threadIdx.x.
    For n=128 this gives: blockIdx.x=1, threadIdx.x=32, serial=1, vectorized=4.
    """

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((n,), dtype), B: T.Buffer((n,), dtype)):
            T.func_attr({"tirx.noalias": True})
            for i0_0 in T.thread_binding(1, thread="blockIdx.x"):
                for i0_1_0 in T.thread_binding(32, thread="threadIdx.x"):
                    for i0_1_1_0 in range(1):
                        for i0_1_1_1 in T.vectorized(4):
                            with T.sblock("B"):
                                v_i0 = T.axis.spatial(n, i0_1_0 * 4 + i0_1_1_0 * 4 + i0_1_1_1)
                                T.reads(A[v_i0])
                                T.writes(B[v_i0])
                                B[v_i0] = compute_fn(A[v_i0])

    return tvm.compile(Module, target="cuda")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vectorized_intrin1():
    test_funcs = [
        (tvm.tirx.floor, lambda x: np.floor(x)),
        (tvm.tirx.ceil, lambda x: np.ceil(x)),
        (tvm.tirx.trunc, lambda x: np.trunc(x)),
        (tvm.tirx.abs, lambda x: np.fabs(x)),
        (tvm.tirx.round, lambda x: np.round(x)),
        (tvm.tirx.exp, lambda x: np.exp(x)),
        (tvm.tirx.exp2, lambda x: np.exp2(x)),
        (tvm.tirx.exp10, lambda x: np.power(10, x)),
        (tvm.tirx.log, lambda x: np.log(x)),
        (tvm.tirx.log2, lambda x: np.log2(x)),
        (tvm.tirx.log10, lambda x: np.log10(x)),
        (tvm.tirx.tan, lambda x: np.tan(x)),
        (tvm.tirx.cos, lambda x: np.cos(x)),
        (tvm.tirx.cosh, lambda x: np.cosh(x)),
        (tvm.tirx.sin, lambda x: np.sin(x)),
        (tvm.tirx.sinh, lambda x: np.sinh(x)),
        (tvm.tirx.atan, lambda x: np.arctan(x)),
        (tvm.tirx.tanh, lambda x: np.tanh(x)),
        (tvm.tirx.sqrt, lambda x: np.sqrt(x)),
    ]

    def run_test(tvm_intrin, np_func, dtype):
        if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
            print("Skip because gpu does not have fp16 support")
            return
        # set of intrinsics does not support fp16 yet.
        skip_set = {
            tvm.tirx.abs,
            tvm.tirx.round,
            tvm.tirx.tan,
            tvm.tirx.atan,
            tvm.tirx.tanh,
            tvm.tirx.cosh,
            tvm.tirx.sinh,
        }
        if dtype == "float16" and tvm_intrin in skip_set:
            print(f"Skip because '{tvm_intrin.__name__}' does not support fp16 yet")
            return

        n = 128
        f = sched(tvm_intrin, dtype, n)

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev)
            b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
            f(a, b)
            tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)

    for func in test_funcs:
        run_test(*func, "float32")
        run_test(*func, "float16")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vectorized_intrin2(dtype="float32"):
    c2 = tvm.tirx.const(2, dtype=dtype)
    test_funcs = [
        (tvm.tirx.power, lambda x: np.power(x, 2.0)),
        (tvm.tirx.fmod, lambda x: np.fmod(x, 2.0)),
    ]

    def run_test(tvm_intrin, np_func):
        n = 128
        f = sched(lambda x: tvm_intrin(x, c2), dtype, n)

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev)
            b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
            f(a, b)
            tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)

        tvm.testing.run_with_gpu_lock(run_and_check)

    for func in test_funcs:
        run_test(*func)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vectorized_popcount():
    def ref_popcount(x):
        cnt = 0
        while x:
            x -= x & -x
            cnt += 1
        return cnt

    def run_test(dtype):
        n = 128
        f = sched(lambda x: tvm.tirx.popcount(x), dtype, n)

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.tensor(np.random.randint(0, 100000, size=n).astype(dtype), dev)
            b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
            f(a, b)
            ref = np.vectorize(ref_popcount)(a.numpy())
            tvm.testing.assert_allclose(b.numpy(), ref)

        tvm.testing.run_with_gpu_lock(run_and_check)

    run_test("uint32")
    run_test("uint64")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_load_permute_pad():
    def check_cuda(dtype, n, l, padding, lanes):
        if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
            print("Skip because gpu does not have fp16 support")
            return

        zero = tvm.tirx.const(0, dtype)
        dim0 = n // lanes
        dim1 = l + 2 * padding

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n, l), dtype), B: T.Buffer((dim0, dim1, lanes), dtype)):
                T.func_attr({"tirx.noalias": True})
                for i in T.thread_binding(dim0, thread="blockIdx.x"):
                    for j in T.thread_binding(dim1, thread="threadIdx.x"):
                        for k in T.vectorized(lanes):
                            with T.sblock("B"):
                                v_i, v_j, v_k = T.axis.remap("SSS", [i, j, k])
                                T.reads(A[v_i * lanes + v_k, v_j - padding])
                                T.writes(B[v_i, v_j, v_k])
                                B[v_i, v_j, v_k] = T.if_then_else(
                                    v_j < padding or l + padding <= v_j,
                                    zero,
                                    A[v_i * lanes + v_k, v_j - padding],
                                )

        fun = tvm.compile(Module, target="cuda")

        np_a = np.random.randint(low=-128, high=127, size=(n, l)).astype(dtype)
        np_a_reshape = np_a.reshape(n // lanes, lanes, l).transpose(0, 2, 1)
        ref = np.pad(
            np_a_reshape, ((0, 0), (padding, padding), (0, 0)), mode="constant", constant_values=0
        )

        def run_and_check():
            dev = tvm.cuda(0)
            a = tvm.runtime.empty((n, l), dtype, dev).copyfrom(np_a)
            b = tvm.runtime.empty((dim0, dim1, lanes), dtype, dev)
            fun(a, b)
            tvm.testing.assert_allclose(b.numpy(), ref)

        tvm.testing.run_with_gpu_lock(run_and_check)

    check_cuda("int8", 64, 16, 3, 2)
    check_cuda("uint8", 64, 16, 3, 2)
    check_cuda("int8", 64, 16, 3, 4)
    check_cuda("uint8", 64, 16, 3, 4)
    check_cuda("int32", 64, 16, 3, 4)
    check_cuda("float16", 64, 16, 3, 4)
    check_cuda("float32", 64, 16, 3, 4)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_try_unaligned_vector_load():
    def build(N, C_N, offset):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((N,), "float16"), C: T.Buffer((C_N,), "float16")):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(C_N // 2, thread="threadIdx.x"):
                    for i_1 in T.vectorized(2):
                        with T.sblock("C"):
                            v_i = T.axis.spatial(C_N, i_0 * 2 + i_1)
                            T.reads(A[v_i + offset])
                            T.writes(C[v_i])
                            C[v_i] = A[v_i + offset]

        f = tvm.tirx.build(Module, target="cuda")

        kernel_source = f.imports[0].inspect_source()
        a_data = np.arange(0, N).astype("float16")

        def run_and_check():
            dev = tvm.cuda()
            a = tvm.runtime.tensor(a_data, dev)
            c = tvm.runtime.tensor(np.zeros(C_N, dtype="float16"), dev)
            f(a, c)
            return c.numpy()

        c = tvm.testing.run_with_gpu_lock(run_and_check)
        return a_data, c, kernel_source

    # Unaligned case: N=3, C_N=2, offset=1
    a_data, c, kernel_source = build(3, 2, 1)
    # (uint1*)(A + (1)) is invalid
    assert "A_ptr + (1)" not in kernel_source

    expected = a_data[1 : 2 + 1]
    assert np.allclose(c, expected), f"expected={expected}\nactual={c}"

    # Aligned case: N=4, C_N=2, offset=2
    a_data, c, kernel_source = build(4, 2, 2)
    # (uint1*)(A + (2)) is a valid vector load
    assert "A_ptr + 2" in kernel_source

    expected = a_data[2 : 2 + 2]
    assert np.allclose(c, expected), f"expected={expected}\nactual={c}"


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_thread_sync_inside_condition():
    @T.prim_func(s_tir=True)
    def func1(A: T.Buffer((4, 4), "float32")) -> None:
        A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
        for bx in T.thread_binding(1, "blockIdx.x"):
            for tx in T.thread_binding(32, "threadIdx.x"):
                if A[0, 0] > 1.0:
                    for i, j in T.grid(4, 4):
                        A_shared[i, j] = A[i, j]
                    for i, j in T.grid(4, 4):
                        A[i, j] = A_shared[i, j] + 1.0

    @T.prim_func(s_tir=True)
    def func2(A: T.Buffer((4, 4), "float32")) -> None:
        A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
        for bx in T.thread_binding(1, "blockIdx.x"):
            for tx in T.thread_binding(32, "threadIdx.x"):
                if T.tvm_thread_invariant(A[0, 0] > 1.0):
                    for i, j in T.grid(4, 4):
                        A_shared[i, j] = A[i, j]
                    for i, j in T.grid(4, 4):
                        A[i, j] = A_shared[i, j] + 1.0

    @T.prim_func(s_tir=True)
    def func3(A: T.Buffer((4, 4), "float32")) -> None:
        A_shared = T.sblock_alloc_buffer((4, 4), "float32", scope="shared")
        for bx in T.thread_binding(1, "blockIdx.x"):
            for tx in T.thread_binding(32, "threadIdx.x"):
                while T.tvm_thread_invariant(A[0, 0] > 1.0):
                    for i, j in T.grid(4, 4):
                        A_shared[i, j] = A[i, j]
                    for i, j in T.grid(4, 4):
                        A[i, j] = A_shared[i, j] + 1.0

    mod = tvm.IRModule({"main": func1})
    with pytest.raises(tvm.error.InternalError):
        tvm.compile(mod, target="cuda")

    mod = tvm.IRModule({"main": func2})
    tvm.compile(mod, target="cuda")

    mod = tvm.IRModule({"main": func3})
    tvm.compile(mod, target="cuda")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_invalid_reinterpret():
    @T.prim_func(s_tir=True)
    def func(A: T.Buffer((4,), "uint32"), B: T.Buffer((4,), "uint8")) -> None:
        for tx in T.thread_binding(4, "threadIdx.x"):
            B[tx] = T.call_intrin("uint8", "tirx.reinterpret", A[tx])

    with pytest.raises(RuntimeError):
        tvm.compile(func, target="cuda")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
def test_cuda_tensormap():
    # fmt: off
    @T.prim_func(s_tir=True)
    def main(A_ptr: T.handle):
        A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)

        A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
        T.call_packed("runtime.cuTensorMapInit", A_map, "float32", 2, A.data,
                      16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0)

        for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
            for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
                if threadIdx == 0:
                    A[0, 0] = T.reinterpret("float64", A_map)
    # fmt: on

    mod = tvm.IRModule({"main": main})
    mod = tvm.compile(mod, target="cuda")
    assert (
        """
extern "C" __global__ void __launch_bounds__(128) main_kernel(const __grid_constant__ CUtensorMap A_map, float* __restrict__ A_ptr) {
  if (((int)threadIdx.x) == 0) {
    A_ptr[0] = ((float)(*(double *)(&(A_map))));
  }
}""".strip()
        in mod.mod.imports[0].inspect_source()
    )


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_device_func_call():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def add(a: T.float32, b: T.float32) -> T.float32:
            return a + b

        @T.prim_func(s_tir=True)
        def main(
            A: T.Buffer((1024, 1024), "float32"),
            B: T.Buffer((1024, 1024), "float32"),
            C: T.Buffer((1024, 1024), "float32"),
        ):
            for bx in T.thread_binding(1024, "blockIdx.x"):
                for tx in T.thread_binding(1024, "threadIdx.x"):
                    C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])

    lib = tvm.compile(Module, target="cuda")
    cuda_code = lib.mod.imports[0].inspect_source()
    assert 'extern "C" __device__ float add(float a, float b) {\n  return (a + b);\n}' in cuda_code


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_float_const_hex_format():
    """Test that float constants are emitted in hexadecimal format for precision"""

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(
            A: T.Buffer((1024, 1024), "float32"),
        ):
            for bx in T.thread_binding(1024, "blockIdx.x"):
                for tx in T.thread_binding(1024, "threadIdx.x"):
                    A[bx, tx] = T.float32(1 / 27)

    lib = tvm.compile(Module, target="cuda")
    cuda_code = lib.mod.imports[0].inspect_source()
    assert "0x1.2f684bda12f68p-5f" in cuda_code


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_device_host_call_same_func():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(private=True, s_tir=True)
        def add(a: T.int32, b: T.int32) -> T.int32:
            return a + b

        @T.prim_func(s_tir=True)
        def main(
            A: T.Buffer((128, 128), "int32"),
            B: T.Buffer((128, 128), "int32"),
            C: T.Buffer((128, 128), "int32"),
        ):
            length: T.let[T.int32] = Module.add(64, 64)  # Call from host
            for bx in T.thread_binding(length, "blockIdx.x"):
                for tx in T.thread_binding(length, "threadIdx.x"):
                    C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])  # Call from device

    # 1. If we set host to llvm, it will raise an error of
    #    "the tirx.ret should be transformed to return zero before the llvm code generation."
    #    Need to revisit this.
    # 2. We set a dummy mcpu value for testing purpose,
    #    in order to avoid checking a function is host or device based on the "cpu" substring.
    target = tvm.target.Target({"kind": "cuda", "mcpu": "dummy_mcpu"}, host="c")
    lib = tvm.compile(Module, target=target)
    cuda_code = lib.mod.imports[0].inspect_source()
    assert 'extern "C" __device__ int add(int a, int b) {\n  return (a + b);\n}' in cuda_code

    # Run a simple test
    a_np = np.random.randint(0, 10, (128, 128), dtype="int32")
    b_np = np.random.randint(0, 10, (128, 128), dtype="int32")

    def run_and_check():
        dev = tvm.cuda(0)
        a_tvm = tvm.runtime.tensor(a_np, device=dev)
        b_tvm = tvm.runtime.tensor(b_np, device=dev)
        c_tvm = tvm.runtime.empty((128, 128), dtype="int32", device=dev)
        lib["main"](a_tvm, b_tvm, c_tvm)
        tvm.testing.assert_allclose(c_tvm.numpy(), a_np + b_np)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_thread_return():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")):
            for bx in T.thread_binding(32, "blockIdx.x"):
                for tx in T.thread_binding(32, "threadIdx.x"):
                    if bx >= 16 or tx >= 16:
                        T.thread_return()
                    B[bx, tx] = A[bx, tx]

    lib = tvm.compile(Module, target="cuda")
    cuda_code = lib.mod.imports[0].inspect_source()
    assert "return;" in cuda_code


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_loop_step():
    @T.prim_func(s_tir=True)
    def cuda_loop_step(
        A: T.Buffer((1024,), "float32"),
        B: T.Buffer((1024,), "float32"),
        C: T.Buffer((1024,), "float32"),
    ):
        # Each thread computes a strided subset of the i loop: start = tx*3, step = 96 (3 * 32 threads)
        for bx in T.thread_binding(1, "blockIdx.x"):
            for tx in T.thread_binding(96, "threadIdx.x"):
                for i in T.serial(tx, 1024, step=96):
                    C[i] = A[i] + B[i]

    target = tvm.target.Target({"kind": "cuda"})
    with tvm.transform.PassContext(disabled_pass=["s_tir.CanonicalizeLoop"]):
        lib = tvm.compile(cuda_loop_step, target=target)

    cuda_src = lib.mod.imports[0].inspect_source()
    assert "i += 96" in cuda_src
    a_np = np.random.uniform(1, 100, (1024,)).astype("float32")
    b_np = np.random.uniform(1, 100, (1024,)).astype("float32")
    c_np = np.zeros((1024,), dtype="float32")

    def run_and_check():
        dev = tvm.cuda(0)
        a_nd = tvm.runtime.tensor(a_np, dev)
        b_nd = tvm.runtime.tensor(b_np, dev)
        c_nd = tvm.runtime.tensor(c_np, dev)
        lib["main"](a_nd, b_nd, c_nd)
        tvm.testing.assert_allclose(c_nd.numpy(), a_np + b_np)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_export_load_with_fallback(monkeypatch, tmp_path):
    """Force the codegen wrapper into the fallback branch, then export+load+run."""
    n = 1024

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
                    with T.sblock("B"):
                        v_i = T.axis.spatial(n, i_0 * 32 + i_1)
                        T.reads(A[v_i])
                        T.writes(B[v_i])
                        B[v_i] = A[v_i] + 1.0

    monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
    host_lib = tvm.compile(Module, target="cuda")
    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")

    lib_path = str(tmp_path / "lib.so")
    host_lib.export_library(lib_path)
    reloaded = tvm.runtime.load_module(lib_path)

    a_np = np.random.uniform(size=(n,)).astype("float32")
    b_np = np.zeros((n,), dtype="float32")

    def run_and_check():
        dev = tvm.cuda(0)
        a = tvm.runtime.tensor(a_np, dev)
        b = tvm.runtime.tensor(b_np, dev)
        reloaded["main"](a, b)
        np.testing.assert_allclose(b.numpy(), a_np + 1.0, rtol=1e-5)

    tvm.testing.run_with_gpu_lock(run_and_check)


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