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

import tvm
import tvm.testing
from tvm.testing import env

pytest.importorskip("scipy")  # tvm.topi.testing imports scipy

import tvm.topi.testing
from tvm import relax
from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
from tvm.contrib.pickle_memoize import memoize
from tvm.relax.backend.cuda.cutlass import partition_for_cutlass
from tvm.relax.testing import (
    get_relax_attention_module,
    get_relax_matmul_module,
    get_relax_stacked_attention_module,
)
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import relax as relax_builder


@pytest.fixture(autouse=True)
def reset_seed():
    np.random.seed(0)


@tvm.script.ir_module
class Conv2dBiasReLU:
    @R.function
    def main(
        data: R.Tensor((16, 32, 32, 16), "float16"),
        weight: R.Tensor((32, 3, 3, 16), "float16"),
        bias: R.Tensor((1, 1, 1, 32), "float16"),
    ):
        with R.dataflow():
            conv1 = R.nn.relu(
                R.nn.conv2d(data, weight, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI")
                + bias,
            )
            R.output(conv1)

        return conv1


@tvm.script.ir_module
class Conv2dx2:
    @R.function
    def main(
        data: R.Tensor((16, 32, 32, 8), "float16"),
        weight1: R.Tensor((8, 3, 3, 8), "float16"),
        weight2: R.Tensor((8, 3, 3, 8), "float16"),
    ):
        with R.dataflow():
            conv1 = relax.op.nn.conv2d(
                data, weight1, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
            )
            conv2 = relax.op.nn.conv2d(
                conv1, weight2, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
            )
            R.output(conv2)

        return conv2


pytestmark = [
    pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass"),
]


def build_and_run(mod, inputs_np, target, legalize=True, cuda_graph=False):
    with tvm.transform.PassContext(
        config={
            "relax.backend.use_cuda_graph": cuda_graph,
            "relax.transform.apply_legalize_ops": legalize,
        }
    ):
        ex = tvm.compile(mod, target)

    def run_and_check():
        dev = tvm.device(target, 0)
        vm = relax.VirtualMachine(ex, dev)
        f = vm["main"]
        inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]

        # For cuda graph, run the compiled function twice to make sure that we can launch the
        # cached graph on the second run.
        if cuda_graph:
            f(*inputs)

        return f(*inputs).numpy()

    return tvm.testing.run_with_gpu_lock(run_and_check)


def build_cutlass(mod, assert_all_bindings_fused=True, num_final_bindings=1):
    mod = partition_for_cutlass(mod)

    if assert_all_bindings_fused:
        assert len(mod["main"].body.blocks[0].bindings) == num_final_bindings, (
            "Not all bindings are fused. " + str(mod["main"])
        )

    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80, "find_first_valid": True}})
    mod = codegen_pass(mod)
    return mod


def get_result_with_relax_cutlass_offload(
    mod, *args, assert_all_bindings_fused=True, num_final_bindings=1
):
    mod = build_cutlass(mod, assert_all_bindings_fused, num_final_bindings)
    return build_and_run(mod, args, "cuda")


def test_kernel_sharing():
    low, high = -1, 1
    data_np = np.random.randint(low, high, size=(16, 32, 32, 8)).astype("float16")
    weight1_np = np.random.randint(low, high, size=(8, 3, 3, 8)).astype("float16")
    weight2_np = np.random.randint(low, high, size=(8, 3, 3, 8)).astype("float16")

    out = get_result_with_relax_cutlass_offload(
        Conv2dx2, data_np, weight1_np, weight2_np, assert_all_bindings_fused=False
    )
    ref = build_and_run(Conv2dx2, [data_np, weight1_np, weight2_np], "llvm")

    np.testing.assert_equal(out, ref)


def get_relax_conv2d_module(
    data_shape,
    weight_shape,
    dtype,
    with_bias=False,
    activation=None,
    residual_bin_op=None,
    residual_activation=None,
):
    with IRBuilder() as builder:
        with relax_builder.function():
            R.func_name("main")
            data = R.arg("data", R.Tensor(data_shape, dtype))
            weight = R.arg("weight", R.Tensor(weight_shape, dtype))
            if with_bias:
                bias = R.arg("bias", R.Tensor((1, 1, 1, weight_shape[0]), dtype))

            with R.dataflow() as frame:
                output = R.emit(
                    R.nn.conv2d(
                        data,
                        weight,
                        out_dtype=dtype,
                        padding=(1, 1),
                        data_layout="NHWC",
                        kernel_layout="OHWI",
                    )
                )
                if with_bias:
                    output = R.emit(output + bias)
                if activation is not None:
                    output = R.emit(activation(output))
                if residual_bin_op is not None:
                    output = R.emit(residual_bin_op(output, data))
                    if residual_activation is not None:
                        output = R.emit(residual_activation(output))
                R.output(output)

            R.func_ret_value(frame.output_vars[0])

    func = builder.get()
    return tvm.IRModule({"main": func})


def _to_concrete_shape(symbolic_shape, var_table=None):
    if var_table is None:
        var_table = {}

    result = []
    for dim in symbolic_shape:
        if isinstance(dim, tuple):
            result.append(_to_concrete_shape(dim, var_table))
            continue

        if not isinstance(dim, tvm.tirx.expr.Var):
            result.append(dim)
            continue

        if dim not in var_table:
            var_table[dim] = np.random.randint(10, 50)
        result.append(var_table[dim])

    return tuple(result)


_vars = {
    "a": tvm.tirx.expr.Var("a", "int64"),
    "b": tvm.tirx.expr.Var("b", "int64"),
}


_epilogue_table = {
    "none": (False, None),
    "bias": (True, None),
    "relu": (True, R.nn.relu),
    "gelu": (True, R.nn.gelu),
    "silu": (True, R.nn.silu),
}


_residual_block_table = {
    "none": (None, None),
    "add_relu": (R.add, R.nn.relu),
    "mul_relu": (R.multiply, R.nn.relu),
    "add": (R.add, None),
    "mul": (R.multiply, None),
}


@pytest.mark.parametrize(
    "data_shape, weight_shape, dtype, epilogue, residual_block",
    [
        # Regular
        ((16, 32, 32, 16), (32, 3, 3, 16), "float16", "none", "none"),
        ((40, 128, 50, 16), (16, 2, 2, 16), "float16", "bias", "none"),
        ((3, 64, 64, 128), (32, 1, 1, 128), "float16", "relu", "none"),
        ((12, 32, 32, 16), (45, 5, 5, 16), "float16", "silu", "none"),
        # residual block
        ((3, 64, 64, 16), (16, 3, 3, 16), "float16", "relu", "add"),
        ((16, 32, 32, 16), (16, 3, 3, 16), "float16", "relu", "mul_relu"),
        ((40, 128, 50, 16), (16, 3, 3, 16), "float16", "bias", "add_relu"),
        ((128, 32, 32, 16), (16, 3, 3, 16), "float16", "silu", "mul"),
    ],
)
def test_conv2d_offload(data_shape, weight_shape, dtype, epilogue, residual_block):
    low, high = -1, 1
    data = np.random.randint(low, high, size=data_shape).astype(dtype)
    weight = np.random.randint(low, high, size=weight_shape).astype(dtype)
    bias = np.random.randint(low, high, size=(1, 1, 1, weight_shape[0])).astype(dtype)

    with_bias, activation = _epilogue_table[epilogue]
    residual_bin_op, residual_activation = _residual_block_table[residual_block]

    if with_bias:
        args = (data, weight, bias)
    else:
        args = (data, weight)

    mod = get_relax_conv2d_module(
        data_shape,
        weight_shape,
        dtype,
        with_bias=with_bias,
        activation=activation,
        residual_bin_op=residual_bin_op,
        residual_activation=residual_activation,
    )
    out = get_result_with_relax_cutlass_offload(mod, *args)

    ref = build_and_run(mod, args, "llvm")

    tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)


@pytest.mark.parametrize(
    "data_shape, weight_shape, dtype",
    [
        # batch dynamism
        ((T.Var("n", "int64"), 32, 32, 16), (32, 3, 3, 16), "float16"),
        # channel dynamism
        ((16, 32, 32, T.Var("c", "int64")), (32, 3, 3, T.Var("c", "int64")), "float16"),
    ],
)
def test_conv2d_dynamic(data_shape, weight_shape, dtype):
    # Create dynamic conv2d module.
    mod = get_relax_conv2d_module(
        data_shape,
        weight_shape,
        dtype,
    )
    # Attempt to offload to cutlass, should run without an error
    # but not offload due to incompatibility.
    mod = build_cutlass(mod)
    # Check that no cutlass call is introduced (until we support dynamism).
    assert "call_dps" not in str(mod.__repr__())


def test_cutlass_partition_conv2d_residual_blocked():
    @tvm.script.ir_module
    class Conv2dReLU:
        """
        This conv2d should not be fused as conv2d residual block, because both lhs and rhs of
        the last R.add depends on the result of conv2d.
        """

        @R.function
        def main(
            data: R.Tensor((32, 3, 3, 16), "float32"),
            weight: R.Tensor((16, 3, 3, 16), "float32"),
            bias: R.Tensor((1, 1, 1, 16), "float32"),
        ):
            with R.dataflow():
                conv1 = R.nn.conv2d(
                    data,
                    weight,
                    padding=(1, 1),
                    data_layout="NHWC",
                    kernel_layout="OHWI",
                )
                out = R.nn.relu(conv1 + bias)
                # residual depends on conv result, which cannot be handled in cutlass
                result = out + out
                R.output(result)

            return result

    mod = partition_for_cutlass(Conv2dReLU, annotate_codegen=False)
    for f_var in mod.functions:
        func = mod[f_var]
        if "Composite" in func.attrs:
            # verify that the function is not fused as residual block
            assert func.attrs["Composite"] == "cutlass.conv2d_bias_relu"


