# Copyright (c) 2023 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.

import os
import shutil
import subprocess
from packaging.version import parse, Version

import paddle
from paddle.utils.cpp_extension import CUDAExtension, setup

sm_version = int(os.getenv("CUDA_SM_VERSION", "0"))

def get_nvcc_cuda_version(cuda_dir: str) -> Version:
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
    nvcc_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
                                          universal_newlines=True)
    output = nvcc_output.split()
    release_idx = output.index("release") + 1
    nvcc_cuda_version = parse(output[release_idx].split(",")[0])
    return nvcc_cuda_version

def update_git_submodule():
    try:
        subprocess.run(["git", "submodule", "update", "--init"], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error occurred while updating git submodule: {str(e)}")
        raise


def find_end_files(directory, end_str):
    gen_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(end_str):
                gen_files.append(os.path.join(root, file))
    return gen_files


def get_sm_version():
    if sm_version > 0:
        return sm_version
    else:
        prop = paddle.device.cuda.get_device_properties()
        cc = prop.major * 10 + prop.minor
        return cc


def strtobool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise ValueError(
            f"Truthy value expected: got {v} but expected one of yes/no, true/false, t/f, y/n, 1/0 (case insensitive)."
        )


def get_gencode_flags():
    if not strtobool(os.getenv("FLAG_LLM_PDC", "False")):
        cc = get_sm_version()
        if cc == 90:
            cc = f"{cc}a"
        return ["-gencode", "arch=compute_{0},code=sm_{0}".format(cc)]
    else:
        # support more cuda archs
        return [
            "-gencode",
            "arch=compute_80,code=sm_80",
            "-gencode",
            "arch=compute_75,code=sm_75",
            "-gencode",
            "arch=compute_70,code=sm_70",
        ]


gencode_flags = get_gencode_flags()
library_path = [os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")]

sources = [
    "./gpu/save_with_output.cc",
    "./gpu/set_value_by_flags.cu",
    "./gpu/token_penalty_multi_scores.cu",
    "./gpu/token_penalty_multi_scores_v2.cu",
    "./gpu/stop_generation_multi_ends.cu",
    "./gpu/fused_get_rope.cu",
    "./gpu/get_padding_offset.cu",
    "./gpu/qkv_transpose_split.cu",
    "./gpu/rebuild_padding.cu",
    "./gpu/transpose_removing_padding.cu",
    "./gpu/write_cache_kv.cu",
    "./gpu/encode_rotary_qk.cu",
    "./gpu/get_padding_offset_v2.cu",
    "./gpu/rebuild_padding_v2.cu",
    "./gpu/set_value_by_flags_v2.cu",
    "./gpu/stop_generation_multi_ends_v2.cu",
    "./gpu/get_output.cc",
    "./gpu/save_with_output_msg.cc",
    "./gpu/write_int8_cache_kv.cu",
    "./gpu/step.cu",
    "./gpu/quant_int8.cu",
    "./gpu/dequant_int8.cu",
    "./gpu/moe/preprocess_for_moe.cu",
    "./gpu/get_position_ids_and_mask_encoder_batch.cu",
    "./gpu/fused_rotary_position_encoding.cu",
    "./gpu/flash_attn_bwd.cc",
    "./gpu/tune_cublaslt_gemm.cu",
    "./gpu/sample_kernels/top_p_sampling_reject.cu",
    "./gpu/update_inputs_v2.cu",
    "./gpu/noaux_tc.cu",
    "./gpu/set_preids_token_penalty_multi_scores.cu",
    "./gpu/speculate_decoding_kernels/ngram_match.cc",
    "./gpu/speculate_decoding_kernels/speculate_save_output.cc",
    "./gpu/speculate_decoding_kernels/speculate_get_output.cc",
    "./gpu/save_output_dygraph.cu",
    "./gpu/all_reduce.cu",
    "./gpu/quantization/per_token_group_quant.cu",
    "./gpu/quantization/per_tensor_quant_fp8.cu",
]
sources += find_end_files("./gpu/speculate_decoding_kernels", ".cu")

nvcc_compile_args = gencode_flags
update_git_submodule()
nvcc_compile_args += [
    "-O3",
    "-DNDEBUG",
    "-U__CUDA_NO_HALF_OPERATORS__",
    "-U__CUDA_NO_HALF_CONVERSIONS__",
    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
    "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
]

include_dirs = [
    "./gpu",
    "./gpu/cutlass_kernels",
    "./gpu/fp8_gemm_with_cutlass",
    "./gpu/cutlass_kernels/fp8_gemm_fused/autogen",
    "./third_party/cutlass/include",
    "./third_party/cutlass/tools/util/include",
    "./third_party/nlohmann_json/single_include",
    "./gpu/sample_kernels",
    "./gpu/moe/fused_moe",
]
cc = get_sm_version()
cuda_version = float(paddle.version.cuda())
nvcc_version = get_nvcc_cuda_version(os.environ.get("CUDA_HOME", "/usr/local/cuda"))

if cc >= 80:
    sources += ["gpu/int8_gemm_with_cutlass/gemm_dequant.cu"]

    sources += ["./gpu/append_attention.cu", "./gpu/multi_head_latent_attention.cu"]

    sources += find_end_files("./gpu/append_attn", ".cu")
    sources += find_end_files("./gpu/append_attn/template_instantiation", ".cu")
    sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
    sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
    sources += "./gpu/cpp_extensions.cu",


fp8_auto_gen_directory = "gpu/cutlass_kernels/fp8_gemm_fused/autogen"
if os.path.isdir(fp8_auto_gen_directory):
    shutil.rmtree(fp8_auto_gen_directory)


if cc == 89 and cuda_version >= 12.4:
    os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels.py --cuda_arch 89")
    os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels.py --cuda_arch 89")
    sources += find_end_files(fp8_auto_gen_directory, ".cu")
    sources += [
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
    ]

if cc >= 80 and nvcc_version >= Version("12.4"):
    os.environ.pop('PADDLE_CUDA_ARCH_LIST', None)
    nvcc_compile_args += [
        "-std=c++17",
        "--use_fast_math",
        "--threads=8",
        "-D_GLIBCXX_USE_CXX11_ABI=1",
    ]
    sources += ["./gpu/sage_attn_kernels/sageattn_fused.cu"]
    if cc >= 80 and cc < 89:
        sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f16_kernel_sm80.cu"]
        nvcc_compile_args += ["-gencode", "arch=compute_80,code=compute_80"]
    elif cc >= 89 and cc < 90:
        sources += ["./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu"]
        nvcc_compile_args += ["-gencode", "arch=compute_89,code=compute_89"]
    elif cc >= 90:
        sources += [
            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm90.cu",
            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_dsk_kernel_sm90.cu",
        ]
        nvcc_compile_args += ["-gencode", "arch=compute_90a,code=compute_90a"]

if cc >= 90 and cuda_version >= 12.0:
    os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py --cuda_arch 90")
    os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_ptr_scale_sm90.py --cuda_arch 90")
    os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels_sm90.py --cuda_arch 90")
    os.system("python utils/auto_gen_fp8_fp8_block_gemm_fused_kernels_sm90.py --cuda_arch 90")
    sources += find_end_files(fp8_auto_gen_directory, ".cu")
    sources += [
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu",
        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm_ptr_scale.cu",
    ]
    sources += find_end_files("./gpu/mla_attn", ".cu")

ops_name = f"paddlenlp_ops_{sm_version}" if sm_version != 0 else "paddlenlp_ops"

setup(
    name=ops_name,
    ext_modules=CUDAExtension(
        sources=sources,
        extra_compile_args={
            "cxx": ["-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"],
            "nvcc": nvcc_compile_args,
        },
        libraries=["cublasLt"],
        library_dirs=library_path,
        include_dirs=include_dirs,
    ),
)
