sources = []
headers = []
defines = []
with_cuda = False
if torch.cuda.is_available():
print('Including CUDA code.')
sources += ['src/sync_bn.c']
headers += ['src/sync_bn.h']
defines += [('WITH_CUDA', None)]
with_cuda = True
assert with_cuda
current_dir = osp.dirname(osp.realpath(__file__))
extra_objects = ['src/cuda/sync_bn_kernel.o']
extra_objects = [osp.join(current_dir, fname) for fname in extra_objects]
ffi = create_extension(
'_ext.sync_bn_lib',
headers=headers,
sources=sources,
define_macros=defines,
relative_to=__file__,
with_cuda=with_cuda,
extra_objects=extra_objects,
)
if __name__ == '__main__':
ffi.build()
if torch.cuda.is_available():
print('Including CUDA code.')
sources += ['src/dcn_v2_cuda.c']
headers += ['src/dcn_v2_cuda.h']
defines += [('WITH_CUDA', None)]
extra_objects += ['src/cuda/dcn_v2_im2col_cuda.cu.o']
extra_objects += ['src/cuda/dcn_v2_psroi_pooling_cuda.cu.o']
with_cuda = True
else:
raise ValueError('CUDA is not available')
extra_compile_args = ['-fopenmp', '-std=c99']
this_file = os.path.dirname(os.path.realpath(__file__))
print(this_file)
sources = [os.path.join(this_file, fname) for fname in sources]
headers = [os.path.join(this_file, fname) for fname in headers]
extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]
ffi = create_extension('_ext.dcn_v2',
headers=headers,
sources=sources,
define_macros=defines,
relative_to=__file__,
with_cuda=with_cuda,
extra_objects=extra_objects,
extra_compile_args=extra_compile_args)
if __name__ == '__main__':
ffi.build()
this_file = os.path.dirname(__file__)
sources = []
headers = []
defines = []
with_cuda = False
if torch.cuda.is_available():
print('Including CUDA code.')
sources += ['pytorch_fft/src/th_fft_cuda.c']
headers += ['pytorch_fft/src/th_fft_cuda.h']
defines += [('WITH_CUDA', None)]
with_cuda = True
ffi = create_extension(
'pytorch_fft._ext.th_fft',
package=True,
headers=headers,
sources=sources,
define_macros=defines,
relative_to=__file__,
with_cuda=with_cuda,
include_dirs=[os.getcwd() + '/pytorch_fft/src'],
library_dirs=['/usr/local/cuda/lib64'],
libraries=['cufft']
)
if __name__ == '__main__':
ffi.build()