@pytest.mark.parametrize(
    "x_shape, y_shape, transpose_y, epilogue, residual_block",
    [
        # Regular
        ((32, 6), (6, 16), False, "none", "none"),
        ((_vars["a"], 6), (6, 16), False, "bias", "none"),
        # Transposed
        ((4, 16), (16, 128), True, "relu", "none"),
        ((35, 8), (8, 8), True, "gelu", "none"),
        # 3D x 3D
        ((6, 32, 8), (6, 8, 10), False, "bias", "none"),
        ((6, 32, 8), (6, 8, 10), True, "none", "none"),
        ((_vars["a"], 32, 8), (_vars["a"], 8, 10), True, "gelu", "none"),
        # 3D x 2D
        ((6, 32, 8), (8, 10), False, "none", "none"),
        ((_vars["a"], 32, 8), (8, 10), False, "bias", "none"),
        ((10, 16, 8), (8, 10), True, "relu", "none"),
        # 2D x 3D
        ((32, 8), (10, 8, 10), False, "relu", "none"),
        ((32, 8), (_vars["a"], 8, 10), True, "gelu", "none"),
        # ND x 2D
        ((3, 6, 32, 8), (8, 10), False, "bias", "none"),
        ((_vars["a"], _vars["b"], 6, 32, 8), (8, 10), False, "none", "none"),
        # 2D x ND
        ((32, 8), (5, 3, 8, 10), False, "gelu", "none"),
        # ND x ND
        ((5, 3, 32, 8), (5, 3, 8, 10), True, "relu", "none"),
        ((3, 2, 4, 16, 15), (1, 1, 15, 2), True, "gelu", "none"),
        ((1, 1, 16, 15), (3, 2, _vars["a"], 15, 2), False, "none", "none"),
        # Residual
        ((32, 8), (8, 8), False, "bias", "add"),
        ((4, 16), (16, 16), True, "relu", "add_relu"),
        ((8, 32, 8), (8, 8, 8), False, "bias", "add"),
        ((5, 3, 32, 8), (8, 8), True, "relu", "add"),
        # Residual fusion without bias - this is supported via the matmul + bias pattern
        # where bias == residual input
        ((4, 16), (16, 16), False, "none", "add"),
    ],
)
@pytest.mark.parametrize(
    "dtype",
    [
        "float16",
    ],
)
def test_matmul_offload(
    x_shape,
    y_shape,
    transpose_y,
    epilogue,
    residual_block,
    dtype,
):
    with_bias, activation = _epilogue_table[epilogue]
    var_table = {}
    concrete_x_shape = _to_concrete_shape(x_shape, var_table)
    concrete_y_shape = _to_concrete_shape(y_shape, var_table)
    x = np.random.randn(*concrete_x_shape).astype(dtype)
    y = np.random.randn(*concrete_y_shape).astype(dtype)

    if transpose_y:
        y = np.swapaxes(y, -2, -1)
        y_shape = (*y_shape[:-2], y_shape[-1], y_shape[-2])

    if with_bias:
        bias = np.random.randn(concrete_y_shape[-1]).astype(dtype)
        args = (x, y, bias)
    else:
        bias = None
        args = (x, y)

    residual_bin_op, residual_activation = _residual_block_table[residual_block]

    mod = get_relax_matmul_module(
        x_shape,
        y_shape,
        dtype,
        bias_shape=bias.shape if with_bias else None,
        transposed_y=transpose_y,
        activation=activation,
        residual_bin_op=residual_bin_op,
        residual_activation=residual_activation,
    )
    out = get_result_with_relax_cutlass_offload(mod, *args)
    ref = build_and_run(mod, args, "llvm")

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_matmul_with_3d_bias_offload():
    x_shape = (1, 4, 8)
    y_shape = (1, 8, 16)
    dtype = "float16"

    x = np.random.randn(*x_shape).astype(dtype)
    y = np.random.randn(*y_shape).astype(dtype)
    bias = np.random.randn(1, x_shape[-2], y_shape[-1]).astype(dtype)
    args = (x, y, bias)

    @tvm.script.ir_module
    class Mod:
        @R.function
        def main(
            x: R.Tensor((1, 4, 8), "float16"),
            y: R.Tensor((1, 8, 16), "float16"),
            bias: R.Tensor((1, 4, 16), "float16"),
        ):
            with R.dataflow():
                lv1 = R.matmul(x, y)
                gv1 = lv1 + bias
                R.output(gv1)

            return gv1

    out = get_result_with_relax_cutlass_offload(Mod, *args)
    ref = build_and_run(Mod, args, "llvm", legalize=True)

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@pytest.mark.parametrize(
    "x_shape, y_shape, expected",
    [
        # Regular matmul
        ((3, 4), (4, 5), True),
        # Batch matmul without stretching
        ((3, 16, 15), (3, 15, 2), True),
        ((_vars["a"], 16, 15), (_vars["a"], 15, 2), True),
        # Broadcast 2D to 3D
        ((3, 16, 15), (15, 2), True),
        ((_vars["a"], 16, 15), (15, 2), True),
        ((16, 15), (3, 15, 2), True),
        # Broadcast one-length dimension
        ((1, 16, 15), (3, 15, 2), True),
        ((3, 16, 15), (1, 15, 2), True),
        ((1, 1, 16, 15), (3, 2, 4, 15, 2), True),
        ((1, 1, 16, 15), (3, _vars["a"], 4, 15, 2), True),
        # ND x ND
        ((3, 2, 4, 16, 15), (3, 2, 4, 15, 2), True),
        ((_vars["a"], 2, 4, 16, 15), (_vars["a"], 2, 4, 15, 2), True),
        (
            (_vars["a"], _vars["b"], 4, 16, 15),
            (_vars["a"], _vars["b"], 4, 15, 2),
            True,
        ),
        # ND x ND with one-length dimension
        ((1, 2, 4, 16, 15), (1, 2, 4, 15, 2), True),
        ((3, 2, 1, 16, 15), (3, 2, 1, 15, 2), True),
        # Extra one-length dimension doesn't block broadcasting
        ((3, 2, 1, 16, 15), (1, 1, 3, 2, 1, 15, 2), True),
        # Not broadcasting all dims. Cannot be computed by stride-based batch gemm
        ((3, 1, 1, 16, 15), (3, 2, 4, 15, 2), False),
        ((3, 2, 4, 16, 15), (2, 4, 15, 2), False),
        # Different shape
        ((3, 4, 16, 15), (3, 2, 15, 2), False),
        ((3, _vars["a"], 16, 15), (3, _vars["b"], 15, 2), False),
        # Cannot prove that broadcast dimensions are equal
        ((_vars["a"], 16, 15), (3, 15, 2), False),
        ((3, _vars["a"], 1, 16, 15), (1, 1, 3, 2, 1, 15, 2), False),
        # Reduction axis must be constant
        ((3, _vars["a"]), (_vars["a"], 5), False),
    ],
)
def test_is_shape_valid_for_cutlass_matmul(x_shape, y_shape, expected):
    assert is_shape_valid_for_cutlass_matmul(x_shape, y_shape) == expected


@pytest.mark.parametrize(
    "x_shape, y_shape, transpose_y, dtype",
    [
        # Not broadcasting all dims. Cannot be computed by stride-based batch gemm
        ((3, 1, 1, 16, 15), (3, 2, 4, 15, 2), False, "float16"),
        ((3, 2, _vars["a"], 16, 15), (3, 2, 4, 15, 2), False, "float16"),
        ((1, 2, 1, 16, 15), (2, 1, 4, 15, 2), False, "float16"),
        ((3, 2, 4, 16, 15), (2, 4, 15, 2), True, "float16"),
        ((3, 16, 15), (2, 1, 3, 15, 2), True, "float16"),
        ((3, 16, 15), (_vars["a"], 1, 3, 15, 2), True, "float16"),
        ((_vars["a"], 1, 3, 16, 15), (_vars["b"], 1, 3, 15, 2), True, "float16"),
        ((_vars["a"], _vars["b"], 3, 16, 15), (_vars["a"], 1, 3, 15, 2), True, "float16"),
    ],
)
def test_cutlass_partition_matmul_blocked(x_shape, y_shape, transpose_y, dtype):
    if transpose_y:
        y_shape = (*y_shape[:-2], y_shape[-1], y_shape[-2])

    mod = get_relax_matmul_module(x_shape, y_shape, dtype, transposed_y=transpose_y)
    mod = partition_for_cutlass(mod)

    assert len(mod.functions) == 1


def test_cutlass_partition_matmul_tuple_return_blocked():
    @tvm.script.ir_module
    class TransposedMatmul:
        @R.function
        def main(
            x: R.Tensor((4, 4), "float32"),
            y: R.Tensor((4, 4), "float32"),
        ):
            with R.dataflow():
                lv1 = R.permute_dims(y)
                # Because lv1 is used by both lv2 and out, it should stay out of
                # the fused function. Otherwise the fused function will return
                # tuple output, which isn't possible in cutlass, e.g.
                # @R.function
                # def fused_relax_permute_dims_relax_matmul(...):
                #     R.func_attr({"Composite": "cutlass.matmul_transposed", "Primitive": True})
                #     with R.dataflow():
                #         gv: R.Tensor((4, 4), dtype="float32") = R.permute_dims(y, axes=None)
                #         gv1: R.Tensor((4, 4), dtype="float32") = R.matmul(x, gv)
                #         R.output(gv, gv1)
                #     return (gv, gv1)  # Cannot get `gv` if dispatch to cutlass kernel.
                lv2 = R.matmul(x, lv1)
                out = R.matmul(lv1, lv2)
                R.output(out)

            return out

    mod = partition_for_cutlass(TransposedMatmul, annotate_codegen=False)
    for f_var in mod.functions:
        func = mod[f_var]
        if "Composite" in func.attrs:
            # verify that the function is not fused as transposed matmul
            assert func.attrs["Composite"] == "cutlass.matmul"


def test_cutlass_partition_matmul_cyclic_dependency_blocked():
    @tvm.script.ir_module
    class Module:
        @R.function
        def main(x: R.Tensor((128, 128), "float16"), w: R.Tensor((128, 128), "float16")):
            with R.dataflow():
                # Because lv1 depends on lv, this block should be fused as matmul instead of matmul_bias.
                lv = R.matmul(x, w)
                lv1 = R.power(lv, R.const(2.0, "float16"))
                lv2 = R.add(lv, lv1)
                R.output(lv2)
            return lv2

    mod = partition_for_cutlass(Module, annotate_codegen=False)
    for f_var in mod.functions:
        func = mod[f_var]
        if "Composite" in func.attrs:
            assert func.attrs["Composite"] == "cutlass.matmul"


@pytest.fixture(params=["float16", "float32"])
def attention_dtype(request):
    return request.param


@pytest.fixture(
    params=[
        # B, S, N, H
        (32, (_vars["a"], 8), 16, (8, 8)),
        (32, (8, 8), 16, (8, 8)),
        (4, (16, 8), 32, (8, 8)),  # s != s_kv
        (4, (16, 8), 32, (8, 16)),  # h != h_v
        (32, (8, 8), 16, (4, 4)),  # h is not aligned
        (2, (8, 8), 8, (256, 256)),  # needs output accumulator buffer
    ]
)
def attention_size(request):
    return request.param


def get_numpy_attention_ref(
    b,
    s,
    s_kv,
    n,
    h,
    h_v,
    bias_shape,
    qk_scale,
    causal,
    dtype,
    window_size=None,
    num_kv_head=None,
):
    num_kv_head = num_kv_head or n
    q = np.random.randn(b, s, n, h).astype(dtype)
    k = np.random.randn(b, s_kv, num_kv_head, h).astype(dtype)
    v = np.random.randn(b, s_kv, num_kv_head, h_v).astype(dtype)
    if bias_shape == "none":
        bias = None
    else:
        bias = np.random.randn(*bias_shape).astype(dtype)

    ref = tvm.topi.testing.attention_python(
        q, k, v, bias, qk_scale, causal=causal, window_size=window_size, layout="BSNH"
    )

    return q, k, v, bias, ref


