# 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
# ruff: noqa: F821
"""Common functions and classes for CUTLASS GEMM and Conv2d geneator."""

import logging
import math
import multiprocessing
import os
import re
import subprocess
import tempfile

import tvm_ffi

from tvm.runtime import Object
from tvm.tirx import IntImm

from . import _ffi_api as ffi
from .attention_operation import (
    instantiate_attention_template,
    instantiate_flash_attention_template,
    instantiate_flash_attention_var_len_template,
)
from .conv2d_operation import instantiate_conv2d_template
from .gemm_operation import emit_fp16A_intB_matmul, instantiate_gemm_template
from .layer_norm_operation import instantiate_layer_norm_template
from .library import (
    DataType,
    DataTypeSize,
    DataTypeTag,
    EpilogueFunctor,
    MathInstruction,
    MathOperation,
    OpcodeClass,
    TileDescription,
)
from .rms_norm_operation import instantiate_rms_norm_template

logger = logging.getLogger("cutlass")


dtype_map = {
    "int8": DataType.s8,
    "uint8": DataType.u8,
    "int32": DataType.s32,
    "float32": DataType.f32,
    "float16": DataType.f16,
}


def generate_tensor_op_common(
    math_instructions, alignment_constraints, get_tile_descriptions, op_creator
):
    """Common kernel generator to be used by archtecture specific generators."""
    ops = []
    for math_inst in math_instructions:
        tile_descriptions = get_tile_descriptions(math_inst)
        data_type = [
            math_inst.element_a,
            math_inst.element_b,
            math_inst.element_c,
            math_inst.element_accumulator,
        ]

        out = op_creator(tile_descriptions, data_type, alignment_constraints)

        ops.extend(out)

    return ops


def generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumulator_dtype="float32"):
    """Gemerate GEMM or Conv2D SIMT kernels"""
    # pylint: disable=unused-argument
    min_cc = 50
    max_cc = 1024
    if arg0_dtype == "float32" and arg1_dtype == "float32":
        assert out_dtype == "float32" and accumulator_dtype == "float32"
        math_instructions = [
            MathInstruction(
                [1, 1, 1],
                DataType.f32,
                DataType.f32,
                DataType.f32,
                DataType.f32,
                OpcodeClass.Simt,
                MathOperation.multiply_add,
            )
        ]
        alignment_constraints = [1]
        tile_descriptions = [
            ([128, 128, 8], 2, [4, 2, 1], min_cc, max_cc),
            ([128, 64, 8], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 128, 8], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 64, 8], 2, [2, 1, 1], min_cc, max_cc),
            ([128, 32, 8], 2, [2, 1, 1], min_cc, max_cc),
            ([32, 128, 8], 2, [1, 2, 1], min_cc, max_cc),
        ]

        def get_tile_descriptions(math_inst):
            return [
                TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
                for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
            ]

        return generate_tensor_op_common(
            math_instructions, alignment_constraints, get_tile_descriptions, op_creator
        )
    else:
        raise NotImplementedError()


