import enum
import os
from itertools import chain, product

from .cutlass_library import (
    enum_auto,
    DataTypeNames,
    DataTypeSize,
    DataType,
    DataTypeTag,
    GemmKind,
    GemmKindNames,
    KernelScheduleType,
    KernelScheduleTag,
    KernelScheduleSuffixes,
    EpilogueScheduleType,
    EpilogueScheduleTag,
    EpilogueScheduleSuffixes,
)


################################################################################
# Epilogue Tag enum and string utils
class TrtLlm_EpilogueTag(enum.Enum):
    epilogue_op_default = enum_auto()
    epilogue_op_bias = enum_auto()
    epilogue_op_silu = enum_auto()
    epilogue_op_gelu = enum_auto()


class TrtLlm_EpilogueFusion(enum.Enum):
    epilogue_fusion_none = enum_auto()
    epilogue_fusion_finalize = enum_auto()


EpiTagNames = {
    TrtLlm_EpilogueTag.epilogue_op_default: "lc",  # linear combination
    TrtLlm_EpilogueTag.epilogue_op_bias: "lc_bias",  # linear combination with bias addition
    TrtLlm_EpilogueTag.epilogue_op_silu: "silu",  # silu or swiglu
    TrtLlm_EpilogueTag.epilogue_op_gelu: "gelu",  # gelu or geglu
}

EpiTag = {
    TrtLlm_EpilogueTag.epilogue_op_default: "tensorrt_llm::cutlass_extensions::EpilogueOpDefault",
    TrtLlm_EpilogueTag.epilogue_op_bias: "tensorrt_llm::cutlass_extensions::EpilogueOpBias",
    TrtLlm_EpilogueTag.epilogue_op_silu: "tensorrt_llm::cutlass_extensions::EpilogueOpDefaultSilu",
    TrtLlm_EpilogueTag.epilogue_op_gelu: "tensorrt_llm::cutlass_extensions::EpilogueOpDefaultFtGelu",
}

EpiFusion = {
    TrtLlm_EpilogueFusion.epilogue_fusion_none: "tensorrt_llm::TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE",
    TrtLlm_EpilogueFusion.epilogue_fusion_finalize: "tensorrt_llm::TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE",
}

EpiFusionSuffixes = {
    None: "",
    TrtLlm_EpilogueFusion.epilogue_fusion_none: "EpilogueFusion_NONE",
    TrtLlm_EpilogueFusion.epilogue_fusion_finalize: "EpilogueFusion_FINALIZE",
}


################################################################################
# Quantization Operation and string utils
class TrtLlm_QuantOp(enum.Enum):
    per_column_scale_only = enum_auto()
    finegrained_scale_only = enum_auto()
    finegrained_scale_and_zeros = enum_auto()
    none = enum_auto()


QuantOpNames = {
    TrtLlm_QuantOp.per_column_scale_only: "cs",
    TrtLlm_QuantOp.finegrained_scale_only: "fgs",
    TrtLlm_QuantOp.finegrained_scale_and_zeros: "fgsz",
    TrtLlm_QuantOp.none: "noquant",
}

QuantOpTag = {
    TrtLlm_QuantOp.per_column_scale_only: "cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY",
    TrtLlm_QuantOp.finegrained_scale_only: "cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY",
    TrtLlm_QuantOp.finegrained_scale_and_zeros: "cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_AND_ZEROS",
    TrtLlm_QuantOp.none: "void",
}

################################################################################
# The activations, biases, scales and zeros are instantiated using CUDA types,
# not CUTLASS types. This map materializes the name of the CUDA type.


class e2m1_type:  # WAR until we have upgraded everything to a supported version
    pass


e2m1 = e2m1_type()


def GetDataTypeBits(type):
    if isinstance(type, e2m1_type):
        return 4
    return DataTypeSize[type]


def GetDataTypeNames(type, is_mx_fpx=None):
    mxprefix = ""
    if is_mx_fpx is not None:
        mxprefix = "mx_" if is_mx_fpx else "nv_"
    if isinstance(type, e2m1_type):
        return mxprefix + "e2m1"
    return mxprefix + DataTypeNames[type]


CudaTypeName = {
    e2m1: "SafeFP4",
    DataType.e4m3: "__nv_fp8_e4m3",
    DataType.bf16: "__nv_bfloat16",
    DataType.f16: "half",
    DataType.f32: "float",
    DataType.e2m1: "__nv_fp4_e2m1",
    DataType.ue8m0: "cutlass::float_ue8m0_t",
    DataType.u4: "cutlass::uint4b_t",
}

MixedInputScaleModeTag = {
    "post_mma": "cutlass::gemm::collective::MixedInputScaleMode::kPostMma",
    "pre_mma_e8m0": "cutlass::gemm::collective::MixedInputScaleMode::kPreMmaE8M0",
}

Sm90MixedInputKernelScheduleTag = {
    "single_warpgroup_prefill": "tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_PREFILL",
    "single_warpgroup_rolling": "tensorrt_llm::cutlass_extensions::MainloopScheduleType::SINGLE_WARPGROUP_ROLLING",
}


