# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
from pathlib import Path

import jinja2

FILE_HEAD = """
// auto generated by generate_kernels.py
// clang-format off

#include "kernel.h"
#include "marlin_template.h"

namespace MARLIN_NAMESPACE_NAME {
""".strip()

FILE_TAIL = "}\n"

TEMPLATE = ("template __global__ void Marlin<"
            "{{scalar_t}}, "
            "{{w_type_id}}, "
            "{{s_type_id}}, "
            "{{threads}}, "
            "{{thread_m_blocks}}, "
            "{{thread_n_blocks}}, "
            "{{thread_k_blocks}}, "
            "{{'true' if m_block_size_8 else 'false'}}, "
            "{{stages}}, "
            "{{group_blocks}}, "
            "{{'true' if is_zp_float else 'false'}}>"
            "( MARLIN_KERNEL_PARAMS );")

# int8 with zero point case (vllm::kU8) is also supported,
# we don't add it to reduce wheel size.
SCALAR_TYPES = [
    "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn",
    "vllm::kFE2M1f"
]
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128),
                  (128, 64, 128)]

THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
# group_blocks:
#   = 0 : act order case
#   = -1 : channelwise quantization
#   > 0 : group_size=16*group_blocks
GROUP_BLOCKS = [0, 1, -1, 2, 4, 8]
DTYPES = ["fp16", "bf16"]


def remove_old_kernels() -> None:
    root = Path(__file__).parent
    for path in root.glob("kernel_*.cu"):
        path.unlink(missing_ok=True)

def _write_kernel_file(scalar_type: str, dtype: str, templates: list[str]) -> Path:
    root = Path(__file__).parent
    scalar_suffix = scalar_type.split("::", 1)[1].lower() if "::" in scalar_type else scalar_type.lower()
    output_path = root / f"kernel_{dtype}_{scalar_suffix}.cu"

    lines = [FILE_HEAD, "", f"// Instantiations for dtype={dtype}, weight={scalar_type}", ""]
    lines.append("\n".join(templates))
    lines.append("")
    lines.append(FILE_TAIL)

    output_path.write_text("\n".join(lines), encoding="utf-8")
    return output_path


def render_templates_for_combo(scalar_type: str, dtype: str) -> list[str]:
    results: list[str] = []
    for group_blocks, m_blocks, thread_configs in itertools.product(
            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):

        # act order case only support gptq-int4 and gptq-int8
        if group_blocks == 0 and scalar_type not in [
                "vllm::kU4B8", "vllm::kU8B128"
        ]:
            continue
        if thread_configs[2] == 256:
            # for small batch (m_blocks == 1), we only need (128, 128, 256)
            # for large batch (m_blocks > 1), we only need (64, 256, 256)
            if m_blocks <= 1 and thread_configs[0] != 128:
                continue
            if m_blocks > 1 and thread_configs[0] != 64:
                continue

        # we only support channelwise quantization and group_size == 128
        # for fp8
        if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
            continue
        # nvfp4 only supports group_size == 16
        # mxfp4 only supports group_size == 32
        if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
            continue
        # other quantization methods don't support group_size = 16
        if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
            continue

        k_blocks = thread_configs[0] // 16
        n_blocks = thread_configs[1] // 16
        threads = thread_configs[2]

        c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"

        is_zp_float_list = [False]
        if dtype == "fp16" and scalar_type == "vllm::kU4" and \
                group_blocks == 4:
            # HQQ (is_zp_float = true) only supports
            # 4bit quantization and fp16
            is_zp_float_list.append(True)

        if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
            s_type = "vllm::kFE4M3fn"
        elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
            s_type = "vllm::kFE8M0fnu"
            if dtype == "fp16":
                # we cannot safely dequantize e8m0 to fp16, so skip this
                continue
        elif dtype == "fp16":
            s_type = "vllm::kFloat16"
        elif dtype == "bf16":
            s_type = "vllm::kBFloat16"

        for is_zp_float in is_zp_float_list:
            template_str = jinja2.Template(TEMPLATE).render(
                scalar_t=c_dtype,
                w_type_id=scalar_type + ".id()",
                s_type_id=s_type + ".id()",
                threads=threads,
                thread_m_blocks=max(m_blocks, 1),
                thread_n_blocks=n_blocks,
                thread_k_blocks=k_blocks,
                m_block_size_8=m_blocks == 0.5,
                stages="pipe_stages",
                group_blocks=group_blocks,
                is_zp_float=is_zp_float,
            )

            results.append(template_str)

    return results


def generate_new_kernels() -> None:
    emitted = False
    for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
        templates = render_templates_for_combo(scalar_type, dtype)
        if not templates:
            continue

        _write_kernel_file(scalar_type, dtype, templates)
        emitted = True

    if not emitted:
        raise RuntimeError("No marlin kernels were generated; check template configuration.")


if __name__ == "__main__":
    remove_old_kernels()
    generate_new_kernels()
