def get_extensions():
extensions = []
extensions_dir = osp.join('csrc')
main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
for main, suffix in product(main_files, suffices):
define_macros = []
extra_compile_args = {'cxx': ['-O2']}
if not os.name == 'nt': # Not on Windows:
extra_compile_args['cxx'] += ['-Wno-sign-compare']
extra_link_args = ['-s']
info = parallel_info()
if ('backend: OpenMP' in info and 'OpenMP not found' not in info
and sys.platform != 'darwin'):
extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
if sys.platform == 'win32':
extra_compile_args['cxx'] += ['/openmp']
else:
extra_compile_args['cxx'] += ['-fopenmp']
else:
print('Compiling without OpenMP...')
# Compile for mac arm64
if (sys.platform == 'darwin' and platform.machine() == 'arm64'):
extra_compile_args['cxx'] += ['-arch', 'arm64']
extra_link_args += ['-arch', 'arm64']
if suffix == 'cuda':
define_macros += [('WITH_CUDA', None)]
nvcc_flags = os.getenv('NVCC_FLAGS', '')
nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
nvcc_flags += ['--expt-relaxed-constexpr', '-O2']
extra_compile_args['nvcc'] = nvcc_flags
name = main.split(os.sep)[-1][:-4]
sources = [main]
path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
if osp.exists(path):
sources += [path]
path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
if suffix == 'cuda' and osp.exists(path):
sources += [path]
Extension = CppExtension if suffix == 'cpu' else CUDAExtension
extension = Extension(
f'torch_spline_conv._{name}_{suffix}',
sources,
include_dirs=[extensions_dir],
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
)
extensions += [extension]
return extensions
def parallel_backend():
# https://github.com/suphoff/pytorch_parallel_extension_cpp/blob/master/setup.py
from torch.__config__ import parallel_info
parallel_info_string = parallel_info()
parallel_info_array = parallel_info_string.splitlines()
backend_lines = [line for line in parallel_info_array if line.startswith('ATen parallel backend:')]
if len(backend_lines) != 1:
return None
backend = backend_lines[0].rsplit(': ')[1]
return backend
def parallel_backend():
parallel_info_string = parallel_info()
parallel_info_array = parallel_info_string.splitlines()
backend_lines = [
line for line in parallel_info_array
if line.startswith('ATen parallel backend:')
]
if len(backend_lines) is not 1:
return None
backend = backend_lines[0].rsplit(': ')[1]
return backend
def get_extensions():
Extension = CppExtension
define_macros = []
extra_compile_args = {'cxx': []}
extra_link_args = []
info = parallel_info()
if 'parallel backend: OpenMP' in info and 'OpenMP not found' not in info:
extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
if sys.platform == 'win32':
extra_compile_args['cxx'] += ['/openmp']
else:
extra_compile_args['cxx'] += ['-fopenmp']
else:
print('Compiling without OpenMP...')
if WITH_CUDA:
Extension = CUDAExtension
define_macros += [('WITH_CUDA', None)]
nvcc_flags = os.getenv('NVCC_FLAGS', '')
nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
nvcc_flags += ['-arch=sm_35', '--expt-relaxed-constexpr']
extra_compile_args['nvcc'] = nvcc_flags
extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
extensions = []
for main in main_files:
name = main.split(os.sep)[-1][:-4]
sources = [main]
path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
if osp.exists(path):
sources += [path]
path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
if WITH_CUDA and osp.exists(path):
sources += [path]
extension = Extension(
'torch_cluster._' + name,
sources,
include_dirs=[extensions_dir],
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
)
extensions += [extension]
return extensions
def get_extensions():
Extension = CppExtension
define_macros = []
extra_compile_args = {"cxx": []}
extra_link_args = []
info = parallel_info()
if "parallel backend: OpenMP" in info and "OpenMP not found" not in info:
extra_compile_args["cxx"] += ["-DAT_PARALLEL_OPENMP"]
if sys.platform == "win32":
extra_compile_args["cxx"] += ["/openmp"]
else:
extra_compile_args["cxx"] += ["-fopenmp"]
else:
print("Compiling without OpenMP...")
if WITH_CUDA:
Extension = CUDAExtension
define_macros += [("WITH_CUDA", None)]
nvcc_flags = os.getenv("NVCC_FLAGS", "")
nvcc_flags = [] if nvcc_flags == "" else nvcc_flags.split(" ")
nvcc_flags += ["-arch=sm_35", "--expt-relaxed-constexpr"]
extra_compile_args["nvcc"] = nvcc_flags
extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), "csrc")
main_files = glob.glob(osp.join(extensions_dir, "*.cpp"))
extensions = []
for main in main_files:
name = main.split(os.sep)[-1][:-4]
sources = [main]
path = osp.join(extensions_dir, "cpu", f"{name}_cpu.cpp")
if osp.exists(path):
sources += [path]
path = osp.join(extensions_dir, "cuda", f"{name}_cuda.cu")
if WITH_CUDA and osp.exists(path):
sources += [path]
extension = Extension(
"torch_isin._" + name,
sources,
include_dirs=[extensions_dir],
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
)
extensions += [extension]
return extensions
def get_extensions():
extensions = []
extensions_dir = osp.join('csrc')
main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
for main, suffix in product(main_files, suffices):
define_macros = []
libraries = []
if WITH_METIS:
define_macros += [('WITH_METIS', None)]
libraries += ['metis']
if WITH_MTMETIS:
define_macros += [('WITH_MTMETIS', None)]
define_macros += [('MTMETIS_64BIT_VERTICES', None)]
define_macros += [('MTMETIS_64BIT_EDGES', None)]
define_macros += [('MTMETIS_64BIT_WEIGHTS', None)]
define_macros += [('MTMETIS_64BIT_PARTITIONS', None)]
libraries += ['mtmetis', 'wildriver']
extra_compile_args = {'cxx': ['-O2']}
extra_link_args = ['-s']
info = parallel_info()
if ('backend: OpenMP' in info and 'OpenMP not found' not in info
and sys.platform != 'darwin'):
extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
if sys.platform == 'win32':
extra_compile_args['cxx'] += ['/openmp']
else:
extra_compile_args['cxx'] += ['-fopenmp']
else:
print('Compiling without OpenMP...')
if suffix == 'cuda':
define_macros += [('WITH_CUDA', None)]
nvcc_flags = os.getenv('NVCC_FLAGS', '')
nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
nvcc_flags += ['--expt-relaxed-constexpr', '-O2']
extra_compile_args['nvcc'] = nvcc_flags
if sys.platform == 'win32':
extra_link_args += ['cusparse.lib']
else:
extra_link_args += ['-lcusparse', '-l', 'cusparse']
name = main.split(os.sep)[-1][:-4]
sources = [main]
path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
if osp.exists(path):
sources += [path]
path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
if suffix == 'cuda' and osp.exists(path):
sources += [path]
Extension = CppExtension if suffix == 'cpu' else CUDAExtension
extension = Extension(
f'torch_sparse._{name}_{suffix}',
sources,
include_dirs=[extensions_dir],
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
libraries=libraries,
)
extensions += [extension]
return extensions
def get_extensions():
extensions = []
extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), "csrc")
main_files = glob.glob(osp.join(extensions_dir, "*.cpp"))
for main, suffix in product(main_files, suffices):
define_macros = []
libraries = []
extra_compile_args = {"cxx": ["-O2"]}
if not os.name == "nt": # Not on Windows:
extra_compile_args["cxx"] += ["-Wno-sign-compare"]
extra_link_args = ["-s"]
info = parallel_info()
if "backend: OpenMP" in info and "OpenMP not found" not in info:
extra_compile_args["cxx"] += ["-DAT_PARALLEL_OPENMP"]
if sys.platform == "win32":
extra_compile_args["cxx"] += ["/openmp"]
else:
extra_compile_args["cxx"] += ["-fopenmp"]
else:
print("Compiling without OpenMP...")
if suffix == "cuda":
define_macros += [("WITH_CUDA", None)]
nvcc_flags = os.getenv("NVCC_FLAGS", "")
nvcc_flags = [] if nvcc_flags == "" else nvcc_flags.split(" ")
nvcc_flags += ["--expt-relaxed-constexpr", "-O2"]
extra_compile_args["nvcc"] = nvcc_flags
if sys.platform == "win32":
extra_link_args += ["cusparse.lib"]
else:
extra_link_args += ["-lcusparse", "-l", "cusparse"]
name = main.split(os.sep)[-1][:-4]
sources = [main]
path = osp.join(extensions_dir, "cpu", f"{name}_cpu.cpp")
if osp.exists(path):
sources += [path]
path = osp.join(extensions_dir, "cuda", f"{name}_cuda.cu")
if suffix == "cuda" and osp.exists(path):
sources += [path]
Extension = CppExtension if suffix == "cpu" else CUDAExtension
extension = Extension(
f"thgsp._{name}_{suffix}",
sources,
include_dirs=[extensions_dir],
define_macros=define_macros,
extra_compile_args=extra_compile_args,
extra_link_args=extra_link_args,
libraries=libraries,
)
extensions += [extension]
return extensions
