Repository navigation
Expand file tree
/
Copy pathutils.py
More file actions
140 lines (121 loc) · 5.15 KB
/
Copy pathutils.py
File metadata and controls
140 lines (121 loc) · 5.15 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0
import os
import sys
import warnings
from setuptools.command.build_ext import build_ext as _build_ext
def detect_cuda_paths():
# The search logic below is designed to support all use cases:
# - build against pip wheels (skipped via --no-build-isolation)
# - build against custom CUDA location (CUDA_PATH or CUDA_HOME)
# - build against default Linux CUDA Location
# The bindings can be built against any recent CUDA version, and we don't really
# need any version/library detection here. But, we do need basic CUDA driver/runtime
# headers, and in the wheel case they are scattered in two wheels. When build
# isolation is on, the build prefix is added to sys.path, but this is the only
# implementation detail that we rely on.
# TODO: move to cuda.pathfinder.
potential_build_prefixes = [os.path.join(p, "nvidia/cu13") for p in sys.path] + [
os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME", "")),
"/usr/local/cuda",
]
# Filter out empty paths, e.g. when neither CUDA_PATH nor CUDA_HOME is set
potential_build_prefixes = [p for p in potential_build_prefixes if p]
cuda_paths = []
def check_path(header):
for prefix in potential_build_prefixes:
cuda_h = os.path.join(prefix, "include", header)
if os.path.isfile(cuda_h):
if prefix not in cuda_paths:
cuda_paths.append(prefix)
break
else:
searched_paths = "\n ".join(potential_build_prefixes)
warnings.warn(
f"include/{header} not found in any of these paths:\n "
f"{searched_paths}\n"
"Compilation of nvmath-python may fail. Set CUDA_PATH to suppress this warning."
)
check_path("cuda.h")
check_path("crt/host_defines.h")
return cuda_paths
def decide_lib_name(ext_name):
# TODO: move the record of the supported lib list elsewhere?
for lib in (
"cublasMp",
"cublas",
"cusolver",
"cufftMp",
"cufft",
"cusparse",
"curand",
"nvpl",
"nvshmem",
"mathdx",
"cudss",
"cutensor",
):
if lib in ext_name:
return lib
else:
return None
class build_ext(_build_ext):
def __init__(self, *args, **kwargs):
self._nvmath_cuda_paths = detect_cuda_paths()
print("\n" + "*" * 80)
for p in self._nvmath_cuda_paths:
print("CUDA path(s):", p)
print("*" * 80 + "\n")
super().__init__(*args, **kwargs)
def _prep_includes_libs_rpaths(self, lib_name):
"""
Set cuda_incl_dir and extra_linker_flags.
"""
cuda_incl_dir = [os.path.join(p, "include") for p in self._nvmath_cuda_paths]
extra_linker_flags = []
site_packages = ["$ORIGIN/../../.."]
if self.editable_mode:
import site
site_packages = site.getsitepackages()
else:
# strip binaries to remove debug symbols which significantly
# increase wheel size
extra_linker_flags = ["-Wl,--strip-all"]
nvpl_site_packages = [f"{p}/.." for p in site_packages]
# Note: soname = library major version
# We need to be able to search for cuBLAS/cuSOLVER/... at run time, in case they
# are installed via pip wheels.
# The rpaths must be adjusted given the following full-wheel installation:
# - $ORIGIN: site-packages/nvmath/bindings/_internal/
# - cublas: site-packages/nvidia/cublas/lib/
# - cusolver: site-packages/nvidia/cusolver/lib/
# - ... ...
if lib_name is None:
return cuda_incl_dir, extra_linker_flags
ldflag = "-Wl,--disable-new-dtags"
if lib_name == "nvpl":
# 1. the nvpl bindings land in
# site-packages/nvmath/bindings/nvpl/_internal/ as opposed to other
# packages that have their bindings in
# site-packages/nvmath/bindings/_internal/, so we need one extra
# `..` to get into `site-packages` and then the lib_name=nvpl is not
# in nvidia dir but directly in the site-packages.
# 2. mkl lib is placed directly in the python `lib` directory, not
# in python{ver}/site-packages
rpath = ":".join([f"{pth}/{lib_name}/lib:{pth}/../../" for pth in nvpl_site_packages])
ldflag += f",-rpath,{rpath}"
extra_linker_flags.append(ldflag)
return cuda_incl_dir, extra_linker_flags
def build_extension(self, ext):
lib_name = decide_lib_name(ext.name)
ext.include_dirs, ext.extra_link_args = self._prep_includes_libs_rpaths(lib_name)
if ext.name.endswith("cusparse"):
# too noisy
ext.define_macros = [
("DISABLE_CUSPARSE_DEPRECATED", None),
]
super().build_extension(ext)
def build_extensions(self):
self.parallel = os.cpu_count()
super().build_extensions()