#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

CLUSTER_K = 1
OLD_GENERATED_PREFIX = "moe_gemm_tma_ws_sm90_mixed_fp4"
NEW_GENERATED_PREFIX = "moe_gemm_tma_ws_sm90_fp4"


@dataclass(frozen=True, order=True)
class Instantiation:
    type_name: str
    cpp_type: str
    m: int
    n: int
    k: int
    cluster_m: int
    cluster_n: int
    schedule: str
    fusion: str = "none"  # "none" or "finalize"
    quick_build: bool = False

    @property
    def file_name(self) -> str:
        fusion_suffix = "" if self.fusion == "none" else f"_{self.fusion}"
        return (
            f"{NEW_GENERATED_PREFIX}_{self.type_name}_m{self.m}_n{self.n}_k{self.k}"
            f"_cm{self.cluster_m}_cn{self.cluster_n}_{self.schedule}{fusion_suffix}.generated.cu"
        )


def add_full_build_configs(instantiations: set[Instantiation], type_name: str, cpp_type: str) -> None:
    cluster_shapes = ((1, 1), (2, 1), (1, 2), (2, 2))
    k_tiles = (128, 256)
    fusions = ("none", "finalize")

    for k in k_tiles:
        for fusion in fusions:
            for m in (64,):
                for n in (16, 32, 64):
                    for cluster_m, cluster_n in cluster_shapes:
                        instantiations.add(
                            Instantiation(type_name, cpp_type, m, n, k, cluster_m, cluster_n, "pp", fusion)
                        )

            for n in (16, 32, 64):
                for cluster_m, cluster_n in cluster_shapes:
                    instantiations.add(
                        Instantiation(type_name, cpp_type, 128, n, k, cluster_m, cluster_n, "pp", fusion)
                    )
                    instantiations.add(
                        Instantiation(type_name, cpp_type, 128, n, k, cluster_m, cluster_n, "co", fusion)
                    )

            for cluster_m, cluster_n in cluster_shapes:
                instantiations.add(Instantiation(type_name, cpp_type, 128, 128, k, cluster_m, cluster_n, "pp", fusion))


def get_instantiations() -> list[Instantiation]:
    instantiations: set[Instantiation] = set()

    add_full_build_configs(instantiations, "fp16", "half")
    add_full_build_configs(instantiations, "bf16", "__nv_bfloat16")

    # Quick-build: a subset of configs for faster compilation
    for k in (128, 256):
        for n in (16, 32, 64, 128):
            instantiations.add(Instantiation("fp16", "half", 128, n, k, 1, 1, "pp", "none", quick_build=True))
            instantiations.add(Instantiation("fp16", "half", 128, n, k, 1, 1, "pp", "finalize", quick_build=True))

    return sorted(instantiations)


def render(instantiation: Instantiation) -> str:
    bf16_open = "#ifdef ENABLE_BF16\n" if instantiation.type_name == "bf16" else ""
    bf16_close = "#endif  // ENABLE_BF16\n" if instantiation.type_name == "bf16" else ""
    quick_open = "" if instantiation.quick_build else "#ifndef ORT_QUICK_BUILD\n"
    quick_close = "" if instantiation.quick_build else "#endif  // !ORT_QUICK_BUILD\n"

    if instantiation.fusion == "finalize":
        if instantiation.schedule == "pp":
            macro = "ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_PP_FINALIZE"
        else:
            macro = "ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_CO_FINALIZE"
    else:
        macro = (
            "ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_PP"
            if instantiation.schedule == "pp"
            else "ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_CO"
        )

    return f"""/*
 * Copyright (c) Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License.
 *
 * Generated by generate_moe_gemm_tma_ws_sm90_fp4.py. Do not edit manually.
 */

#ifdef COMPILE_HOPPER_TMA_GROUPED_GEMMS
#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE)
{quick_open}{bf16_open}#include "contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_sm90_fp4_instantiation.cuh"

namespace onnxruntime::llm::kernels::cutlass_kernels {{

{macro}({instantiation.cpp_type}, {instantiation.m}, {instantiation.n}, {instantiation.k}, {instantiation.cluster_m}, {instantiation.cluster_n}, {CLUSTER_K});

}}  // namespace onnxruntime::llm::kernels::cutlass_kernels

{bf16_close}{quick_close}#endif  // ENABLE_FP4 && USE_FP4_QMOE
#endif  // COMPILE_HOPPER_TMA_GROUPED_GEMMS
"""


def main() -> None:
    launcher_dir = Path(__file__).resolve().parent
    generated_names = {instantiation.file_name for instantiation in get_instantiations()}

    for generated_file in launcher_dir.glob(f"{OLD_GENERATED_PREFIX}*.generated.cu"):
        generated_file.unlink()
    for generated_file in launcher_dir.glob(f"{NEW_GENERATED_PREFIX}_*.generated.cu"):
        if generated_file.name not in generated_names:
            generated_file.unlink()

    for instantiation in get_instantiations():
        (launcher_dir / instantiation.file_name).write_text(render(instantiation), encoding="utf-8")


if __name__ == "__main__":
    main()