################################################################################
# A data structure holding all info to instantiate gemm launchers in TRT LLM.
class TrtLlm_GemmLauncher:
    def __init__(
        self,
        gemm_kind,
        arch,
        act_type,
        weight_type,
        scalezero_type,
        bias_type,
        output_type,
        quant_op,
        epi_tag,
        cta_shape,
        warp_shape,
        stages,
        cga_shape,
        mainloop_schedule,
        epi_schedule,
        epi_fusion=None,
        is_mx_fpx=False,
        mixed_input_scale_mode="post_mma",
        sm90_mixed_input_kernel_type="warp_specialized",
        dynamic_cga=False,
        swap_ab=False,
    ):
        self.gemm_kind = gemm_kind
        self.arch = arch
        self.act_type = act_type
        self.weight_type = weight_type
        self.scalezero_type = scalezero_type
        self.bias_type = bias_type
        self.output_type = output_type
        self.quant_op = quant_op
        self.epi_tag = epi_tag
        self.cta_shape = cta_shape
        self.warp_shape = warp_shape
        self.stages = stages
        self.cga_shape = cga_shape
        self.dynamic_cga = dynamic_cga
        self.mainloop_schedule = mainloop_schedule
        self.epi_schedule = epi_schedule
        self.epi_fusion = epi_fusion
        self.is_mx_fpx = is_mx_fpx
        self.mixed_input_scale_mode = mixed_input_scale_mode
        self.sm90_mixed_input_kernel_type = sm90_mixed_input_kernel_type
        self.swap_ab = swap_ab

    def __repr__(self):
        kernel_prefix = "{}_sm{}_{}_{}_{}_{}_{}_{}_{}_{}x{}x{}_{}x{}x{}_{}".format(
            GemmKindNames[self.gemm_kind],
            self.arch,
            GetDataTypeNames(self.act_type, self.is_mx_fpx),
            GetDataTypeNames(self.weight_type, self.is_mx_fpx),
            GetDataTypeNames(self.scalezero_type),
            GetDataTypeNames(self.bias_type),
            GetDataTypeNames(self.output_type),
            QuantOpNames[self.quant_op],
            EpiTagNames[self.epi_tag],
            self.cta_shape[0],
            self.cta_shape[1],
            self.cta_shape[2],
            self.warp_shape[0],
            self.warp_shape[1],
            self.warp_shape[2],
            self.stages,
        )

        hopper_suffix = "_{}x{}x{}{}{}{}{}{}".format(
            self.cga_shape[0],
            self.cga_shape[1],
            self.cga_shape[2],
            KernelScheduleSuffixes[self.mainloop_schedule],
            EpilogueScheduleSuffixes[self.epi_schedule],
            EpiFusionSuffixes[self.epi_fusion],
            "_mxfpx_" if self.is_mx_fpx else "",
            "_swap_ab" if self.swap_ab else "",
        )
        if self.mixed_input_scale_mode != "post_mma":
            hopper_suffix += f"_{self.mixed_input_scale_mode}"
        if self.sm90_mixed_input_kernel_type != "warp_specialized":
            hopper_suffix += f"_{self.sm90_mixed_input_kernel_type}"

        if self.arch >= 90:
            return kernel_prefix + hopper_suffix
        elif self.arch > 100:
            raise ValueError(f"SM{self.arch} not supported yet.")
        return kernel_prefix


################################################################################
def tuple_to_cute_shape(shape):
    return f"cute::Shape<cute::Int<{shape[0]}>, cute::Int<{shape[1]}>, cute::Int<{shape[2]}>>"


