# 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, F841

import re

import numpy as np
import pytest

import tvm
import tvm.testing
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import ir as I_builder
from tvm.script.ir_builder import tirx as T_builder
from tvm.testing import env

dtype = tvm.testing.parameter("float32", "int32", "float16", "int8")
fuzz_seed = tvm.testing.parameter(range(25))


# Explicitly specify a target, as this test is looking at the
# generated shader code, and is not running on an actual device.
@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled(
        {
            "kind": "vulkan",
            "supports_int8": 1,
            "supports_8bit_buffer": 1,
            "supports_storage_buffer_storage_class": 1,
            "supports_float16": 1,
            "supports_16bit_buffer": 1,
        }
    ),
    reason="vulkan not enabled",
)
def test_vector_comparison(dtype):
    target = {
        "kind": "vulkan",
        "supports_int8": 1,
        "supports_8bit_buffer": 1,
        "supports_storage_buffer_storage_class": 1,
        "supports_float16": 1,
        "supports_16bit_buffer": 1,
    }
    target = tvm.target.Target(target)
    zero = tvm.tirx.const(0, dtype)
    one = tvm.tirx.const(1, dtype)

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((1024,), dtype), B: T.Buffer((1024,), dtype)):
            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
                for i_1 in T.thread_binding(32, thread="threadIdx.x"):
                    for i_2 in T.vectorized(4):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(1024, i_0 * 128 + i_1 * 4 + i_2)
                            B[v_i] = T.Select(A[v_i] >= zero, A[v_i] + one, zero)

    # Build
    f = tvm.tirx.build(Module, target=target)

    # Verify we generate the boolx4 type declaration and the OpSelect
    # v4{float,half,int} instruction
    assembly = f.imports[0].inspect_source()
    matches = re.findall("%v4bool = OpTypeVector %bool 4", assembly)
    assert len(matches) == 1
    matches = re.findall("OpSelect %v4.*", assembly)
    assert len(matches) == 1


@pytest.mark.parametrize(
    "target",
    [
        "llvm",
        pytest.param("cuda", marks=pytest.mark.gpu),
        pytest.param("rocm", marks=pytest.mark.gpu),
        pytest.param("vulkan", marks=pytest.mark.gpu),
        pytest.param("metal", marks=pytest.mark.gpu),
        pytest.param("opencl", marks=pytest.mark.gpu),
        pytest.param("nvptx", marks=pytest.mark.gpu),
    ],
)
def test_array_copy(target, dtype, fuzz_seed):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")
    np.random.seed(fuzz_seed)

    log_arr_size = np.random.uniform(low=np.log(1), high=np.log(32768))
    arr_size = np.exp(log_arr_size).astype(int)
    a_np = np.random.uniform(size=(arr_size,)).astype(dtype)

    def run_and_check():
        dev = tvm.device(target)
        a = tvm.runtime.empty((arr_size,), dtype, dev).copyfrom(a_np)
        tvm.testing.assert_allclose(a_np, a.numpy())

    if target == "llvm":
        run_and_check()
    else:
        tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_array_vectorize_add(dtype):
    target = {"kind": "vulkan", "from_device": 0}
    target = tvm.target.Target(target)
    arr_size = 64
    lanes = 2

    vec_dtype = f"{dtype}x{lanes}"
    one = tvm.tirx.const(1, vec_dtype)

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((64,), vec_dtype), B: T.Buffer((64,), vec_dtype)):
            for i_0 in T.thread_binding(16, thread="blockIdx.x"):
                for i_1 in T.thread_binding(4, thread="threadIdx.x"):
                    with T.sblock("B"):
                        v_i = T.axis.spatial(64, i_0 * 4 + i_1)
                        B[v_i] = A[v_i] + one

    f = tvm.compile(Module, target=target)

    def run_and_check():
        dev = tvm.device(target.kind.name)
        a = tvm.runtime.empty((arr_size,), vec_dtype, dev).copyfrom(
            np.random.uniform(size=(arr_size, lanes))
        )
        c = tvm.runtime.empty((arr_size,), vec_dtype, dev)
        f(a, c)
        tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vulkan_bool_load():
    target = {"kind": "vulkan", "from_device": 0}
    target = tvm.target.Target(target)
    arr_size = 1024

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((1024,), "bool"), B: T.Buffer((1024,), "int32")):
            for i_0 in T.thread_binding(8, thread="blockIdx.x"):
                for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                    with T.sblock("B"):
                        v_i = T.axis.spatial(1024, i_0 * 128 + i_1)
                        B[v_i] = T.Cast("int32", A[v_i])

    f = tvm.compile(Module, target=target)

    a_np = np.random.uniform(size=arr_size) > 0.5
    b_np = np.zeros((arr_size,), dtype="int32")
    ref = a_np.astype(np.int32)

    def run_and_check():
        dev = tvm.device(target.kind.name)
        a = tvm.runtime.tensor(a_np, dev)
        b = tvm.runtime.tensor(b_np, dev)
        f(a, b)
        tvm.testing.assert_allclose(b.numpy(), ref)

    tvm.testing.run_with_gpu_lock(run_and_check)


