# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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.
"""generate gemm_fused_kernel code."""

import argparse
import os
import re


def get_candidate_tiles():
    """
    获取候选的tile配置列表。

    Args:
        无参数。

    Returns:
        List[Tuple[str, str, str]]: 包含tile配置的三元组列表，每个三元组中的字符串表示tile的形状"。

    """
    base_configs = [("<64, 64, 64>", "<32, 32, 64>", "<16, 8, 32>")]
    base_configs.extend(
        [
            ("<16, 32, 64>", "<16, 32, 64>", "<16, 8, 32>"),
            ("<16, 64, 64>", "<16, 32, 64>", "<16, 8, 32>"),
            ("<32, 128, 64>", "<32, 32, 64>", "<16, 8, 32>"),
            ("<64, 128, 64>", "<32, 64, 64>", "<16, 8, 32>"),
            ("<64, 64, 128>", "<32, 64, 64>", "<16, 8, 32>"),
            ("<64, 128, 64>", "<64, 32, 64>", "<16, 8, 32>"),
            ("<128, 64, 64>", "<64, 32, 64>", "<16, 8, 32>"),
            ("<128, 128, 64>", "<64, 32, 64>", "<16, 8, 32>"),
            ("<128, 128, 64>", "<64, 64, 64>", "<16, 8, 32>"),
            ("<128, 128, 64>", "<128, 32, 64>", "<16, 8, 32>"),
            ("<128, 256, 64>", "<64, 64, 64>", "<16, 8, 32>"),
            ("<256, 128, 64>", "<64, 64, 64>", "<16, 8, 32>"),
            ("<128, 64, 128>", "<64, 32, 128>", "<16, 8, 32>"),
            ("<16, 256, 128>", "<16, 64, 128>", "<16, 8, 32>"),
        ]
    )

    return base_configs


def get_candidate_configs(sm, min_stages, max_stages):
    """
    获取候选的gemm算子配置列表。

    Args:
        sm (str): 计算能力，如"70"。
    """
    tiles = get_candidate_tiles()
    candidate_configs = list()

    stages = tuple(i for i in range(min_stages, max_stages + 1, 1))
    hasbias = ("false", "true")

    candidate_configs.extend([(stages, tiles, hasbias)])

    return candidate_configs


# this is a file's header part
CommonHead = """// Generated by generate_code_gemm_fused_kernels.py - Do not edit.

#pragma once

#include "fp8_gemm_fused/visitor_fp8_gemm_fused_template.h"
"""

CommonTail = """

"""

GemmDeclare = """
template<>
bool dispatch_visitor_fuse_gemm<phi::dtype::{input_type}, phi::dtype::{output_type},
                        cutlass::gemm::GemmShape{thread_block_shape}, cutlass::gemm::GemmShape{warp_shape},
                        cutlass::gemm::GemmShape{mma_shape}, {num_stages}, {hasbias}, {SM}>(GemmEpilogueAllParams);


"""

LaunchGemmHead = """
#pragma once

#include "fp8_gemm_fused/visitor_fp8_gemm_fused.h"

"""

LaunchGemmDeclare = """
bool launch_visitor_gemm_fused_kernel_{gemm_config}(const int type_id, GemmEpilogueAllParams params);
"""

LaunchGemmPart0 = """
#pragma once

#include "fp8_gemm_fused/visitor_fp8_gemm_fused_template.h"

bool launch_visitor_gemm_fused_kernel_{gemm_config}(const int type_id, GemmEpilogueAllParams params){
    switch (type_id) {"""

LaunchGemmPart1 = """
        case {type_id}:
            return dispatch_visitor_fuse_gemm<phi::dtype::{input_type}, phi::dtype::{output_type},
                                    cutlass::gemm::GemmShape{thread_block_shape}, cutlass::gemm::GemmShape{warp_shape},
                                    cutlass::gemm::GemmShape{mma_shape}, {num_stages}, {hasbias}, {SM}>(params);
            break;"""

LaunchGemmPart2 = """
        default:
            throw std::runtime_error("cutlass gemm config is invalid.");
            break;
    }
    return false;
}
"""

code_part0 = """// Generated by generate_code_gemm_fused_kernels.py - Do not edit.

#include <map>
#include <regex>
#include <limits>
#include "helper.h"
#include "visitor_fp8_gemm_fused.h"
#include "launch_visitor_gemm_fused_kernel.h"

COMMON_DECLARE_string(use_cutlass_device_best_config_path);

std::map<std::string, int> per_channel_gemm_type_map{"""

code_part1 = """
    {"{input_type}_{output_type}_{hasbias}",   {type_id}}, """

code_part2 = """
};

std::map<std::string, int> per_channel_gemm_config_map{
"""

code_part3 = """    {"{thread_block_shape}, {warp_shape}, {mma_shape}, {num_stages}", {tile_id}},
"""