def instantiate_operation_tma_warp_specialized(operation):
    act_tag = CudaTypeName[operation.act_type]
    scale_zero_tag = CudaTypeName[operation.scalezero_type]
    bias_tag = CudaTypeName[operation.bias_type]
    out_tag = CudaTypeName[operation.output_type]

    quant_op = QuantOpTag[operation.quant_op]
    epi_tag = EpiTag[operation.epi_tag]

    cute_cta_shape = tuple_to_cute_shape(operation.cta_shape)
    cute_cga_shape = tuple_to_cute_shape(operation.cga_shape)

    kernel_sched = KernelScheduleTag[operation.mainloop_schedule]
    epi_sched = "void"
    if operation.epi_schedule is not None:
        epi_sched = EpilogueScheduleTag[operation.epi_schedule]

    if operation.gemm_kind == GemmKind.Gemm:
        weight_tag = DataTypeTag[operation.weight_type]
        instantiation = f"""
template void sm90_generic_mixed_gemm_kernelLauncher<{act_tag}, {weight_tag}, {scale_zero_tag}, {bias_tag}, {out_tag},
{quant_op}, {epi_tag},
{cute_cta_shape}, {cute_cga_shape},
{kernel_sched}, {epi_sched}> (
const {act_tag}*, const {weight_tag}*, const {scale_zero_tag}*, const {scale_zero_tag}*, const {bias_tag}*, const float,
{out_tag}*, int, int, int, const int, tensorrt_llm::cutlass_extensions::CutlassGemmConfig, char*, size_t, cudaStream_t, int*
);"""
    elif operation.gemm_kind == GemmKind.Grouped:
        if operation.arch == 90 and operation.act_type != operation.weight_type:
            # Mixed MoE GEMM
            weight_tag = CudaTypeName[operation.weight_type]
            optional_template_args = []
            if operation.mixed_input_scale_mode != "post_mma":
                optional_template_args.append(
                    MixedInputScaleModeTag[operation.mixed_input_scale_mode]
                )
            if operation.sm90_mixed_input_kernel_type != "warp_specialized":
                if operation.mixed_input_scale_mode == "post_mma":
                    optional_template_args.append(
                        MixedInputScaleModeTag[operation.mixed_input_scale_mode]
                    )
                optional_template_args.append(
                    Sm90MixedInputKernelScheduleTag[
                        operation.sm90_mixed_input_kernel_type
                    ]
                )
            optional_template_args = "".join(
                f", {arg}" for arg in optional_template_args
            )
            launcher_name = (
                "sm90_generic_mixed_moe_gemm_kernelLauncher"
                if operation.sm90_mixed_input_kernel_type == "warp_specialized"
                else "sm90_generic_mixed_moe_small_k_kernelLauncher"
            )
            instantiation = f"""
template void {launcher_name}<{act_tag}, {weight_tag}, {out_tag},
{epi_tag}, {cute_cta_shape}, {cute_cga_shape}, {kernel_sched}, {epi_sched}, {quant_op}{optional_template_args}> (
GroupedGemmInput<{act_tag}, {weight_tag}, {out_tag}, {out_tag}>inputs, TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size);
"""
        else:
            # Similar to MixedInput above, we must modify the tags for grouped gemm as CUTLASS library does not have the updated schedules
            assert operation.mainloop_schedule in [
                KernelScheduleType.TmaWarpSpecializedCooperative,
                KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum,
            ]
            kernel_sched.replace("::Kernel", "::KernelGrouped")
            # epi_sched += "Grouped"
            # arch_tag = f"cutlass::arch::Sm{operation.arch}"
            arch_tag = f"Sm{operation.arch}"
            weight_tag = CudaTypeName[operation.weight_type]
            assert operation.epi_fusion is not None
            epi_fusion = EpiFusion[operation.epi_fusion]

            # We need to remove the '::' because this will break the instantiation macro
            epi_fusion = epi_fusion.split(":")[-1]
            epi_tag = epi_tag.split(":")[-1]
            epi_sched = epi_sched.split(":")[-1]
            epi_sched = epi_sched.replace(
                "1Sm", ""
            )  # Hack to WAR missing `PtrArrayTmaWarpSpecialized` type

            guard_map = {
                e2m1: "defined(ENABLE_FP4)",
                DataType.e4m3: "defined(ENABLE_FP8)",
                DataType.bf16: "defined(ENABLE_BF16)",
            }
            guard_act = guard_map.get(operation.act_type, "1")
            guard_weight = guard_map.get(operation.weight_type, "1")

            is_mx_fpx = str(operation.is_mx_fpx).lower()
            use_dynamic_cga = str(operation.dynamic_cga).lower()
            use_bias = str(False).lower()
            swap_ab = str(operation.swap_ab).lower()
            # TODO Revert this once compiler bug is fixed so we can use template instead of macro again
            #         instantiation = f"""
            #         template void tma_warp_specialized_generic_moe_gemm_kernelLauncher<{arch_tag}, {act_tag}, {weight_tag}, {out_tag},
            #                 {epi_tag}, {epi_fusion}, {cute_cta_shape}, {cute_cga_shape}, false>
            #                 (TmaWarpSpecializedGroupedGemmInput, int, int, cudaStream_t, int*, size_t*);
            # """
            instantiation = f"""
#if {guard_act} && {guard_weight}
        INSTANTIATE_TMA_WARP_SPECIALIZED_MOE_GEMM({arch_tag}, {act_tag}, {weight_tag}, {out_tag},
        {epi_sched}, {epi_tag}, {epi_fusion},
        {operation.cta_shape[0]}, {operation.cta_shape[1]}, {operation.cta_shape[2]}, {operation.cga_shape[0]}, {operation.cga_shape[1]}, {operation.cga_shape[2]},
        {is_mx_fpx}, {use_dynamic_cga}, {use_bias}, {swap_ab});
#endif"""
    return instantiation


def instantiate_operation_sm80(operation):
    act_tag = DataTypeTag[operation.dtype]
    weight_tag = DataTypeTag[operation.dtype]
    epi_tag = EpiTag[operation.epi_tag]

    instantiation = f"""
            template void sm80_generic_fused_moe_gemm_kernelLauncher<{act_tag}, {weight_tag}, {operation.cta_shape[0]}, {operation.cta_shape[1]}, {operation.cta_shape[2]}, {operation.stage}, {epi_tag}>
                    ({act_tag} const* A, {weight_tag} const* B, {act_tag} const* biases, bool bias_is_broadcast, {act_tag}* C, int64_t const* total_tokens_including_expert, int64_t num_rows, int64_t gemm_n, int64_t gemm_k, int num_experts, int multi_processor_count, cudaStream_t stream, int* kernel_occupancy);"""
    return instantiation


def instantiate_operation(operation):
    if operation.arch == 80:
        return instantiate_operation_sm80(operation)
    elif operation.arch >= 90:
        return instantiate_operation_tma_warp_specialized(operation)


def get_file_content(launcher_inl_files, operations):
    assert operations
    include_list = list()
    for file in launcher_inl_files:
        include_list.append(f'#include "{file}"')
    includes = "\n".join(include_list)

    insts_list = list()
    for op in operations:
        insts_list.append(instantiate_operation(op))
    instantiations = "\n".join(insts_list)

    file_content = f"""{includes}
namespace tensorrt_llm
{{
namespace kernels
{{
namespace cutlass_kernels_oss
{{

{instantiations}

}} // namespace cutlass_kernels_oss
}} // namespace kernels
}} // namespace tensorrt_llm
"""
    return file_content