def test_attention_offload(attention_size, attention_dtype):
    b, (s, s_kv), n, (h, h_v) = attention_size
    concrete_s, concrete_s_kv = _to_concrete_shape((s, s_kv))
    q, k, v, _, ref = get_numpy_attention_ref(
        b, concrete_s, concrete_s_kv, n, h, h_v, "none", "none", "none", attention_dtype
    )

    q_shape = (b, s, n, h)
    k_shape = (b, s_kv, n, h)
    v_shape = (b, s_kv, n, h_v)

    mod = get_relax_attention_module(q_shape, k_shape, v_shape, dtype=attention_dtype)
    out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2)

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@pytest.fixture(
    params=[
        # B, S, N, H, bias_shape
        (4, (16, 8), 32, (8, 16), (4, 32, 16, 8)),
        (4, (16, 8), 32, (8, 16), (4, 1, 16, 8)),
        (4, (16, 8), 32, (8, 16), (4, 32, 1, 8)),
        (4, (16, 8), 32, (8, 16), (4, 1, 1, 8)),
        (4, (16, 8), 32, (8, 16), (1, 32, 16, 8)),
        (4, (16, 8), 32, (8, 16), (1, 1, 16, 8)),
        (4, (16, 8), 32, (8, 16), (1, 32, 1, 8)),
        (4, (16, 8), 32, (8, 16), (1, 1, 1, 8)),
    ]
)
def attention_bias_size(request):
    return request.param


def test_attention_bias_offload(attention_bias_size):
    b, (s, s_kv), n, (h, h_v), bias_shape = attention_bias_size
    concrete_s, concrete_s_kv, concrete_bias_shape = _to_concrete_shape((s, s_kv, bias_shape))

    q, k, v, bias, ref = get_numpy_attention_ref(
        b, concrete_s, concrete_s_kv, n, h, h_v, concrete_bias_shape, "none", "none", "float32"
    )

    q_shape = (b, s, n, h)
    k_shape = (b, s_kv, n, h)
    v_shape = (b, s_kv, n, h_v)

    mod = get_relax_attention_module(
        q_shape, k_shape, v_shape, bias_shape=bias_shape, dtype="float32"
    )
    out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2)

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@pytest.fixture(
    params=[
        # B, S, N, H, bias_shape
        (4, (16, 8), 32, (8, 16), (4, 32, 16, 8)),
        (4, (16, 8), 32, (8, 16), "none"),
    ]
)
def attention_scale_size(request):
    return request.param


@pytest.fixture(params=[0.01, 1e-8, -0.5, 1.23])
def attention_scale(request):
    return request.param


def test_attention_scale_offload(attention_scale_size, attention_scale):
    b, (s, s_kv), n, (h, h_v), bias_shape = attention_scale_size
    q, k, v, bias, ref = get_numpy_attention_ref(
        b, s, s_kv, n, h, h_v, bias_shape, attention_scale, "none", "float32"
    )

    q_shape = (b, s, n, h)
    k_shape = (b, s_kv, n, h)
    v_shape = (b, s_kv, n, h_v)

    mod = get_relax_attention_module(
        q_shape, k_shape, v_shape, dtype="float32", bias_shape=bias_shape, qk_scale=attention_scale
    )
    if bias is None:
        out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2)
    else:
        out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2)
    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@pytest.fixture(
    params=[
        # B, S, N, H, bias_shape
        (2, (16, 8), 4, (8, 16), "none"),
        (2, (8, 16), 4, (8, 16), "none"),
        (2, (16, 8), 4, (8, 16), (2, 4, 16, 8)),
    ]
)
def attention_causal_size(request):
    return request.param


@pytest.fixture(params=["TopLeft", "BottomRight"])
def attention_causal(request):
    return request.param


def test_attention_causal_offload(attention_causal_size, attention_causal):
    b, (s, s_kv), n, (h, h_v), bias_shape = attention_causal_size
    q, k, v, bias, ref = get_numpy_attention_ref(
        b, s, s_kv, n, h, h_v, bias_shape, "none", attention_causal, "float16"
    )

    q_shape = (b, s, n, h)
    k_shape = (b, s_kv, n, h)
    v_shape = (b, s_kv, n, h_v)

    mod = get_relax_attention_module(
        q_shape,
        k_shape,
        v_shape,
        dtype="float16",
        bias_shape=bias_shape,
        causal_mask=attention_causal,
    )

    if bias is None:
        out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2)
    else:
        out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2)
    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@memoize("topi.tests.test_codegen_cutlass.test_stacked_attention_offload")
def get_numpy_stacked_attention_ref(b, s, n, h, h_v, bias_shape, qk_scale, dtype):
    qkv = np.random.randn(b, s, n * h + n * h + n * h_v).astype(dtype)
    split_qkv = np.split(qkv, [n * h, n * h * 2], axis=2)
    q = np.reshape(split_qkv[0], (b, s, n, h))
    k = np.reshape(split_qkv[1], (b, s, n, h))
    v = np.reshape(split_qkv[2], (b, s, n, h_v))
    if not bias_shape == "none":
        bias = np.random.randn(*bias_shape).astype(dtype)
    else:
        bias = None
    ref = tvm.topi.testing.attention_python(
        q, k, v, bias, qk_scale, causal="none", window_size=None, layout="BSNH"
    )
    return qkv, bias, ref


@pytest.fixture(
    params=[
        # B, S, N, H, bias_shape, scale, single_shape
        (4, 8, 32, (64, 32), "none", "none", False),
        (4, 8, 32, (64, 32), (4, 32, 8, 8), 0.5, False),
        (4, 8, 32, (64, 64), "none", "none", True),
    ]
)
def stacked_attention_size(request):
    return request.param


def test_stacked_attention_split_offload(stacked_attention_size):
    b, s, n, (h, h_v), bias_shape, scale, single_shape = stacked_attention_size
    qkv, bias, ref = get_numpy_stacked_attention_ref(b, s, n, h, h_v, bias_shape, scale, "float16")
    if scale == "none":
        mod = get_relax_stacked_attention_module(
            qkv,
            b,
            s,
            n,
            h,
            h_v,
            "split",
            bias,
            single_shape=single_shape,
            layout="BS3NH",
        )
    else:
        mod = get_relax_stacked_attention_module(
            qkv,
            b,
            s,
            n,
            h,
            h_v,
            "split",
            bias,
            scale,
            single_shape=single_shape,
            layout="BS3NH",
        )

    if bias is None:
        out = get_result_with_relax_cutlass_offload(mod, qkv, num_final_bindings=2)
    else:
        out = get_result_with_relax_cutlass_offload(mod, qkv, bias, num_final_bindings=2)
    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_stacked_attention_strided_slice_offload(stacked_attention_size):
    b, s, n, (h, h_v), bias_shape, scale, single_shape = stacked_attention_size
    qkv, bias, ref = get_numpy_stacked_attention_ref(b, s, n, h, h_v, bias_shape, scale, "float32")
    if scale == "none":
        mod = get_relax_stacked_attention_module(
            qkv,
            b,
            s,
            n,
            h,
            h_v,
            "strided_slice",
            bias,
            single_shape=single_shape,
            layout="BS3NH",
        )
    else:
        mod = get_relax_stacked_attention_module(
            qkv,
            b,
            s,
            n,
            h,
            h_v,
            "strided_slice",
            bias,
            scale,
            single_shape=single_shape,
            layout="BS3NH",
        )
    if bias is None:
        out = get_result_with_relax_cutlass_offload(mod, qkv, num_final_bindings=2)
    else:
        out = get_result_with_relax_cutlass_offload(mod, qkv, bias, num_final_bindings=2)
    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


@pytest.fixture(
    params=[
        # B, S, N, H, bias_shape, scale
        (4, (16, 8), 32, (8, 16), "none", 0.5),
        (4, (16, 8), 32, (8, 16), (4, 32, 16, 8), 0.5),
        (4, (16, 8), "none", (8, 16), "none", 0.5),
        (4, (16, 8), "none", (8, 16), (4, 32, 16, 8), 0.5),
    ]
)
def attention_rewrite_size(request):
    return request.param


def get_relax_attention_rewrite_module(
    q_shape, k_shape, v_shape, out_shape, dtype, bias_shape=None, scale=None
):
    from tvm.script.ir_builder import IRBuilder
    from tvm.script.ir_builder import relax as relax_builder
    from tvm.script.ir_builder import tirx as T

    with IRBuilder() as builder:
        with relax_builder.function():
            R.func_name("main")
            q = R.arg("q", R.Tensor(q_shape, dtype))
            k = R.arg("k", R.Tensor(k_shape, dtype))
            v = R.arg("v", R.Tensor(v_shape, dtype))
            if bias_shape is not None:
                bias = R.arg("bias", R.Tensor(bias_shape, dtype))
            with R.dataflow() as frame:
                if len(q_shape) == 4:
                    q = R.emit(R.permute_dims(q, axes=[0, 2, 1, 3]))
                    q = R.emit(R.reshape(q, [q_shape[0] * q_shape[2], q_shape[1], q_shape[3]]))

                if len(k_shape) == 4:
                    k = R.emit(R.permute_dims(k, axes=[0, 2, 1, 3]))
                    k = R.emit(R.reshape(k, [k_shape[0] * k_shape[2], k_shape[1], k_shape[3]]))

                if len(v_shape) == 4:
                    v = R.emit(R.permute_dims(v, axes=[0, 2, 1, 3]))
                    v = R.emit(R.reshape(v, [v_shape[0] * v_shape[2], v_shape[1], v_shape[3]]))

                k = R.emit(R.permute_dims(k, axes=[0, 2, 1]))
                qk = R.emit(R.matmul(q, k))
                qk_scaled = R.emit(R.multiply(qk, R.const(scale, "float32")))
                if bias_shape is not None:
                    if len(bias_shape) == 4:
                        bias = R.emit(
                            R.reshape(bias, [bias_shape[0] * bias_shape[1], *bias_shape[2:]])
                        )
                    qk_added = R.emit(R.add(qk_scaled, bias))
                    softmax = R.emit(R.nn.softmax(qk_added, axis=-1))
                else:
                    softmax = R.emit(R.nn.softmax(qk_scaled, axis=-1))
                out = R.emit(R.matmul(softmax, v))

                if len(out_shape) == 4:
                    out = R.emit(
                        R.reshape(
                            out,
                            [out_shape[0], out_shape[2], out_shape[1], out_shape[3]],
                        )
                    )
                    out = R.emit(R.permute_dims(out, axes=[0, 2, 1, 3]))
                R.output(out)

            R.func_ret_value(frame.output_vars[0])

    original_func = builder.get()

    if scale is not None:
        scale = T.FloatImm("float32", scale)

    with IRBuilder() as builder:
        with relax_builder.function():
            R.func_name("main")
            q = R.arg("q", R.Tensor(q_shape, dtype))
            k = R.arg("k", R.Tensor(k_shape, dtype))
            v = R.arg("v", R.Tensor(v_shape, dtype))
            if bias_shape is not None:
                bias = R.arg("bias", R.Tensor(bias_shape, dtype))
            with R.dataflow() as frame:
                if len(q_shape) == 3:
                    q = R.emit(R.reshape(q, [q_shape[0], q_shape[1], 1, q_shape[2]]))

                if len(k_shape) == 3:
                    k = R.emit(R.reshape(k, [k_shape[0], k_shape[1], 1, k_shape[2]]))

                if len(v_shape) == 3:
                    v = R.emit(R.reshape(v, [v_shape[0], v_shape[1], 1, v_shape[2]]))

                if bias_shape is not None:
                    if len(bias_shape) == 4:
                        bias = R.emit(
                            R.reshape(
                                bias,
                                [
                                    bias_shape[0] * bias_shape[1],
                                    bias_shape[2],
                                    bias_shape[3],
                                ],
                            )
                        )
                        bias = R.emit(
                            R.reshape(
                                bias,
                                [
                                    bias_shape[0],
                                    bias_shape[1],
                                    bias_shape[2],
                                    bias_shape[3],
                                ],
                            )
                        )
                    elif len(bias_shape) == 3:
                        bias = R.emit(
                            R.reshape(bias, [bias_shape[0], 1, bias_shape[1], bias_shape[2]])
                        )
                else:
                    bias = None
                out = R.emit(R.nn.attention(q, k, v, bias, scale))

                if len(out_shape) == 3:
                    out = R.emit(R.reshape(out, [out_shape[0], out_shape[1], out_shape[2]]))
                R.output(out)

            R.func_ret_value(frame.output_vars[0])

    expected_func = builder.get()
    return tvm.IRModule({"main": original_func}), tvm.IRModule({"main": expected_func})


