1111WITH_CUDA = torch .cuda .is_available () and CUDA_HOME is not None
1212if os .getenv ('FORCE_CUDA' , '0' ) == '1' :
1313 WITH_CUDA = True
14- if os .getenv ('FORCE_NON_CUDA ' , '0' ) == '1' :
14+ if os .getenv ('FORCE_CPU ' , '0' ) == '1' :
1515 WITH_CUDA = False
1616
1717BUILD_DOCS = os .getenv ('BUILD_DOCS' , '0' ) == '1'
@@ -23,39 +23,37 @@ def get_extensions():
2323 extra_compile_args = {'cxx' : [], 'nvcc' : []}
2424 extra_link_args = []
2525
26- # Windows users: Edit both of these to contain your VS include path, i.e.:
27- # extra_compile_args['cxx'] += ['-I{VISUAL_STUDIO_DIR}\\include']
28- # extra_compile_args['nvcc'] += ['-I{VISUAL_STUDIO_DIR}\\include']
29-
3026 if WITH_CUDA :
3127 Extension = CUDAExtension
3228 define_macros += [('WITH_CUDA' , None )]
3329 nvcc_flags = os .getenv ('NVCC_FLAGS' , '' )
3430 nvcc_flags = [] if nvcc_flags == '' else nvcc_flags .split (' ' )
3531 nvcc_flags += ['-arch=sm_35' , '--expt-relaxed-constexpr' ]
36- extra_compile_args ['cxx' ] += ['-O0' ]
3732 extra_compile_args ['nvcc' ] += nvcc_flags
38- if sys .platform == 'win32' :
3933
34+ if sys .platform == 'win32' :
4035 extra_link_args = ['cusparse.lib' ]
4136 else :
4237 extra_link_args = ['-lcusparse' , '-l' , 'cusparse' ]
4338
44- if sys .platform == 'win32' :
45- extra_compile_args ['cxx' ] += ['/MP' ]
46-
4739 extensions_dir = osp .join (osp .dirname (osp .abspath (__file__ )), 'csrc' )
4840 main_files = glob .glob (osp .join (extensions_dir , '*.cpp' ))
4941 extensions = []
5042 for main in main_files :
5143 name = main .split (os .sep )[- 1 ][:- 4 ]
5244
53- sources = [main , osp .join (extensions_dir , 'cpu' , f'{ name } _cpu.cpp' )]
54- if WITH_CUDA :
55- sources += [osp .join (extensions_dir , 'cuda' , f'{ name } _cuda.cu' )]
45+ sources = [main ]
46+
47+ path = osp .join (extensions_dir , 'cpu' , name + '_cpu.cpp' )
48+ if osp .exists (path ):
49+ sources += [path ]
50+
51+ path = osp .join (extensions_dir , 'cuda' , name + '_cuda.cu' )
52+ if WITH_CUDA and osp .exists (path ):
53+ sources += [path ]
5654
5755 extension = Extension (
58- f 'torch_sparse._{ name } ' ,
56+ 'torch_sparse._' + name ,
5957 sources ,
6058 include_dirs = [extensions_dir ],
6159 define_macros = define_macros ,
@@ -73,14 +71,15 @@ def get_extensions():
7371
7472setup (
7573 name = 'torch_sparse' ,
76- version = '0.5.0 ' ,
74+ version = '0.5.1 ' ,
7775 author = 'Matthias Fey' ,
78767977 url = 'https://github.com/rusty1s/pytorch_sparse' ,
8078 description = ('PyTorch Extension Library of Optimized Autograd Sparse '
8179 'Matrix Operations' ),
8280 keywords = ['pytorch' , 'sparse' , 'sparse-matrices' , 'autograd' ],
8381 license = 'MIT' ,
82+ python_requires = '>=3.5' ,
8483 install_requires = install_requires ,
8584 setup_requires = setup_requires ,
8685 tests_require = tests_require ,
0 commit comments