def clean_leftover_files(output_dir, generated_files):
    """Remove leftover generated files that weren't created in this run."""
    for root, _dirs, files in os.walk(output_dir):
        for file in files:
            file_path = os.path.join(root, file)
            if file_path not in generated_files:
                os.remove(file_path)


def write_file(launcher_inl_files, operations, output_file):
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    # Avoid changing modified time if file content is up to date
    content = get_file_content(launcher_inl_files, operations)
    try:
        with open(output_file, mode="r") as f:
            if f.read() == content:
                return
    except FileNotFoundError:
        pass
    with open(output_file, mode="w") as f:
        f.write(content)


def is_gemm_op_valid_sm100(op):
    # TODO These are much more restricted than theory dictates, investigate if more can be enabled in future
    tile_m, tile_n, _ = op.cta_shape
    cga_m, cga_n, cga_k = op.cga_shape

    if (
        op.epi_fusion == TrtLlm_EpilogueFusion.epilogue_fusion_finalize
        and op.epi_schedule != EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm
    ):
        return False

    # We use a runtime cluster shape for SM100, so we only use cluster shapes to distinguish between 1SM and 2SM variants.
    if cga_m > 2 or cga_n != 1 or cga_k != 1:
        return False

    if op.arch == 103:
        return (
            op.act_type == e2m1
            and op.weight_type == e2m1
            and tile_m == 128
            and tile_n in [128, 256]
        )

    # Default shapes
    # This is epilogue tile size. For two CTA this is actually size 128/256 for the MMA
    if tile_m not in [64, 128]:
        return False

    # FP4 Has some much more limited sizes
    if op.act_type == e2m1 or op.weight_type == e2m1:
        if tile_n not in [64, 128, 256] or tile_m != 128:
            return False
        # TODO Revert this once cutlass adds support for blockscaled + no smem
        if (
            op.arch == 100
            and op.epi_schedule == EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm
        ):
            return False

    # MXFP block-scaled paths currently follow the same shape/schedule limits as FP4 block scaling.
    if op.is_mx_fpx:
        if tile_n not in [64, 128, 256] or tile_m != 128:
            return False
        if (
            op.arch == 100
            and op.epi_schedule == EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm
        ):
            return False

    # Shapes for fp8 small N shapes
    if (
        (op.act_type == DataType.e4m3)
        and (tile_n == 16 or tile_n == 8)
        and (cga_m == 1 and cga_n == 1)
    ):
        # todo: double check why tile_n = 8 is disabled in CUTLASS backend. @yuhan
        return tile_m != 128 or tile_n % 16 == 0

    # Default alignment requirements
    if tile_n % 32 != 0 or tile_n < 32 or tile_n > 256:
        return False

    # Two CTA mode needs bigger tile n alignment
    if cga_m % 2 == 0 and tile_n % 64 != 0:
        return False

    return True


def is_gemm_op_valid(op):
    tile_m, tile_n, _ = op.cta_shape
    cga_m, cga_n, _ = op.cga_shape

    if cga_m == 1 and cga_n == 1:
        return True

    if cga_m == 2 and cga_n == 1 and tile_m >= 128:
        return True

    if cga_m == 1 and cga_n == 2 and tile_n >= 128:
        return True

    if cga_m == 2 and cga_n == 2 and tile_m >= 128 and tile_n >= 128:
        return True

    return False


def is_grouped_gemm_op_valid(op):
    if not is_gemm_op_valid(op):
        return False

    if op.epi_tag != TrtLlm_EpilogueTag.epilogue_op_default:
        return False

    if (
        op.epi_schedule is not None
        and op.epi_schedule != EpilogueScheduleType.NoSmemWarpSpecialized
    ):
        return False

    if op.mainloop_schedule not in [
        KernelScheduleType.TmaWarpSpecializedCooperative,
        KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum,
    ]:
        return False

    return True


def is_op_valid(op):
    if op.arch >= 100:
        return is_gemm_op_valid_sm100(op)

    if op.gemm_kind == GemmKind.Gemm:
        return is_gemm_op_valid(op)
    if op.gemm_kind == GemmKind.Grouped:
        return is_grouped_gemm_op_valid(op)


