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