# 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.
# pylint: disable=invalid-name, dangerous-default-value, arguments-differ
# ruff: noqa: F821
"""Driver for partitioning and building a Relax module for CUTLASS offload."""

import itertools
import logging
import multiprocessing
import operator
import os
from collections.abc import Sequence
from functools import reduce

from tvm_ffi import register_global_func

import tvm
from tvm import relax, runtime
from tvm.support.nvcc import get_cuda_version
from tvm.topi.utils import get_const_tuple

from .gen_conv2d import CutlassConv2DProfiler
from .gen_gemm import CutlassGemmProfiler
from .library import ConvKind, LayoutType

logger = logging.getLogger("cutlass")


def has_cutlass():
    """Returns true if the CUTLASS custom codegen is available"""
    return tvm.get_global_func("relax.ext.cutlass", True) is not None


def _get_cutlass_path():
    invalid_paths = []
    for rel in ["../../../../", "../../../", "../../"]:
        tvm_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), rel)
        cutlass_path = os.path.join(tvm_root, "3rdparty/cutlass")
        if os.path.exists(cutlass_path):
            return cutlass_path
        invalid_paths.append(cutlass_path)
    raise AssertionError(f"The CUTLASS root directory not found in: {invalid_paths}")


def _get_cutlass_compile_options(sm, threads, use_fast_math=False):
    cutlass_root = _get_cutlass_path()
    cutlass_include = os.path.join(cutlass_root, "include")
    cutlass_util_include = os.path.join(cutlass_root, "tools/util/include")
    cutlass_attention_include = os.path.join(cutlass_root, "examples/41_fused_multi_head_attention")
    cutlass_fpA_intB_gemm_include = os.path.join(cutlass_root, "../cutlass_fpA_intB_gemm")
    flash_attn_include = os.path.join(cutlass_root, "../libflash_attn/include")

    kwargs = {}
    kwargs["cc"] = "nvcc"
    kwargs["options"] = [
        "-c",
        "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
        f"-gencode=arch=compute_{sm},code=[sm_{sm},compute_{sm}]",
        "-DNDEBUG",
        "-Xcompiler=-fPIC",
        "-Xcompiler=-Wconversion",
        "-Xcompiler=-fno-strict-aliasing",
        "-Xcompiler=-fvisibility=hidden",
        "-O3",
        "-std=c++17",
        f"-I{cutlass_include}",
        f"-I{cutlass_util_include}",
        f"-I{cutlass_attention_include}",
        f"-I{cutlass_fpA_intB_gemm_include}",
        f"-I{flash_attn_include}",
    ]
    if use_fast_math:
        kwargs["options"].append("-DCUTLASS_USE_TANH_FOR_SIGMOID")
    cuda_ver = get_cuda_version()
    if cuda_ver >= (11, 2):
        ncpu = multiprocessing.cpu_count() if threads < 0 else threads
        kwargs["options"].append(f"-t {ncpu}")
    return kwargs


def select_gemm_kernel(
    cutlass_profiler,
    op_type,
    MM,
    KK,
    NN,
    out_dtype,
    arg0_dtype,
    arg1_dtype,
    use_3xtf32,
    batched,
    find_first_valid,
    use_multiprocessing,
):
    """Run CUTLASS profiler to select the best kernel, or return the default one for dynamic
    workloads."""
    if any(isinstance(s, tvm.tirx.Any) for s in [MM, KK, NN]):
        out = cutlass_profiler.get_default(
            op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32, batched=batched
        )
        name, cutlass_op_def = out["name"], out["opdef"]
        logger.info("Picked the default kernel %s", name)
    else:
        name, cutlass_op_def, _ = cutlass_profiler.profile(
            op_type,
            MM,
            NN,
            KK,
            out_dtype,
            arg0_dtype,
            arg1_dtype,
            use_3xtf32,
            batched=batched,
            find_first_valid=find_first_valid,
            use_multiprocessing=use_multiprocessing,
        )
        if not find_first_valid:
            logger.info("The best kernel is %s", name)
        else:
            logger.info("Picked the first kernel found %s", name)

    return name, cutlass_op_def


