# 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 numpy as np
import pytest
import tvm_ffi

import tvm
import tvm.testing
from tvm import s_tir
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.testing import env


def test_cp_async_raw_dtype_round_trips():
    # The raw cp.async form emitted by InjectPTXAsyncCopy carries the element
    # dtype in Call.dtype and must survive a TVMScript print -> parse round-trip
    # (it prints dtype-first via tirx.ptx.cp_async_raw). Guards the regression
    # where the element dtype was dropped after the flat op was phased out.
    @T.prim_func
    def f(A: T.Buffer((128,), "float16"), B: T.Buffer((128,), "float16")):
        T.func_attr({"global_symbol": "f"})
        for i in T.serial(8):
            T.ptx.cp_async("float16", B.data, i * 16, A.data, i * 16, 16)

    reparsed = tvm.script.from_source(f.script())
    tvm.ir.assert_structural_equal(f, reparsed)


def count_cp_async(stmt):
    num_alloc = [0]

    def verify(n):
        if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.ptx.cp_async_raw":
            num_alloc[0] += 1

    tvm.tirx.stmt_functor.post_order_visit(stmt, verify)
    return num_alloc[0]


def generate_global_to_shared_vectorized_copy(dtype, vector_size):
    num_iters = 128 // vector_size
    vector_size_expr = tvm.runtime.convert(vector_size)

    @T.prim_func(s_tir=True)
    def ptx_global_to_shared_copy(
        A: T.Buffer((32, 128), dtype), B: T.Buffer((32, 128), dtype)
    ) -> None:
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        bx = T.env_thread("blockIdx.x")
        tx = T.env_thread("threadIdx.x")
        T.launch_thread(bx, 1)
        T.launch_thread(tx, 32)
        with T.sblock():
            A_shared = T.sblock_alloc_buffer([32, 128], dtype, scope="shared")
            T.reads(A[0:32, 0:128])
            T.writes(B[0:32, 0:128])

            T.attr("default", "async_scope", 1)
            for i in T.serial(num_iters):
                for j in T.vectorized(vector_size):
                    A_shared[tx, i * vector_size_expr + j] = A[tx, i * vector_size_expr + j]

            T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
            T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))

            for i in range(128):
                B[tx, i] = A_shared[tx, i]

    return ptx_global_to_shared_copy


@T.prim_func(s_tir=True)
def ptx_global_to_shared_copy_fp32x1(
    A: T.Buffer((32, 128), "float32"), B: T.Buffer((32, 128), "float32")
) -> None:
    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
    bx = T.env_thread("blockIdx.x")
    tx = T.env_thread("threadIdx.x")
    T.launch_thread(bx, 1)
    T.launch_thread(tx, 32)
    with T.sblock():
        A_shared = T.sblock_alloc_buffer([32, 128], "float32", scope="shared")
        T.reads(A[0:32, 0:128])
        T.writes(B[0:32, 0:128])

        T.attr("default", "async_scope", 1)
        for i in T.serial(128):
            A_shared[tx, i] = A[tx, i]

        T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
        T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))

        for i in range(128):
            B[tx, i] = A_shared[tx, i]


