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)