def generate_sm75_tensor_op_1688(
    out_dtype,
    arg0_dtype,
    arg1_dtype,
    op_creator,
    check_align,
    _,
    profile_all_alignments=False,
    accumlator_dtype="float32",
):
    """Generate GEMM or Conv2D kernels for Turing."""
    assert out_dtype in ["float32", "float16", "int32"]
    min_cc = 75
    max_cc = 1024

    if arg0_dtype == "float16" and arg1_dtype == "float16":
        math_instructions = [
            MathInstruction(
                [16, 8, 8],
                DataType.f16,
                DataType.f16,
                dtype_map[out_dtype],
                dtype_map[accumlator_dtype],
                OpcodeClass.TensorOp,
                MathOperation.multiply_add,
            )
        ]
        alignment_constraints = [8, 4, 2, 1]
        tile_descriptions = [
            ([256, 128, 32], 2, [4, 2, 1], min_cc, max_cc),
            ([128, 256, 32], 2, [2, 4, 1], min_cc, max_cc),
            ([128, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
            ([128, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 128, 64], 2, [1, 2, 2], min_cc, max_cc),
        ]

    elif "int8" in arg0_dtype and "int8" in arg1_dtype:
        assert out_dtype == "int32"
        math_instructions = [
            MathInstruction(
                [8, 8, 16],
                dtype_map[arg0_dtype],
                dtype_map[arg1_dtype],
                DataType.s32,
                DataType.s32,
                OpcodeClass.TensorOp,
                MathOperation.multiply_add_saturate,
            )
        ]
        alignment_constraints = [16, 8, 4, 2, 1]
        tile_descriptions = [
            ([256, 128, 64], 2, [4, 2, 1], min_cc, max_cc),
            ([128, 256, 64], 2, [2, 4, 1], min_cc, max_cc),
            ([128, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 256, 64], 2, [1, 4, 1], min_cc, max_cc),
            ([256, 64, 64], 2, [4, 1, 1], min_cc, max_cc),
            ([64, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
            ([128, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
            ([64, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
        ]
    elif arg0_dtype == "float32" and arg1_dtype == "float32" and out_dtype == "float32":
        return generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumlator_dtype)
    else:
        raise NotImplementedError()

    alignment_constraints = [align for align in alignment_constraints if check_align(align)]
    assert len(alignment_constraints) > 0

    if not profile_all_alignments:
        alignment_constraints = [alignment_constraints[0]]

    def get_tile_descriptions(math_inst):
        return [
            TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
            for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
        ]

    return generate_tensor_op_common(
        math_instructions, alignment_constraints, get_tile_descriptions, op_creator
    )


def generate_sm80_tensor_op_16816(
    out_dtype,
    arg0_dtype,
    arg1_dtype,
    op_creator,
    check_align,
    use_3xtf32=True,
    profile_all_alignments=False,
    accumlator_dtype="float32",
):
    """Generate GEMM or Conv2D kernels for Ampere."""
    min_cc = 80
    max_cc = 1024
    max_cc_smem_limited = 80

    def get_default_tile_descriptions(block_k_factor):
        return [
            ([128, 256, int(32 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc),
            ([256, 128, int(32 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc),
            ([256, 64, int(32 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
            ([256, 64, int(32 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc),
            ([64, 256, int(32 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc),
            ([128, 128, int(32 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
            ([128, 128, int(32 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
            ([128, 128, int(32 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
            ([128, 64, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
            ([64, 128, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
            ([64, 64, int(32 * block_k_factor)], 10, [2, 2, 1], min_cc, max_cc),
            ([256, 128, int(64 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
            ([128, 256, int(64 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
            ([256, 64, int(64 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc_smem_limited),
            ([64, 256, int(64 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc_smem_limited),
            ([128, 128, int(64 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
            ([256, 64, int(64 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
            ([64, 256, int(64 * block_k_factor)], 3, [1, 4, 1], min_cc, max_cc),
            ([128, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
            ([128, 64, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
            ([64, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
            ([64, 64, int(64 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
        ]

    if arg0_dtype == "float16" and arg1_dtype == "float16":
        math_instructions = [
            MathInstruction(
                [16, 8, 16],
                DataType.f16,
                DataType.f16,
                dtype_map[out_dtype],
                dtype_map[accumlator_dtype],
                OpcodeClass.TensorOp,
                MathOperation.multiply_add,
            )
        ]
        alignment_constraints = [8, 4, 2]
        tile_descriptions = get_default_tile_descriptions(1)
    elif arg0_dtype == "float32" and arg1_dtype == "float32":
        math_instructions = [
            MathInstruction(
                [16, 8, 8],
                DataType.f32,
                DataType.f32,
                DataType.f32,
                DataType.f32,
                OpcodeClass.TensorOp,
                MathOperation.multiply_add_fast_f32 if use_3xtf32 else MathOperation.multiply_add,
            )
        ]
        alignment_constraints = [4, 2, 1]

        if use_3xtf32:
            # tf32
            tile_descriptions = [
                ([128, 128, 16], 4, [4, 2, 1], min_cc, max_cc),
                ([128, 128, 16], 3, [4, 2, 1], min_cc, max_cc),
                ([256, 64, 16], 3, [4, 2, 1], min_cc, max_cc),
                ([64, 256, 16], 3, [2, 4, 1], min_cc, max_cc),
                ([128, 64, 16], 4, [2, 2, 1], min_cc, max_cc),
                ([64, 128, 16], 4, [2, 2, 1], min_cc, max_cc),
                ([64, 64, 16], 3, [2, 2, 1], min_cc, max_cc),
                ([128, 128, 32], 3, [4, 2, 1], min_cc, max_cc),
                ([256, 64, 32], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
                ([64, 256, 32], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
                ([128, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
                ([64, 128, 32], 3, [2, 2, 1], min_cc, max_cc),
                ([64, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
            ]
        else:
            tile_descriptions = get_default_tile_descriptions(0.5)
    else:
        assert out_dtype == "int32"
        math_instructions = [
            MathInstruction(
                [16, 8, 32],
                dtype_map[arg0_dtype],
                dtype_map[arg1_dtype],
                DataType.s32,
                DataType.s32,
                OpcodeClass.TensorOp,
                MathOperation.multiply_add_saturate,
            )
        ]
        alignment_constraints = [16, 8, 4]
        tile_descriptions = get_default_tile_descriptions(2)

    def get_tile_descriptions(math_inst):
        return [
            TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
            for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
        ]

    alignment_constraints = [align for align in alignment_constraints if check_align(align)]

    if len(alignment_constraints) > 0 and not profile_all_alignments:
        alignment_constraints = [alignment_constraints[0]]

    if arg0_dtype != "float32" and arg1_dtype != "float32":
        sm75_kernels = generate_sm75_tensor_op_1688(
            out_dtype,
            arg0_dtype,
            arg1_dtype,
            op_creator,
            check_align,
            False,
            profile_all_alignments,
            accumlator_dtype=accumlator_dtype,
        )
    else:
        # TF32 (float32 + float32 case) is only supported on sm80
        sm75_kernels = []

    if len(alignment_constraints) > 0:
        sm80_kernels = generate_tensor_op_common(
            math_instructions, alignment_constraints, get_tile_descriptions, op_creator
        )
    else:
        sm80_kernels = []

    # TODO(masahi): For int8 kernels, The CUTLASS generator modifies the output tensor alignment
    # after ops are created. Revisit how important this modification is.
    # for op in operations:
    #     if op.tile_description.threadblock_shape[1] >= 128:
    #         op.C.alignment = 16
    #     else:
    #         op.C.alignment = 8

    return sm75_kernels + sm80_kernels


GENERATOR_FUNC_TABLE = {75: generate_sm75_tensor_op_1688, 80: generate_sm80_tensor_op_16816}


# (Epilogue functor name, no_beta_scaling)
EPILOGUE_MAP = {
    "cutlass.dense": (EpilogueFunctor.LinearCombination, False),
    "cutlass.dense_bias": (EpilogueFunctor.LinearCombinationBias, True),
    "cutlass.dense_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
    "cutlass.dense_bias_gelu_fp16": (EpilogueFunctor.LinearCombinationGelu, False),
    "cutlass.dense_bias_gelu_fp32": (EpilogueFunctor.LinearCombinationGelu, False),
    "cutlass.matmul": (EpilogueFunctor.LinearCombination, False),
    "cutlass.matmul_bias": (EpilogueFunctor.LinearCombinationBias, True),
    "cutlass.matmul_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
    "cutlass.matmul_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
    "cutlass.matmul_transposed": (EpilogueFunctor.LinearCombination, False),
    "cutlass.matmul_transposed_bias": (EpilogueFunctor.LinearCombinationBias, True),
    "cutlass.matmul_transposed_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
    "cutlass.matmul_transposed_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
    "cutlass.batch_matmul": (EpilogueFunctor.LinearCombination, False),
    "cutlass.conv2d_bias_hardswish": (EpilogueFunctor.LinearCombinationHardSwish, False),
    "cutlass.conv2d_bias_silu": (EpilogueFunctor.LinearCombinationSilu, False),
    "cutlass.conv2d_bias_sigmoid": (EpilogueFunctor.LinearCombinationSigmoid, False),
    "cutlass.conv2d_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
    "cutlass.conv2d_bias": (EpilogueFunctor.LinearCombinationBias, True),
    "cutlass.conv2d": (EpilogueFunctor.LinearCombination, False),
    "cutlass.conv2d_transpose": (EpilogueFunctor.LinearCombination, False),
    "cutlass.conv2d_backward_weight": (EpilogueFunctor.LinearCombination, False),
}


class ProfilerEngine:
    """Compile and run a given profiler executable."""

    def __init__(self, cuda_arch, cutlass_path, binary_prefix):
        self.cuda_arch = cuda_arch
        self.binary_prefix = binary_prefix
        self.cutlass = cutlass_path
        self.cflags = f"-I{cutlass_path}/include -I{cutlass_path}/tools/util/include -O3 -std=c++17"
        self.cflags += " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
        self.cflags += (
            f" -gencode=arch=compute_{cuda_arch},code=[sm_{cuda_arch},compute_{cuda_arch}]"
        )
        self.cflags += " -Xcompiler=-Wconversion -Xcompiler=-fno-strict-aliasing"
        self.cmd = "nvcc {cflags} {src} -o {output}"

    def _compile(self, op):
        os.makedirs(self.binary_prefix, exist_ok=True)
        opath = os.path.join(self.binary_prefix, op["name"])
        if os.path.exists(opath):
            return
        fi = tempfile.NamedTemporaryFile("w", delete=False, prefix=self.binary_prefix, suffix=".cu")
        fi.write(op["src"])
        fi.close()
        cmd = self.cmd.format(cflags=self.cflags, src=fi.name, output=opath)
        logger.info("invoking compilation %s", cmd)
        os.system(cmd)
        os.unlink(fi.name)

    def compile_all(self, ops, use_multiprocessing=False):
        """Compile all profiler executables."""
        if use_multiprocessing:
            pool = multiprocessing.Pool(multiprocessing.cpu_count())
            pool.map(self._compile, ops)
        else:
            for op in ops:
                self._compile(op)

    def evaluate(self, op, args):
        """Run the profiler executable corresponding to op_name with args."""
        op_name = op["name"]
        opath = os.path.join(self.binary_prefix, op_name)
        if not os.path.exists(opath):
            self._compile(op)
        if not os.path.exists(opath):
            # Bail out if compilation fails for a whatever reason (e.g. static assert failure)
            return float("inf")
        cmd = [opath]
        for arg in args:
            cmd.append(str(arg))
        try:
            logger.info("invoking evaluation %s", cmd)
            sp = subprocess.run(cmd, capture_output=True, check=True)
            rt = float(sp.stdout)
            if rt == 0.0:
                # This seems to happen with split-k using invalid split-k-slices
                rt = float("inf")
            logger.info("%s, %f", op_name, rt)
        except subprocess.CalledProcessError:
            rt = float("inf")
        return rt


class CodegenResult(Object):
    """The holder for the generated code and required headers."""

    def __init__(self, code, headers):
        self.__init_handle_by_constructor__(ffi.CodegenResult, code, headers)


def _get_optional_int_annotation(annotations, key, default=None):
    value = annotations.get(key, default)
    if value is None:
        return default
    return int(value)


@tvm_ffi.register_global_func("contrib.cutlass.instantiate_template")
def instantiate_template(func_name, annotations, func_args):
    """Return CUTLASS host code based on a template and the provided annotations.

    Parameters
    ----------
    func_name: str
        A string to identify the type of the kernel (dense/matmul, batched_matmul, or conv2d).

    annotations: tvm_ffi.Map
        Key and value pairs annotated during kernel selection.

    func_args: list
        Names of the function arguments.

    Returns
    -------
    codegen_result : CodegenResult
        Generated CUTLASS host code and required header-file names.
    """
    attrs = {}

    for k in ["lda", "ldb", "ldc", "cutlass_op_def", "cutlass_op_name", "op_type"]:
        if k in annotations:
            attrs[k] = annotations[k]

    headers = ["tvm/ffi/function.h", "tvm/ffi/extra/c_env_api.h"]

    if "relu" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h")
    elif "gelu" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_gelu.h")
    elif "sigmoid" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_sigmoid.h")
    elif "silu" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_silu.h")
    elif "hardswish" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_hardswish.h")
    else:
        headers.append("cutlass/epilogue/thread/linear_combination.h")

    if "residual" in func_name:
        headers.append("cutlass/epilogue/thread/linear_combination_residual_block.h")

    def get_dim(shape_annot, var_name, axis_idx, batched_offset=0):
        if isinstance(shape_annot, IntImm):
            return str(int(shape_annot))
        return f"{var_name}->shape[{batched_offset + axis_idx}]"

    def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_idx):
        if isinstance(stride_annot, IntImm):
            return str(int(stride_annot))
        dim1 = func_args[arg0_idx] + f"->shape[{arg0_axis_idx}]"
        dim2 = func_args[arg1_idx] + f"->shape[{arg1_axis_idx}]"
        return dim1 + " * " + dim2

    def get_flattened_batch_dim(arg_name, batch_rank):
        return " * ".join([f"{arg_name}->shape[{i}]" for i in range(batch_rank)])

    if "decode_matmul" in func_name:
        headers.append("cutlass_kernels/fpA_intB_gemm.h")
        lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
        rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
        scales_arg_idx = _get_optional_int_annotation(annotations, "scales_arg_idx", 2)
        bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)

        attrs["A_arg"] = func_args[lhs_arg_idx]
        attrs["B_arg"] = func_args[rhs_arg_idx]
        attrs["scales_arg"] = func_args[scales_arg_idx]
        attrs["activation"] = annotations.get("activation", "identity")
        attrs["bias_stride"] = annotations["bias_stride"]
        attrs["M"] = annotations["M"]
        attrs["group_size"] = annotations["group_size"]

        if not isinstance(attrs["M"], tvm.tirx.IntImm):
            attrs["M"] = get_flattened_batch_dim(
                func_args[lhs_arg_idx], int(annotations["batch_rank"])
            )

        if bias_arg_idx is not None:
            attrs["bias_arg"] = func_args[bias_arg_idx]

        if residual_arg_idx is not None:
            attrs["residual_arg"] = func_args[residual_arg_idx]
            attrs["binary_op"] = annotations["binary_op"]
            attrs["unary_op"] = annotations["unary_op"]

        if annotations["weight_nbit"] == 4:
            attrs["weight_dtype"] = "cutlass::uint4b_t"
            attrs["float_per_int"] = 2
        else:
            assert annotations["weight_nbit"] == 8
            attrs["weight_dtype"] = "uint8_t"
            attrs["float_per_int"] = 1

        code = emit_fp16A_intB_matmul(attrs)
        return CodegenResult(code, headers)

    elif "dense" in func_name or "matmul" in func_name:
        batched = "batch" in annotations
        # dense is equal to transposed_matmul
        transposed = "transposed" in func_name or "dense" in func_name
        lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
        rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
        if "bias" in func_name:
            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", 2)
        else:
            bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)

        lhs_arg = func_args[lhs_arg_idx]
        rhs_arg = func_args[rhs_arg_idx]
        lhs_shape = annotations[f"arg{lhs_arg_idx}_shape"]
        rhs_shape = annotations[f"arg{rhs_arg_idx}_shape"]
        lhs_batched_offset = len(lhs_shape) - 2
        rhs_batched_offset = len(rhs_shape) - 2

        attrs["lhs_arg"] = lhs_arg
        attrs["rhs_arg"] = rhs_arg

        if bias_arg_idx is not None:
            attrs["bias_arg"] = func_args[bias_arg_idx]
        if residual_arg_idx is not None:
            attrs["residual_arg"] = func_args[residual_arg_idx]

        attrs["ElementInputA"] = DataTypeTag[dtype_map[annotations[f"arg{lhs_arg_idx}_dtype"]]]
        attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
        attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]

        attrs["K"] = lhs_shape[lhs_batched_offset + 1]
        attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)

        if transposed:
            attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
        else:
            attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)

        if batched:
            headers.append("cutlass/gemm/device/gemm_batched.h")

            def get_batch_on_arg(arg_name, arg_shape):
                return " * ".join(f"{arg_name}->shape[{i}]" for i in range(len(arg_shape) - 2))

            if isinstance(annotations["batch"], IntImm):
                attrs["batch"] = str(int(annotations["batch"]))
            elif annotations["batch_stride_A"] == 0:
                # 2D x ND
                attrs["batch"] = get_batch_on_arg(rhs_arg, rhs_shape)
            else:
                # ND x 2D or ND x ND
                attrs["batch"] = get_batch_on_arg(lhs_arg, lhs_shape)

            attrs["batch_stride_A"] = get_batch_stride(
                annotations["batch_stride_A"],
                lhs_arg_idx,
                lhs_arg_idx,
                lhs_batched_offset,
                lhs_batched_offset + 1,
            )
            attrs["batch_stride_B"] = get_batch_stride(
                annotations["batch_stride_B"],
                rhs_arg_idx,
                rhs_arg_idx,
                rhs_batched_offset,
                rhs_batched_offset + 1,
            )

            if transposed:
                attrs["batch_stride_C"] = get_batch_stride(
                    annotations["batch_stride_C"],
                    lhs_arg_idx,
                    rhs_arg_idx,
                    lhs_batched_offset,
                    rhs_batched_offset,
                )
            else:
                attrs["batch_stride_C"] = get_batch_stride(
                    annotations["batch_stride_C"],
                    lhs_arg_idx,
                    rhs_arg_idx,
                    lhs_batched_offset,
                    rhs_batched_offset + 1,
                )
        else:
            headers.append("cutlass/gemm/device/gemm.h")

        if "residual" in func_name:
            headers.append("cutlass/gemm/device/gemm_universal_with_broadcast.h")

        code = instantiate_gemm_template(attrs)
        return CodegenResult(code, headers)

    elif "conv2d" in func_name:
        data_arg_idx = _get_optional_int_annotation(annotations, "data_arg_idx", 0)
        weight_arg_idx = _get_optional_int_annotation(annotations, "weight_arg_idx", 1)
        bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
        residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)

        attrs["data_arg"] = func_args[data_arg_idx]
        attrs["weight_arg"] = func_args[weight_arg_idx]

        if bias_arg_idx is not None:
            attrs["bias_arg"] = func_args[bias_arg_idx]
        if residual_arg_idx is not None:
            attrs["residual_arg"] = func_args[residual_arg_idx]

        activation_shape = annotations[f"arg{data_arg_idx}_shape"]
        weight_shape = annotations[f"arg{weight_arg_idx}_shape"]
        output_shape = annotations["ret_shape"]

        if "conv2d_transpose" in func_name:
            headers.append("cutlass/conv/kernel/default_conv2d_dgrad.h")
            activation_shape = output_shape
            output_shape = annotations["arg0_shape"]
        elif "backward" in func_name:
            headers.append("cutlass/conv/kernel/default_conv2d_wgrad.h")
            activation_shape = annotations["arg1_shape"]
            weight_shape = output_shape
            output_shape = annotations["arg0_shape"]
        elif "residual" in func_name:
            headers.append("cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h")
        else:
            headers.append("cutlass/conv/kernel/default_conv2d_fprop.h")

        headers.append("cutlass/conv/device/implicit_gemm_convolution.h")

        op_name = attrs["cutlass_op_name"]

        if "splitk" in op_name:
            headers += [
                "cutlass/reduction/device/reduce_split_k.h",
                "cutlass/reduction/thread/reduction_operators.h",
            ]

        data_arg = attrs["data_arg"]
        attrs["N"] = get_dim(activation_shape[0], data_arg, 0)
        attrs["H"] = get_dim(activation_shape[1], data_arg, 1)
        attrs["W"] = get_dim(activation_shape[2], data_arg, 2)
        attrs["C"] = activation_shape[3]
        attrs["P"] = get_dim(output_shape[1], "out0", 1)
        attrs["Q"] = get_dim(output_shape[2], "out0", 2)
        attrs["K"] = output_shape[3]
        attrs["R"] = weight_shape[1]
        attrs["S"] = weight_shape[2]
        attrs["pad_h"] = annotations["padding"][0]
        attrs["pad_w"] = annotations["padding"][1]
        attrs["stride_h"] = annotations["strides"][0]
        attrs["stride_w"] = annotations["strides"][1]
        attrs["dilation_h"] = annotations["dilation"][0]
        attrs["dilation_w"] = annotations["dilation"][1]

        if "splitk" in op_name:
            attrs["split_k_mode"] = "kParallel"
            attrs["split_k_slices"] = str(re.search(r"splitk(\d+)", op_name).group(1))
        else:
            attrs["split_k_mode"] = "kSerial"
            attrs["split_k_slices"] = 1

        if "residual_shape" in annotations:
            attrs["residual_shape"] = annotations["residual_shape"]

        code = instantiate_conv2d_template(attrs)
        return CodegenResult(code, headers)

    elif "attention" in func_name:
        is_var_len = "var_len" in func_name
        data_type = dtype_map[annotations["arg0_dtype"]]

        attrs["qkv_layout"] = annotations["qkv_layout"]
        if attrs["qkv_layout"] == "default":
            attrs["query"] = func_args[0]
            attrs["key"] = func_args[1]
            attrs["value"] = func_args[2]
            attrs["num_queries"] = s = get_dim(annotations["num_queries"], func_args[0], 1)
            attrs["num_keys"] = get_dim(annotations["num_keys"], func_args[1], 1)
            if len(func_args) > 4 and not is_var_len:  # +1 for workspace, the last arg
                attrs["bias"] = func_args[3]
        elif attrs["qkv_layout"] == "qkv_stacked":
            attrs["qkv"] = func_args[0]
            attrs["num_queries"] = s = annotations["num_queries"]
            attrs["num_keys"] = annotations["num_keys"]
            if len(func_args) > 2 and not is_var_len:  # +1 for workspace, the last arg
                attrs["bias"] = func_args[1]
        else:
            raise NotImplementedError()

        attrs["data_type"] = DataTypeTag[data_type]
        attrs["num_batches"] = b = annotations["num_batches"]
        attrs["head_dim"] = h = annotations["head_dim"]
        attrs["head_dim_value"] = h_v = annotations["head_dim_value"]
        attrs["kMaxK"] = max(int(attrs["head_dim"]), int(attrs["head_dim_value"]))
        attrs["scale"] = (
            float(1 / math.sqrt(h.value)) if annotations["scale"] is None else annotations["scale"]
        )

        if is_var_len:
            attrs["seqstart_q"] = func_args[int(annotations["seqstart_q_idx"])]
            attrs["seqstart_k"] = func_args[int(annotations["seqstart_k_idx"])]
            attrs["max_seqlen_q"] = func_args[int(annotations["max_seqlen_q_idx"])]
            attrs["max_seqlen_k"] = func_args[int(annotations["max_seqlen_k_idx"])]

        is_mqa = annotations["num_q_heads"] != annotations["num_kv_heads"]

        use_flash = (
            annotations["ret_dtype"] == "float16"
            and "bias" not in attrs
            and int(attrs["head_dim"]) <= 256
            and int(attrs["head_dim"]) % 8 == 0
            and int(attrs["head_dim"]) == int(attrs["head_dim_value"])
            # For the causal case (custom mask = "BottomRight"), only use flash for multi-query
            # attention workloads. Otherwise, CUTLASS fMHA seems faster for causal attention
            # with a single query.
            # In addition, sliding-window attention is only supported by flash.
            and (
                int(annotations["custom_mask_type"]) == 0
                or (int(annotations["custom_mask_type"]) == 2 and is_mqa)
                or (int(annotations["custom_mask_type"]) == 2 and "window_size" in annotations)
            )
            # Flash v2 is currently not supported for sm < 80
            and int(annotations["arch"]) >= 80
        )

        # See https://github.com/Dao-AILab/flash-attention/blob/
        # 92dd5703ecdb99aa4a4aee9817f28557907403a2/csrc/flash_attn/flash_api.cpp#L111-L116
        if "window_size" in annotations:
            assert use_flash, "Sliding-window attention is supported only by Flash Attention."
            assert int(annotations["custom_mask_type"]) == 2, (
                "Sliding-window attention is only supported for causal with bottom right mask."
            )
            attrs["window_size_left"] = int(annotations["window_size"]) - 1
            attrs["window_size_right"] = 0
            attrs["is_causal"] = False
        else:
            if int(annotations["custom_mask_type"]) == 2:
                attrs["window_size_left"] = attrs["num_keys"]
                attrs["window_size_right"] = 0
                attrs["is_causal"] = True
            else:
                attrs["window_size_left"] = -1
                attrs["window_size_right"] = -1
                attrs["is_causal"] = False

        if use_flash:
            headers.append("flash.h")
            attrs["num_q_heads"] = annotations["num_q_heads"]
            attrs["num_kv_heads"] = annotations["num_kv_heads"]

            if is_var_len:
                code = instantiate_flash_attention_var_len_template(attrs)
            else:
                code = instantiate_flash_attention_template(attrs)
        else:
            headers.append("kernel_forward.h")

            assert not is_mqa, (
                "The number of query and KV heads need to be the same for CUTLASS fMHA."
            )

            attrs["num_heads"] = n = annotations["num_q_heads"]

            data_type_size = DataTypeSize[data_type]
            if (data_type_size * h // 8) % 16 == 0 and (data_type_size * h_v // 8) % 16 == 0:
                attrs["kIsAligned"] = True
            elif (h % 4 == 0) and (h_v % 4 == 0):
                attrs["kIsAligned"] = False
            else:
                raise NotImplementedError()
            if h_v > 64:
                attrs["kQueriesPerBlock"] = 32
                attrs["kKeysPerBlock"] = 128
                attrs["kSingleValueIteration"] = h_v <= 128
            else:
                attrs["kQueriesPerBlock"] = 64
                attrs["kKeysPerBlock"] = 64
                attrs["kSingleValueIteration"] = True

            assert attrs["scale"] > 0 or attrs["scale"] < 0, (
                "Cutlass may generate nan occasionally when scale == 0.0"
            )
            attrs["arch"] = "cutlass::arch::Sm{}".format(annotations["arch"])
            attrs["kSupportsDropout"] = False

            attrs["output_size"] = f"{b} * {s} * {n} * {h_v}"

            attrs["custom_mask_type"] = annotations["custom_mask_type"]

            for arg in func_args:
                if "workspace" in arg:
                    attrs["workspace"] = arg
            if "bias" in attrs:
                attrs["kSupportsBias"] = True
                if len(annotations["bias_shape"]) == 4:
                    strides = "p.num_keys"
                    if annotations["bias_shape"][2] == 1:
                        attrs["bias_strideM"] = 0
                    else:
                        attrs["bias_strideM"] = strides
                        strides = f"p.num_queries * {strides}"
                    if annotations["bias_shape"][1] == 1:
                        attrs["bias_strideH"] = 0
                    else:
                        attrs["bias_strideH"] = strides
                        strides = f"p.num_heads * {strides}"
                    if annotations["bias_shape"][0] == 1:
                        attrs["bias_strideB"] = 0
                    else:
                        attrs["bias_strideB"] = strides
                else:
                    raise NotImplementedError()
            else:
                # To support negative scale in current Cutlass implementation,
                # kSupportsBias should be set true, or there are nan's as result.
                attrs["kSupportsBias"] = attrs["scale"] < 0

            code = instantiate_attention_template(attrs)

        return CodegenResult(code, headers)
    elif "layer_norm" in func_name:
        headers.append("cutlass/util/device_layernorm.h")
        headers.append("cutlass/layout/matrix.h")
        attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
        attrs.update(dict(annotations))

        if not isinstance(attrs["M"], tvm.tirx.IntImm):
            attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))

        code = instantiate_layer_norm_template(attrs)
        return CodegenResult(code, headers)
    elif "rms_norm" in func_name:
        headers.append("cutlass/util/device_rmsnorm.h")
        headers.append("cutlass/layout/matrix.h")
        attrs = {"input": func_args[0], "weight": func_args[1]}
        attrs.update(dict(annotations))

        if not isinstance(attrs["M"], tvm.tirx.IntImm):
            attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))

        code = instantiate_rms_norm_template(attrs)
        return CodegenResult(code, headers)

    raise ValueError(f"Do not have a template for {func_name}")