def handle_batch_matmul(
    cutlass_profiler,
    op_type,
    arg0_shape,
    arg1_shape,
    out_dtype,
    arg0_dtype,
    arg1_dtype,
    use_3xtf32,
    find_first_valid,
    use_multiprocessing,
):
    """Profile and select a kernel for batch_matmul op workload."""
    MM = arg0_shape[1]
    KK = arg0_shape[2]
    NN = arg1_shape[1]

    name, cutlass_op_def = select_gemm_kernel(
        cutlass_profiler,
        op_type,
        MM,
        KK,
        NN,
        out_dtype,
        arg0_dtype,
        arg1_dtype,
        use_3xtf32,
        True,
        find_first_valid,
        use_multiprocessing,
    )

    return {
        "batch": arg0_shape[0],
        "batch_stride_A": arg0_shape[1] * arg0_shape[2],
        "batch_stride_B": arg1_shape[1] * arg1_shape[2],
        "batch_stride_C": arg0_shape[1] * arg1_shape[1],
        "cutlass_op_def": cutlass_op_def,
        "cutlass_op_name": name,
        "lda": "K",
        "ldb": "K",
        "ldc": "N",
    }


def handle_dense(
    cutlass_profiler,
    op_type,
    arg0_shape,
    arg1_shape,
    out_dtype,
    arg0_dtype,
    arg1_dtype,
    use_3xtf32,
    find_first_valid,
    use_multiprocessing,
):
    """Profile and select a kernel for dense op workload."""
    MM = arg0_shape[0]
    KK = arg0_shape[1]
    NN = arg1_shape[0]

    name, cutlass_op_def = select_gemm_kernel(
        cutlass_profiler,
        op_type,
        MM,
        KK,
        NN,
        out_dtype,
        arg0_dtype,
        arg1_dtype,
        use_3xtf32,
        False,
        find_first_valid,
        use_multiprocessing,
    )

    assert "tn_align" in name, "Only supports (row_major, col_major) input layout for now."

    return {
        "cutlass_op_def": cutlass_op_def,
        "cutlass_op_name": name,
        "lda": "K",
        "ldb": "K",
        "ldc": "N",
    }


def handle_conv2d(
    cutlass_profiler,
    op_type,
    d_shape,
    w_shape,
    padding,
    strides,
    dilation,
    out_dtype,
    data_dtype,
    weight_dtype,
    use_3xtf32,
    split_k_slices,
    profile_all_alignments,
    find_first_valid,
    use_multiprocessing,
):
    """Profile and select a kernel for conv2d op workload."""
    if "conv2d_transpose" in op_type:
        conv_kind = ConvKind.Dgrad
    elif "backward_weight" in op_type:
        conv_kind = ConvKind.Wgrad
    else:
        conv_kind = ConvKind.Fprop

    if any(isinstance(s, tvm.tirx.Any) for s in d_shape):
        out = cutlass_profiler.get_default(
            op_type, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, strides
        )
        name, cutlass_op_def = out["name"], out["opdef"]
        logger.info("Picked the default kernel %s", name)
    else:
        name, cutlass_op_def, _ = cutlass_profiler.profile(
            op_type,
            d_shape,
            w_shape,
            padding,
            strides,
            dilation,
            out_dtype,
            data_dtype,
            weight_dtype,
            use_3xtf32,
            conv_kind,
            split_k_slices,
            profile_all_alignments,
            find_first_valid=find_first_valid,
            use_multiprocessing=use_multiprocessing,
        )
        if not find_first_valid:
            logger.info("The best kernel is %s", name)
        else:
            logger.info("Picked the first kernel found %s", name)

    return {"cutlass_op_def": cutlass_op_def, "cutlass_op_name": name}


def num_cutlass_partitions(mod):
    return sum([(1 if "cutlass" in var.name_hint else 0) for var in mod.get_global_vars()])