code_part4 = """};

bool launch_visitor_gemm_fused_kernel(const int type_id, const int kernel_id, GemmEpilogueAllParams params){
    switch (kernel_id) {"""

code_part5 = """
        case {tile_id}:
            return launch_visitor_gemm_fused_kernel_{gemm_config}(type_id, params);
            break;"""

code_part6 = """
        default:
            throw std::runtime_error("fp8_fp8_bf16_gemm_fused Config is invalid.");
            break;
    }
    return false;
}



bool fp8_visitor_gemm_fused(GemmEpilogueAllParams params) {
  if (per_channel_gemm_type_map.find(params.fuse_gemm_config) == per_channel_gemm_type_map.end()) {
    throw std::runtime_error("fp8_visitor_gemm_fused config is invalid.");
  }

  int type_id = per_channel_gemm_type_map[params.fuse_gemm_config];
  int M = (params.M+31)/32 *32;
  int N = params.N;
  int K = params.K;

  std::string mnk_string = "per_channel_gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
  auto encoded_mnk_string = base64_encode(mnk_string);

  int kernel_id;
  std::string best_config;
  CutlassGemmConfigMannager& best_config_mannager = CutlassGemmConfigMannager::getInstance();
  if(getenv("FLAGS_use_cutlass_device_best_config_path")){ // run kernel
    std::string config_file_path = getenv("FLAGS_use_cutlass_device_best_config_path");
    nlohmann::json* config_json = new nlohmann::json();
    if (config_file_path != "default") {
        config_json = best_config_mannager.get_gemm_best_configs(config_file_path);
    }

    best_config = get_relative_best<std::string>(config_json, encoded_mnk_string, "<64, 64, 64>, <32, 32, 64>, <16, 8, 32>, 3");

    if (per_channel_gemm_config_map.find(best_config) == per_channel_gemm_config_map.end()) {
        throw std::runtime_error("This config'kernel not be generate, please check auto_gen_visitor_fp8_gemm_fused_kernels.py and re-generate.");
    } else {
        kernel_id = per_channel_gemm_config_map[best_config];
    }
    return launch_visitor_gemm_fused_kernel(type_id, kernel_id, params);
  } else { // tune kernel
    int warm_up_times = 5;
    int tune_times = 10;
    std::string best_kernel_id = "";
    float duratation = 1000000.f;
    // tune all kernel_id kernels
    for(const auto& config_pair : per_channel_gemm_config_map){
        bool is_valid = true;
        // warm up
        for(int num_time = 0; num_time < warm_up_times; ++num_time){
            if(!launch_visitor_gemm_fused_kernel(type_id, config_pair.second, params)){
                is_valid = false;
                break;
            }
        }
        if(!is_valid){
            continue;
        }
        cudaEvent_t start, stop;
        cudaEventCreate(&start);
        cudaEventCreate(&stop);
        cudaStreamSynchronize(params.stream);
        cudaEventRecord(start, params.stream);
        for(int num_time = 0; num_time < tune_times; ++num_time){
            if(!launch_visitor_gemm_fused_kernel(type_id, config_pair.second, params)){
                is_valid = false;
                break;
            };
        }
        cudaEventRecord(stop, params.stream);
        cudaEventSynchronize(stop);
        float elapsedTime;
        if(is_valid){
            cudaEventElapsedTime(&elapsedTime, start, stop);
        } else {
            continue;
        }
        cudaEventDestroy(start);
        cudaEventDestroy(stop);
        if(elapsedTime < duratation){
            best_kernel_id = config_pair.first;
            duratation = elapsedTime;
        }
    }


    nlohmann::json new_json;
    new_json[encoded_mnk_string] = best_kernel_id;
    best_config_mannager.up_date_configs(new_json);
    std::cout <<"Gemm tune result for " << mnk_string<< ": best config is: "<< best_kernel_id << std::endl;
    return true;
  }
}
"""


def SubstituteTemplate(template, values):
    """
    生成函数模板
    """
    text = template
    changed = True
    while changed:
        changed = False
        for key, value in values.items():
            regex = f"\\{{{key}\\}}"
            newtext = re.sub(regex, value, text)
            if newtext != text:
                changed = True
            text = newtext
    return text


def parse_args():
    """
    代码参数解析
    """
    parser = argparse.ArgumentParser(
        description="The argument for generating the generic_mixed_gemm_kernelLauncher instance."
    )
    parser.add_argument(
        "--cuda_arch",
        type=str,
        nargs="+",
        default=["89"],
        help="The CUDA architecture to be generated.",
    )

    parser.add_argument(
        "--min_stages",
        type=int,
        default=2,
        help="The min stages for the gemm kernel.",
    )

    parser.add_argument(
        "--max_stages",
        type=int,
        default=8,
        help="The max stages for the gemm kernel.",
    )

    args = parser.parse_args()
    return args


