import errno

import os

import re

import sys

import shutil

import platform

import subprocess

from contextlib import contextmanager

from setuptools import setup

def set_rpath(lib, rpath):

command = "patchelf --set-rpath '{}' {}".format(rpath, lib)

if os.system(command) != 0:

raise Exception("patch {} failed, command: {}".format(lib, command))

def git_commit():

try:

cmd = ['git', 'rev-parse', 'HEAD']

git_commit = subprocess.Popen(cmd, stdout = subprocess.PIPE,

cwd="${PROJECT_SOURCE_DIR}").communicate()[0].strip()

except:

git_commit = b'Unknown'

git_commit = git_commit.decode()

return str(git_commit)

def _get_version_detail(idx):

assert idx < 3, "vesion info consists of %(major)d.%(minor)d.%(patch)d, \

so detail index must less than 3"

if re.match('${TAG_VERSION_REGEX}', '${PADDLE_VERSION}'):

version_details = '${PADDLE_VERSION}'.split('.')

if len(version_details) >= 3:

return version_details[idx]

return 0

def get_major():

return int(_get_version_detail(0))

def get_minor():

return int(_get_version_detail(1))

def get_patch():

return str(_get_version_detail(2))

def get_cuda_version():

if '${WITH_GPU}' == 'ON':

return '${CUDA_VERSION}'

else:

return 'False'

def get_cudnn_version():

if '${WITH_GPU}' == 'ON':

temp_cudnn_version = ''

if '${CUDNN_MAJOR_VERSION}':

temp_cudnn_version += '${CUDNN_MAJOR_VERSION}'

if '${CUDNN_MINOR_VERSION}':

temp_cudnn_version += '.${CUDNN_MINOR_VERSION}'

if '${CUDNN_PATCHLEVEL_VERSION}':

temp_cudnn_version += '.${CUDNN_PATCHLEVEL_VERSION}'

return temp_cudnn_version

else:

return 'False'

def is_tagged():

try:

cmd = ['git', 'describe', '--exact-match', '--tags', 'HEAD', '2>/dev/null']

git_tag = subprocess.Popen(cmd, stdout = subprocess.PIPE, cwd="${PROJECT_SOURCE_DIR}").communicate()[0].strip()

git_tag = git_tag.decode()

except:

return False

if str(git_tag).replace('v', '') == '${CINN_VERSION}':

return True

else:

return False

def write_version_py(filename='cinn/version/info.py'):

cnt = '''# THIS FILE IS GENERATED FROM CINN SETUP.PY

#

full_version = '%(major)d.%(minor)d.%(patch)s'

major = '%(major)d'

minor = '%(minor)d'

patch = '%(patch)s'

cuda_version = '%(cuda)s'

cudnn_version = '%(cudnn)s'

is_tagged = %(is_tagged)s

commit = '%(commit)s'

with_mkl = '%(with_mkl)s'

'''

commit = git_commit()

dirname = os.path.dirname(filename)

try:

os.makedirs(dirname)

except OSError as e:

if e.errno != errno.EEXIST:

raise

with open(filename, 'w') as f:

f.write(cnt % {

'major': get_major(),

'minor': get_minor(),

'patch': get_patch(),

'version': '${CINN_VERSION}',

'cuda': get_cuda_version(),

'cudnn': get_cudnn_version(),

'commit': commit,

'is_tagged': is_tagged(),

'with_mkl': '${WITH_MKL}'})

write_version_py(filename='${CMAKE_BINARY_DIR}/python/cinn/version/info.py')

if sys.platform != 'win32':

@contextmanager

def redirect_stdout():

f_log = open('${SETUP_LOG_FILE}', 'w')

origin_stdout = sys.stdout

sys.stdout = f_log

yield

f_log = sys.stdout

sys.stdout = origin_stdout

f_log.close()

else:

@contextmanager

def redirect_stdout():

yield

libs_path = '${CMAKE_BINARY_DIR}/python/cinn/libs'

os.makedirs(libs_path, exist_ok=True)

cinnlibs = []

package_data = {'cinn': ['core_api.so'], 'cinn.libs': []}

if '${WITH_MKL}' == 'ON':

cinnlibs.append('${MKLML_LIB}')

cinnlibs.append('${MKLML_IOMP_LIB}')

if '${WITH_ONEDNN}' == 'ON':

cinnlibs.append('${ONEDNN_SHARED_LIB}')

cinnlibs.append('${PHI_LIB}')

cinnlibs.append('${PHI_CORE_LIB}')

if '${WITH_GPU}' == 'ON' or '${WITH_ROCM}' == 'ON':

cinnlibs.append('${PHI_GPU_LIB}')

