From cb73882b7f6b48f4ba73426b140e77d0d1d97468 Mon Sep 17 00:00:00 2001 From: sniklaus Date: Sat, 9 Sep 2017 22:59:59 -0700 Subject: no message --- install.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 install.py (limited to 'install.py') diff --git a/install.py b/install.py new file mode 100644 index 0000000..42847fa --- /dev/null +++ b/install.py @@ -0,0 +1,33 @@ +import os +import torch +import torch.utils.ffi + +strBasepath = os.path.split(os.path.abspath(__file__))[0] + '/' +strHeaders = [] +strSources = [] +strDefines = [] +strObjects = [] + +if torch.cuda.is_available() == True: + strHeaders += ['src/SeparableConvolution_cuda.h'] + strSources += ['src/SeparableConvolution_cuda.c'] + strDefines += [('WITH_CUDA', None)] + strObjects += ['src/SeparableConvolution_kernel.o'] +# end + +objectExtension = torch.utils.ffi.create_extension( + name='_ext.cunnex', + headers=strHeaders, + sources=strSources, + verbose=False, + with_cuda=any(strDefine[0] == 'WITH_CUDA' for strDefine in strDefines), + package=False, + relative_to=strBasepath, + include_dirs=[os.path.expandvars('$CUDA_HOME') + '/include'], + define_macros=strDefines, + extra_objects=[os.path.join(strBasepath, strObject) for strObject in strObjects] +) + +if __name__ == '__main__': + objectExtension.build() +# end \ No newline at end of file -- cgit v1.2.3-70-g09d2