1*4c07ec22SJames Wrightfrom pathlib import Path 2*4c07ec22SJames Wrightimport torch 3*4c07ec22SJames Wrightimport torch.utils.cpp_extension as C 4*4c07ec22SJames Wrightimport torch.utils as tutils 5*4c07ec22SJames Wrightimport re 6*4c07ec22SJames Wright 7*4c07ec22SJames Wrightbuild_dir = Path('./build') 8*4c07ec22SJames Wrightpkgconfig_path = build_dir / 'libtorch.pc' 9*4c07ec22SJames Wright 10*4c07ec22SJames Wrightvariables = {} 11*4c07ec22SJames Wrightkeywords = {} 12*4c07ec22SJames Wright 13*4c07ec22SJames Wright 14*4c07ec22SJames Wrightdef add_variable(file, variable, value): 15*4c07ec22SJames Wright file.write(f"{variable}={value}\n") 16*4c07ec22SJames Wright 17*4c07ec22SJames Wright 18*4c07ec22SJames Wrightdef add_keyword(file, key, value): 19*4c07ec22SJames Wright file.write(f"{key}: {value}\n") 20*4c07ec22SJames Wright 21*4c07ec22SJames Wright 22*4c07ec22SJames Wrightvariables['prefix'] = Path(C.library_paths()[0]).parent.as_posix() 23*4c07ec22SJames Wright 24*4c07ec22SJames Wrightkeywords['Name'] = 'libTorch' 25*4c07ec22SJames Wrightkeywords['Description'] = 'Custom made PC for PyTorch' 26*4c07ec22SJames Wrightkeywords['Version'] = torch.__version__ 27*4c07ec22SJames Wright 28*4c07ec22SJames Wrightkeywords['Cflags'] = '' 29*4c07ec22SJames Wrightfor include_path in C.include_paths(): 30*4c07ec22SJames Wright keywords['Cflags'] += f'-I{include_path} ' 31*4c07ec22SJames Wright 32*4c07ec22SJames Wright# Need to search the CMake file to see whether the library was compiled with the CXX11 ABI standard 33*4c07ec22SJames Wrightregex_ABI = re.compile(r'"(\S*GLIBCXX_USE_CXX11_ABI\S*)"') 34*4c07ec22SJames WrighttorchCMakePath = Path(tutils.cmake_prefix_path) / 'Torch/TorchConfig.cmake' 35*4c07ec22SJames Wrightabi_flag = '' 36*4c07ec22SJames Wrightwith torchCMakePath.open('r') as f: 37*4c07ec22SJames Wright for line in f: 38*4c07ec22SJames Wright regex_result = regex_ABI.search(line) 39*4c07ec22SJames Wright if regex_result: 40*4c07ec22SJames Wright abi_flag = regex_result[1] 41*4c07ec22SJames Wright 42*4c07ec22SJames Wrightkeywords['Cflags'] += abi_flag 43*4c07ec22SJames Wright 44*4c07ec22SJames Wrightkeywords['Libs'] = '' 45*4c07ec22SJames Wrightfor lib_path in C.library_paths(): 46*4c07ec22SJames Wright keywords['Libs'] += f'-L{lib_path} ' 47*4c07ec22SJames Wrightkeywords['Libs'] += '-lc10 -ltorch_cpu ' 48*4c07ec22SJames Wrightif torch.cuda.is_available(): 49*4c07ec22SJames Wright keywords['Libs'] += '-lc10_cuda -ltorch_cuda ' 50*4c07ec22SJames Wright # Need to force linking with libtorch_cuda.so, so find path and specify linking flag to force it 51*4c07ec22SJames Wright # This flag might be of limited portability 52*4c07ec22SJames Wright for lib_path in C.library_paths(): 53*4c07ec22SJames Wright torch_cuda_path = Path(lib_path) / 'libtorch_cuda.so' 54*4c07ec22SJames Wright if torch_cuda_path.exists(): 55*4c07ec22SJames Wright variables['torch_cuda_path'] = torch_cuda_path.as_posix() 56*4c07ec22SJames Wright keywords['Libs'] += f'-Wl,--no-as-needed,"{torch_cuda_path.as_posix()}" ' 57*4c07ec22SJames Wrightkeywords['Libs'] += '-ltorch ' 58*4c07ec22SJames Wrightkeywords['Libs.private'] = '' 59*4c07ec22SJames Wright 60*4c07ec22SJames Wrightwith pkgconfig_path.open('w') as file: 61*4c07ec22SJames Wright for variable, value in variables.items(): 62*4c07ec22SJames Wright add_variable(file, variable, value) 63*4c07ec22SJames Wright 64*4c07ec22SJames Wright file.write('\n') 65*4c07ec22SJames Wright 66*4c07ec22SJames Wright for keyword, value in keywords.items(): 67*4c07ec22SJames Wright add_keyword(file, keyword, value) 68*4c07ec22SJames Wright 69*4c07ec22SJames Wrightprint(pkgconfig_path.absolute()) 70