def get_numpy_attention_input(q_shape, k_shape, v_shape, bias_shape, dtype):
    q = np.random.randn(*q_shape).astype(dtype)
    k = np.random.randn(*k_shape).astype(dtype)
    v = np.random.randn(*v_shape).astype(dtype)
    if not bias_shape == "none":
        bias = np.random.randn(*bias_shape).astype(dtype)
    else:
        bias = None
    return q, k, v, bias


def test_attention_rewrite_offload(attention_rewrite_size):
    b, (s, s_kv), n, (h, h_v), bias_shape, scale = attention_rewrite_size
    q_shape = [b, s, n, h] if n != "none" else [b, s, h]
    k_shape = [b, s_kv, n, h] if n != "none" else [b, s_kv, h]
    v_shape = [b, s_kv, n, h_v] if n != "none" else [b, s_kv, h_v]
    out_shape = [b, s, n, h_v] if n != "none" else [b, s, h_v]
    bias_shape = [b, n, s, s_kv] if n != "none" else [b, s, s_kv]
    q, k, v, bias = get_numpy_attention_input(q_shape, k_shape, v_shape, bias_shape, "float32")
    original_mod, expected_mod = get_relax_attention_rewrite_module(
        q_shape, k_shape, v_shape, out_shape, "float32", bias_shape, scale
    )
    original_mod = partition_for_cutlass(original_mod, True)
    expected_mod = partition_for_cutlass(expected_mod, True)
    tvm.ir.assert_structural_equal(original_mod, expected_mod, True)

    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80, "find_first_valid": True}})
    original_mod = codegen_pass(original_mod)
    expected_mod = codegen_pass(expected_mod)
    if bias is None:
        original_out = build_and_run(original_mod, [q, k, v], "cuda")
        expected_out = build_and_run(expected_mod, [q, k, v], "cuda")
        tvm.testing.assert_allclose(original_out, expected_out, rtol=1e-5, atol=1e-5)
    else:
        original_out = build_and_run(original_mod, [q, k, v, bias], "cuda", legalize=False)
        expected_out = build_and_run(expected_mod, [q, k, v, bias], "cuda", legalize=False)
        tvm.testing.assert_allclose(original_out, expected_out, rtol=1e-5, atol=1e-5)


def test_conv2d_residual_broadcast():
    data_shape = (2, 64, 64, 8)
    weight_shape = (8, 3, 3, 8)
    dtype = "float16"

    def get_mod(residual_batch):
        with IRBuilder() as builder:
            with relax_builder.function():
                R.func_name("main")
                data = R.arg("data", R.Tensor(data_shape, dtype))
                weight = R.arg("weight", R.Tensor(weight_shape, dtype))
                bias = R.arg("bias", R.Tensor((1, 1, weight_shape[0]), dtype))
                residual = R.arg(
                    "residual", R.Tensor((residual_batch, 1, 1, weight_shape[0]), dtype)
                )

                with R.dataflow() as frame:
                    output = R.emit(
                        R.nn.conv2d(
                            data,
                            weight,
                            out_dtype=dtype,
                            padding=(1, 1),
                            data_layout="NHWC",
                            kernel_layout="OHWI",
                        )
                    )
                    output = R.emit(output + bias)
                    output = R.emit(R.nn.relu(output))
                    output = R.emit(R.add(output, residual))
                    R.output(output)

                R.func_ret_value(frame.output_vars[0])

        func = builder.get()
        return tvm.IRModule({"main": func})

    low = -1
    high = 1

    residual_batch = 1
    mod = get_mod(residual_batch)
    data = np.random.randint(low, high, size=data_shape).astype(dtype)
    weight = np.random.randint(low, high, size=weight_shape).astype(dtype)
    bias = np.random.randint(low, high, size=(1, 1, weight_shape[0])).astype(dtype)
    bias2 = np.random.randint(low, high, size=(residual_batch, 1, 1, weight_shape[0])).astype(dtype)

    args = [data, weight, bias, bias2]
    out = get_result_with_relax_cutlass_offload(mod, *args)
    ref = build_and_run(mod, args, "llvm")
    tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)


@pytest.mark.parametrize(
    "data_shape, dtype, axes",
    [
        ((2, 128, 64), "float16", [-1]),
        ((128, 30), "float32", [-1]),
        ((2, 128, 64), "float32", [1]),
        ((2, 128, 64), "float32", [1, 2]),
    ],
)
def test_layer_norm(data_shape, dtype, axes):
    def get_mod(data_shape, dtype, axes):
        reduced_shape = [data_shape[axis] for axis in axes]
        with IRBuilder() as builder:
            with relax_builder.function():
                R.func_name("main")
                inp = R.arg("input", R.Tensor(data_shape, dtype))
                gamma = R.arg("gamma", R.Tensor(reduced_shape, dtype))
                beta = R.arg("beta", R.Tensor(reduced_shape, dtype))

                with R.dataflow() as frame:
                    output = R.emit(R.nn.layer_norm(inp, gamma, beta, axes))
                    R.output(output)

                R.func_ret_value(frame.output_vars[0])

        func = builder.get()
        return tvm.IRModule({"main": func})

    Module = get_mod(data_shape, dtype, axes)
    mod = partition_for_cutlass(Module)

    if len(axes) != 1 or (axes[0] != -1 and axes[0] != len(data_shape) - 1):
        tvm.ir.assert_structural_equal(mod, Module)
        return

    mod = relax.transform.RunCodegen()(mod)

    inp = np.random.randn(*data_shape).astype(dtype)
    gamma = np.random.randn(data_shape[-1]).astype(dtype)
    beta = np.random.randn(data_shape[-1]).astype(dtype)
    out = build_and_run(mod, [inp, gamma, beta], "cuda")
    ref = build_and_run(Module, [inp, gamma, beta], "llvm")

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_attention_rewrite_fp16():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            q: R.Tensor((4, 16, 32, 8), dtype="float16"),
            k: R.Tensor((4, 8, 32, 8), dtype="float16"),
            v: R.Tensor((4, 8, 32, 16), dtype="float16"),
            bias: R.Tensor((4, 32, 16, 8), dtype="float16"),
        ) -> R.Tensor((4, 16, 32, 16), dtype="float16"):
            R.func_attr({"num_input": 4})
            with R.dataflow():
                lv = R.permute_dims(q, axes=[0, 2, 1, 3])
                lv1 = R.reshape(lv, R.shape([128, 16, 8]))
                lv2 = R.permute_dims(k, axes=[0, 2, 1, 3])
                lv3 = R.reshape(lv2, R.shape([128, 8, 8]))
                lv4 = R.permute_dims(v, axes=[0, 2, 1, 3])
                lv5 = R.reshape(lv4, R.shape([128, 8, 16]))
                lv6 = R.permute_dims(lv3, axes=[0, 2, 1])
                lv7 = R.matmul(lv1, lv6, out_dtype="float16")
                lv3_1 = R.astype(R.const(0.5, "float32"), dtype="float16")
                lv8 = R.multiply(lv7, lv3_1)
                lv9 = R.reshape(bias, R.shape([128, 16, 8]))
                lv10 = R.add(lv8, lv9)
                lv10_fp16 = R.astype(lv10, dtype="float16")
                lv11 = R.nn.softmax(lv10_fp16, axis=2)
                lv5_1 = R.astype(lv11, dtype="float16")
                lv12 = R.matmul(lv5_1, lv5, out_dtype="float16")
                lv13 = R.reshape(lv12, R.shape([4, 32, 16, 16]))
                lv6_1 = R.permute_dims(lv13, axes=[0, 2, 1, 3])
                lv14 = R.astype(lv6_1, dtype="float32")
                R.output(lv14)
            return lv14

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def fused_relax_nn_attention_bias_cutlass1(
            q: R.Tensor((4, 16, 32, 8), dtype="float16"),
            k: R.Tensor((4, 8, 32, 8), dtype="float16"),
            v: R.Tensor((4, 8, 32, 16), dtype="float16"),
            lv1: R.Tensor((4, 32, 16, 8), dtype="float16"),
            workspace: R.Tensor((65536,), dtype="uint8"),
        ) -> R.Tensor((4, 16, 32, 16), dtype="float16"):
            R.func_attr(
                {
                    "Codegen": "cutlass",
                    "WorkspaceSize": T.int64(65536),
                    "global_symbol": "fused_relax_nn_attention_bias_cutlass1",
                }
            )

            @R.function
            def gv_1(
                q_1: R.Tensor((4, 16, 32, 8), dtype="float16"),
                k_1: R.Tensor((4, 8, 32, 8), dtype="float16"),
                v_1: R.Tensor((4, 8, 32, 16), dtype="float16"),
                lv1_1: R.Tensor((4, 32, 16, 8), dtype="float16"),
                workspace_1: R.Tensor((65536,), dtype="uint8"),
            ) -> R.Tensor((4, 16, 32, 16), dtype="float16"):
                R.func_attr(
                    {
                        "Composite": "cutlass.attention_bias",
                        "WorkspaceSize": T.int64(65536),
                    }
                )
                with R.dataflow():
                    gv_2 = R.nn.attention(
                        q_1, k_1, v_1, lv1_1, scale=T.float32(0.5), causal_mask=None
                    )
                    R.output(gv_2)
                return gv_2

            gv1: R.Tensor((4, 16, 32, 16), dtype="float16") = gv_1(q, k, v, lv1, workspace)
            return gv1

        @R.function
        def main(
            q: R.Tensor((4, 16, 32, 8), dtype="float16"),
            k: R.Tensor((4, 8, 32, 8), dtype="float16"),
            v: R.Tensor((4, 8, 32, 16), dtype="float16"),
            bias: R.Tensor((4, 32, 16, 8), dtype="float16"),
        ) -> R.Tensor((4, 16, 32, 16), dtype="float32"):
            R.func_attr({"num_input": 4})
            cls = Expected
            with R.dataflow():
                workspace_main = R.builtin.alloc_tensor(
                    R.shape([65536]), R.dtype("uint8"), R.prim_value(0)
                )
                lv_1 = R.reshape(bias, R.shape([128, 16, 8]))
                lv1 = R.reshape(lv_1, R.shape([4, 32, 16, 8]))
                lv_2 = cls.fused_relax_nn_attention_bias_cutlass1(q, k, v, lv1, workspace_main)
                lv14 = R.astype(lv_2, dtype="float32")
                R.output(lv14)
            return lv14

    mod = partition_for_cutlass(Module)
    tvm.ir.assert_structural_equal(mod, Expected)