# generate source .cu
def generate_source_cu(
    inputs_type: str,
    outputs_type: str,
    stages: int,
    tiles: str,
    hasbiases: str,
    sm: str,
):
    """
    生成.cu源文件
    """
    all_code = SubstituteTemplate(CommonHead, {})

    for input_type in inputs_type:
        for output_type in outputs_type:
            for stage in stages:
                for hasbias in hasbiases:
                    for tile_config in tiles:
                        value_dict = {
                            "input_type": input_type,
                            "output_type": output_type,
                            "thread_block_shape": tile_config[0],
                            "warp_shape": tile_config[1],
                            "mma_shape": tile_config[2],
                            "num_stages": str(stage),
                            "hasbias": hasbias,
                            "SM": sm,
                        }
                        all_code += SubstituteTemplate(GemmDeclare, value_dict)

    all_code += CommonTail
    return all_code


# generate gemm launch .cu
def generate_launch_gemm_cus(
    generate_dir: str,
    inputs_type: str,
    outputs_type: str,
    stages: int,
    tiles: str,
    hasbiases: str,
    sm: str,
):
    """
    生成含有CUDA执行器的launch_visitor_gemm_fused_kernel.h和launch_visitor_gemm_fused_kernel_*.cu文件。

    Args:
        generate_dir (str): 生成文件的目录路径。
        inputs_type (str): 输入类型，可以是"float", "half"或者"bfloat16"中的一种。
        outputs_type (str): 输出类型，可以是"float", "half"或者"bfloat16"中的一种。
        stages (int): Gemm算子的阶段数。
        tiles (str): 包含三个元素的列表，每个元素都是包含三个元素的列表，分别代表线程块形状、整体线程形状和MMA形状。例如：["32,8,4","16,8,4"]。
        hasbiases (str): 是否包含偏置量，可以是"true"或者"false"中的一种。
        sm (str): GPU的SM大小，可以是"70"或者"80"中的一种。

    Returns:
        tuple (str, dict):
            - str (head_all_code) - 所有头部代码的字符串。
            - dict (code_map) - 包含每个Gemm配置对应的源代码的字典，格式为{"gemm_config": source_code}。
    """
    code_map = {}
    head_path = os.path.join(generate_dir, "launch_visitor_gemm_fused_kernel.h")
    head_all_code = LaunchGemmHead
    for tile in tiles:
        blocks, warps, mmas = [s.replace(" ", "").strip("<>").split(",") for s in tile]
        gemm_config = f"block{blocks[0]}x{blocks[1]}x{blocks[2]}_warp{warps[0]}x{warps[1]}x{warps[2]}_mma{mmas[0]}x{mmas[1]}x{mmas[2]}"
        for stage in stages:
            gemm_config_str = gemm_config + f"_stage{stage}"
            value_dict = {
                "gemm_config": gemm_config_str,
            }
            head_all_code += SubstituteTemplate(LaunchGemmDeclare, value_dict)
    os.makedirs(generate_dir, exist_ok=True)
    with open(head_path, "w") as f:
        f.write(head_all_code)
        f.close()

    for tile in tiles:
        blocks, warps, mmas = [s.replace(" ", "").strip("<>").split(",") for s in tile]
        gemm_config = f"block{blocks[0]}x{blocks[1]}x{blocks[2]}_warp{warps[0]}x{warps[1]}x{warps[2]}_mma{mmas[0]}x{mmas[1]}x{mmas[2]}"
        for stage in stages:
            gemm_config_str = gemm_config + f"_stage{stage}"
            value_dict = {
                "gemm_config": gemm_config_str,
            }
            source_all_code = SubstituteTemplate(LaunchGemmPart0, value_dict)
            type_id = 0
            for input_type in inputs_type:
                for output_type in outputs_type:
                    for hasbias in hasbiases:
                        value_dict = {
                            "input_type": input_type,
                            "output_type": output_type,
                            "hasbias": hasbias,
                            "type_id": str(type_id),
                            "thread_block_shape": tile[0],
                            "warp_shape": tile[1],
                            "mma_shape": tile[2],
                            "num_stages": str(stage),
                            "SM": sm,
                        }
                        source_all_code += SubstituteTemplate(LaunchGemmPart1, value_dict)
                        type_id += 1
            source_all_code += LaunchGemmPart2
            code_map[gemm_config_str] = source_all_code
            source_path = os.path.join(
                generate_dir,
                f"launch_visitor_gemm_fused_kernel_{gemm_config_str}.cu",
            )
            with open(source_path, "w") as f:
                f.write(source_all_code)
                f.close()

    return head_all_code, code_map


