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