def tune_cutlass_kernels(
    mod,
    sm,
    use_3xtf32=True,
    split_k_slices=[1],
    profile_all_alignments=False,
    find_first_valid=False,
    use_multiprocessing=False,
    tmp_dir="./tmp",
):
    """Given a module partitioned for CUTLASS offloading, profile each workload to select which
    kernels to emit.

    Parameters
    ----------
    mod : IRModule
        The IRModule with cutlass partitions.

    sm : int
        An integer specifying the compute capability. For example, 75 for Turing and
        80 or 86 for Ampere.

    use_3xtf32 : bool
        Wheter or not use slower but very accurate (compared to tf32) 3xtf32 mode for
        fp32 inputs on tensorcore.

    split_k_slices : list of int
        Split factor candidates for split-K GEMM. If split-K > 1, the GEMM K-loop is computed in
        parallel across split-K blocks, and a separate global reduction kernel is launched to
        accumulate partial reductions. The profiler will pick the best split-k factor from the
        given candidate list. Note that the larger split-K factor requires a larger workspace.
        Currently, parallel split-k has been tested only for wgrad. For GEMM and other conv2d
        kinds, split_k_slices is ignored.

    profile_all_alignments : bool
        When True, profile all kernal variants with smaller alignments than the largest possible.

    find_first_valid : bool
        Whether or not profile all candidate kernels, or stop profiling after
        the first applicable kernel is found.

    use_multiprocessing : bool
        Whether or not compile profiler executables for different kernels in parallel.

    tmp_dir : string, optional
        A temporary directory where intermediate compiled artifacts will be stored.

    Returns
    -------
    mod : IRModule
        The updated module annotated with cutlass profiling information.

    num_cutlass_partition : int
        The number of partitioned functions created for CUTLASS.
    """
    gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)
    conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
    num_cutlass_partition = 0
    for var in mod.get_global_vars():
        fun_name = var.name_hint
        func = mod[fun_name]
        if "cutlass" in fun_name:
            num_cutlass_partition += 1
            new_func = tune_cutlass_function(
                func,
                use_3xtf32,
                split_k_slices,
                profile_all_alignments,
                find_first_valid,
                use_multiprocessing,
                gemm_profiler,
                conv2d_profiler,
            )
            mod.update_func(var, new_func)

    return mod, num_cutlass_partition


def _get_call_node(expr: relax.Expr, op_name: str) -> relax.Call | None:
    node = None

    def fvisit(e):
        nonlocal node
        if isinstance(e, relax.Call) and e.op.name == op_name:
            node = e

    relax.analysis.post_order_visit(expr, fvisit)
    return node


def _extract_relax_function_signature(f):
    signature = {}

    for i, arg in enumerate(f.params):
        ty = arg.ty
        if isinstance(ty, relax.TensorType):
            signature[f"arg{i}_shape"] = get_const_tuple(ty.shape)
            signature[f"arg{i}_dtype"] = ty.dtype
        elif isinstance(ty, relax.ShapeType):
            signature[f"arg{i}_shape"] = get_const_tuple(ty.values)
        else:
            raise NotImplementedError()

    ret_ty = f.ret_ty
    if ret_ty.shape is not None:
        signature["ret_shape"] = get_const_tuple(ret_ty.shape)
    else:
        signature["ret_shape"] = None
    signature["ret_dtype"] = ret_ty.dtype

    return signature


def _extract_arg_idx(pattern_name, f):
    extract_func = tvm.get_global_func("relax.contrib.extract_arg_idx")
    arg_indices = extract_func(pattern_name, f)
    return {k: int(v) for k, v in arg_indices.items()}


def is_shape_valid_for_cutlass_matmul(
    lhs_shape: Sequence[tvm.ir.Expr],
    rhs_shape: Sequence[tvm.ir.Expr],
) -> bool:
    """
    Check whether the shape of inputs can be handled by CUTLASS GEMM.

    The stride-based batch matmul in CUTLASS cannot handle cases that some of
    the batch dimensions need to be stretched while others don't. This means
    it can only handle ND x ND whose batch dimensions match exactly on both side,
    as well as ND x 2D and 2D x ND. For example, it cannot handle matmul with shape
    (2, 1, 4, 8) x (2, 3, 8, 16), because the batch stride of lhs is not constant.
    """
    if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
        # Reduction axis must be constant
        return False

    lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
    rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
    if lhs_batches == 1 or rhs_batches == 1:
        # This could be regular matmul or batch matmul with shape ND x 2D or 2D x ND
        return True

    analyzer = tvm.arith.Analyzer()
    # If one side has less dimensions, use 1 to fill the gap
    batch_dim_pairs = list(
        itertools.zip_longest(
            list(lhs_shape)[-3::-1],  # Remove the last two dimensions and reverse
            list(rhs_shape)[-3::-1],
            fillvalue=1,
        )
    )
    return all(analyzer.can_prove_equal(p[0], p[1]) for p in batch_dim_pairs)


