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

"""Generate SM120 FP4-based MoE GEMM kernel instantiations.

SM120 (Blackwell SM120) supports native FP4 grouped GEMM via CUTLASS TMA
warp-specialized kernels.  This script generates one .cu file per
(output_type, tile_shape) combination, each containing the
INSTANTIATE_TMA_WARP_SPECIALIZED_MOE_GEMM macro call.

Supported combinations:
  - FP4xFP4: activation=FP4, weight=FP4, output=fp16/bf16
    Tile shapes (element counts): 128x128x128, 128x128x256, 128x256x128, 256x128x128
    Uses NV-native nv_float4_t with ue4m3 SF (IsMXFPX=false).

  - FP8xFP4: activation=FP8 (e4m3), weight=FP4, output=fp16/bf16
    Tile shapes (element counts): 128x128x128
    Uses MX-format tuple with ue8m0 SF (IsMXFPX=true).
    Larger tiles (128x256, 256x128) exceed shared memory for >= 2 pipeline stages.
    K=64 is not supported (TMA tile size constraint for MXF8F6F4 on SM120).

Cluster shape: 1x1x1 only.
Fusion: NONE only (isValidSM120MOESpecialisation constraint).
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

GENERATED_PREFIX = "moe_gemm_tma_ws_sm120_fp4"
CLUSTER_M, CLUSTER_N, CLUSTER_K = 1, 1, 1


@dataclass(frozen=True, order=True)
class Instantiation:
    act_type_name: str  # "fp4" or "fp8"
    act_cpp_type: str  # "SafeFP4" or "SafeFP8"
    wt_cpp_type: str  # "SafeFP4"
    out_type_name: str  # "fp16" or "bf16"
    out_cpp_type: str  # "half" or "SafeBF16"
    m: int
    n: int
    k: int
    is_mxfpx: bool  # True for FP8xFP4

    @property
    def file_name(self) -> str:
        return (
            f"{GENERATED_PREFIX}_{self.act_type_name}_{self.out_type_name}_m{self.m}_n{self.n}_k{self.k}.generated.cu"
        )


def get_instantiations() -> list[Instantiation]:
    out_types = [
        ("fp16", "half"),
        ("bf16", "SafeBF16"),
    ]

    instantiations: set[Instantiation] = set()

    # FP4xFP4: all four tile shapes (K in FP4 elements)
    fp4_tile_shapes = [
        (128, 128, 128),
        (128, 128, 256),
        (128, 256, 128),
        (256, 128, 128),
    ]
    for out_name, out_cpp in out_types:
        for m, n, k in fp4_tile_shapes:
            instantiations.add(Instantiation("fp4", "SafeFP4", "SafeFP4", out_name, out_cpp, m, n, k, False))

    # FP8xFP4: only 128x128x128 fits in smem with >= 2 pipeline stages
    fp8_fp4_tile_shapes = [
        (128, 128, 128),
    ]
    for out_name, out_cpp in out_types:
        for m, n, k in fp8_fp4_tile_shapes:
            instantiations.add(Instantiation("fp8", "SafeFP8", "SafeFP4", out_name, out_cpp, m, n, k, True))

    return sorted(instantiations)


def render(inst: Instantiation) -> str:
    bf16_open = "#ifdef ENABLE_BF16\n" if inst.out_type_name == "bf16" else ""
    bf16_close = "#endif  // ENABLE_BF16\n" if inst.out_type_name == "bf16" else ""
    mxfpx_str = "true" if inst.is_mxfpx else "false"
    is_fp8 = inst.act_type_name == "fp8"
    fp8_open = "#ifdef ENABLE_FP8\n" if is_fp8 else ""
    fp8_close = "#endif  // ENABLE_FP8\n" if is_fp8 else ""

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

#ifndef EXCLUDE_SM_120
#ifdef COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS
#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE)
{fp8_open}{bf16_open}
#include "contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl"

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

INSTANTIATE_TMA_WARP_SPECIALIZED_MOE_GEMM(Sm120, {inst.act_cpp_type}, {inst.wt_cpp_type}, {inst.out_cpp_type}, EpilogueOpDefault, NONE, {inst.m}, {inst.n}, {inst.k}, {CLUSTER_M}, {CLUSTER_N}, {CLUSTER_K}, {mxfpx_str}, false)

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

{bf16_close}{fp8_close}#endif  // ENABLE_FP4 && USE_FP4_QMOE
#endif  // COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS
#endif  // EXCLUDE_SM_120
"""


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

    # Clean up stale generated files
    for generated_file in launcher_dir.glob(f"{GENERATED_PREFIX}_*.generated.cu"):
        if generated_file.name not in generated_names:
            generated_file.unlink()

    for inst in get_instantiations():
        path = launcher_dir / inst.file_name
        content = render(inst)
        if path.exists() and path.read_text(encoding="utf-8") == content:
            print(f"File {path.name} is up to date")
            continue
        path.write_text(content, encoding="utf-8")
        print(f"Generated {path.name}")


if __name__ == "__main__":
    main()
