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