@relax.expr_functor.mutator
class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
    """A Relax function mutator that tunes and annotates CUTLASS composite functions
    with shape, dtype and generated templates.
    """

    def __init__(
        self,
        mod,
        conv2d_profiler: CutlassConv2DProfiler,
        gemm_profiler: CutlassGemmProfiler,
        options,
    ):
        super().__init__(mod)
        self.options = options
        self.conv2d_profiler = conv2d_profiler
        self.gemm_profiler = gemm_profiler

    def handle_conv2d(self, f, op_type):
        """Tune and annotate a conv2d op."""
        signature = _extract_relax_function_signature(f)
        arg_idx = _extract_arg_idx(op_type, f)
        op_attrs = _get_call_node(f.body, "relax.nn.conv2d").attrs

        data_arg = f"arg{arg_idx['lhs']}"
        weight_arg = f"arg{arg_idx['rhs']}"

        d_shape = signature[f"{data_arg}_shape"]
        w_shape = signature[f"{weight_arg}_shape"]
        out_shape = signature["ret_shape"]
        data_dtype = signature[f"{data_arg}_dtype"]
        weight_dtype = signature[f"{weight_arg}_dtype"]
        out_dtype = signature["ret_dtype"]
        padding = op_attrs["padding"]
        strides = op_attrs["strides"]
        dilation = op_attrs["dilation"]
        conv_kind = ConvKind.Fprop

        use_3xtf32 = self.options.get("use_3xtf32", False)
        profile_all_alignments = self.options.get("profile_all_alignments", False)
        find_first_valid = self.options.get("find_first_valid", True)
        use_multiprocessing = self.options.get("use_multiprocessing", True)
        split_k_slices = self.options.get("split_k_slices", [1])

        op_name, op_def, _ = self.conv2d_profiler.profile(
            op_type,
            d_shape,
            w_shape,
            padding,
            strides,
            dilation,
            out_dtype,
            data_dtype,
            weight_dtype,
            use_3xtf32,
            conv_kind,
            split_k_slices,
            profile_all_alignments,
            find_first_valid=find_first_valid,
            use_multiprocessing=use_multiprocessing,
        )

        attrs = {
            "op_type": op_type,
            "data_arg_idx": arg_idx["lhs"],
            "weight_arg_idx": arg_idx["rhs"],
            "bias_arg_idx": arg_idx.get("bias"),
            "residual_arg_idx": arg_idx.get("residual"),
            "arg0_dtype": data_dtype,
            "arg1_dtype": weight_dtype,
            "ret_dtype": out_dtype,
            "arg0_shape": d_shape,
            "arg1_shape": w_shape,
            "ret_shape": out_shape,
            "strides": strides,
            "padding": padding,
            "dilation": dilation,
            "cutlass_op_name": op_name,
            "cutlass_op_def": op_def,
        }

        residual_arg = arg_idx.get("residual")

        if residual_arg:
            residual_shape = signature[f"arg{residual_arg}_shape"]
            attrs["residual_shape"] = residual_shape
        elif "residual" in op_type:
            attrs["residual_shape"] = d_shape

        return f.with_attrs(attrs)

    def handle_decode_matmul(self, f, op_type):
        """Annotate a decode -> matmul op."""
        arg_idx = _extract_arg_idx(op_type, f)
        signature = _extract_relax_function_signature(f)
        lhs_arg = f"arg{arg_idx['lhs']}"
        rhs_arg = f"arg{arg_idx['w_encoded']}"
        lhs_shape = signature[f"{lhs_arg}_shape"]
        rhs_shape = signature[f"{rhs_arg}_shape"]
        ret_shape = signature["ret_shape"]
        scale_arg = f"arg{arg_idx['scales']}"
        scale_shape = signature[f"{scale_arg}_shape"]
        N = ret_shape[-1]

        attrs = {
            "op_type": op_type,
            "lhs_arg_idx": arg_idx["lhs"],
            "rhs_arg_idx": arg_idx["w_encoded"],
            "scales_arg_idx": arg_idx["scales"],
            "bias_arg_idx": arg_idx.get("bias"),
            "activation": "identity",
        }
        # TODO(wuwei): find a better way to get group size
        attrs["group_size"] = 64 if len(scale_shape) == 2 and scale_shape[0] != 1 else -1

        attrs["batch_rank"] = len(lhs_shape[:-1])
        attrs["M"] = reduce(operator.mul, lhs_shape[:-1], 1)

        attrs["bias_stride"] = 0

        if "bias" in arg_idx:
            bias_shape = signature[f"arg{arg_idx['bias']}_shape"]
            bias_shape_1d = reduce(operator.mul, bias_shape, 1)
            if bias_shape_1d != bias_shape[-1]:
                attrs["bias_stride"] = bias_shape[-1]

        if N == rhs_shape[1]:
            attrs["weight_nbit"] = 8
        else:
            assert N == rhs_shape[1] * 2
            attrs["weight_nbit"] = 4

        if "residual" in op_type:
            residual_pos = op_type.find("residual_")
            postfix = op_type[residual_pos + len("residual_") :]

            if postfix.startswith("multiply"):
                binary_op = "multiply"
            else:
                binary_op = "plus"

            if "relu" in postfix:
                unary_op = "relu"
            else:
                unary_op = "identity"

            activation = "identity"

            for act in ["relu", "silu", "gelu"]:
                if act in op_type[op_type.find("matmul_") + len("matmul_") : residual_pos]:
                    activation = act
                    break

            attrs.update(
                {
                    "unary_op": unary_op,
                    "binary_op": binary_op,
                    "activation": activation,
                    "residual_arg_idx": arg_idx["residual"],
                }
            )
        else:
            for act in ["relu", "silu", "gelu"]:
                if act in op_type:
                    attrs["activation"] = act
                    break

        return f.with_attrs(attrs)

    def handle_matmul(self, f, op_type):
        """Tune and annotate a matmul op."""
        signature = _extract_relax_function_signature(f)
        arg_idx = _extract_arg_idx(op_type, f)

        lhs_arg = f"arg{arg_idx['lhs']}"
        rhs_arg = f"arg{arg_idx['rhs']}"

        lhs_shape = signature[f"{lhs_arg}_shape"]
        rhs_shape = signature[f"{rhs_arg}_shape"]
        out_shape = signature["ret_shape"]
        lhs_dtype = signature[f"{lhs_arg}_dtype"]
        rhs_dtype = signature[f"{rhs_arg}_dtype"]
        out_dtype = signature["ret_dtype"]

        if not is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape):
            raise ValueError(f"Cannot handle the input shapes, lhs: {lhs_shape}, rhs: {rhs_shape}")

        MM = lhs_shape[-2]
        KK = lhs_shape[-1]
        if "transposed" in op_type:
            NN = rhs_shape[-2]
            ldb = "K"
            layout_b = LayoutType.ColumnMajor
        else:
            NN = rhs_shape[-1]
            ldb = "N"
            layout_b = LayoutType.RowMajor

        lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
        rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
        if lhs_batches == 1 and rhs_batches == 1:
            # Regular matmul
            is_batched = False
            batch_attrs = {}
        else:
            is_batched = True
            batch_attrs = {
                # If both lhs_batches and rhs_batches are greater than 1,
                # they must be equal. This is checked by is_shape_valid_for_cutlass_matmul.
                "batch": lhs_batches if rhs_batches == 1 else rhs_batches,
                "batch_stride_A": 0 if lhs_batches == 1 else MM * KK,
                "batch_stride_B": 0 if rhs_batches == 1 else KK * NN,
                "batch_stride_C": MM * NN,
            }

        use_3xtf32 = self.options.get("use_3xtf32", False)
        find_first_valid = self.options.get("find_first_valid", True)
        use_multiprocessing = self.options.get("use_multiprocessing", True)

        op_name, op_def, _ = self.gemm_profiler.profile(
            op_type,
            MM,
            NN,
            KK,
            out_dtype,
            lhs_dtype,
            rhs_dtype,
            use_3xtf32,
            batched=is_batched,
            find_first_valid=find_first_valid,
            use_multiprocessing=use_multiprocessing,
            layout_b=layout_b,
        )

        return f.with_attrs(
            {
                "op_type": op_type,
                "lhs_arg_idx": arg_idx["lhs"],
                "rhs_arg_idx": arg_idx["rhs"],
                "residual_arg_idx": arg_idx.get("residual"),
                "bias_arg_idx": arg_idx.get("bias"),
                "arg0_dtype": signature["arg0_dtype"],
                "arg1_dtype": signature["arg1_dtype"],
                "ret_dtype": out_dtype,
                "arg0_shape": signature["arg0_shape"],
                "arg1_shape": signature["arg1_shape"],
                "ret_shape": out_shape,
                "lda": "K",
                "ldb": ldb,
                "ldc": "N",
                "cutlass_op_name": op_name,
                "cutlass_op_def": op_def,
                **batch_attrs,
            }
        )

    def handle_attention(self, f, op_type):
        """Annotate an attention op."""
        signature = _extract_relax_function_signature(f)

        if _get_call_node(f.body, "relax.nn.attention") is not None:
            attention_node = _get_call_node(f.body, "relax.nn.attention")
            op_attrs = attention_node.attrs
        elif _get_call_node(f.body, "relax.nn.attention_bias") is not None:
            attention_node = _get_call_node(f.body, "relax.nn.attention_bias")
            op_attrs = attention_node.attrs
        elif _get_call_node(f.body, "relax.nn.attention_var_len") is not None:
            attention_node = _get_call_node(f.body, "relax.nn.attention_var_len")
            op_attrs = attention_node.attrs
        else:
            raise ValueError("Cannot find call node for attention")
        arg = {}

        if "stacked_attention" in op_type:
            arg["arg0_dtype"] = signature["arg0_dtype"]
            q_shape = get_const_tuple(attention_node.args[0].ty.shape)
            k_shape = get_const_tuple(attention_node.args[1].ty.shape)
            v_shape = get_const_tuple(attention_node.args[2].ty.shape)
            if len(attention_node.args) == 4:
                arg["bias_shape"] = get_const_tuple(attention_node.args[3].ty.shape)
                arg["bias_dtype"] = attention_node.args[3].ty.dtype

            qkv_layout = "qkv_stacked"
        else:
            # arg0: q, arg1: k, arg2: v,  arg3: bias, arg4: workspace
            arg["arg0_shape"] = q_shape = signature["arg0_shape"]
            arg["arg1_shape"] = k_shape = signature["arg1_shape"]
            arg["arg2_shape"] = v_shape = signature["arg2_shape"]
            arg["arg0_dtype"] = signature["arg0_dtype"]
            arg["arg1_dtype"] = signature["arg1_dtype"]
            arg["arg2_dtype"] = signature["arg2_dtype"]

            if "arg4_dtype" in signature:
                arg["bias_dtype"] = signature["arg3_dtype"]
            if "arg4_shape" in signature:
                arg["bias_shape"] = signature["arg3_shape"]

            qkv_layout = "default"

        out_shape = signature["ret_shape"]
        out_dtype = signature["ret_dtype"]
        num_batches, num_queries, num_q_heads, head_dim = q_shape
        _, num_keys, num_kv_heads, _ = k_shape
        _, _, _, head_dim_value = v_shape
        scale = op_attrs.scale

        if op_attrs.causal_mask is None:
            custom_mask_type = 0
        elif op_attrs.causal_mask == "TopLeft":
            custom_mask_type = 1
        elif op_attrs.causal_mask == "BottomRight":
            custom_mask_type = 2
        else:
            raise NotImplementedError()

        attrs = {
            "op_type": op_type,
            "ret_dtype": out_dtype,
            "ret_shape": out_shape,
            "num_batches": num_batches,
            "num_queries": num_queries,
            "num_keys": num_keys,
            "num_q_heads": num_q_heads,
            "num_kv_heads": num_kv_heads,
            "head_dim": head_dim,
            "head_dim_value": head_dim_value,
            "scale": scale,
            "arch": self.options["sm"],
            "qkv_layout": qkv_layout,
            "custom_mask_type": custom_mask_type,
            **arg,
        }

        if "var_len" in op_type:
            arg_idx = _extract_arg_idx(op_type, f)
            for arg in ["seqstart_q", "seqstart_k", "max_seqlen_q", "max_seqlen_k"]:
                if arg in arg_idx:
                    attrs[arg + "_idx"] = arg_idx[arg]

        if op_attrs.window_size:
            attrs["window_size"] = op_attrs.window_size

        return f.with_attrs(attrs)

    def handle_norm(self, f, op_type):
        """Annotate a layer or rms norm op."""
        signature = _extract_relax_function_signature(f)
        attrs = {}
        attrs["batch_rank"] = len(signature["arg0_shape"][:-1])
        attrs["M"] = reduce(operator.mul, signature["arg0_shape"][:-1], 1)
        attrs["N"] = signature["arg0_shape"][-1]
        dtype = signature["arg0_dtype"]
        attrs["data_type"] = {"float32": "float", "float16": "cutlass::half_t"}[str(dtype)]

        if "rms" in op_type:
            attrs["rms_eps"] = self.options.get("rms_eps", 1e-5)
        else:
            attrs["layer_norm_eps"] = self.options.get("layer_nrom_eps", 1e-5)

        return f.with_attrs(attrs)

    def visit_function_(self, f):
        if "Composite" not in f.attrs:
            body = super().visit_expr(f.body)
            return relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)

        op_type = f.attrs["Composite"]

        if "conv2d" in op_type:
            return self.handle_conv2d(f, op_type)
        elif "decode" in op_type:
            return self.handle_decode_matmul(f, op_type)
        elif "matmul" in op_type:
            return self.handle_matmul(f, op_type)
        elif "attention" in op_type:
            return self.handle_attention(f, op_type)
        elif "layer_norm" in op_type or "rms_norm" in op_type:
            return self.handle_norm(f, op_type)

        raise ValueError(f"Unsupported composite {op_type}")

    def visit_span(self, span):
        return span


