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