def build(args):
extra_objects = args.objs
extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')]
print(extra_objects)
# ffi = create_extension(
ffi = BuildExtension(
'_ext.pointnet2',
headers=[a for a in glob.glob("cinclude/*_wrapper.h")],
sources=[a for a in glob.glob("csrc/*.c")],
define_macros=[('WITH_CUDA', None)],
relative_to=__file__,
with_cuda=True,
extra_objects=extra_objects,
include_dirs=[osp.join(base_dir, 'cinclude')],
verbose=False,
package=False)
ffi.build()
with_cuda = False
if torch.cuda.is_available():
with_cuda = True
headers += ['src/prroi_pooling_gpu.h']
sources += ['src/prroi_pooling_gpu.c']
defines += [('WITH_CUDA', None)]
this_file = os.path.dirname(os.path.realpath(__file__))
extra_objects_cuda = ['src/prroi_pooling_gpu_impl.cu.o']
extra_objects_cuda = [
os.path.join(this_file, fname) for fname in extra_objects_cuda
]
extra_objects.extend(extra_objects_cuda)
else:
# TODO(Jiayuan Mao @ 07/13): remove this restriction after we support the cpu implementation.
raise NotImplementedError(
'Precise RoI Pooling only supports GPU (cuda) implememtations.')
ffi = BuildExtension('_prroi_pooling',
headers=headers,
sources=sources,
define_macros=defines,
relative_to=__file__,
with_cuda=with_cuda,
extra_objects=extra_objects)
if __name__ == '__main__':
ffi.build()