################################################################################
def generate_sm90_mixed_gemm_operations():
    arch = 90

    # For legacy reasons, we use unsigned types for the weights. The instanitated template
    # will remap those back to the signed type.
    # Takes the form (activation_type, weight_type, scalezero_type, bias_type, output_type)
    supported_dtypes = [
        (DataType.e4m3, DataType.u4, DataType.f16, DataType.f16, DataType.f16),
        (DataType.e4m3, DataType.u4, DataType.f16, DataType.bf16, DataType.bf16),
        (DataType.f16, DataType.u4, DataType.f16, DataType.f16, DataType.f16),
        (DataType.bf16, DataType.u4, DataType.bf16, DataType.bf16, DataType.bf16),
        (DataType.f16, DataType.u8, DataType.f16, DataType.f16, DataType.f16),
        (DataType.bf16, DataType.u8, DataType.bf16, DataType.bf16, DataType.bf16),
    ]

    quant_ops = [
        TrtLlm_QuantOp.per_column_scale_only,
        TrtLlm_QuantOp.finegrained_scale_only,
        TrtLlm_QuantOp.finegrained_scale_and_zeros,
    ]

    epi_tags = [TrtLlm_EpilogueTag.epilogue_op_bias]

    M_TILES = [64, 128]
    N_TILES = [16, 32, 64, 128, 256]
    cta_shapes_mn = product(M_TILES, N_TILES)

    warp_shape = [4, 1, 1]
    stages = 0  # auto

    cga_shapes = product([1, 2], [1, 2], [1])

    partial_args = product(
        supported_dtypes, quant_ops, epi_tags, cta_shapes_mn, cga_shapes
    )

    operations = list()
    for dtype_combo, quant_op, epi_tag, cta_shape_mn, cga_shape in partial_args:
        max_k_bits = 128 * 8
        cta_shape_k = max_k_bits // GetDataTypeBits(dtype_combo[0])
        cta_shape_mnk = cta_shape_mn + (cta_shape_k,)

        use_coop = cta_shape_mn[0] == 128
        mainloop_schedule = (
            KernelScheduleType.TmaWarpSpecializedCooperative
            if use_coop
            else KernelScheduleType.TmaWarpSpecializedPingpong
        )
        epi_schedule = (
            EpilogueScheduleType.TmaWarpSpecializedCooperative
            if use_coop
            else EpilogueScheduleType.TmaWarpSpecialized
        )

        fpA_intB_operation = TrtLlm_GemmLauncher(
            GemmKind.Gemm,
            arch,
            *dtype_combo,
            quant_op,
            epi_tag,
            cta_shape_mnk,
            warp_shape,
            stages,
            cga_shape,
            mainloop_schedule,
            epi_schedule,
        )

        if is_op_valid(fpA_intB_operation):
            operations.append(fpA_intB_operation)

    return operations


def generate_sm90_grouped_gemm_operations(is_arch_enabled):
    if not is_arch_enabled:
        return []
    arch = 90
    supported_dtypes = [DataType.f16, DataType.bf16, DataType.f32, DataType.e4m3]
    quant_ops = [TrtLlm_QuantOp.none]
    epi_tags = [TrtLlm_EpilogueTag.epilogue_op_default]
    M_TILES = [128]  # Currently M tile must be 128 for Grouped GEMM
    N_TILES = [16, 32, 64, 128, 256]
    cta_shapes_mn = list(product(M_TILES, N_TILES)) + [(256, 128)]

    warp_shape = [0, 0, 0]  # ignored except for naming
    stages = 0  # auto

    epi_fusions = [
        TrtLlm_EpilogueFusion.epilogue_fusion_none,
        TrtLlm_EpilogueFusion.epilogue_fusion_finalize,
    ]

    swap_ab = [True, False]

    cga_shapes = product([1, 2], [1, 2], [1])

    partial_args = product(
        supported_dtypes,
        quant_ops,
        epi_tags,
        epi_fusions,
        cta_shapes_mn,
        cga_shapes,
        swap_ab,
    )

    operations = list()
    for (
        dtype,
        quant_op,
        epi_tag,
        epi_fusion,
        cta_shape_mn,
        cga_shape,
        swap_ab,
    ) in partial_args:
        max_k_bits = 128 * 8
        cta_shape_k = max_k_bits // GetDataTypeBits(dtype)
        cta_shape_mnk = cta_shape_mn + (cta_shape_k,)

        mainloop_schedule = (
            KernelScheduleType.TmaWarpSpecializedCooperative
            if dtype != DataType.e4m3
            else KernelScheduleType.TmaWarpSpecializedCooperativeFP8FastAccum
        )
        epi_schedule = None

        otypes = [dtype]
        if dtype == DataType.e4m3:
            otypes = [DataType.f16, DataType.bf16]

        for otype in otypes:
            moe_gemm_operation = TrtLlm_GemmLauncher(
                GemmKind.Grouped,
                arch,
                dtype,
                dtype,
                dtype,
                dtype,
                otype,
                quant_op,
                epi_tag,
                cta_shape_mnk,
                warp_shape,
                stages,
                cga_shape,
                mainloop_schedule,
                epi_schedule,
                epi_fusion,
                swap_ab=swap_ab,
            )

            if is_op_valid(moe_gemm_operation):
                operations.append(moe_gemm_operation)
    return operations