vulkan_parameter_impl = tvm.testing.parameter("push_constants", "ubo")
vulkan_parameter_dtype = tvm.testing.parameter("int32", "float32", "int64")


# Only run on vulkan because extremely large numbers of input
# parameters can crash cuda/llvm compiler.
@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
    target = {"kind": "vulkan", "from_device": 0}
    target = tvm.target.Target(target)
    dtype = vulkan_parameter_dtype

    if not target.attrs.get("supports_int64", False):
        pytest.xfail("Vulkan target does not support Int64 variables")

    # f_add has 3+num_int_params scalar parameters.  The other three
    # are length_n, stride1, and stride2.
    if vulkan_parameter_impl == "push_constants":
        # 4 params, 32 bytes.  Within 128-byte spec-guaranteed size of
        # push constants.  Uses push constants.
        num_int_params = 1
    else:
        # 24 params, 192 bytes.  May be above spec-guaranteed size of 128
        # bytes for push constants.  Uses either push constants or UBO,
        # depending on the device.
        max_push_constants_size = int(target.attrs.get("max_push_constants_size", 128))
        max_int_params_in_push = max_push_constants_size // 8 - 3
        num_int_params = max_int_params_in_push + 1

    # Build IRModule programmatically since num_int_params is dynamic
    with IRBuilder() as ib:
        with I_builder.ir_module():
            with T_builder.prim_func():
                T_builder.func_name("main")
                scalar_vars = []
                for i in range(num_int_params):
                    v = T_builder.arg(f"scale{i}", tvm.tirx.Var("", dtype))
                    scalar_vars.append(v)
                var_A = T_builder.arg("var_A", T_builder.handle())
                var_B = T_builder.arg("var_B", T_builder.handle())
                T_builder.func_attr({"tirx.noalias": True})
                n_var = T_builder.int32()
                A = T_builder.match_buffer(var_A, (n_var,), dtype)
                B = T_builder.match_buffer(var_B, (n_var,), dtype)
                scalar_sum = scalar_vars[0]
                for s in scalar_vars[1:]:
                    scalar_sum = scalar_sum + s
                with T_builder.thread_binding(
                    tvm.tirx.ceildiv(n_var, 64), thread="blockIdx.x"
                ) as i_0:
                    with T_builder.thread_binding(64, thread="threadIdx.x") as i_1:
                        with T_builder.sblock("B"):
                            v_i = T_builder.axis.spatial(n_var, i_0 * 64 + i_1)
                            T_builder.where(i_0 * 64 + i_1 < n_var)
                            T_builder.reads(A[v_i])
                            T_builder.writes(B[v_i])
                            T_builder.buffer_store(B, scalar_sum + A[v_i], [v_i])
    mod = ib.get()
    f_add = tvm.compile(mod, target=target)

    n = 1024
    scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)

    def run_and_check():
        dev = tvm.device(target.kind.name)
        a = tvm.runtime.tensor(np.random.uniform(size=n).astype(dtype), dev)
        b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
        f_add(*scalars, a, b)
        tvm.testing.assert_allclose(a.numpy() + sum(scalars), b.numpy())

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vulkan_while_if():
    target = {"kind": "vulkan", "from_device": 0}
    target = tvm.target.Target(target)
    n = 1
    dtype = "int32"

    @T.prim_func(s_tir=True)
    def while_if_gpu(A: T.Buffer((1,), "int32"), B: T.Buffer((1,), "int32")):
        for bx in T.thread_binding(1, thread="blockIdx.x"):
            iterations = T.decl_buffer((1,), "int32", scope="local")
            iterations[0] = 0
            B[0] = 0
            while iterations[0] < T.if_then_else(A[0] > 0, 10, 20):
                iterations[0] = iterations[0] + 1
                B[0] = B[0] + iterations[0]

    mod = tvm.IRModule.from_expr(while_if_gpu.with_attr("target", target))
    compiled_func = tvm.compile(mod, target=target)

    def run_and_check():
        dev = tvm.device(target.kind.name)
        for input_value, expected in [(5, [55]), (-5, [210])]:
            a = tvm.runtime.tensor(np.array([input_value], dtype=dtype), dev)
            b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
            compiled_func(a, b)
            tvm.testing.assert_allclose(b.numpy(), expected)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vulkan_local_threadidx():
    target = {"kind": "vulkan", "from_device": 0}
    target = tvm.target.Target(target)
    n = 32

    @T.prim_func(s_tir=True)
    def local_threadidx_func(A: T.Buffer((32,), "int32"), B: T.Buffer((32,), "int32")):
        # First block with thread extent 16
        for _ in range(1):
            for tx in T.thread_binding(16, thread="threadIdx.x"):
                B[tx + 0] = A[tx + 0]
        # Second block with thread extent 16
        for _ in range(1):
            for tx in T.thread_binding(16, thread="threadIdx.x"):
                B[tx + 16] = A[tx + 16]

    mod = tvm.IRModule.from_expr(local_threadidx_func)
    func = tvm.compile(mod, target=target)

    a_np = np.arange(n).astype(dtype="int32")
    b_np = np.zeros((n,), dtype="int32")

    def run_and_check():
        dev = tvm.device(target.kind.name)
        a = tvm.runtime.tensor(a_np, dev)
        b = tvm.runtime.tensor(b_np, dev)
        func(a, b)
        tvm.testing.assert_allclose(b.numpy(), a_np)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vectorized_index_ramp():
    """Test vectorized copy with ramp indices (load N values, write to N locations)"""
    target = {"kind": "vulkan", "from_device": 0}
    n = 4
    ramp_index = tvm.tirx.Ramp(0, 1, 4)

    @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})
            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
            with T.sblock("compute"):
                T.reads()
                T.writes()
                bx = T.launch_thread("blockIdx.x", 1)
                B[ramp_index] = A[ramp_index]

    f = tvm.compile(Module, target=target)

    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
    b_np = np.zeros(n, dtype="int32")

    def run_and_check():
        dev = tvm.device(target["kind"])
        a = tvm.runtime.tensor(a_np, dev)
        b = tvm.runtime.tensor(b_np, dev)
        f(a, b)
        tvm.testing.assert_allclose(b.numpy(), a_np)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_vectorized_index_broadcast():
    """Test broadcast index (load 1 value, write to N locations)"""
    target = {"kind": "vulkan", "from_device": 0}
    n = 4
    broadcast_index = tvm.tirx.Broadcast(0, 4)
    ramp_index = tvm.tirx.Ramp(0, 1, 4)

    @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})
            A = T.match_buffer(var_A, (n,), "int32", offset_factor=1)
            B = T.match_buffer(var_B, (n,), "int32", offset_factor=1)
            with T.sblock("compute"):
                T.reads()
                T.writes()
                bx = T.launch_thread("blockIdx.x", 1)
                # Load from broadcast index (single element), store to ramp index
                B[ramp_index] = A[broadcast_index]

    f = tvm.compile(Module, target=target)

    a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
    b_np = np.zeros(n, dtype="int32")

    def run_and_check():
        dev = tvm.device(target["kind"])
        a = tvm.runtime.tensor(a_np, dev)
        b = tvm.runtime.tensor(b_np, dev)
        f(a, b)
        # All elements of b should be a[0] (broadcast load)
        tvm.testing.assert_allclose(b.numpy(), np.full(n, a_np[0]))

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(
    not tvm.testing.device_enabled({"kind": "vulkan", "from_device": 0}),
    reason="vulkan not enabled",
)
def test_negative_operand_divmod():
    """Test handling of negative offsets to floormod/floordiv

    Even though the SPIR-V spec states that OpSRem and OpSMod can give
    the signed modulo, the Vulkan spec states that any use of negative
    operands is undefined behavior.  This test starts with negative
    operands to floordiv, validating that they are simplified into the
    corresponding positive operands, such that the final TIR can be
    expressed using only positive operands.

    SPIR-V: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpSRem
    Vulkan: https://registry.khronos.org/vulkan/specs/1.3/html/chap37.html#spirvenv-op-prec
    """
    target = {"kind": "vulkan", "from_device": 0}

    N = 32
    offset = 16
    divisor = 5

    @T.prim_func(s_tir=True)
    def func(A: T.Buffer((N, 2), "int32")):
        for i in T.thread_binding(N, thread="threadIdx.x"):
            with T.sblock("A"):
                v_i = T.axis.spatial(N, i)
                A[v_i, 0] = T.floordiv(v_i - offset, divisor)
                A[v_i, 1] = T.floormod(v_i - offset, divisor)

    built = tvm.compile(func, target=target)

    def run_and_check():
        dev = tvm.device(target["kind"])
        a_dev = tvm.runtime.empty([N, 2], "int32", dev)
        built(a_dev)
        a = a_dev.numpy()
        np.testing.assert_array_equal(a[:, 0], (np.arange(N) - offset) // divisor)
        np.testing.assert_array_equal(a[:, 1], (np.arange(N) - offset) % divisor)

    tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.parametrize("out_dtype", ["float32", "float16"])
def test_cooperative_matrix(out_dtype):
    M, N, K = 16, 16, 32

    # fmt: off
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(X: T.Buffer((16, 32), "float16"), W: T.Buffer((32, 16), "float16"), compute: T.Buffer((16, 16), out_dtype)):
            T.func_attr({"tirx.noalias": True})
            X_shared = T.sblock_alloc_buffer((16, 32), "float16", scope="shared")
            W_shared = T.sblock_alloc_buffer((32, 16), "float16", scope="shared")
            X_shared_wmma_matrix_a = T.sblock_alloc_buffer((16, 32), "float16", scope="wmma.matrix_a")
            W_shared_wmma_matrix_b = T.sblock_alloc_buffer((32, 16), "float16", scope="wmma.matrix_b")
            compute_wmma_accumulator = T.sblock_alloc_buffer((16, 16), out_dtype, scope="wmma.accumulator")
            for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
                with T.sblock("compute_init_o"):
                    v_i_o = T.axis.spatial(1, 0)
                    v_j_o = T.axis.spatial(1, 0)
                    T.reads()
                    T.writes(compute_wmma_accumulator[0:16, 0:16])
                    C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
                    T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0.0))
                for k_0 in range(2):
                    for ax0_ax1_fused_0 in range(2):
                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
                            for ax0_ax1_fused_2 in T.vectorized(4):
                                with T.sblock("X_shared"):
                                    v0 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
                                    v1 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
                                    T.reads(X[v0, v1])
                                    T.writes(X_shared[v0, v1])
                                    X_shared[v0, v1] = X[v0, v1]
                    for ax0_ax1_fused_0 in range(2):
                        for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
                            for ax0_ax1_fused_2 in T.vectorized(4):
                                with T.sblock("W_shared"):
                                    v0 = T.axis.spatial(32, k_0 * 16 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 16)
                                    v1 = T.axis.spatial(16, (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 16)
                                    T.reads(W[v0, v1])
                                    T.writes(W_shared[v0, v1])
                                    W_shared[v0, v1] = W[v0, v1]
                    for ax0_0 in T.unroll(1):
                        for ax1_0 in T.unroll(1):
                            with T.sblock("X_shared_wmma.matrix_a_o"):
                                v0_o = T.axis.spatial(1, ax0_0)
                                v1_o = T.axis.spatial(2, k_0 + ax1_0)
                                T.reads(X_shared[0:16, v1_o * 16:v1_o * 16 + 16])
                                T.writes(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16])
                                A = T.match_buffer(X_shared[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
                                C = T.match_buffer(X_shared_wmma_matrix_a[0:16, v1_o * 16:v1_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
                    for ax0_0 in T.unroll(1):
                        for ax1_0 in T.unroll(1):
                            with T.sblock("W_shared_wmma.matrix_b_o"):
                                v0_o = T.axis.spatial(2, k_0 + ax0_0)
                                v1_o = T.axis.spatial(1, ax1_0)
                                T.reads(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16])
                                T.writes(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16])
                                A = T.match_buffer(W_shared[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=16)
                                C = T.match_buffer(W_shared_wmma_matrix_b[v0_o * 16:v0_o * 16 + 16, 0:16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
                                T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
                    with T.sblock("compute_update_o"):
                        v_i_o = T.axis.spatial(1, 0)
                        v_j_o = T.axis.spatial(1, 0)
                        v_k_o = T.axis.reduce(2, k_0)
                        T.reads(compute_wmma_accumulator[0:16, 0:16], X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16])
                        T.writes(compute_wmma_accumulator[0:16, 0:16])
                        A = T.match_buffer(X_shared_wmma_matrix_a[0:16, v_k_o * 16:v_k_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
                        B = T.match_buffer(W_shared_wmma_matrix_b[v_k_o * 16:v_k_o * 16 + 16, 0:16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
                        C = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
                        T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
                with T.sblock("compute_wmma.accumulator_o"):
                    v0_o = T.axis.spatial(1, 0)
                    v1_o = T.axis.spatial(1, 0)
                    T.reads(compute_wmma_accumulator[0:16, 0:16])
                    T.writes(compute[0:16, 0:16])
                    A = T.match_buffer(compute_wmma_accumulator[0:16, 0:16], (16, 16), out_dtype, strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
                    C = T.match_buffer(compute[0:16, 0:16], (16, 16), out_dtype, strides=("C_s0", "C_s1"), offset_factor=16)
                    T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation(out_dtype), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
    # fmt: on

    target = {"kind": "vulkan", "from_device": 0}
    tgt_attrs = tvm.target.Target(target).attrs

    if tgt_attrs.get("supports_cooperative_matrix"):
        f = tvm.compile(Module, target=target)

        def run_and_check():
            dev = tvm.device("vulkan", 0)
            A = tvm.runtime.tensor(np.random.randn(M, K).astype("float16"), dev)
            B = tvm.runtime.tensor(np.random.randn(K, N).astype("float16"), dev)
            C = tvm.runtime.tensor(np.random.randn(M, N).astype(out_dtype), dev)
            f(A, B, C)
            A_np = A.numpy()
            B_np = B.numpy()
            ref = np.dot(A_np.astype("float32"), B_np.astype("float32"))
            tvm.testing.assert_allclose(C.numpy(), ref, rtol=1e-2, atol=1e-2)

        tvm.testing.run_with_gpu_lock(run_and_check)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
def test_codegen_decl_buffer():
    """The codegen should accept DeclBuffer nodes in its input"""

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def kernel():
            T.func_attr({"calling_conv": 2, "global_symbol": "kernel", "tirx.noalias": True})
            A = T.alloc_buffer((256,), dtype="float32", scope="local")
            A_buf = T.decl_buffer([256], dtype="float32", scope="local", data=A.data)

    target = tvm.target.Target("vulkan")
    vulkan_codegen = tvm.get_global_func("target.build.vulkan")
    vulkan_codegen(Module, target)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
def test_codegen_static_shared_memory():
    """The codegen should accept static shared/workgroup allocations."""

    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
            A_shared = T.alloc_buffer((128,), dtype="float32", scope="shared")

            for bx in T.thread_binding(1, thread="blockIdx.x"):
                for tx in T.thread_binding(128, thread="threadIdx.x"):
                    A_shared[tx] = A[tx]
                    B[tx] = A_shared[tx]

    tvm.compile(Module, target="vulkan")


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
def test_unary():
    test_funcs = [
        (tvm.tirx.sin, lambda x: np.sin(x)),
        (tvm.tirx.cos, lambda x: np.cos(x)),
        (tvm.tirx.tan, lambda x: np.tan(x)),
        (tvm.tirx.sinh, lambda x: np.sinh(x)),
        (tvm.tirx.cosh, lambda x: np.cosh(x)),
        (tvm.tirx.tanh, lambda x: np.tanh(x)),
        (tvm.tirx.asin, lambda x: np.arcsin(x)),
        (tvm.tirx.acos, lambda x: np.arccos(x)),
        (tvm.tirx.atan, lambda x: np.arctan(x)),
        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
    ]

    def run_test(tvm_intrin, np_func):
        n = 16

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(var_A: T.handle, var_B: T.handle):
                m = T.int32()
                A = T.match_buffer(var_A, (m,), "float32")
                B = T.match_buffer(var_B, (m,), "float32")
                for i_0 in T.thread_binding((m + 63) // 64, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(64, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(m, i_0 * 64 + i_1)
                            T.where(i_0 * 64 + i_1 < m)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = tvm_intrin(A[v_i])

        target = tvm.target.Target("vulkan")
        func = tvm.compile(Module, target=target)

        if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
            data = np.random.uniform(-1.0, 1.0, size=n)
        elif tvm_intrin == tvm.tirx.atanh:
            data = np.random.uniform(-0.999, 0.999, size=n)
        elif tvm_intrin == tvm.tirx.acosh:
            data = np.random.uniform(1.0, 5.0, size=n)
        else:
            data = np.random.uniform(0.1, 0.9, size=n)

        def run_and_check():
            dev = tvm.device(target.kind.name, 0)
            a = tvm.runtime.tensor(data.astype("float32"), dev)
            b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
            func(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_vulkan(), reason="need vulkan")
def test_export_load_with_fallback(monkeypatch, tmp_path):
    """Force the codegen wrapper into the fallback branch, then export."""
    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="vulkan")
    monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")

    lib_path = str(tmp_path / "lib.so")
    host_lib.export_library(lib_path)


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