def build_torch_extension(build_ext, options, torch_version): pytorch_compile_flags = [ "-std=c++14" if flag == "-std=c++11" else flag for flag in options['COMPILE_FLAGS'] ] have_cuda = is_torch_cuda(build_ext, include_dirs=options['INCLUDES'], extra_compile_args=pytorch_compile_flags) if not have_cuda and check_macro(options['MACROS'], 'HAVE_CUDA'): raise DistutilsPlatformError( 'byteps build with GPU support was requested, but this PyTorch ' 'installation does not support CUDA.') # Update HAVE_CUDA to mean that PyTorch supports CUDA. updated_macros = set_macro(options['MACROS'], 'HAVE_CUDA', str(int(have_cuda))) # Export TORCH_VERSION equal to our representation of torch.__version__. Internally it's # used for backwards compatibility checks. updated_macros = set_macro(updated_macros, 'TORCH_VERSION', str(torch_version)) # Always set _GLIBCXX_USE_CXX11_ABI, since PyTorch can only detect whether it was set to 1. import torch updated_macros = set_macro(updated_macros, '_GLIBCXX_USE_CXX11_ABI', str(int(torch.compiled_with_cxx11_abi()))) # PyTorch requires -DTORCH_API_INCLUDE_EXTENSION_H updated_macros = set_macro(updated_macros, 'TORCH_API_INCLUDE_EXTENSION_H', '1') if have_cuda: from torch.utils.cpp_extension import CUDAExtension as TorchExtension else: # CUDAExtension fails with `ld: library not found for -lcudart` if CUDA is not present from torch.utils.cpp_extension import CppExtension as TorchExtension ext = TorchExtension( pytorch_lib.name, define_macros=updated_macros, include_dirs=options['INCLUDES'], sources=options['SOURCES'] + [ 'byteps/torch/ops.cc', 'byteps/torch/ready_event.cc', 'byteps/torch/cuda_util.cc', 'byteps/torch/adapter.cc', 'byteps/torch/handle_manager.cc' ], extra_compile_args=pytorch_compile_flags, extra_link_args=options['LINK_FLAGS'], extra_objects=options['EXTRA_OBJECTS'], library_dirs=options['LIBRARY_DIRS'], libraries=options['LIBRARIES']) # Patch an existing pytorch_lib extension object. for k, v in ext.__dict__.items(): pytorch_lib.__dict__[k] = v build_ext.build_extension(pytorch_lib)
def build_torch_extension_v2(build_ext, options, torch_version): have_cuda = is_torch_cuda_v2(build_ext, include_dirs=options['INCLUDES'], extra_compile_args=options['COMPILE_FLAGS']) if not have_cuda and check_macro(options['MACROS'], 'HAVE_CUDA'): raise DistutilsPlatformError( 'Horovod build with GPU support was requested, but this PyTorch ' 'installation does not support CUDA.') # Update HAVE_CUDA to mean that PyTorch supports CUDA. Internally, we will be checking # HOROVOD_GPU_(ALLREDUCE|ALLGATHER|BROADCAST) to decide whether we should use GPU # version or transfer tensors to CPU memory for those operations. updated_macros = set_macro(options['MACROS'], 'HAVE_CUDA', str(int(have_cuda))) # Export TORCH_VERSION equal to our representation of torch.__version__. Internally it's # used for backwards compatibility checks. updated_macros = set_macro(updated_macros, 'TORCH_VERSION', str(torch_version)) # Always set _GLIBCXX_USE_CXX11_ABI, since PyTorch can only detect whether it was set to 1. import torch updated_macros = set_macro(updated_macros, '_GLIBCXX_USE_CXX11_ABI', str(int(torch.compiled_with_cxx11_abi()))) # PyTorch requires -DTORCH_API_INCLUDE_EXTENSION_H updated_macros = set_macro(updated_macros, 'TORCH_API_INCLUDE_EXTENSION_H', '1') if have_cuda: from torch.utils.cpp_extension import CUDAExtension as TorchExtension else: # CUDAExtension fails with `ld: library not found for -lcudart` if CUDA is not present from torch.utils.cpp_extension import CppExtension as TorchExtension ext = TorchExtension( torch_mpi_lib_v2.name, define_macros=updated_macros, include_dirs=options['INCLUDES'], sources=options['SOURCES'] + [ 'horovod/torch/mpi_ops_v2.cc', 'horovod/torch/handle_manager.cc', 'horovod/torch/ready_event.cc', 'horovod/torch/cuda_util.cc', 'horovod/torch/adapter_v2.cc' ], extra_compile_args=options['COMPILE_FLAGS'], extra_link_args=options['LINK_FLAGS'], library_dirs=options['LIBRARY_DIRS'], libraries=options['LIBRARIES']) # Patch an existing torch_mpi_lib_v2 extension object. for k, v in ext.__dict__.items(): torch_mpi_lib_v2.__dict__[k] = v build_ext.build_extension(torch_mpi_lib_v2)
def build_torch_extension(build_ext, global_options, torch_version): # Backup the options, preventing other plugins access libs that # compiled with compiler of this plugin import torch is_cxx11_abi = torch.compiled_with_cxx11_abi() options = copy.deepcopy(global_options) have_cuda = is_torch_cuda(build_ext, include_dirs=options['INCLUDES'], extra_compile_args=options['COMPILE_FLAGS']) have_cuda = have_cuda and torch.cuda.is_available() if have_cuda: cuda_include_dirs, cuda_lib_dirs = get_cuda_dirs( build_ext, options['COMPILE_FLAGS']) nvcc_cmd = get_nvcc_cmd() cuda_extra_objects = build_nvcc_extra_objects(nvcc_cmd, is_cxx11_abi) options['EXTRA_OBJECTS'] += cuda_extra_objects options['INCLUDES'] += cuda_include_dirs options['LIBRARY_DIRS'] += cuda_lib_dirs options['LIBRARIES'] += ['cudart'] print('INFO: Try PyTorch extension with CUDA.') # Update HAVE_CUDA to mean that PyTorch supports CUDA. updated_macros = set_macro(options['MACROS'], 'HAVE_CUDA', str(int(have_cuda))) # TODO(ybc) make this into common options? have_nccl = os.getenv('BLUEFOG_WITH_NCCL', '0') assert have_nccl in ['0', '1'], "BLUEFOG_WITH_NCCL has to be either 0 or 1" if have_cuda and have_nccl == '1': nccl_include_dirs, nccl_lib_dirs, nccl_lib = get_nccl_dirs( build_ext, cuda_include_dirs, cuda_lib_dirs, options['COMPILE_FLAGS']) options['INCLUDES'] += nccl_include_dirs options['LIBRARY_DIRS'] += nccl_lib_dirs options['LIBRARIES'] += nccl_lib options['SOURCES'] += [ "bluefog/common/nccl_controller.cc", "bluefog/common/nccl_win.cc" ] print('INFO: Try PyTorch extension with NCCL.') updated_macros = set_macro(updated_macros, 'HAVE_NCCL', str(int(have_nccl))) updated_macros = set_macro(updated_macros, 'TORCH_VERSION', str(torch_version)) # Always set _GLIBCXX_USE_CXX11_ABI, since PyTorch can only detect whether it was set to 1. updated_macros = set_macro(updated_macros, '_GLIBCXX_USE_CXX11_ABI', str(int(is_cxx11_abi))) # PyTorch requires -DTORCH_API_INCLUDE_EXTENSION_H updated_macros = set_macro(updated_macros, 'TORCH_API_INCLUDE_EXTENSION_H', '1') if have_cuda: from torch.utils.cpp_extension import CUDAExtension as TorchExtension else: # CUDAExtension fails with `ld: library not found for -lcudart` if CUDA is not present from torch.utils.cpp_extension import CppExtension as TorchExtension ext = TorchExtension( bluefog_torch_mpi_lib.name, define_macros=updated_macros, include_dirs=options['INCLUDES'], sources=options['SOURCES'] + [ "bluefog/torch/adapter.cc", "bluefog/torch/handle_manager.cc", "bluefog/torch/mpi_ops.cc", "bluefog/torch/mpi_win_ops.cc" ], extra_compile_args=options['COMPILE_FLAGS'], extra_link_args=options['LINK_FLAGS'], library_dirs=options['LIBRARY_DIRS'], extra_objects=options['EXTRA_OBJECTS'], libraries=options['LIBRARIES']) # Patch an existing bluefog_torch_mpi_lib extension object. for k, v in ext.__dict__.items(): bluefog_torch_mpi_lib.__dict__[k] = v build_ext.build_extension(bluefog_torch_mpi_lib)