1013a5551SJames Wrightfrom pathlib import Path 2013a5551SJames Wrightimport torch 3013a5551SJames Wrightimport torch.utils.cpp_extension as C 4013a5551SJames Wrightimport torch.utils as tutils 5013a5551SJames Wrightimport re 6013a5551SJames Wright 7013a5551SJames Wrightbuild_dir = Path('./build') 8*637c7b11SJames Wrightif not build_dir.is_dir(): 9*637c7b11SJames Wright build_dir.mkdir() 10013a5551SJames Wrightpkgconfig_path = build_dir / 'libtorch.pc' 11013a5551SJames Wright 12013a5551SJames Wrightvariables = {} 13013a5551SJames Wrightkeywords = {} 14013a5551SJames Wright 15013a5551SJames Wright 16013a5551SJames Wrightdef add_variable(file, variable, value): 17013a5551SJames Wright file.write(f"{variable}={value}\n") 18013a5551SJames Wright 19013a5551SJames Wright 20013a5551SJames Wrightdef add_keyword(file, key, value): 21013a5551SJames Wright file.write(f"{key}: {value}\n") 22013a5551SJames Wright 23013a5551SJames Wright 24013a5551SJames Wrightvariables['prefix'] = Path(C.library_paths()[0]).parent.as_posix() 25013a5551SJames Wright 26013a5551SJames Wrightkeywords['Name'] = 'libTorch' 27013a5551SJames Wrightkeywords['Description'] = 'Custom made PC for PyTorch' 28013a5551SJames Wrightkeywords['Version'] = torch.__version__ 29013a5551SJames Wright 30013a5551SJames Wrightkeywords['Cflags'] = '' 31013a5551SJames Wrightfor include_path in C.include_paths(): 32013a5551SJames Wright keywords['Cflags'] += f'-I{include_path} ' 33013a5551SJames Wright 34013a5551SJames Wright# Need to search the CMake file to see whether the library was compiled with the CXX11 ABI standard 35013a5551SJames Wrightregex_ABI = re.compile(r'"(\S*GLIBCXX_USE_CXX11_ABI\S*)"') 36013a5551SJames WrighttorchCMakePath = Path(tutils.cmake_prefix_path) / 'Torch/TorchConfig.cmake' 37013a5551SJames Wrightabi_flag = '' 38013a5551SJames Wrightwith torchCMakePath.open('r') as f: 39013a5551SJames Wright for line in f: 40013a5551SJames Wright regex_result = regex_ABI.search(line) 41013a5551SJames Wright if regex_result: 42013a5551SJames Wright abi_flag = regex_result[1] 43013a5551SJames Wright 44013a5551SJames Wrightkeywords['Cflags'] += abi_flag 45013a5551SJames Wright 46013a5551SJames Wrightkeywords['Libs'] = '' 47013a5551SJames Wrightfor lib_path in C.library_paths(): 48013a5551SJames Wright keywords['Libs'] += f'-L{lib_path} ' 49013a5551SJames Wrightkeywords['Libs'] += '-lc10 -ltorch_cpu ' 50013a5551SJames Wrightif torch.cuda.is_available(): 51013a5551SJames Wright keywords['Libs'] += '-lc10_cuda -ltorch_cuda ' 52013a5551SJames Wright # Need to force linking with libtorch_cuda.so, so find path and specify linking flag to force it 53013a5551SJames Wright # This flag might be of limited portability 54013a5551SJames Wright for lib_path in C.library_paths(): 55013a5551SJames Wright torch_cuda_path = Path(lib_path) / 'libtorch_cuda.so' 56013a5551SJames Wright if torch_cuda_path.exists(): 57013a5551SJames Wright variables['torch_cuda_path'] = torch_cuda_path.as_posix() 58013a5551SJames Wright keywords['Libs'] += f'-Wl,--no-as-needed,"{torch_cuda_path.as_posix()}" ' 59013a5551SJames Wrightkeywords['Libs'] += '-ltorch ' 60013a5551SJames Wrightkeywords['Libs.private'] = '' 61013a5551SJames Wright 62013a5551SJames Wrightwith pkgconfig_path.open('w') as file: 63013a5551SJames Wright for variable, value in variables.items(): 64013a5551SJames Wright add_variable(file, variable, value) 65013a5551SJames Wright 66013a5551SJames Wright file.write('\n') 67013a5551SJames Wright 68013a5551SJames Wright for keyword, value in keywords.items(): 69013a5551SJames Wright add_keyword(file, keyword, value) 70013a5551SJames Wright 71013a5551SJames Wrightprint(pkgconfig_path.absolute()) 72