cinnlibs.append('${IR_LIB}')

cinnlibs.append('${COMMON_LIB}')

if '${WITH_GPU}' == 'ON':

cinnlibs.append('${CMAKE_BINARY_DIR}/dist/cinn/include/paddle/cinn/runtime/cuda/cinn_cuda_runtime_source.cuh')

cinnlibs.append('${CMAKE_BINARY_DIR}/dist/cinn/include/paddle/cinn/runtime/cuda/float16.h')

cinnlibs.append('${CMAKE_BINARY_DIR}/dist/cinn/include/paddle/cinn/runtime/cuda/bfloat16.h')

if '${WITH_ROCM}' == 'ON':

cinnlibs.append('${CMAKE_BINARY_DIR}/dist/cinn/include/paddle/cinn/runtime/hip/cinn_hip_runtime_source.h')

for lib in cinnlibs:

shutil.copy(lib, libs_path)

libname = os.path.basename(lib)

if lib.endswith('so'):

set_rpath(os.path.join(libs_path, libname) , '$ORIGIN/')

package_data['cinn.libs'].append(libname)

set_rpath('${CMAKE_BINARY_DIR}/python/cinn/core_api.so', '$ORIGIN/../nvidia/cuda_runtime/lib:$ORIGIN/../nvidia/cuda_nvrtc/lib:$ORIGIN/../nvidia/cudnn/lib:$ORIGIN/../nvidia/nvtx/lib:$ORIGIN/../nvidia/cublas/lib:$ORIGIN/../nvidia/curand/lib:$ORIGIN/../nvidia/cusolver/lib:$ORIGIN/libs/')

def git_commit():

try:

cmd = ['git', 'rev-parse', 'HEAD']

git_commit = subprocess.Popen(cmd, stdout = subprocess.PIPE,

cwd="@PADDLE_SOURCE_DIR@").communicate()[0].strip()

except:

git_commit = 'Unknown'

git_commit = git_commit.decode()

return str(git_commit)

packages = ["cinn",

"cinn.auto_schedule",

"cinn.auto_schedule.cost_model",

"cinn.ir",

"cinn.libs",

"cinn.version",

"cinn.runtime"

]

install_requires=[]

if platform.system() == 'Linux' and platform.machine() == 'x86_64':

paddle_cuda_install_requirements = os.getenv(

"PADDLE_CUDA_INSTALL_REQUIREMENTS", None

)

if paddle_cuda_install_requirements == "ON":

PADDLE_CUDA_INSTALL_REQUIREMENTS = {

"V11": (

"nvidia-cuda-runtime-cu11==11.8.89; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cuda-cupti-cu11==11.8.87; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cudnn-cu11==8.7.0.84; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cublas-cu11==11.11.3.6; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cufft-cu11==10.9.0.58; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-curand-cu11==10.3.0.86; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cusolver-cu11==11.4.1.48; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cusparse-cu11==11.7.5.86; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-nccl-cu11==2.19.3; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-nvtx-cu11==11.8.86; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cuda-nvrtc-cu11==11.8.89; platform_system == 'Linux' and platform_machine == 'x86_64'"

),

"V12": (

"nvidia-cuda-runtime-cu12==12.3.101; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cuda-cupti-cu12==12.3.101; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cudnn-cu12==9.0.0.312; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cublas-cu12==12.3.4.1; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cufft-cu12==11.2.1.3; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-curand-cu12==10.3.5.147; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cusolver-cu12==11.6.1.9; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cusparse-cu12==12.3.1.170; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-nccl-cu12==2.19.3; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-nvtx-cu12==12.4.127; platform_system == 'Linux' and platform_machine == 'x86_64' | "

"nvidia-cuda-nvrtc-cu12==12.3.107; platform_system == 'Linux' and platform_machine == 'x86_64'"

),

}

try:

output = subprocess.check_output(['nvcc', '--version']).decode('utf-8')

version_line = [line for line in output.split('\n') if 'release' in line][0]

version = version_line.split(' ')[-1].split(',')[0]

cuda_major_version = version.split('.')[0]

except Exception as e:

raise ValueError("CUDA not found")

install_requires.append(PADDLE_CUDA_INSTALL_REQUIREMENTS[cuda_major_version].split("|"))

with redirect_stdout():

setup(

name='${PACKAGE_NAME}',

version='${CINN_VERSION}',

description='CINN: a Compiler Infrastructure for Neural Networks',

maintainer="PaddlePaddle",

maintainer_email="[email protected]",

url='https://github.com/PaddlePaddle/Paddle',

license='Apache Software License',

packages=packages,

install_requires=install_requires,

package_data=package_data

)