import argparse
import copy
from distutils import ccompiler
from distutils import errors
from distutils import msvccompiler
from distutils import sysconfig
from distutils import unixccompiler
import os
from os import path
import shutil
import sys
import pkg_resources
import setuptools
from setuptools.command import build_ext
from setuptools.command import sdist
from install import build
from install.build import PLATFORM_LINUX
from install.build import PLATFORM_WIN32
required_cython_version = pkg_resources.parse_version('0.28.0')
ignore_cython_versions = [
]
use_hip = bool(int(os.environ.get('CUPY_INSTALL_USE_HIP', '0')))
# The value of the key 'file' is a list that contains extension names
# or tuples of an extension name and a list of other souces files
# required to build the extension such as .cpp files and .cu files.
#
# | (, a list of )
#
# The extension name is also interpreted as the name of the Cython
# source file required to build the extension with appending '.pyx'
# file extension.
MODULES = []
cuda_files = [
'cupy_backends.cuda.api.driver',
'cupy_backends.cuda.api.runtime',
'cupy_backends.cuda.libs.cublas',
'cupy_backends.cuda.libs.curand',
'cupy_backends.cuda.libs.cusparse',
'cupy_backends.cuda.libs.nvrtc',
'cupy_backends.cuda.libs.profiler',
'cupy_backends.cuda.stream',
'cupy.core._accelerator',
'cupy.core._carray',
'cupy.core._cub_reduction',
'cupy.core._dtype',
'cupy.core._fusion_kernel',
'cupy.core._fusion_thread_local',
'cupy.core._fusion_trace',
'cupy.core._fusion_variable',
'cupy.core._kernel',
'cupy.core._memory_range',
'cupy.core._optimize_config',
'cupy.core._reduction',
'cupy.core._routines_binary',
'cupy.core._routines_indexing',
'cupy.core._routines_linalg',
'cupy.core._routines_logic',
'cupy.core._routines_manipulation',
'cupy.core._routines_math',
'cupy.core._routines_sorting',
'cupy.core._routines_statistics',
'cupy.core._scalar',
'cupy.core.core',
'cupy.core.dlpack',
'cupy.core.flags',
'cupy.core.internal',
'cupy.core.fusion',
'cupy.core.new_fusion',
'cupy.core.raw',
'cupy.cuda.common',
'cupy.cuda.cufft',
'cupy.cuda.device',
'cupy.cuda.memory',
'cupy.cuda.memory_hook',
'cupy.cuda.pinned_memory',
'cupy.cuda.function',
'cupy.cuda.stream',
'cupy.cuda.texture',
'cupy.fft._cache',
'cupy.fft._callback',
'cupy.lib._polynomial',
'cupy._util'
]
if use_hip:
# We handle nvtx (and likely any other future support) here, because
# the HIP stubs (hip/cupy_*.h) would cause many symbols
# to leak into all these modules even if unused. It's easier for all of
# them to link to the same set of shared libraries.
MODULES.append({
# TODO(leofang): call this "rocm" or "hip" to avoid confusion?
'name': 'cuda',
'file': cuda_files + [
'cupy.cuda.nvtx',
'cupy_backends.cuda.libs.cusolver',
],
'include': [
'hip/hip_runtime_api.h',
'hip/hiprtc.h',
'hipblas.h',
'hiprand/hiprand.h',
'hipfft.h',
'roctx.h',
'rocsolver.h',
],
'libraries': [
'amdhip64', # was hiprtc and hip_hcc before ROCm 3.8.0
'hipblas',
'hiprand',
'rocfft',
'roctx64',
'rocblas',
'rocsolver',
],
'check_method': build.check_hip_version,
'version_method': build.get_hip_version,
})
else:
MODULES.append({
'name': 'cuda',
'file': cuda_files,
'include': [
'cublas_v2.h',
'cuda.h',
'cuda_profiler_api.h',
'cuda_runtime.h',
'cufft.h',
'curand.h',
'cusparse.h',
'nvrtc.h',
],
'libraries': [
'cublas',
'cuda',
'cudart',
'cufft',
'curand',
'cusparse',
'nvrtc',
],
'check_method': build.check_cuda_version,
'version_method': build.get_cuda_version,
})
if not use_hip:
MODULES.append({
'name': 'cusolver',
'file': [
'cupy_backends.cuda.libs.cusolver',
],
'include': [
'cusolverDn.h',
],
'libraries': [
'cusolver',
],
'check_method': build.check_cuda_version,
})
if not use_hip:
MODULES.append({
'name': 'cudnn',
'file': [
'cupy_backends.cuda.libs.cudnn',
'cupy.cudnn',
],
'include': [
'cudnn.h',
],
'libraries': [
'cudnn',
],
'check_method': build.check_cudnn_version,
'version_method': build.get_cudnn_version,
})
MODULES.append({
'name': 'nccl',
'file': [
'cupy.cuda.nccl',
],
'include': [
'nccl.h',
],
'libraries': [
'nccl',
],
'check_method': build.check_nccl_version,
'version_method': build.get_nccl_version,
})
MODULES.append({
'name': 'nvtx',
'file': [
'cupy.cuda.nvtx',
],
'include': [
'nvToolsExt.h',
],
'libraries': [
'nvToolsExt' if not PLATFORM_WIN32 else 'nvToolsExt64_1',
],
'check_method': build.check_nvtx,
})
MODULES.append({
'name': 'cutensor',
'file': [
'cupy_backends.cuda.libs.cutensor',
'cupy.cutensor',
],
'include': [
'cutensor.h',
],
'libraries': [
'cutensor',
'cublas',
],
'check_method': build.check_cutensor_version,
'version_method': build.get_cutensor_version,
})
MODULES.append({
'name': 'cub',
'file': [
('cupy.cuda.cub', ['cupy/cuda/cupy_cub.cu']),
],
'include': [
'cub/util_namespace.cuh', # dummy
],
'libraries': [
'cudart',
],
'check_method': build.check_cub_version,
'version_method': build.get_cub_version,
})
MODULES.append({
'name': 'jitify',
'file': [
'cupy.cuda.jitify',
],
'include': [
'cuda.h',
'cuda_runtime.h',
'nvrtc.h',
],
'libraries': [
'cuda',
'cudart',