import os import re import sys import shutil import errno 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_taged(): 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' istaged = %(istaged)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, 'istaged': is_taged(), 'with_mkl': '${WITH_MKL}'}) write_version_py(filename='${CINN_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' cinnlibs = [] package_data = {'cinn': ['core_api.so'], 'cinn.libs': []} if '${WITH_MKL}' == 'ON': cinnlibs.append('${MKLML_LIB}') cinnlibs.append('${MKLML_IOMP_LIB}') if '${WITH_MKLDNN}' == 'ON': cinnlibs.append('${MKLDNN_SHARED_LIB_2}') 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') 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/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" ] with redirect_stdout(): setup( name='${PACKAGE_NAME}', version='${CINN_VERSION}', description='CINN: a Compiler Infrastructure for Neural Networks', maintainer="PaddlePaddle", maintainer_email="Paddle-better@baidu.com", url='https://github.com/PaddlePaddle/Paddle', license='Apache Software License', packages=packages, package_data=package_data )