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