def split_transform_deploy_mod(mod):
    mod_transform = tvm.IRModule()
    mod_deploy = tvm.IRModule().with_attrs(mod.attrs)

    transform_func_name = None

    for gv, func in mod.functions.items():
        if "transform_params" in gv.name_hint:
            transform_func_name = gv.name_hint
            mod_transform[gv] = func
        elif isinstance(func, tvm.tirx.PrimFunc):
            mod_transform[gv] = func
        else:
            mod_deploy[gv] = func

    assert transform_func_name is not None
    return mod_transform, mod_deploy, transform_func_name


def test_fp16A_int4B_gemm():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def decode(
            A: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            B: T.Buffer((T.int64(128),), "float16"),
            decode_1: T.Buffer((T.int64(64), T.int64(128)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i, j in T.grid(T.int64(64), T.int64(128)):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(A[v_i, v_j // T.int64(2)], B[v_j])
                    T.writes(decode_1[v_i, v_j])
                    decode_1[v_i, v_j] = (
                        T.Cast(
                            "float16",
                            T.shift_right(
                                T.shift_left(
                                    T.bitwise_and(
                                        T.shift_right(
                                            T.Cast("int32", A[v_i, v_j // T.int64(2)]),
                                            T.Cast("int32", v_j % T.int64(2)) * 4,
                                        ),
                                        15,
                                    ),
                                    28,
                                ),
                                28,
                            ),
                        )
                        * B[v_j]
                    )

        @T.prim_func(s_tir=True)
        def encode(
            A: T.Buffer((T.int64(128), T.int64(64)), "float16"),
            w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            compute: T.Buffer((T.int64(128),), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            max_abs_value = T.sblock_alloc_buffer((T.int64(128),), "float16")
            scale = T.sblock_alloc_buffer((T.int64(128),))
            for i, k in T.grid(T.int64(128), T.int64(64)):
                with T.sblock("max_abs_value"):
                    v_i, v_k = T.axis.remap("SR", [i, k])
                    T.reads(A[v_i, v_k])
                    T.writes(max_abs_value[v_i])
                    with T.init():
                        max_abs_value[v_i] = T.float16(-65504)
                    max_abs_value[v_i] = T.max(max_abs_value[v_i], T.fabs(A[v_i, v_k]))
            for i in range(T.int64(128)):
                with T.sblock("scale"):
                    v_i = T.axis.spatial(T.int64(128), i)
                    T.reads(max_abs_value[v_i])
                    T.writes(scale[v_i])
                    scale[v_i] = T.max(
                        T.Cast("float32", max_abs_value[v_i]), T.float32(0.0001)
                    ) * T.float32(0.125)
            for j, i, k in T.grid(T.int64(64), T.int64(64), T.int64(2)):
                with T.sblock("w_gathered"):
                    v_j, v_i, v_k = T.axis.remap("SSR", [j, i, k])
                    T.reads(A[v_i * T.int64(2) + v_k, v_j], scale[v_i * T.int64(2) + v_k])
                    T.writes(w_gathered[v_j, v_i])
                    with T.init():
                        w_gathered[v_j, v_i] = T.int8(0)
                    w_gathered[v_j, v_i] = T.bitwise_or(
                        w_gathered[v_j, v_i],
                        T.if_then_else(
                            v_i * T.int64(2) + v_k < T.int64(128),
                            T.shift_left(
                                T.bitwise_and(
                                    T.Cast(
                                        "int8",
                                        T.min(
                                            T.max(
                                                T.round(
                                                    T.Cast(
                                                        "float32", A[v_i * T.int64(2) + v_k, v_j]
                                                    )
                                                    / scale[v_i * T.int64(2) + v_k]
                                                ),
                                                T.float32(-8),
                                            ),
                                            T.float32(7),
                                        ),
                                    ),
                                    T.int8(15),
                                ),
                                T.Cast("int8", v_k) * T.int8(4),
                            ),
                            T.int8(0),
                        ),
                    )
            for i0 in range(T.int64(128)):
                with T.sblock("compute"):
                    v_i0 = T.axis.spatial(T.int64(128), i0)
                    T.reads(scale[v_i0])
                    T.writes(compute[v_i0])
                    compute[v_i0] = T.Cast("float16", scale[v_i0])

        @R.function
        def main_bias(
            x: R.Tensor((64, 64), dtype="float16"),
            y: R.Tensor((128, 64), dtype="float16"),
            bias: R.Tensor((1, 128), dtype="float16"),
        ) -> R.Tensor((64, 128), dtype="float16"):
            R.func_attr({"num_input": 1})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[R.Tensor((64, 64), dtype="int8"), R.Tensor((128,), dtype="float16")],
                )
                lv1 = lv[0]
                lv2 = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    80,
                    True,
                    ty_args=(R.Tensor((64, 64), dtype="int8"),),
                )
                lv3: R.Tensor((128,), dtype="float16") = lv[1]
                lv6 = R.call_tir(
                    cls.decode, (lv2, lv3), out_ty=R.Tensor((64, 128), dtype="float16")
                )
                lv1_1: R.Tensor((64, 128), dtype="float16") = R.matmul(x, lv6, out_dtype="float16")
                lv2_1: R.Tensor((64, 128), dtype="float16") = R.add(lv1_1, bias)
                R.output(lv2_1)
            return lv2_1

        @R.function
        def main_cast_bias(
            x: R.Tensor((64, 64), dtype="float16"),
            y: R.Tensor((128, 64), dtype="float16"),
            bias: R.Tensor((1, 128), dtype="float16"),
        ) -> R.Tensor((64, 128), dtype="float16"):
            R.func_attr({"num_input": 1})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[R.Tensor((64, 64), dtype="int8"), R.Tensor((128,), dtype="float16")],
                )
                lv1 = lv[0]
                lv2 = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    80,
                    True,
                    ty_args=(R.Tensor((64, 64), dtype="int8"),),
                )
                lv3: R.Tensor((128,), dtype="float16") = lv[1]
                lv6 = R.call_tir(
                    cls.decode, (lv2, lv3), out_ty=R.Tensor((64, 128), dtype="float16")
                )
                lv1_1: R.Tensor((64, 128), dtype="float32") = R.matmul(x, lv6, out_dtype="float32")
                cast: R.Tensor((64, 128), dtype="float16") = R.astype(lv1_1, dtype="float16")
                lv2_1: R.Tensor((64, 128), dtype="float16") = R.add(cast, bias)
                R.output(lv2_1)
            return lv2_1

        @R.function
        def main_residual(
            x: R.Tensor((64, 64), dtype="float16"),
            residual: R.Tensor((64, 128), dtype="float16"),
            y: R.Tensor((128, 64), dtype="float16"),
            bias: R.Tensor((1, 128), dtype="float16"),
        ) -> R.Tensor((64, 128), dtype="float16"):
            R.func_attr({"num_input": 2})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[R.Tensor((64, 64), dtype="int8"), R.Tensor((128,), dtype="float16")],
                )
                lv1 = lv[0]
                lv2 = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    80,
                    True,
                    ty_args=(R.Tensor((64, 64), dtype="int8"),),
                )
                lv3: R.Tensor((128,), dtype="float16") = lv[1]
                lv6 = R.call_tir(
                    cls.decode, (lv2, lv3), out_ty=R.Tensor((64, 128), dtype="float16")
                )
                lv1_1: R.Tensor((64, 128), dtype="float16") = R.matmul(x, lv6, out_dtype="float16")
                lv2_1: R.Tensor((64, 128), dtype="float16") = R.add(lv1_1, bias)
                lv3_1: R.Tensor((64, 128), dtype="float16") = R.add(lv2_1, residual)
                R.output(lv3_1)
            return lv3_1

    x_shape = (64, 64)
    y_shape = (128, 64)

    mod = partition_for_cutlass(Module)
    func_names = [name.name_hint for (name, _) in mod.functions.items()]
    assert "fused_decode_relax_matmul_relax_add_cutlass" in func_names
    assert "fused_decode_relax_matmul_relax_add_relax_add_cutlass" in func_names
    assert "fused_decode_relax_matmul_relax_astype_relax_add_cutlass" in func_names

    mod = relax.transform.RunCodegen(
        {"cutlass": {"sm": 80, "find_first_valid": False}},
        entry_functions=["main_bias", "main_residual", "main_cast_bias"],
    )(mod)

    x = np.random.randn(*x_shape).astype("float16")
    y = np.random.normal(0, 0.002, size=y_shape).astype("float16")
    bias = np.random.randn(1, y_shape[0]).astype("float16")
    residual = np.random.randn(x_shape[0], y_shape[0]).astype("float16")

    mod = relax.pipeline.get_pipeline()(mod)
    mod = relax.transform.LiftTransformParams()(mod)

    mod_transform, mod_deploy, transform_func_name = split_transform_deploy_mod(mod)

    ex = tvm.compile(mod_transform, target="llvm")
    vm = relax.vm.VirtualMachine(ex, tvm.cpu(0))

    packed_weight, scales, bias_trans = vm[transform_func_name](
        (tvm.runtime.tensor(y), tvm.runtime.tensor(bias))
    )

    ex_cuda = tvm.compile(mod_deploy, target="cuda")

    def run_and_check():
        dev = tvm.device("cuda", 0)
        vm = relax.vm.VirtualMachine(ex_cuda, dev)
        x_nd = tvm.runtime.tensor(x, dev)
        residual_nd = tvm.runtime.tensor(residual, dev)
        params = [packed_weight.copyto(dev), scales.copyto(dev), bias_trans.copyto(dev)]

        for f_name in ["main_bias", "main_cast_bias", "main_residual"]:
            with_residual = "residual" in f_name

            if with_residual:
                inp = [x_nd, residual_nd] + params
            else:
                inp = [x_nd] + params

            out = vm[f_name](*inp).numpy()
            ref = np.dot(x, y.transpose()) + bias
            if with_residual:
                ref += residual
            tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_fp16A_int8B_gemm():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def decode(
            A: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            B: T.Buffer((T.int64(64),), "float16"),
            decode_1: T.Buffer((T.int64(64), T.int64(64)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i, j in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(A[v_i, v_j], B[v_j])
                    T.writes(decode_1[v_i, v_j])
                    decode_1[v_i, v_j] = T.Cast("float16", A[v_i, v_j]) * B[v_j]

        @T.prim_func(s_tir=True)
        def encode(
            A: T.Buffer((T.int64(64), T.int64(64)), "float16"),
            w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            compute: T.Buffer((T.int64(64),), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            max_abs_value = T.sblock_alloc_buffer((T.int64(64),), "float16")
            scale = T.sblock_alloc_buffer((T.int64(64),))
            for i, k in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("max_abs_value"):
                    v_i, v_k = T.axis.remap("SR", [i, k])
                    T.reads(A[v_i, v_k])
                    T.writes(max_abs_value[v_i])
                    with T.init():
                        max_abs_value[v_i] = T.float16(-65504)
                    max_abs_value[v_i] = T.max(max_abs_value[v_i], T.fabs(A[v_i, v_k]))
            for i in range(T.int64(64)):
                with T.sblock("scale"):
                    v_i = T.axis.spatial(T.int64(64), i)
                    T.reads(max_abs_value[v_i])
                    T.writes(scale[v_i])
                    scale[v_i] = T.max(
                        T.Cast("float32", max_abs_value[v_i]), T.float32(0.0001)
                    ) * T.float32(0.0078125)
            for j, i in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("w_gathered"):
                    v_j, v_i = T.axis.remap("SS", [j, i])
                    T.reads(A[v_i, v_j], scale[v_i])
                    T.writes(w_gathered[v_j, v_i])
                    w_gathered[v_j, v_i] = T.Cast(
                        "int8",
                        T.min(
                            T.max(
                                T.round(T.Cast("float32", A[v_i, v_j]) / scale[v_i]),
                                T.float32(-128),
                            ),
                            T.float32(127),
                        ),
                    )
            for i0 in range(T.int64(64)):
                with T.sblock("compute"):
                    v_i0 = T.axis.spatial(T.int64(64), i0)
                    T.reads(scale[v_i0])
                    T.writes(compute[v_i0])
                    compute[v_i0] = T.Cast("float16", scale[v_i0])

        @R.function
        def main(
            x: R.Tensor((64, 64), dtype="float16"),
            y: R.Tensor((64, 64), dtype="float16"),
            bias: R.Tensor((64, 64), dtype="float16"),
        ) -> R.Tensor((64, 64), dtype="float16"):
            R.func_attr({"num_input": 1})
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[R.Tensor((64, 64), dtype="int8"), R.Tensor((64,), dtype="float16")],
                )
                lv1: R.Tensor((64, 64), dtype="int8") = lv[0]
                lv2: R.Tensor((64, 64), dtype="int8") = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    R.prim_value(80),
                    R.prim_value(0),
                    ty_args=(R.Tensor((64, 64), dtype="int8"),),
                )
                lv3: R.Tensor((64,), dtype="float16") = lv[1]
                lv4: R.Tensor((64, 64), dtype="int8") = R.builtin.stop_lift_params(lv2)
                lv5: R.Tensor((64,), dtype="float16") = R.builtin.stop_lift_params(lv3)
                lv6 = R.call_tir(cls.decode, (lv4, lv5), out_ty=R.Tensor((64, 64), dtype="float16"))
                lv1_1: R.Tensor((64, 64), dtype="float16") = R.matmul(x, lv6, out_dtype="float16")
                lv2_1: R.Tensor((64, 128), dtype="float16") = R.add(lv1_1, bias)
                lv2_2: R.Tensor((64, 128), dtype="float16") = R.nn.gelu(lv2_1)
                R.output(lv2_2)
            return lv2_2

    x_shape = (64, 64)
    y_shape = (64, 64)

    mod = partition_for_cutlass(Module)
    func_names = [name.name_hint for (name, _) in mod.functions.items()]
    assert "fused_decode_relax_matmul_relax_add_relax_nn_gelu_cutlass" in func_names

    mod = relax.transform.RunCodegen(
        {"cutlass": {"sm": 80, "find_first_valid": False}},
    )(mod)

    x = np.random.randn(*x_shape).astype("float16")
    y = np.random.normal(0, 0.002, size=y_shape).astype("float16")
    bias = np.random.randn(x_shape[0], y_shape[0]).astype("float16")

    mod = relax.pipeline.get_pipeline()(mod)
    mod = relax.transform.LiftTransformParams()(mod)

    mod_transform, mod_deploy, transform_func_name = split_transform_deploy_mod(mod)

    ex = tvm.compile(mod_transform, target="llvm")
    vm = relax.vm.VirtualMachine(ex, tvm.cpu(0))

    packed_weight, scales, bias_trans = vm[transform_func_name](
        (tvm.runtime.tensor(y), tvm.runtime.tensor(bias))
    )

    ex_cuda = tvm.compile(mod_deploy, target="cuda")

    def gelu_fp16(x):
        erf_inp = x * (0.5**0.5)
        from scipy.special import erf

        erf_out = erf(erf_inp.astype("float32")).astype("float16")
        return x * 0.5 * (1.0 + erf_out)

    def run_and_check():
        dev = tvm.device("cuda", 0)
        vm = relax.vm.VirtualMachine(ex_cuda, dev)
        x_nd = tvm.runtime.tensor(x, dev)
        inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev), bias_trans.copyto(dev)]
        out = vm["main"](*inp).numpy()
        ref = gelu_fp16(np.dot(x, y.transpose()) + bias)
        tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_rms_norm():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def rms_norm(
            A: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"),
            B: T.Buffer((T.int64(4096),), "float16"),
            rms_norm: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            Ared_temp = T.sblock_alloc_buffer((T.int64(1), T.int64(1)))
            for bsz, i, k in T.grid(T.int64(1), T.int64(1), T.int64(4096)):
                with T.sblock("Ared_temp"):
                    v_bsz, v_i, v_k = T.axis.remap("SSR", [bsz, i, k])
                    T.reads(A[v_bsz, v_i, v_k])
                    T.writes(Ared_temp[v_bsz, v_i])
                    with T.init():
                        Ared_temp[v_bsz, v_i] = T.float32(0)
                    Ared_temp[v_bsz, v_i] = Ared_temp[v_bsz, v_i] + T.Cast(
                        "float32", A[v_bsz, v_i, v_k]
                    ) * T.Cast("float32", A[v_bsz, v_i, v_k])
            for bsz, i, k in T.grid(T.int64(1), T.int64(1), T.int64(4096)):
                with T.sblock("rms_norm"):
                    v_bsz, v_i, v_k = T.axis.remap("SSS", [bsz, i, k])
                    T.reads(B[v_k], A[v_bsz, v_i, v_k], Ared_temp[v_bsz, v_i])
                    T.writes(rms_norm[v_bsz, v_i, v_k])
                    rms_norm[v_bsz, v_i, v_k] = T.Cast(
                        "float16",
                        T.Cast("float32", B[v_k])
                        * (
                            T.Cast("float32", A[v_bsz, v_i, v_k])
                            / T.sqrt(
                                Ared_temp[v_bsz, v_i] * T.float32(0.000244140625)
                                + T.float32(9.9999999999999995e-07)
                            )
                        ),
                    )

        @R.function
        def main(
            input: R.Tensor((1, 1, 4096), dtype="float16"),
            weight: R.Tensor((4096,), dtype="float16"),
        ) -> R.Tensor((1, 1, 4096), dtype="float16"):
            cls = Module
            with R.dataflow():
                lv = R.call_tir(
                    cls.rms_norm, (input, weight), out_ty=R.Tensor((1, 1, 4096), dtype="float16")
                )
                R.output(lv)
            return lv

    data_shape = (1, 1, 4096)
    dtype = "float16"
    mod = partition_for_cutlass(Module)

    # TODO(@tvm-team): This is temporary patch.Currently, the remaining packed function triggers error since it is not scheduled.
    # This is because RunCodegen does not support PrimFunc well yet.
    # i.e., it does remove the global symbol of PrimFunc, which would be no longer used,
    # and thus, the following DCE cannot remove this. Revisit when resolved.
    with tvm.target.Target("cuda"):
        mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)

    mod = relax.transform.RunCodegen(
        {"cutlass": {"rms_eps": 1e-6}},
    )(mod)

    inp = np.random.randn(*data_shape).astype(dtype)
    weight = np.random.randn(data_shape[-1]).astype(dtype)
    out = build_and_run(mod, [inp, weight], "cuda")
    ref = build_and_run(Module, [inp, weight], "llvm", legalize=True)

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_conv2d_cuda_graph():
    @tvm.script.ir_module
    class Conv2d:
        @R.function
        def main(
            data: R.Tensor((16, 32, 32, 16), "float16"),
            weight1: R.Tensor((16, 3, 3, 16), "float16"),
            weight2: R.Tensor((16, 3, 3, 16), "float16"),
            weight3: R.Tensor((16, 3, 3, 16), "float16"),
            gamma: R.Tensor((16,), "float16"),
            beta: R.Tensor((16,), "float16"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                conv1 = R.nn.relu(
                    R.nn.conv2d(
                        data, weight1, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )
                conv2 = R.nn.relu(
                    R.nn.conv2d(
                        conv1, weight2, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )
                ln = R.nn.layer_norm(conv2, gamma, beta, axes=[-1])
                conv3 = R.nn.relu(
                    R.nn.conv2d(
                        ln, weight3, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )
                R.output(conv3)

            return conv3

    low, high = -1, 1
    data_shape = (16, 32, 32, 16)
    weight_shape = (16, 3, 3, 16)
    dtype = "float16"
    data = np.random.randint(low, high, size=data_shape).astype(dtype)
    weight1 = np.random.randint(low, high, size=weight_shape).astype(dtype)
    weight2 = np.random.randint(low, high, size=weight_shape).astype(dtype)
    weight3 = np.random.randint(low, high, size=weight_shape).astype(dtype)
    gamma = np.random.randint(low, high, size=(weight_shape[0],)).astype(dtype)
    beta = np.random.randint(low, high, size=(weight_shape[0],)).astype(dtype)
    inputs = [data, weight1, weight2, weight3, gamma, beta]

    mod = partition_for_cutlass(Conv2d)
    mod = relax.transform.RunCodegen({"cutlass": {"sm": 80, "find_first_valid": True}})(mod)
    mod = relax.pipeline.get_pipeline()(mod)  # pylint: disable=no-value-for-parameter

    with tvm.target.Target("cuda"):
        mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)

    out = build_and_run(mod, inputs, "cuda", cuda_graph=True)
    ref = build_and_run(Conv2d, inputs, "llvm", legalize=True)
    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_fp16A_int8B_gemm_batched():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def decode(
            A: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            B: T.Buffer((T.int64(64),), "float16"),
            decode_1: T.Buffer((T.int64(64), T.int64(64)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i, j in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(A[v_i, v_j], B[v_j])
                    T.writes(decode_1[v_i, v_j])
                    decode_1[v_i, v_j] = T.Cast("float16", A[v_i, v_j]) * B[v_j]

        @T.prim_func(s_tir=True)
        def encode(
            A: T.Buffer((T.int64(64), T.int64(64)), "float16"),
            w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            compute: T.Buffer((T.int64(64),), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            max_abs_value = T.sblock_alloc_buffer((T.int64(64),), "float16")
            scale = T.sblock_alloc_buffer((T.int64(64),))
            for i, k in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("max_abs_value"):
                    v_i, v_k = T.axis.remap("SR", [i, k])
                    T.reads(A[v_i, v_k])
                    T.writes(max_abs_value[v_i])
                    with T.init():
                        max_abs_value[v_i] = T.float16(-65504)
                    max_abs_value[v_i] = T.max(max_abs_value[v_i], T.fabs(A[v_i, v_k]))
            for i in range(T.int64(64)):
                with T.sblock("scale"):
                    v_i = T.axis.spatial(T.int64(64), i)
                    T.reads(max_abs_value[v_i])
                    T.writes(scale[v_i])
                    scale[v_i] = T.max(
                        T.Cast("float32", max_abs_value[v_i]), T.float32(0.0001)
                    ) * T.float32(0.0078125)
            for j, i in T.grid(T.int64(64), T.int64(64)):
                with T.sblock("w_gathered"):
                    v_j, v_i = T.axis.remap("SS", [j, i])
                    T.reads(A[v_i, v_j], scale[v_i])
                    T.writes(w_gathered[v_j, v_i])
                    w_gathered[v_j, v_i] = T.Cast(
                        "int8",
                        T.min(
                            T.max(
                                T.round(T.Cast("float32", A[v_i, v_j]) / scale[v_i]),
                                T.float32(-128),
                            ),
                            T.float32(127),
                        ),
                    )
            for i0 in range(T.int64(64)):
                with T.sblock("compute"):
                    v_i0 = T.axis.spatial(T.int64(64), i0)
                    T.reads(scale[v_i0])
                    T.writes(compute[v_i0])
                    compute[v_i0] = T.Cast("float16", scale[v_i0])

        @R.function
        def main(
            x: R.Tensor(("b", 64, 64), dtype="float16"),
            y: R.Tensor((64, 64), dtype="float16"),
        ) -> R.Tensor(("b", 64, 64), dtype="float16"):
            R.func_attr({"num_input": 1})
            cls = Module
            b = T.int64()
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[R.Tensor((64, 64), dtype="int8"), R.Tensor((64,), dtype="float16")],
                )
                lv1: R.Tensor((64, 64), dtype="int8") = lv[0]
                lv2: R.Tensor((64, 64), dtype="int8") = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    R.prim_value(80),
                    R.prim_value(0),
                    ty_args=(R.Tensor((64, 64), dtype="int8"),),
                )
                lv3: R.Tensor((64,), dtype="float16") = lv[1]
                lv4: R.Tensor((64, 64), dtype="int8") = R.builtin.stop_lift_params(lv2)
                lv5: R.Tensor((64,), dtype="float16") = R.builtin.stop_lift_params(lv3)
                lv6 = R.call_tir(cls.decode, (lv4, lv5), out_ty=R.Tensor((64, 64), dtype="float16"))
                lv1_1: R.Tensor((b, 64, 64), dtype="float16") = R.matmul(
                    x, lv6, out_dtype="float16"
                )
                R.output(lv1_1)
            return lv1_1

    x_shape = (4, 64, 64)
    y_shape = (64, 64)

    mod = partition_for_cutlass(Module)

    mod = relax.transform.RunCodegen(
        {"cutlass": {"sm": 80, "find_first_valid": False}},
    )(mod)

    x = np.random.randn(*x_shape).astype("float16")
    y = np.random.normal(0, 0.002, size=y_shape).astype("float16")

    mod = relax.pipeline.get_pipeline()(mod)
    mod = relax.transform.LiftTransformParams()(mod)

    mod_transform, mod_deploy, transform_func_name = split_transform_deploy_mod(mod)

    ex = tvm.compile(mod_transform, target="llvm")
    vm = relax.vm.VirtualMachine(ex, tvm.cpu(0))

    packed_weight, scales = vm[transform_func_name]((tvm.runtime.tensor(y),))

    ex_cuda = tvm.compile(mod_deploy, target="cuda")

    def run_and_check():
        dev = tvm.device("cuda", 0)
        vm = relax.vm.VirtualMachine(ex_cuda, dev)
        x_nd = tvm.runtime.tensor(x, dev)
        inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev)]
        out = vm["main"](*inp).numpy()
        ref = np.dot(x, y.transpose())
        tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_fp16A_int8B_gemm_batched_finegrained():
    @I.ir_module(s_tir=True)
    class Module:
        @T.prim_func(s_tir=True)
        def decode(
            A: T.Buffer((T.int64(128), T.int64(128)), "int8"),
            B: T.Buffer((T.int64(2), T.int64(128)), "float16"),
            decode_1: T.Buffer((T.int64(128), T.int64(128)), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            for i, j in T.grid(T.int64(128), T.int64(128)):
                with T.sblock("decode"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(A[v_i, v_j], B[v_i // T.int64(64), v_j])
                    T.writes(decode_1[v_i, v_j])
                    decode_1[v_i, v_j] = T.Cast("float16", A[v_i, v_j]) * B[v_i // T.int64(64), v_j]

        @T.prim_func(s_tir=True)
        def encode(
            A: T.Buffer((T.int64(128), T.int64(128)), "float16"),
            w_gathered: T.Buffer((T.int64(128), T.int64(128)), "int8"),
            compute: T.Buffer(
                (
                    T.int64(2),
                    T.int64(128),
                ),
                "float16",
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            max_abs_value = T.sblock_alloc_buffer(
                (
                    T.int64(2),
                    T.int64(128),
                ),
                "float16",
            )
            scale = T.sblock_alloc_buffer(
                (
                    T.int64(2),
                    T.int64(128),
                )
            )
            for i, j, k in T.grid(T.int64(2), T.int64(128), T.int64(64)):
                with T.sblock("max_abs_value"):
                    v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
                    T.reads(A[v_j, v_i * T.int64(64) + v_k])
                    T.writes(max_abs_value[v_i, v_j])
                    with T.init():
                        max_abs_value[v_i, v_j] = T.float16(-65504)
                    max_abs_value[v_i, v_j] = T.max(
                        max_abs_value[v_i, v_j], T.fabs(A[v_j, v_i * T.int64(64) + v_k])
                    )
            for i, j in T.grid(T.int64(2), T.int64(128)):
                with T.sblock("scale"):
                    v_i, v_j = T.axis.remap("SS", [i, j])
                    T.reads(max_abs_value[v_i, v_j])
                    T.writes(scale[v_i, v_j])
                    scale[v_i, v_j] = T.max(
                        T.Cast("float32", max_abs_value[v_i, v_j]), T.float32(0.0001)
                    ) * T.float32(0.0078125)
            for j, i in T.grid(T.int64(128), T.int64(128)):
                with T.sblock("w_gathered"):
                    v_j, v_i = T.axis.remap("SS", [j, i])
                    T.reads(A[v_i, v_j], scale[v_j // T.int64(64), v_i])
                    T.writes(w_gathered[v_j, v_i])
                    w_gathered[v_j, v_i] = T.Cast(
                        "int8",
                        T.min(
                            T.max(
                                T.round(
                                    T.Cast("float32", A[v_i, v_j]) / scale[v_j // T.int64(64), v_i]
                                ),
                                T.float32(-128),
                            ),
                            T.float32(127),
                        ),
                    )
            for i0, i1 in T.grid(T.int64(2), T.int64(128)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(scale[v_i0, v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.Cast("float16", scale[v_i0, v_i1])

        @R.function
        def main(
            x: R.Tensor(("b", 128, 128), dtype="float16"),
            y: R.Tensor((128, 128), dtype="float16"),
        ) -> R.Tensor(("b", 128, 128), dtype="float16"):
            R.func_attr({"num_input": 1})
            cls = Module
            b = T.int64()
            with R.dataflow():
                lv = R.call_tir(
                    cls.encode,
                    (y,),
                    out_ty=[
                        R.Tensor((128, 128), dtype="int8"),
                        R.Tensor((2, 128), dtype="float16"),
                    ],
                )
                lv1: R.Tensor((128, 128), dtype="int8") = lv[0]
                lv2: R.Tensor((128, 128), dtype="int8") = R.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    lv1,
                    R.prim_value(80),
                    R.prim_value(0),
                    ty_args=(R.Tensor((128, 128), dtype="int8"),),
                )
                lv3: R.Tensor((2, 128), dtype="float16") = lv[1]
                lv4: R.Tensor((128, 128), dtype="int8") = R.builtin.stop_lift_params(lv2)
                lv5: R.Tensor((2, 128), dtype="float16") = R.builtin.stop_lift_params(lv3)
                lv6 = R.call_tir(
                    cls.decode, (lv4, lv5), out_ty=R.Tensor((128, 128), dtype="float16")
                )
                lv1_1: R.Tensor((b, 128, 128), dtype="float16") = R.matmul(
                    x, lv6, out_dtype="float16"
                )
                R.output(lv1_1)
            return lv1_1

    x_shape = (4, 128, 128)
    y_shape = (128, 128)

    mod = partition_for_cutlass(Module)

    mod = relax.transform.RunCodegen(
        {"cutlass": {"sm": 80, "find_first_valid": False}},
    )(mod)

    x = np.random.randn(*x_shape).astype("float16")
    y = np.random.normal(0, 0.002, size=y_shape).astype("float16")

    mod = relax.pipeline.get_pipeline()(mod)
    mod = relax.transform.LiftTransformParams()(mod)

    mod_transform, mod_deploy, transform_func_name = split_transform_deploy_mod(mod)

    ex = tvm.compile(mod_transform, target="llvm")
    vm = relax.vm.VirtualMachine(ex, tvm.cpu(0))

    packed_weight, scales = vm[transform_func_name]((tvm.runtime.tensor(y),))

    ex_cuda = tvm.compile(mod_deploy, target="cuda")

    def run_and_check():
        dev = tvm.device("cuda", 0)
        vm = relax.vm.VirtualMachine(ex_cuda, dev)
        x_nd = tvm.runtime.tensor(x, dev)
        inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev)]
        out = vm["main"](*inp).numpy()
        ref = np.dot(x, y.transpose())
        tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_attention_rewrite_multi_query():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            q: R.Tensor((4, 16, 32, 16), dtype="float16"),
            k_single: R.Tensor((4, 16, 1, 16), dtype="float16"),
            v_single: R.Tensor((4, 16, 1, 16), dtype="float16"),
        ) -> R.Tensor((4, 16, 32, 8), dtype="float16"):
            with R.dataflow():
                k = R.repeat(k_single, 32, axis=2)
                v = R.repeat(v_single, 32, axis=2)

                lv = R.permute_dims(q, axes=[0, 2, 1, 3])
                lv1 = R.reshape(lv, R.shape([128, 16, 16]))
                lv2 = R.permute_dims(k, axes=[0, 2, 1, 3])
                lv3 = R.reshape(lv2, R.shape([128, 16, 16]))
                lv4 = R.permute_dims(v, axes=[0, 2, 1, 3])
                lv5 = R.reshape(lv4, R.shape([128, 16, 16]))

                lv6 = R.permute_dims(lv3, axes=[0, 2, 1])
                lv7 = R.matmul(lv1, lv6, out_dtype="float16")
                lv3_1 = R.astype(R.const(0.25, "float32"), "float16")
                lv8 = R.multiply(lv7, lv3_1)
                lv11 = R.astype(R.nn.softmax(R.astype(lv8, "float32"), axis=2), "float16")
                lv12 = R.matmul(lv11, lv5, out_dtype="float16")
                lv13 = R.reshape(lv12, R.shape([4, 32, 16, 16]))
                lv6_1 = R.permute_dims(lv13, axes=[0, 2, 1, 3])
                R.output(lv6_1)
            return lv6_1

    q_np = np.random.randn(4, 16, 32, 16).astype("float16")
    k_np = np.random.randn(4, 16, 1, 16).astype("float16")
    v_np = np.random.randn(4, 16, 1, 16).astype("float16")
    args = [q_np, k_np, v_np]
    ref = build_and_run(Module, args, "llvm", legalize=True)

    mod = partition_for_cutlass(Module, use_flash_mqa=True)
    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80}})
    mod = codegen_pass(mod)

    out = build_and_run(mod, args, "cuda")

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def _test_batched_var_len_attention(
    mod, seq_lens, num_head, num_kv_head, head_size, window_size=None
):
    if not tvm.get_global_func("tvm.contrib.thrust.sum_scan", True):
        return

    hidden_size = num_head * head_size

    batched_queries = []
    batched_keys = []
    batched_values = []
    batched_refs = []

    for s in seq_lens:
        q, k, v, _, ref = get_numpy_attention_ref(
            1,
            s,
            s,
            num_head,
            head_size,
            head_size,
            "none",
            "none",
            "BottomRight",
            "float16",
            num_kv_head=num_kv_head,
            window_size=window_size,
        )
        batched_queries.append(np.reshape(q, [-1, hidden_size]))
        batched_keys.append(np.reshape(k, [-1, num_kv_head * head_size]))
        batched_values.append(np.reshape(v, [-1, num_kv_head * head_size]))
        batched_refs.append(np.reshape(ref, [-1, hidden_size]))

    batched_queries = np.vstack(batched_queries)
    batched_keys = np.vstack(batched_keys)
    batched_values = np.vstack(batched_values)
    ref = np.vstack(batched_refs)

    mod = partition_for_cutlass(mod)
    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80}})
    mod = codegen_pass(mod)

    with tvm.target.Target("cuda"):
        mod = relax.transform.LegalizeOps()(mod)
        mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)

    out = build_and_run(
        mod,
        [
            batched_queries,
            batched_keys,
            batched_values,
            np.array(seq_lens, dtype="int32"),
        ],
        "cuda",
    )

    ############# xformer reference for verification #############

    # attn_bias = BlockDiagonalCausalMask.from_seqlens(seq_lens)

    # queries = torch.from_numpy(np.reshape(batched_queries, [1, -1, num_head, head_size])).to("cuda")
    # keys = torch.from_numpy(np.reshape(batched_keys, [1, -1, num_head, head_size])).to("cuda")
    # values = torch.from_numpy(np.reshape(batched_values, [1, -1, num_head, head_size])).to("cuda")

    # out = xops.memory_efficient_attention_forward(
    #     queries, keys, values,
    #     attn_bias=attn_bias,
    # ).cpu().numpy()[0]
    # out = np.reshape(out, [-1, hidden_size])

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_batched_var_len_attention():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                ]
            }
        )

        @R.function
        def main(
            queries: R.Tensor(("num_tokens", 4096), dtype="float16"),
            keys: R.Tensor(("num_tokens", 4096), dtype="float16"),
            values: R.Tensor(("num_tokens", 4096), dtype="float16"),
            seq_lens: R.Tensor(("num_seq",), dtype="int32"),
        ) -> R.Tensor(("num_tokens", 4096), dtype="float16"):
            R.func_attr({"num_input": 4})
            cls = Module
            num_tokens = T.int64()
            num_seq = T.int64()

            with R.dataflow():
                # TODO(masahi): Workaround for the broken Relax cumsum op on GPU.
                # https://github.com/apache/tvm/issues/15851
                cumsum = R.call_dps_packed(
                    "tvm.contrib.thrust.sum_scan", seq_lens, out_ty=seq_lens.ty
                )
                max_seqlen_q = R.to_vdevice(R.max(seq_lens), "llvm:0")
                seqstart_q = R.concat([R.zeros((1,), "int32"), cumsum])
                q = R.reshape(queries, R.shape([1, num_tokens, 128, 32]))
                k = R.reshape(keys, R.shape([1, num_tokens, 128, 32]))
                v = R.reshape(values, R.shape([1, num_tokens, 128, 32]))
                attn_out = R.nn.attention_var_len(
                    q,
                    k,
                    v,
                    seqstart_q,
                    max_seqlen_q,
                    causal_mask="BottomRight",
                )
                out = R.reshape(attn_out, R.shape([num_tokens, 4096]))
                R.output(out)
            return out

    seq_lens = [5, 3, 8]
    num_head = 128
    head_size = 32

    _test_batched_var_len_attention(Module, seq_lens, num_head, num_head, head_size)


def test_batched_var_len_multi_query_attention():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                ]
            }
        )

        @R.function
        def main(
            queries: R.Tensor(("num_tokens", 4096), dtype="float16"),
            keys: R.Tensor(("num_tokens", 512), dtype="float16"),
            values: R.Tensor(("num_tokens", 512), dtype="float16"),
            seq_lens: R.Tensor(("num_seq",), dtype="int32"),
        ) -> R.Tensor(("num_tokens", 4096), dtype="float16"):
            R.func_attr({"num_input": 4})
            cls = Module
            num_tokens = T.int64()
            num_seq = T.int64()

            with R.dataflow():
                # TODO(masahi): Workaround for the broken Relax cumsum op on GPU.
                # https://github.com/apache/tvm/issues/15851
                cumsum = R.call_dps_packed(
                    "tvm.contrib.thrust.sum_scan", seq_lens, out_ty=seq_lens.ty
                )
                max_seqlen_q = R.to_vdevice(R.max(seq_lens), "llvm:0")
                seqstart_q = R.concat([R.zeros((1,), "int32"), cumsum])
                q = R.reshape(queries, R.shape([1, num_tokens, 128, 32]))
                k = R.reshape(keys, R.shape([1, num_tokens, 16, 32]))
                v = R.reshape(values, R.shape([1, num_tokens, 16, 32]))
                attn_out = R.nn.attention_var_len(
                    q,
                    k,
                    v,
                    seqstart_q,
                    max_seqlen_q,
                    causal_mask="BottomRight",
                )
                out = R.reshape(attn_out, R.shape([num_tokens, 4096]))
                R.output(out)
            return out

    seq_lens = [5, 3, 8]
    num_head = 128
    num_kv_head = 16
    head_size = 32

    _test_batched_var_len_attention(Module, seq_lens, num_head, num_kv_head, head_size)


def test_sliding_window():
    q_shape = (1, 64, 16, 8)
    k_shape = v_shape = q_shape
    window_size = 8
    causal = "BottomRight"

    mod = get_relax_attention_module(
        q_shape,
        k_shape,
        v_shape,
        dtype="float16",
        causal_mask=causal,
        window_size=window_size,
    )

    q, k, v, _, ref = get_numpy_attention_ref(
        1, 64, 64, 16, 8, 8, "none", "none", causal, "float16", window_size=window_size
    )

    out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2)

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)

    ############# xformer reference for verification #############

    # attn_bias = BlockDiagonalCausalMask.from_seqlens([64])

    # if window_size > 0:
    #     attn_bias = attn_bias.make_local_attention(window_size)

    # query = torch.from_numpy(q).to("cuda")
    # key = torch.from_numpy(k).to("cuda")
    # value = torch.from_numpy(v).to("cuda")

    # ref = xops.memory_efficient_attention_forward(
    #     query, key, value, attn_bias=attn_bias,
    # ).cpu().numpy()

    # tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def test_batched_var_len_sliding_window():
    @I.ir_module(s_tir=True)
    class Module:
        I.module_global_infos(
            {
                "vdevice": [
                    I.vdevice("llvm"),
                ]
            }
        )

        @R.function
        def main(
            queries: R.Tensor(("num_tokens", 4096), dtype="float16"),
            keys: R.Tensor(("num_tokens", 4096), dtype="float16"),
            values: R.Tensor(("num_tokens", 4096), dtype="float16"),
            seq_lens: R.Tensor(("num_seq",), dtype="int32"),
        ) -> R.Tensor(("num_tokens", 4096), dtype="float16"):
            R.func_attr({"num_input": 4})
            cls = Module
            num_tokens = T.int64()
            num_seq = T.int64()

            with R.dataflow():
                # TODO(masahi): Workaround for the broken Relax cumsum op on GPU.
                # https://github.com/apache/tvm/issues/15851
                cumsum = R.call_dps_packed(
                    "tvm.contrib.thrust.sum_scan", seq_lens, out_ty=seq_lens.ty
                )
                max_seqlen_q = R.to_vdevice(R.max(seq_lens), "llvm:0")
                seqstart_q = R.concat([R.zeros((1,), "int32"), cumsum])
                q = R.reshape(queries, R.shape([1, num_tokens, 128, 32]))
                k = R.reshape(keys, R.shape([1, num_tokens, 128, 32]))
                v = R.reshape(values, R.shape([1, num_tokens, 128, 32]))
                attn_out = R.nn.attention_var_len(
                    q,
                    k,
                    v,
                    seqstart_q,
                    max_seqlen_q,
                    causal_mask="BottomRight",
                    window_size=T.IntImm("int32", 8),
                )
                out = R.reshape(attn_out, R.shape([num_tokens, 4096]))
                R.output(out)
            return out

    seq_lens = [64, 64, 64]
    num_head = 128
    num_kv_head = 128
    head_size = 32
    window_size = 8

    _test_batched_var_len_attention(Module, seq_lens, num_head, num_kv_head, head_size, window_size)


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