def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled):
    if not is_arch_enabled:
        return []
    arch = 90

    # act_type, weight_type, scalezero_type, bias_type, output_type
    supported_dtypes_int4 = [
        (DataType.e4m3, DataType.u4, DataType.f16, DataType.f16, DataType.f16),
        (DataType.e4m3, DataType.u4, DataType.bf16, DataType.bf16, DataType.bf16),
    ]

    supported_dtypes_fp4 = [
        (
            DataType.bf16,
            e2m1,
            DataType.ue8m0,
            DataType.bf16,
            DataType.bf16,
        ),
        (
            DataType.e4m3,
            e2m1,
            DataType.ue8m0,
            DataType.f16,
            DataType.f16,
        ),
        (
            DataType.e4m3,
            e2m1,
            DataType.ue8m0,
            DataType.bf16,
            DataType.bf16,
        ),
    ]

    quant_ops = [TrtLlm_QuantOp.finegrained_scale_only]

    epi_tags = [TrtLlm_EpilogueTag.epilogue_op_default]

    cta_shapes_mnk_mixed_input = list(product([64], [16, 32, 64, 128], [128, 256, 512]))
    cta_shapes_mnk_mixed_input.extend(
        product([128], [16, 32, 64, 128], [128, 256, 512])
    )
    cta_shapes_mnk_mixed_input.extend([(128, 256, k_tile) for k_tile in [128, 256]])
    cta_shapes_mnk_mixed_input.extend([(256, 128, k_tile) for k_tile in [128, 256]])
    cta_shapes_mnk_mixed_input.append((256, 256, 128))
    cta_shapes_mnk_int4 = list(cta_shapes_mnk_mixed_input)
    cta_shapes_mnk_fp4 = list(cta_shapes_mnk_mixed_input)

    warp_shape = [0, 0, 0]  # ignored except for naming
    stages = 0  # auto

    cga_shapes = list(product([1, 2], [1, 2], [1]))

    partial_args_int4 = product(
        supported_dtypes_int4, quant_ops, epi_tags, cta_shapes_mnk_int4, cga_shapes
    )
    partial_args_fp4 = product(
        supported_dtypes_fp4, quant_ops, epi_tags, cta_shapes_mnk_fp4, cga_shapes
    )
    partial_args = chain(partial_args_int4, partial_args_fp4)

    operations = list()
    for dtype_combo, quant_op, epi_tag, cta_shape_mnk, cga_shape in partial_args:
        is_fp8_mxfp4 = dtype_combo[0] == DataType.e4m3 and dtype_combo[1] == e2m1
        mixed_input_scale_mode = "pre_mma_e8m0" if is_fp8_mxfp4 else "post_mma"

        use_coop = cta_shape_mnk[0] >= 128
        mainloop_schedules = (
            [
                KernelScheduleType.TmaWarpSpecializedCooperative,
                KernelScheduleType.TmaWarpSpecializedPingpong,
            ]
            if use_coop
            else [KernelScheduleType.TmaWarpSpecializedPingpong]
        )
        epi_schedule = EpilogueScheduleType.TmaWarpSpecializedCooperative
        for mainloop_schedule in mainloop_schedules:
            moe_gemm_operation = TrtLlm_GemmLauncher(
                GemmKind.Grouped,
                arch,
                *dtype_combo,
                quant_op,
                epi_tag,
                cta_shape_mnk,
                warp_shape,
                stages,
                cga_shape,
                mainloop_schedule,
                epi_schedule,
                mixed_input_scale_mode=mixed_input_scale_mode,
            )
            operations.append(moe_gemm_operation)

    small_k_shapes = [(128, token_tile, 128) for token_tile in [8, 16, 32, 40]]
    small_k_kernel_types = [
        "single_warpgroup_prefill",
        "single_warpgroup_rolling",
    ]
    for dtype_combo in supported_dtypes_fp4:
        if dtype_combo[0] != DataType.e4m3:
            continue
        for cta_shape_mnk, kernel_type in product(small_k_shapes, small_k_kernel_types):
            operations.append(
                TrtLlm_GemmLauncher(
                    GemmKind.Grouped,
                    arch,
                    *dtype_combo,
                    TrtLlm_QuantOp.finegrained_scale_only,
                    TrtLlm_EpilogueTag.epilogue_op_default,
                    cta_shape_mnk,
                    warp_shape,
                    3,
                    (1, 1, 1),
                    KernelScheduleType.TmaWarpSpecializedPingpong,
                    EpilogueScheduleType.TmaWarpSpecializedCooperative,
                    mixed_input_scale_mode="pre_mma_e8m0",
                    sm90_mixed_input_kernel_type=kernel_type,
                )
            )
    return operations


def generate_sm90_operations(is_arch_enabled):
    operations = generate_sm90_mixed_gemm_operations()
    operations.extend(generate_sm90_grouped_gemm_operations(is_arch_enabled))
    operations.extend(generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled))
    return operations


def calc_shape_mnk_sm100_grouped_gemm(cta_shape_mn, dtype):
    max_k_bits = 128 * 8
    cta_shape_k = max_k_bits // GetDataTypeBits(dtype)
    if dtype == DataType.e4m3 and (cta_shape_mn[1] == 8):
        cta_shape_k = 256
    return cta_shape_mn + (cta_shape_k,)


def generate_sm120_grouped_gemm_operations(is_arch_enabled):
    if not is_arch_enabled:
        return []
    arch = 120
    supported_dtypes = [e2m1, (DataType.e4m3, e2m1)]
    quant_ops = [TrtLlm_QuantOp.none]
    epi_tags = [TrtLlm_EpilogueTag.epilogue_op_default]
    cta_shapes_mnk = [
        [128, 32, 128],
        [128, 32, 256],
        [128, 64, 128],
        [128, 64, 256],
        [128, 128, 128],
        [128, 128, 256],
        [256, 128, 128],
        [128, 256, 128],
    ]

    warp_shape = [0, 0, 0]  # ignored except for naming
    stages = 0  # auto

    epi_fusions = [
        TrtLlm_EpilogueFusion.epilogue_fusion_none,
        TrtLlm_EpilogueFusion.epilogue_fusion_finalize,
    ]

    cga_shapes = [[1, 1, 1]]

    swap_ab = [True, False]

    partial_args = product(
        supported_dtypes,
        quant_ops,
        epi_tags,
        epi_fusions,
        cta_shapes_mnk,
        cga_shapes,
        swap_ab,
    )

    operations = list()
    for (
        dtype,
        quant_op,
        epi_tag,
        epi_fusion,
        cta_shape_mnk,
        cga_shape,
        swap_ab,
    ) in partial_args:
        # Ignored
        mainloop_schedule = KernelScheduleType.TmaWarpSpecializedCooperative
        epi_schedule = None

        if isinstance(dtype, tuple):
            act_type, weight_type = dtype
        else:
            act_type, weight_type = dtype, dtype

        # Minimal filter: for mixed FP8xFP4 on SM120/SM121, only emit 128xNx128
        if act_type == DataType.e4m3 and weight_type == e2m1:
            if cta_shape_mnk not in ([128, 32, 128], [128, 64, 128], [128, 128, 128]):
                continue

        otypes = [act_type]
        if act_type in [DataType.e4m3, e2m1]:
            otypes = [DataType.f16, DataType.bf16]

        for otype in otypes:
            moe_gemm_operation = TrtLlm_GemmLauncher(
                GemmKind.Grouped,
                arch,
                act_type,
                weight_type,
                act_type,
                act_type,
                otype,
                quant_op,
                epi_tag,
                cta_shape_mnk,
                warp_shape,
                stages,
                cga_shape,
                mainloop_schedule,
                epi_schedule,
                epi_fusion,
                is_mx_fpx=(act_type == DataType.e4m3 and weight_type == e2m1),
                swap_ab=swap_ab,
            )

            operations.append(moe_gemm_operation)
    return operations


