wehub-resource-sync
135 lines
4.1 KiB
Python
135 lines
4.1 KiB
Python
# 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.
|
||
|
||
import os
|
||
import site
|
||
import subprocess
|
||
|
||
from paddle.utils.cpp_extension import CppExtension, setup
|
||
|
||
# from setuptools import Extension, setup
|
||
from setuptools.command.build_ext import build_ext
|
||
|
||
|
||
# refer: https://note.qidong.name/2018/03/setup-warning-strict-prototypes
|
||
# Avoid a gcc warning below:
|
||
# cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid
|
||
# for C/ObjC but not for C++
|
||
class BuildExt(build_ext):
|
||
def build_extensions(self):
|
||
if "-Wstrict-prototypes" in self.compiler.compiler_so:
|
||
self.compiler.compiler_so.remove("-Wstrict-prototypes")
|
||
super().build_extensions()
|
||
|
||
|
||
def check_avx512_bf16__support():
|
||
try:
|
||
result = subprocess.run(
|
||
["lscpu", "|", "grep", '"avx512_bf16"'],
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
shell=True,
|
||
)
|
||
|
||
if "avx512_bf16" in result.stdout.lower():
|
||
return True
|
||
else:
|
||
return False
|
||
|
||
except Exception as e:
|
||
print(f"Error checking AVX512 support: {e}")
|
||
return False
|
||
|
||
|
||
paddle_extra_compile_args = [
|
||
"-std=c++17",
|
||
"-shared",
|
||
"-fPIC",
|
||
"-Wno-parentheses",
|
||
"-DPADDLE_WITH_CUSTOM_KERNEL",
|
||
"-mavx512f",
|
||
"-mavx512vl",
|
||
"-fopenmp",
|
||
"-mavx512bw",
|
||
"-mno-mmx",
|
||
"-Wall",
|
||
"-march=skylake-avx512",
|
||
"-O3",
|
||
"-g",
|
||
]
|
||
|
||
if check_avx512_bf16__support():
|
||
paddle_extra_compile_args += [
|
||
"-DAVX512_BF16_WEIGHT_ONLY_BF16=true",
|
||
"-DAVX512_FP16_WEIGHT_ONLY_INT8=true",
|
||
"-DAVX512_FP16_WEIGHT_ONLY_FP16=true",
|
||
]
|
||
else:
|
||
paddle_extra_compile_args += [
|
||
"-DAVX512_FP32_WEIGHT_ONLY_FP16=true",
|
||
"-DAVX512_FP32_WEIGHT_ONLY_INT8=true",
|
||
]
|
||
# include path
|
||
site_packages_path = site.getsitepackages()
|
||
paddle_custom_kernel_include = [os.path.join(path, "paddle", "include") for path in site_packages_path]
|
||
|
||
XFT_INCLUDE_DIR = os.environ["XFT_HEADER_DIR"]
|
||
XFT_LIBRARY_DIR = os.environ["XFT_LIB_DIR"]
|
||
|
||
# include path third_party
|
||
paddle_custom_kernel_include += [
|
||
os.path.join(XFT_INCLUDE_DIR, "include"),
|
||
os.path.join(XFT_INCLUDE_DIR, "src/common"),
|
||
os.path.join(XFT_INCLUDE_DIR, "src/kernel"),
|
||
os.path.join(XFT_INCLUDE_DIR, "src/layers"),
|
||
os.path.join(XFT_INCLUDE_DIR, "src/models"),
|
||
os.path.join(XFT_INCLUDE_DIR, "src/utils"),
|
||
os.path.join(XFT_INCLUDE_DIR, "3rdparty/onednn/include"),
|
||
os.path.join(XFT_INCLUDE_DIR, "3rdparty/onednn/build/include"),
|
||
os.path.join(XFT_INCLUDE_DIR, "3rdparty/xdnn"),
|
||
os.path.join(XFT_INCLUDE_DIR, "3rdparty"),
|
||
os.path.join(XFT_INCLUDE_DIR, "3rdparty/mkl/include"),
|
||
]
|
||
|
||
# libs path
|
||
paddle_custom_kernel_library_dir = [os.path.join(path, "paddle", "base") for path in site_packages_path]
|
||
paddle_custom_kernel_library_dir += [XFT_LIBRARY_DIR]
|
||
|
||
|
||
libs = [":libxfastertransformer.so", ":libxft_comm_helper.so"]
|
||
|
||
custom_kernel_dot_module = CppExtension(
|
||
sources=[
|
||
"../gpu/save_with_output.cc",
|
||
"./src/token_penalty_multi_scores.cc",
|
||
"./src/stop_generation_multi_ends.cc",
|
||
"./src/set_value_by_flags.cc",
|
||
"./src/xft_transformer.cc",
|
||
"./src/avx_weight_only.cc",
|
||
"./src/xft_greedy_search.cc",
|
||
],
|
||
include_dirs=paddle_custom_kernel_include,
|
||
library_dirs=paddle_custom_kernel_library_dir,
|
||
libraries=libs,
|
||
extra_compile_args=paddle_extra_compile_args,
|
||
)
|
||
|
||
setup(
|
||
name="paddlenlp_ops",
|
||
version="1.0",
|
||
description="custom kernel for compiling",
|
||
ext_modules=[custom_kernel_dot_module],
|
||
)
|