@T.prim_func(s_tir=True)
def ptx_global_to_shared_dyn_copy_fp16x8(
    A: T.Buffer((32, 128), "float16"),
    B: T.Buffer((32, 128), "float16"),
    C: T.Buffer((32, 128), "float16"),
) -> None:
    T.func_attr({"global_symbol": "main", "tirx.noalias": True})
    bx = T.env_thread("blockIdx.x")
    tx = T.env_thread("threadIdx.x")
    T.launch_thread(bx, 1)
    T.launch_thread(tx, 32)
    with T.sblock():
        A_shared = T.sblock_alloc_buffer([32, 128], "float16", scope="shared.dyn")
        B_shared = T.sblock_alloc_buffer([32, 128], "float16", scope="shared.dyn")
        T.reads(A[0:32, 0:128], B[0:32, 0:128])
        T.writes(C[0:32, 0:128])

        T.attr("default", "async_scope", 1)
        for i in T.serial(16):
            for j in T.vectorized(8):
                A_shared[tx, i * 8 + j] = A[tx, i * 8 + j]
                B_shared[tx, i * 8 + j] = B[tx, i * 8 + j]

        T.evaluate(T.ptx.cp_async.commit_group(dtype=""))
        T.evaluate(T.ptx.cp_async.wait_group(0, dtype=""))

        for i in range(128):
            C[tx, i] = A_shared[tx, i] + B_shared[tx, i]


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_inject_async_copy():
    for dtype, vec_size in [("float16", 8), ("float16", 4), ("float32", 4), ("float32", 1)]:
        if vec_size == 1:
            f = ptx_global_to_shared_copy_fp32x1
        else:
            f = generate_global_to_shared_vectorized_copy(dtype, vec_size)

        mod = tvm.IRModule.from_expr(f)
        mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
        mod = tvm.tirx.transform.FlattenBuffer()(mod)
        if vec_size > 1:
            mod = tvm.tirx.transform.VectorizeLoop()(mod)
        mod = tvm.s_tir.transform.InjectPTXAsyncCopy()(mod)

        assert count_cp_async(mod["main"].body) == 1

        if not tvm.testing.is_ampere_or_newer():
            continue

        with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
            mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")

        A_np = np.random.rand(32, 128).astype(dtype)
        B_np = np.zeros((32, 128)).astype(dtype)

        def run_and_check():
            dev = tvm.cuda(0)
            A_nd = tvm.runtime.tensor(A_np, device=dev)
            B_nd = tvm.runtime.tensor(B_np, device=dev)
            mod(A_nd, B_nd)
            tvm.testing.assert_allclose(B_nd.numpy(), A_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_inject_async_copy_shared_dyn():
    f = ptx_global_to_shared_dyn_copy_fp16x8

    mod = tvm.IRModule.from_expr(f)
    mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
    mod = tvm.tirx.transform.FlattenBuffer()(mod)
    mod = tvm.tirx.transform.VectorizeLoop()(mod)
    mod = tvm.s_tir.transform.MergeSharedMemoryAllocations()(mod)
    mod = tvm.s_tir.transform.InjectPTXAsyncCopy()(mod)

    assert count_cp_async(mod["main"].body) == 2

    if not tvm.testing.is_ampere_or_newer():
        return

    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
        mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")

    A_np = np.random.rand(32, 128).astype("float16")
    B_np = np.random.rand(32, 128).astype("float16")
    C_np = np.zeros((32, 128)).astype("float16")

    def run_and_check():
        dev = tvm.cuda(0)
        A_nd = tvm.runtime.tensor(A_np, device=dev)
        B_nd = tvm.runtime.tensor(B_np, device=dev)
        C_nd = tvm.runtime.tensor(C_np, device=dev)
        mod(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)


# Note: the test_inject_async_copy_barrier case (and its prim_func helper)
# was removed — it relied on the indexed barrier API
# (`create_barriers`, `init_barrier_thread_count`, `arrive_barrier`,
# `wait_barrier`) which fork does not provide; fork uses the
# `ptx_mbarrier_*` family instead.


expected_cuda_script = r"""#include <cuda.h>
#endif

#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \
     (__CUDACC_VER_MAJOR__ > 11))
#define TVM_ENABLE_L2_PREFETCH 1
#else
#define TVM_ENABLE_L2_PREFETCH 0
#endif

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
#endif

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_0() {
    asm volatile("cp.async.wait_group 0;");
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_1() {
    asm volatile("cp.async.wait_group 1;");
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_2() {
    asm volatile("cp.async.wait_group 2;");
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_wait_group_5() {
    asm volatile("cp.async.wait_group 5;");
}

__forceinline__ __device__ void ptx_cp_async_legacy_pred_ca_4_4_4(void* dst, int dst_off, void* src, int src_off, int predicate) {
  uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
  uint8_t* src_p = (uint8_t*)src + src_off * 4;
  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
  __asm__ __volatile__(
    "{\n"
    " .reg .pred p;\n"
    " setp.eq.u32 p, %3, 1;\n"
    " @p cp.async.ca.shared.global [%0], [%1], %2;\n"
    " @!p st.shared.u32 [%0], {%4};\n"
    "}\n"
    :: "r"(dst_addr), "l"(src_p), "n"(4), "r"(predicate), "r"(0)
  );
}

__forceinline__ __device__ void ptx_cp_async_legacy_ca_4_4_4(void* dst, int dst_off, void* src, int src_off) {
  uint8_t* dst_p = (uint8_t*)dst + dst_off * 4;
  uint8_t* src_p = (uint8_t*)src + src_off * 4;
  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);
  asm volatile("cp.async.ca.shared.global [%0], [%1], %2;"
    :: "r"(dst_addr), "l"(src_p), "n"(4));
}

__forceinline__ __device__ void tvm_builtin_ptx_cp_async_commit_group() {
    asm volatile("cp.async.commit_group;");
}
extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr);
extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr) {
  __shared__ alignas(64) float A_shared_ptr[64];
  __shared__ alignas(64) float B_shared_ptr[64];
  A_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
  B_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/;
  tvm_builtin_ptx_cp_async_commit_group();
  int cse_v1 = (((int)threadIdx.x) * 14);
  int cse_v2 = (((int)threadIdx.x) + 16);
  ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 16), A_ptr, (((int)threadIdx.x) * 14));
  ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 16), B_ptr, (((int)threadIdx.x) * 14));
  tvm_builtin_ptx_cp_async_commit_group();
  int cse_v3 = (((int)threadIdx.x) + 32);
  int cse_v6 = ((((int)threadIdx.x) * 14) + 1);
  ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 32), A_ptr, ((((int)threadIdx.x) * 14) + 1));
  ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 32), B_ptr, ((((int)threadIdx.x) * 14) + 1));
  tvm_builtin_ptx_cp_async_commit_group();
  int cse_v4 = (((int)threadIdx.x) * 16);
  for (int i = 0; i < 13; ++i) {
    int cse_v7 = (((((int)threadIdx.x) * 14) + i) + 2);
    int cse_v9 = ((((i + 3) & 3) * 16) + ((int)threadIdx.x));
    ptx_cp_async_legacy_pred_ca_4_4_4(A_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), A_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
    tvm_builtin_ptx_cp_async_commit_group();
    tvm_builtin_ptx_cp_async_wait_group_5();
    __syncthreads();
    int cse_v8 = (((i & 3) * 16) + ((int)threadIdx.x));
    C_ptr[((((int)threadIdx.x) * 16) + i)] = (A_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))] + B_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))]);
    __syncthreads();
    ptx_cp_async_legacy_pred_ca_4_4_4(B_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12));
    tvm_builtin_ptx_cp_async_commit_group();
  }
  tvm_builtin_ptx_cp_async_wait_group_2();
  __syncthreads();
  C_ptr[((((int)threadIdx.x) * 16) + 13)] = (A_shared_ptr[(((int)threadIdx.x) + 16)] + B_shared_ptr[(((int)threadIdx.x) + 16)]);
  tvm_builtin_ptx_cp_async_wait_group_1();
  __syncthreads();
  C_ptr[((((int)threadIdx.x) * 16) + 14)] = (A_shared_ptr[(((int)threadIdx.x) + 32)] + B_shared_ptr[(((int)threadIdx.x) + 32)]);
  tvm_builtin_ptx_cp_async_wait_group_0();
  __syncthreads();
  int cse_v5 = (((int)threadIdx.x) + 48);
  C_ptr[((((int)threadIdx.x) * 16) + 15)] = (A_shared_ptr[(((int)threadIdx.x) + 48)] + B_shared_ptr[(((int)threadIdx.x) + 48)]);
}

"""


@pytest.fixture
def postproc_if_missing_async_support():
    arch = tvm.support.nvcc.get_target_compute_version()
    major, _ = tvm.support.nvcc.parse_compute_version(arch)
    support_async = major >= 8

    func_name = "tvm_callback_cuda_postproc"
    prev_postproc = tvm.get_global_func(func_name, allow_missing=True)

    # Store the generated code prior to the post-processing.  This
    # way, even though the generated code doesn't compile on platforms
    # that do not support async, the comparison against an expected
    # output can still be performed.  We cannot use
    # `mod.inspect_source()`, as that contains the source after all
    # post-processing.
    original_code = None

    def get_original_code():
        nonlocal original_code
        return original_code

    @tvm.register_global_func(func_name, override=True)
    def tvm_callback_cuda_postproc(code, _):
        nonlocal original_code
        original_code = code
        if support_async:
            return code
        else:
            ret = []
            for line in code.split("\n"):
                ret.append(line)
                ret.append("\n")
                if line.startswith('extern "C" __global__') and line.endswith("{"):
                    break
            ret.append("}")
            return "".join(ret)

    yield get_original_code

    # Restore previous postproc func to avoid impacting other tests
    if prev_postproc is None:
        tvm_ffi.registry.remove_global_func(func_name)
    else:
        tvm.register_global_func(func_name, prev_postproc, override=True)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cp_async_in_if_then_else(postproc_if_missing_async_support):
    @T.prim_func(s_tir=True)
    def simple_compute(
        A: T.Buffer((16, 14), "float32"),
        B: T.Buffer((16, 14), "float32"),
        C: T.Buffer((16, 16), "float32"),
    ):
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
            for i in T.serial(
                16,
                annotations={
                    "software_pipeline_stage": [0, 0, 3],
                    "software_pipeline_order": [0, 2, 1],
                    "software_pipeline_async_stages": [0],
                },
            ):
                with T.sblock("compute"):
                    T.reads(A[tx, i])
                    T.writes(C[tx, i])
                    A_shared = T.sblock_alloc_buffer((16, 1), dtype="float32", scope="shared")
                    B_shared = T.sblock_alloc_buffer((16, 1), dtype="float32", scope="shared")
                    with T.sblock():
                        T.reads(A[tx, i])
                        T.writes(A_shared[tx, 0])
                        A_shared[tx, 0] = T.if_then_else(
                            1 <= i and i < 15, A[tx, i - 1], T.float32(0), dtype="float32"
                        )
                    with T.sblock():
                        T.reads(B[tx, i])
                        T.writes(B_shared[tx, 0])
                        B_shared[tx, 0] = T.if_then_else(
                            1 <= i and i < 15, B[tx, i - 1], T.float32(0), dtype="float32"
                        )
                    with T.sblock():
                        T.reads(A_shared[tx, 0], B_shared[tx, 0])
                        T.writes(C[tx, i])
                        C[tx, i] = A_shared[tx, 0] + B_shared[tx, 0]

    mod = tvm.IRModule.from_expr(simple_compute)
    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
        tvm.compile(mod, target="cuda")
    generated_code = postproc_if_missing_async_support()
    print(generated_code)
    # Fork emits an NVRTC-aware preamble (`#ifdef __CUDACC_RTC__ ... #else ...`
    # block) before the apache-style `#include <cuda.h>`; the body after that
    # block matches the expected snippet, so compare from the kernel-body
    # onwards instead of byte-for-byte from the start.
    marker = "#include <cuda.h>"
    expected_body = expected_cuda_script[expected_cuda_script.index(marker) :]
    actual_body = generated_code[generated_code.index(marker) :]
    assert actual_body == expected_body


@pytest.mark.skip(
    reason="This test fails due to an ordering issue with MergeSharedMemoryAllocations "
    "in device_driver_api.cc. However, fixing this causes failures in MLC. "
    "This bug should be addressed. See discussion in https://github.com/apache/tvm/pull/16769 "
    "and https://github.com/apache/tvm/pull/16569#issuecomment-1992720448"
)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
    @T.prim_func(s_tir=True)
    def complex_compute(
        A: T.Buffer((2, 16, 16, 1280), "float16"),
        W: T.Buffer((1280, 3, 3, 1280), "float16"),
        Conv: T.Buffer((512, 1280), "float16"),
    ):
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        # with T.sblock("root"):
        data_im2col_reindex_shared_dyn = T.sblock_alloc_buffer(
            (512, 11520), "float16", scope="shared.dyn"
        )
        data_im2col_reindex_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer(
            (512, 11520), "float16", scope="wmma.matrix_a"
        )
        weight_flatten_reindex_shared_dyn = T.sblock_alloc_buffer(
            (1280, 11520), "float16", scope="shared.dyn"
        )
        weight_flatten_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer(
            (1280, 11520), "float16", scope="wmma.matrix_b"
        )
        Conv_reindex_wmma_accumulator = T.sblock_alloc_buffer(
            (512, 1280), "float16", scope="wmma.accumulator"
        )
        for x_0_0 in T.thread_binding(8, thread="blockIdx.y"):
            for y_0_0 in T.thread_binding(20, thread="blockIdx.x"):
                for x_0_1 in T.thread_binding(2, thread="threadIdx.y"):
                    for y_0_1 in T.thread_binding(2, thread="threadIdx.z"):
                        for x_0_2_init, y_0_2_init in T.grid(2, 2):
                            with T.sblock("Conv_init_o"):
                                v_x_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + x_0_2_init)
                                v_y_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + y_0_2_init)
                                T.reads()
                                T.writes(
                                    Conv_reindex_wmma_accumulator[
                                        v_x_o * 16 : v_x_o * 16 + 16, v_y_o * 16 : v_y_o * 16 + 16
                                    ]
                                )
                                C_s0 = T.int32()
                                C_s1 = T.int32()
                                C = T.match_buffer(
                                    Conv_reindex_wmma_accumulator[
                                        v_x_o * 16 : v_x_o * 16 + 16, v_y_o * 16 : v_y_o * 16 + 16
                                    ],
                                    (16, 16),
                                    "float16",
                                    strides=(C_s0, C_s1),
                                    scope="wmma.accumulator",
                                    offset_factor=16,
                                )
                                T.tvm_fill_fragment(
                                    C.data,
                                    16,
                                    16,
                                    16,
                                    C.elem_offset // C_s0 // 16 * (C_s0 // 16)
                                    + C.elem_offset % C_s0 // 16,
                                    T.float32(0),
                                )
                        for k_0_0 in T.serial(
                            180,
                            annotations={
                                "software_pipeline_stage": [0, 0, 1],
                                "software_pipeline_order": [0, 1, 2],
                                "software_pipeline_async_stages": [0],
                            },
                        ):
                            for ax0_ax1_0_fused_0 in range(4):
                                for ax0_ax1_0_fused_1 in T.thread_binding(2, thread="threadIdx.z"):
                                    for ax0_ax1_0_fused_2 in T.thread_binding(
                                        2, thread="threadIdx.y"
                                    ):
                                        for ax0_ax1_0_fused_3 in T.thread_binding(
                                            32, thread="threadIdx.x"
                                        ):
                                            with T.sblock("data_im2col_reindex_shared.dyn_o"):
                                                v0 = T.axis.spatial(
                                                    512,
                                                    x_0_0 * 64
                                                    + (
                                                        ax0_ax1_0_fused_0 * 128
                                                        + ax0_ax1_0_fused_1 * 64
                                                        + ax0_ax1_0_fused_2 * 32
                                                        + ax0_ax1_0_fused_3
                                                    )
                                                    // 8,
                                                )
                                                v1_o = T.axis.spatial(
                                                    1440,
                                                    k_0_0 * 8
                                                    + (
                                                        ax0_ax1_0_fused_0 * 128
                                                        + ax0_ax1_0_fused_1 * 64
                                                        + ax0_ax1_0_fused_2 * 32
                                                        + ax0_ax1_0_fused_3
                                                    )
                                                    % 8,
                                                )
                                                T.reads(
                                                    A[
                                                        v0 // 256,
                                                        v1_o // 480 + v0 % 256 // 16 - 1,
                                                        v1_o % 480 // 160 + v0 % 16 - 1,
                                                        v1_o % 160 * 8 : v1_o % 160 * 8 + 8,
                                                    ]
                                                )
                                                T.writes(
                                                    data_im2col_reindex_shared_dyn[
                                                        v0, v1_o * 8 : v1_o * 8 + 8
                                                    ]
                                                )
                                                for ax1_1 in T.vectorized(8):
                                                    with T.sblock("data_im2col_reindex_shared.dyn"):
                                                        v1_i = T.axis.spatial(8, ax1_1)
                                                        T.reads(
                                                            A[
                                                                v0 // 256,
                                                                v1_o // 480 + v0 % 256 // 16 - 1,
                                                                v1_o % 480 // 160 + v0 % 16 - 1,
                                                                v1_o % 160 * 8 + v1_i,
                                                            ]
                                                        )
                                                        T.writes(
                                                            data_im2col_reindex_shared_dyn[
                                                                v0, v1_o * 8 + v1_i
                                                            ]
                                                        )
                                                        T.sblock_attr(
                                                            {"buffer_dim_align": [[0, 0, 32, 8]]}
                                                        )
                                                        data_im2col_reindex_shared_dyn[
                                                            v0, v1_o * 8 + v1_i
                                                        ] = T.if_then_else(
                                                            1 <= v1_o // 480 + v0 % 256 // 16
                                                            and v1_o // 480 + v0 % 256 // 16 < 17
                                                            and 1 <= v1_o % 480 // 160 + v0 % 16
                                                            and v1_o % 480 // 160 + v0 % 16 < 17,
                                                            A[
                                                                v0 // 256,
                                                                v1_o // 480 + v0 % 256 // 16 - 1,
                                                                v1_o % 480 // 160 + v0 % 16 - 1,
                                                                v1_o % 160 * 8 + v1_i,
                                                            ],
                                                            T.float16(0),
                                                        )
                            for ax0_ax1_0_fused_0 in range(4):
                                for ax0_ax1_0_fused_1 in T.thread_binding(2, thread="threadIdx.z"):
                                    for ax0_ax1_0_fused_2 in T.thread_binding(
                                        2, thread="threadIdx.y"
                                    ):
                                        for ax0_ax1_0_fused_3 in T.thread_binding(
                                            32, thread="threadIdx.x"
                                        ):
                                            for ax1_1 in T.vectorized(8):
                                                with T.sblock("weight_flatten_reindex_shared.dyn"):
                                                    v0 = T.axis.spatial(
                                                        1280,
                                                        y_0_0 * 64
                                                        + (
                                                            ax0_ax1_0_fused_0 * 128
                                                            + ax0_ax1_0_fused_1 * 64
                                                            + ax0_ax1_0_fused_2 * 32
                                                            + ax0_ax1_0_fused_3
                                                        )
                                                        // 8,
                                                    )
                                                    v1 = T.axis.spatial(
                                                        11520,
                                                        k_0_0 * 64
                                                        + (
                                                            ax0_ax1_0_fused_0 * 128
                                                            + ax0_ax1_0_fused_1 * 64
                                                            + ax0_ax1_0_fused_2 * 32
                                                            + ax0_ax1_0_fused_3
                                                        )
                                                        % 8
                                                        * 8
                                                        + ax1_1,
                                                    )
                                                    T.reads(
                                                        W[
                                                            v0,
                                                            v1 // 3840,
                                                            v1 % 3840 // 1280,
                                                            v1 % 1280,
                                                        ]
                                                    )
                                                    T.writes(
                                                        weight_flatten_reindex_shared_dyn[v0, v1]
                                                    )
                                                    T.sblock_attr(
                                                        {"buffer_dim_align": [[0, 0, 32, 8]]}
                                                    )
                                                    weight_flatten_reindex_shared_dyn[v0, v1] = W[
                                                        v0,
                                                        v1 // 1280 // 3,
                                                        v1 // 1280 % 3,
                                                        v1 % 1280,
                                                    ]
                            for k_0_1 in range(4):
                                for ax0_0, ax1_0 in T.grid(2, 1):
                                    with T.sblock("data_im2col_reindex_shared.dyn_wmma.matrix_a_o"):
                                        v0_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + ax0_0)
                                        v1_o = T.axis.spatial(720, k_0_0 * 4 + k_0_1 + ax1_0)
                                        T.reads(
                                            data_im2col_reindex_shared_dyn[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ]
                                        )
                                        T.writes(
                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ]
                                        )
                                        A_s0 = T.int32()
                                        A_s1 = T.int32()
                                        A_1 = T.match_buffer(
                                            data_im2col_reindex_shared_dyn[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ],
                                            (16, 16),
                                            "float16",
                                            strides=(A_s0, A_s1),
                                            scope="shared.dyn",
                                            offset_factor=16,
                                        )
                                        C_s0 = T.int32()
                                        C_s1 = T.int32()
                                        C = T.match_buffer(
                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
                                                v0_o * 16 : v0_o * 16 + 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_s0 // 16 * (C_s0 // 16)
                                            + C.elem_offset % C_s0 // 16,
                                            T.tvm_access_ptr(
                                                T.type_annotation("float16"),
                                                A_1.data,
                                                A_1.elem_offset,
                                                A_s0 * 16,
                                                1,
                                            ),
                                            A_s0,
                                            "row_major",
                                        )
                                for ax0_0, ax1_0 in T.grid(2, 1):
                                    with T.sblock(
                                        "weight_flatten_reindex_shared.dyn_wmma.matrix_b_o"
                                    ):
                                        v0_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + ax0_0)
                                        v1_o = T.axis.spatial(720, k_0_0 * 4 + k_0_1 + ax1_0)
                                        T.reads(
                                            weight_flatten_reindex_shared_dyn[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ]
                                        )
                                        T.writes(
                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ]
                                        )
                                        A_s0 = T.int32()
                                        A_s1 = T.int32()
                                        A_1 = T.match_buffer(
                                            weight_flatten_reindex_shared_dyn[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 16,
                                            ],
                                            (16, 16),
                                            "float16",
                                            strides=(A_s0, A_s1),
                                            scope="shared.dyn",
                                            offset_factor=16,
                                        )
                                        C_s0 = T.int32()
                                        C_s1 = T.int32()
                                        C = T.match_buffer(
                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
                                                v0_o * 16 : v0_o * 16 + 16,
                                                v1_o * 16 : v1_o * 16 + 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_s0 // 16 * (C_s0 // 16)
                                            + C.elem_offset % C_s0 // 16,
                                            T.tvm_access_ptr(
                                                T.type_annotation("float16"),
                                                A_1.data,
                                                A_1.elem_offset,
                                                A_s0 * 16,
                                                1,
                                            ),
                                            A_s0,
                                            "col_major",
                                        )
                                for x_0_2, y_0_2 in T.grid(2, 2):
                                    with T.sblock("Conv_update_o"):
                                        v_x_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + x_0_2)
                                        v_y_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + y_0_2)
                                        v_k_o = T.axis.reduce(720, k_0_0 * 4 + k_0_1)
                                        T.reads(
                                            Conv_reindex_wmma_accumulator[
                                                v_x_o * 16 : v_x_o * 16 + 16,
                                                v_y_o * 16 : v_y_o * 16 + 16,
                                            ],
                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
                                                v_x_o * 16 : v_x_o * 16 + 16,
                                                v_k_o * 16 : v_k_o * 16 + 16,
                                            ],
                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
                                                v_y_o * 16 : v_y_o * 16 + 16,
                                                v_k_o * 16 : v_k_o * 16 + 16,
                                            ],
                                        )
                                        T.writes(
                                            Conv_reindex_wmma_accumulator[
                                                v_x_o * 16 : v_x_o * 16 + 16,
                                                v_y_o * 16 : v_y_o * 16 + 16,
                                            ]
                                        )
                                        A_s0 = T.int32()
                                        A_s1 = T.int32()
                                        A_1 = T.match_buffer(
                                            data_im2col_reindex_shared_dyn_wmma_matrix_a[
                                                v_x_o * 16 : v_x_o * 16 + 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_s0 = T.int32()
                                        B_s1 = T.int32()
                                        B = T.match_buffer(
                                            weight_flatten_reindex_shared_dyn_wmma_matrix_b[
                                                v_y_o * 16 : v_y_o * 16 + 16,
                                                v_k_o * 16 : v_k_o * 16 + 16,
                                            ],
                                            (16, 16),
                                            "float16",
                                            strides=(B_s0, B_s1),
                                            scope="wmma.matrix_b",
                                            offset_factor=16,
                                        )
                                        C_s0 = T.int32()
                                        C_s1 = T.int32()
                                        C = T.match_buffer(
                                            Conv_reindex_wmma_accumulator[
                                                v_x_o * 16 : v_x_o * 16 + 16,
                                                v_y_o * 16 : v_y_o * 16 + 16,
                                            ],
                                            (16, 16),
                                            "float16",
                                            strides=(C_s0, C_s1),
                                            scope="wmma.accumulator",
                                            offset_factor=16,
                                        )
                                        T.tvm_mma_sync(
                                            C.data,
                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
                                            + C.elem_offset % C_s0 // 16,
                                            A_1.data,
                                            A_1.elem_offset // A_s0 // 16 * (A_s0 // 16)
                                            + A_1.elem_offset % A_s0 // 16,
                                            B.data,
                                            B.elem_offset // B_s0 // 16 * (B_s0 // 16)
                                            + B.elem_offset % B_s0 // 16,
                                            C.data,
                                            C.elem_offset // C_s0 // 16 * (C_s0 // 16)
                                            + C.elem_offset % C_s0 // 16,
                                        )
                        for ax0_0, ax1_0 in T.grid(2, 2):
                            with T.sblock("Conv_reindex_wmma.accumulator_o"):
                                v0_o = T.axis.spatial(32, x_0_0 * 4 + x_0_1 * 2 + ax0_0)
                                v1_o = T.axis.spatial(80, y_0_0 * 4 + y_0_1 * 2 + ax1_0)
                                T.reads(
                                    Conv_reindex_wmma_accumulator[
                                        v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
                                    ]
                                )
                                T.writes(
                                    Conv[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
                                )
                                A_s0 = T.int32()
                                A_s1 = T.int32()
                                A_1 = T.match_buffer(
                                    Conv_reindex_wmma_accumulator[
                                        v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
                                    ],
                                    (16, 16),
                                    "float16",
                                    strides=(A_s0, A_s1),
                                    scope="wmma.accumulator",
                                    offset_factor=16,
                                )
                                C_s0 = T.int32()
                                C_s1 = T.int32()
                                C = T.match_buffer(
                                    Conv[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16],
                                    (16, 16),
                                    "float16",
                                    strides=(C_s0, C_s1),
                                    offset_factor=16,
                                )
                                T.tvm_store_matrix_sync(
                                    A_1.data,
                                    16,
                                    16,
                                    16,
                                    A_1.elem_offset // A_s0 // 16 * (A_s0 // 16)
                                    + A_1.elem_offset % A_s0 // 16,
                                    T.tvm_access_ptr(
                                        T.type_annotation("float16"),
                                        C.data,
                                        C.elem_offset,
                                        C_s0 * 16,
                                        2,
                                    ),
                                    C_s0,
                                    "row_major",
                                )

    mod = tvm.IRModule.from_expr(complex_compute)
    with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
        tvm.compile(mod, target="cuda")
    generated_code = postproc_if_missing_async_support()
    # generated_code must contain "  setp.ne.b32 p, %0, 0;"
    assert "setp.ne.b32" in generated_code


def test_multiplication_nodes_are_inlined():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((32, 128), "float16")):
            tx = T.launch_thread("threadIdx.x", T.int64(32))
            A_flattened = T.decl_buffer((4096,), "float16", data=A.data)
            A_shared = T.decl_buffer([4096], "float16", scope="shared")

            T.attr("default", "async_scope", 1)
            for i in range(16):
                cse_v1: T.int64 = T.Cast("int64", i)
                A_shared[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8), T.int64(1), 8)] = (
                    A_flattened[T.Ramp(tx * T.int64(128) + cse_v1 * T.int64(8), T.int64(1), 8)]
                )
            T.ptx.cp_async.commit_group()
            T.ptx.cp_async.wait_group(0)

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def main(A: T.Buffer((32, 128), "float16")):
            tx = T.launch_thread("threadIdx.x", T.int64(32))
            A_flattened = T.decl_buffer((4096,), "float16", data=A.data)
            A_shared = T.decl_buffer((4096,), "float16", scope="shared")
            for i in range(16):
                cse_v1: T.int64 = T.Cast("int64", i)
                T.ptx.cp_async(
                    "float16",
                    A_shared.data,
                    tx * T.int64(128) + cse_v1 * T.int64(8),
                    A.data,
                    tx * T.int64(128) + cse_v1 * T.int64(8),
                    16,
                )
            T.ptx.cp_async.commit_group()
            T.ptx.cp_async.wait_group(0)

    After = tvm.s_tir.transform.InjectPTXAsyncCopy()(Before)
    tvm.ir.assert_structural_equal(After, Expected)


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