@register_global_func("contrib.cutlass.tune_relax_function")
def profile_relax_function(functions, options):
    """Tune and annotate CUTLASS composite functions with shape, dtype and generated templates."""
    tmp_dir = options.get("tmp_dir", "./tmp")
    sm = options.get("sm", 80)
    conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
    gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)

    annotated_functions = []

    for f in functions:
        annotator = CutlassRelaxFunctionAnnotator(
            tvm.IRModule.from_expr(f), conv2d_profiler, gemm_profiler, options
        )
        annotated_functions.append(annotator.visit_expr(f))

    return annotated_functions


@register_global_func("contrib.cutlass.compile")
def compile_cutlass_module(c_source_module, options):
    """Compile all CUTLASS kernels in the given C-source module.

    Parameters
    ----------
    c_source_module: runtime.Module
        A C-source module containing CUTLASS kernels.

    options: dict
        Compilation options. Currently recognizes
          "sm": The target architecture (compute capability), for example 75 or 80 (default: 80)
          "threads": The number of threads to use in NVCC parallel compilation (default:
          use all logical cores)
          "use_fast_math": Whether or not to use faster but approximate arithmetic in some
          CUTLASS epilogues (default: False)

    Returns
    -------
    rt_mod : runtime.Module
        A runtime module where all cutlass kernels have been compiled.
    """
    tmp_dir = options.get("tmp_dir", "./tmp")
    defaults = {"sm": 80, "threads": -1, "use_fast_math": False}
    compile_config = {key: options.get(key, val) for key, val in defaults.items()}

    function_names = c_source_module.get_function("get_func_names")()
    compile_options = _get_cutlass_compile_options(**compile_config)
    lib_path = os.path.join(tmp_dir, "cutlass.o")
    logger.info("Compiling generated CUTLASS code")
    c_source_module.export_library(lib_path, workspace_dir=tmp_dir, **compile_options)

    # Recover static library
    return tvm.runtime.load_static_library(lib_path, function_names)


def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"):
    """Returns lib with any C source, LLVM and static library modules complied and linked in ready
    for use by the graph or AOT executors. This method is not specific to CUTLASS, however it does
    assume nvcc will be used for final compilation and linking. It is provided here for
    convenience.

    Parameters
    ----------
    lib : runtime.Module
        The output from build.

    lib_path : string
        The path to a shared library which will be generated as the result of the build process.

    tmp_dir : string
        A temporary directory where intermediate compiled artifacts will be stored.

    Returns
    -------
    updated_lib : runtime.Module
        The updated library with all compilation and linking completed.

    """
    lib_path = os.path.join(tmp_dir, lib_path)
    lib.export_library(lib_path, workspace_dir=tmp_dir, cc="nvcc")
    return runtime.load_module(lib_path)