def generate_sm120_operations(is_arch_enabled):
    operations = generate_sm120_grouped_gemm_operations(is_arch_enabled)
    return operations


def generate_sm100_grouped_gemm_operations(is_arch_enabled, arch):
    if not is_arch_enabled:
        return []
    supported_dtypes = [
        DataType.f16,
        DataType.bf16,
        DataType.f32,
        DataType.e4m3,
        e2m1,
        (DataType.e4m3, e2m1),
    ]
    quant_ops = [TrtLlm_QuantOp.none]
    epi_tags = [TrtLlm_EpilogueTag.epilogue_op_default]
    cta_shapes_m = [64, 128]
    cta_shapes_n = [8, 16, 32, 64, 128, 192, 256]
    cta_shapes_mn = product(cta_shapes_m, cta_shapes_n)

    warp_shape = [0, 0, 0]  # ignored except for naming
    stages = 0  # auto

    epi_fusions = [
        TrtLlm_EpilogueFusion.epilogue_fusion_none,
        TrtLlm_EpilogueFusion.epilogue_fusion_finalize,
    ]

    # Some shapes for SM100 are better with NoSmem, note the kernel will internally map to the 1 or 2 SM variants based on the cga_shape[0]
    epi_schedules = [
        EpilogueScheduleType.PtrArrayNoSmemWarpSpecialized1Sm,
        EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm,
    ]

    # We will use dynamic cluster shapes for SM100, so we only need to indicate if we are using 1 or 2 SM version
    cga_shapes = [(1, 1, 1), (2, 1, 1)]

    swap_ab = [True, False]

    dynamic_cga = [True, False]

    partial_args = product(
        supported_dtypes,
        quant_ops,
        epi_tags,
        epi_fusions,
        cta_shapes_mn,
        cga_shapes,
        epi_schedules,
        dynamic_cga,
        swap_ab,
    )

    operations = list()
    for (
        dtype,
        quant_op,
        epi_tag,
        epi_fusion,
        cta_shape_mn,
        cga_shape,
        epi_schedule,
        dynamic_cga,
        swap_ab,
    ) in partial_args:
        if isinstance(dtype, tuple):
            dtype, weight_type = dtype
        else:
            weight_type = dtype

        cta_shape_mnk = calc_shape_mnk_sm100_grouped_gemm(cta_shape_mn, dtype)

        # Ignored
        mainloop_schedule = KernelScheduleType.TmaWarpSpecializedCooperative

        otypes = [dtype]
        if dtype in [DataType.e4m3, e2m1]:
            otypes = [DataType.f16, DataType.bf16]

        mxfp_modes = [False]
        if dtype == DataType.e4m3 and weight_type == e2m1:
            # MXFP8 x MXFP4 path.
            mxfp_modes = [True]
        elif dtype == DataType.e4m3 and weight_type == DataType.e4m3:
            # Emit both regular FP8xFP8 and MXFP8xMXFP8 variants.
            mxfp_modes = [False, True]

        for otype in otypes:
            for is_mx_fpx in mxfp_modes:
                moe_gemm_operation = TrtLlm_GemmLauncher(
                    GemmKind.Grouped,
                    arch,
                    dtype,
                    weight_type,
                    otype,
                    otype,
                    otype,
                    quant_op,
                    epi_tag,
                    cta_shape_mnk,
                    warp_shape,
                    stages,
                    cga_shape,
                    mainloop_schedule,
                    epi_schedule,
                    epi_fusion,
                    is_mx_fpx=is_mx_fpx,
                    dynamic_cga=dynamic_cga,
                    swap_ab=swap_ab,
                )

                if is_op_valid(moe_gemm_operation):
                    operations.append(moe_gemm_operation)
    return operations


def generate_sm103_operations(is_arch_enabled):
    operations = generate_sm100_grouped_gemm_operations(is_arch_enabled, 103)
    return operations


def generate_sm100_operations(is_arch_enabled):
    operations = generate_sm100_grouped_gemm_operations(is_arch_enabled, 100)
    return operations


class GemmSm80LauncherConfig:
    def __init__(self, gemm_kind, arch, dtype, epi_tag, cta_shape, stage):
        self.gemm_kind = gemm_kind
        self.arch = arch
        self.dtype = dtype
        self.epi_tag = epi_tag
        self.cta_shape = cta_shape
        self.stage = stage