# generate fp8_visitor_gemm_fused.cu
def generate_dispatch_gemm_cu(
    inputs_type: str,
    outputs_type: str,
    stages: int,
    tiles: str,
    hasbiases: str,
    sm: str,
):
    """
    生成调度Gemm的CU代码。

    Args:
        inputs_type (str): 输入类型，字符串格式，多个类型用逗号分隔。可选值为 "float", "half".
        outputs_type (str): 输出类型，字符串格式，多个类型用逗号分隔。可选值为 "float", "half".
        stages (int): Gemm的层数。
        tiles (str): 瓦片形状，字符串格式，多个瓦片形状用逗号分隔。每个瓦片形状由三个整数组成，表示线程块形状、线程块内部warp形状和MMA形状。例如："<8,8,8>,<8,8,8>,<8,8,8>"。
        hasbiases (str): 是否有偏置，字符串格式，多个值用逗号分隔。可选值为 "true", "false".
        sm (str): 使用的SM数量，字符串格式，多个值用逗号分隔。可选值为 "32", "64".

    Returns:
        str: 返回一个包含所有代码的字符串。

    """

    all_code = code_part0
    type_id = 0
    for input_type in inputs_type:
        for output_type in outputs_type:
            for hasbias in hasbiases:
                value_dict = {
                    "input_type": input_type,
                    "output_type": output_type,
                    "hasbias": hasbias,
                    "type_id": str(type_id),
                }
                all_code += SubstituteTemplate(code_part1, value_dict)
                type_id += 1

    all_code += code_part2
    tile_id = 0
    for tile in tiles:
        for stage in stages:
            value_dict = {
                "thread_block_shape": tile[0],
                "warp_shape": tile[1],
                "mma_shape": tile[2],
                "num_stages": str(stage),
                "tile_id": str(tile_id),
            }
            all_code += SubstituteTemplate(code_part3, value_dict)
            tile_id += 1
    all_code += code_part4
    tile_id = 0
    for tile in tiles:
        blocks, warps, mmas = [s.replace(" ", "").strip("<>").split(",") for s in tile]
        gemm_config = f"block{blocks[0]}x{blocks[1]}x{blocks[2]}_warp{warps[0]}x{warps[1]}x{warps[2]}_mma{mmas[0]}x{mmas[1]}x{mmas[2]}"
        for stage in stages:
            gemm_config_str = gemm_config + f"_stage{stage}"
            value_dict = {
                "tile_id": str(tile_id),
                "gemm_config": gemm_config_str,
            }
            all_code += SubstituteTemplate(code_part5, value_dict)
            tile_id += 1
    all_code += SubstituteTemplate(code_part6, value_dict)
    return all_code


if __name__ == "__main__":
    args = parse_args()
    archs = args.cuda_arch
    min_stages = args.min_stages
    max_stages = args.max_stages
    inputs_type = ("float8_e4m3fn", "float8_e5m2")
    outputs_type = ("float16", "bfloat16")
    sm_dict = {"89": "cutlass::arch::Sm89", "90": "cutlass::arch::Sm90"}

    for sm in archs:
        if sm == "89":
            fuse_gemm_configs = get_candidate_configs(sm, min_stages, max_stages)
            for fuse_gemm_config in fuse_gemm_configs:
                file_name = (
                    f"gpu_ops/cutlass_kernels/fp8_gemm_fused/autogen/generic_visitor_gemm_fused_kernel_sm{sm}.cu"
                )
                all_code = generate_source_cu(
                    inputs_type,
                    outputs_type,
                    fuse_gemm_config[0],
                    fuse_gemm_config[1],
                    fuse_gemm_config[2],
                    sm_dict[sm],
                )

                file_dir = os.path.dirname(file_name)
                os.makedirs(file_dir, exist_ok=True)
                with open(file_name, "w") as f:
                    f.write(all_code)
                    f.close()

            fuse_gemm_config = list(fuse_gemm_configs)[0]

            # Compile parallelization
            generate_launch_gemm_cus(
                "gpu_ops/cutlass_kernels/fp8_gemm_fused/autogen",
                inputs_type,
                outputs_type,
                fuse_gemm_config[0],
                fuse_gemm_config[1],
                fuse_gemm_config[2],
                sm_dict[sm],
            )

            file_name = "gpu_ops/cutlass_kernels/fp8_gemm_fused/visitor_fp8_gemm_fused.cu"
            all_code = generate_dispatch_gemm_cu(
                inputs_type,
                outputs_type,
                fuse_gemm_config[0],
                fuse_gemm_config[1],
                fuse_gemm_config[2],
                sm_dict[sm],
            )
            file_dir = os.path.dirname(file_name)
            os.makedirs(file_dir, exist_ok=True)
            with open(file_name, "w") as f:
                f.write(all_code)
                f.close()

        elif sm == 90:
            print("Not supported yet.")
            exit(0)
        else:
            raise ValueError(f"Unsupported SM: {sm}")