def generate_sm80_fused_grouped_gemm_operations():
    arch = 80
    supported_dtypes = [DataType.f16, DataType.bf16]
    epi_tags = [
        TrtLlm_EpilogueTag.epilogue_op_silu,
        TrtLlm_EpilogueTag.epilogue_op_gelu,
    ]
    cta_shapes_mnk = [
        (16, 128, 64),
        (16, 256, 64),
        (32, 128, 64),
        (64, 128, 64),
        (128, 128, 64),
    ]

    stages = [2, 3, 4]

    partial_args = product(supported_dtypes, epi_tags, cta_shapes_mnk, stages)

    operations = list()
    for dtype, epi_tag, cta_shape_mnk, stage in partial_args:
        item = GemmSm80LauncherConfig(
            GemmKind.Grouped, arch, dtype, epi_tag, cta_shape_mnk, stage
        )
        operations.append(item)
    return operations


def generate_sm80_operations(is_arch_enabled):
    operations = generate_sm80_fused_grouped_gemm_operations()
    return operations


def generate_gemm_operations(output_dir, architectures):
    arches = architectures.split(";")
    # Get the absolute path of the provided directory
    output_dir = os.path.abspath(output_dir)

    fpA_intB_inl = "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/launchers/fpA_intB_launcher_sm90.inl"
    moe_gemm_inl = "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl"
    # moe_gemm_inl = "tensorrt_llm/kernels/internal_cutlass_kernels/src/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl"
    moe_mixed_gemm_inl = "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl"
    # moe_mixed_gemm_inl = "tensorrt_llm/kernels/internal_cutlass_kernels/src/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl"
    sm80_moe_gemm_inl = "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl"
    # sm80_moe_gemm_inl = "tensorrt_llm/kernels/internal_cutlass_kernels/src/moe_gemm/launchers/fused_moe_gemm_launcher_sm80.inl"

    inl_map = {
        (GemmKind.Gemm, 90): [fpA_intB_inl],
        (GemmKind.Grouped, 90): [moe_gemm_inl],
        (GemmKind.Grouped, 100): [moe_gemm_inl],
        (GemmKind.Grouped, 103): [moe_gemm_inl],
        (GemmKind.Grouped, 120): [moe_gemm_inl],
        (GemmKind.Grouped, 80): [sm80_moe_gemm_inl],
    }

    def has_arch(sm):
        return (
            f"{sm}" in arches
            or f"{sm}-real" in arches
            or f"{sm}f-real" in arches
            or f"{sm}f" in arches
        )

    # The goal here is to group kernels with common instantiations together in order to reduce template instantiation overheads.
    # Template instantiation dominates the time in a compilation unit, so it is the most important factor to improve.
    operations = []
    operations += generate_sm120_operations(has_arch(120) or has_arch(121))
    operations += generate_sm103_operations(has_arch(103))
    operations += generate_sm100_operations(
        has_arch(100) or has_arch(103) or has_arch(107)
    )
    operations += generate_sm90_operations(has_arch(90))
    operations += generate_sm80_operations(has_arch(80) or has_arch(89))

    def should_skip(op):
        return False  # All kernels have a public implementation

    # SM90 mixed-input grouped GEMMs have a dedicated launcher.
    def is_mixed_dtype_grouped(op):
        if isinstance(op, GemmSm80LauncherConfig):
            return False
        return (
            (op.arch == 90)
            and (op.act_type != op.weight_type)
            and (op.gemm_kind == GemmKind.Grouped)
        )

    # Fix OOM error in CI. If len(operations) is more than GROUP_SIZE, it will be split into multiple sub groups.
    GROUP_SIZE = 8
    op_groups = dict()
    for op in operations:
        if should_skip(op):
            continue
        # This dict key is used to group kernels with common instantiations together
        # Similar implementations should live in the same file so the compiler can share the cutlass state
        # Without this we see significant memory consumption, and separating them also does not reduce the compilation time
        # because most time is spent parsing the same cutlass files
        # We separate by: Architecture, Leading dimension of the CTA shape, FP4 (i.e. block scaled MMA), mixed input
        # TODO Do a more scientific analysis of this
        dict_key = (
            op.gemm_kind,
            op.arch,
            op.cta_shape[0],
            op.arch >= 100 and (op.weight_type == e2m1 or op.is_mx_fpx),
            is_mixed_dtype_grouped(op),
        )
        op_group = op_groups.get(dict_key, [])
        if len(op_group) == 0 or len(op_group[-1]) >= GROUP_SIZE:
            op_group.append([op])
        else:
            op_group[-1].append(op)
        op_groups[dict_key] = op_group

    file_list = []
    for key, value in op_groups.items():
        gemm_kind, arch, m, block_scale, is_mixed = key
        for i, op_sub_group in enumerate(value):
            out_file = os.path.join(
                output_dir,
                GemmKindNames[gemm_kind],
                str(arch),
                f"cutlass_kernel_file_{GemmKindNames[gemm_kind]}_sm{arch}_M{m}{'_BS' if block_scale else ''}{'_Mixed' if is_mixed else ''}_group{i}.generated.cu",
            )
            inl_file = [moe_mixed_gemm_inl] if is_mixed else inl_map[key[:2]]
            write_file(inl_file, op_sub_group, out_file)
            file_list.append(out_file)

    # Clean up any leftover files from previous runs
    clean_leftover_files(output_dir, set(file_list))
