Skip to content

Commit 4fcb7ca

Browse files
Produce BOLT-compatible libraries (#3383)
* Produce BOLT-compatible libraries Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix target arch detection Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * fail fast when failing to detect arch Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * reword env var description in docs Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * add guard for when `CMAKE_SYSTEM_PROCESSOR` does not match CXX target arch Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Rename `build_target_arch` to `build_target_cpu_arch` Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate `-fno-reorder-blocks-and-partitions` on compiler identity Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable BOLT flags for all builds Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enforce default compiler consistency for NCCL EP (needed to keep flags consistent) Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * fix greptile PR comments about compiler probe failures and cross build arch mismatch Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent db759b2 commit 4fcb7ca

5 files changed

Lines changed: 191 additions & 8 deletions

File tree

‎build_tools/jax.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
all_files_in_dir,
1717
cudnn_frontend_include_path,
1818
debug_build_enabled,
19+
get_bolt_build_flags,
1920
setup_mpi_flags,
2021
nccl_include_path,
2122
nccl_lib_path,
@@ -122,6 +123,9 @@ def setup_jax_extension(
122123
else:
123124
cxx_flags.append("-g0")
124125

126+
bolt_cxx_flags, linker_flags = get_bolt_build_flags()
127+
cxx_flags.extend(bolt_cxx_flags)
128+
125129
setup_mpi_flags(include_dirs, cxx_flags)
126130

127131
if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))):
@@ -144,5 +148,6 @@ def setup_jax_extension(
144148
sources=[str(path) for path in sources],
145149
include_dirs=[str(path) for path in include_dirs],
146150
extra_compile_args=cxx_flags,
151+
extra_link_args=linker_flags,
147152
**kwargs,
148153
)

‎build_tools/pytorch.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
cuda_version,
1616
get_cuda_include_dirs,
1717
debug_build_enabled,
18+
get_bolt_build_flags,
1819
nccl_ep_enabled,
1920
setup_mpi_flags,
2021
)
@@ -76,6 +77,9 @@ def setup_pytorch_extension(
7677
else:
7778
cxx_flags.append("-g0")
7879

80+
bolt_cxx_flags, linker_flags = get_bolt_build_flags()
81+
cxx_flags.extend(bolt_cxx_flags)
82+
7983
# Version-dependent CUDA options
8084
try:
8185
version = cuda_version()
@@ -127,6 +131,7 @@ def setup_pytorch_extension(
127131
sources=[str(src) for src in sources],
128132
include_dirs=[str(inc) for inc in include_dirs],
129133
extra_compile_args={"cxx": cxx_flags},
134+
extra_link_args=linker_flags,
130135
libraries=[str(lib) for lib in libraries],
131136
library_dirs=[str(lib_dir) for lib_dir in library_dirs],
132137
)

‎build_tools/utils.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import importlib
1010
import os
1111
import re
12+
import shlex
1213
import shutil
1314
import subprocess
1415
import sys
@@ -43,6 +44,72 @@ def debug_build_enabled() -> bool:
4344
return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0")))
4445

4546

47+
@functools.lru_cache(maxsize=None)
48+
def _cxx_compiler_predefined_macros() -> Tuple[str, ...]:
49+
"""Predefined macros from the configured C++ compiler."""
50+
cxx = shlex.split(os.getenv("CXX", "c++"))
51+
if not cxx:
52+
return ()
53+
try:
54+
result = subprocess.run(
55+
[*cxx, "-dM", "-E", "-x", "c++", "-"],
56+
input="",
57+
capture_output=True,
58+
text=True,
59+
)
60+
except OSError:
61+
return ()
62+
if result.returncode != 0:
63+
return ()
64+
return tuple(
65+
fields[1]
66+
for line in result.stdout.splitlines()
67+
if len(fields := line.split(maxsplit=2)) >= 2 and fields[0] == "#define"
68+
)
69+
70+
71+
@functools.lru_cache(maxsize=None)
72+
def cxx_compiler_supports_flag(flag: str) -> bool:
73+
"""Whether the configured C++ compiler accepts a command-line flag."""
74+
cxx = shlex.split(os.getenv("CXX", "c++"))
75+
if not cxx:
76+
return False
77+
try:
78+
result = subprocess.run(
79+
[*cxx, flag, "-x", "c++", "-fsyntax-only", "-"],
80+
input="",
81+
capture_output=True,
82+
text=True,
83+
)
84+
except OSError:
85+
return False
86+
return result.returncode == 0
87+
88+
89+
def cxx_compiler_targets_arm64() -> bool:
90+
"""Whether the configured C++ compiler targets Arm64."""
91+
macros = _cxx_compiler_predefined_macros()
92+
return any(macro in macros for macro in ("__aarch64__", "__arm64__", "_M_ARM64"))
93+
94+
95+
def get_bolt_build_flags() -> Tuple[List[str], List[str]]:
96+
"""BOLT-compatible host compiler and linker flags."""
97+
compiler_flags = ["-fno-jump-tables"]
98+
if cxx_compiler_supports_flag("-fno-reorder-blocks-and-partition"):
99+
compiler_flags.append("-fno-reorder-blocks-and-partition")
100+
linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"]
101+
if cxx_compiler_targets_arm64():
102+
compiler_flags.extend(
103+
[
104+
"-mno-fix-cortex-a53-835769",
105+
"-mno-fix-cortex-a53-843419",
106+
]
107+
)
108+
# The Cortex-A53 843419 workaround is applied by the linker.
109+
linker_flags.append("-mno-fix-cortex-a53-843419")
110+
return compiler_flags, linker_flags
111+
112+
46113
@functools.lru_cache(maxsize=None)
47114
def get_max_jobs_for_parallel_build() -> int:
48115
"""Number of parallel jobs for Nina build"""

‎setup.py‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
cuda_home_path,
2323
cuda_version,
2424
cudnn_frontend_include_path,
25+
get_bolt_build_flags,
2526
get_frameworks,
2627
remove_dups,
2728
min_python_version_str,
@@ -264,6 +265,9 @@ def build_nccl_ep_submodule() -> str:
264265

265266
nproc = get_max_jobs_for_parallel_build()
266267
env = os.environ.copy()
268+
# get_bolt_build_flags() defaults to `c++` when CXX is unset. Export the
269+
# same default so Make does not independently select its `g++` default.
270+
env.setdefault("CXX", "c++")
267271
if (cuda_home := cuda_home_path()) is not None:
268272
env.setdefault("CUDA_HOME", str(cuda_home))
269273
if (nvcc_bin := nvcc_path()) is not None:
@@ -275,13 +279,34 @@ def build_nccl_ep_submodule() -> str:
275279
env["NCCL_HOME"] = nccl_home
276280
env["NCCL_EP_BUILDDIR"] = str(build_dir)
277281

278-
prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None
279-
if not nccl_ep_shared_lib.exists() or prev_gencode != gencode:
280-
if nccl_ep_shared_lib.exists() and prev_gencode != gencode:
281-
print(
282-
f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); "
283-
"rebuilding NCCL EP libraries"
284-
)
282+
bolt_cxx_flags, bolt_linker_flags = get_bolt_build_flags()
283+
nvcc_host_flags = [f"-Xcompiler={flag}" for flag in bolt_cxx_flags]
284+
nvcc_linker_flags = []
285+
if bolt_linker_flags:
286+
nvcc_linker_flags.extend(["-Xlinker=--emit-relocs", "-Xlinker=-z", "-Xlinker=now"])
287+
if "-mno-fix-cortex-a53-843419" in bolt_linker_flags:
288+
nvcc_linker_flags.append("-Xlinker=--no-fix-cortex-a53-843419")
289+
290+
def append_env_flags(name: str, flags: List[str]) -> None:
291+
if flags:
292+
env[name] = " ".join([env.get(name, ""), *flags]).strip()
293+
294+
append_env_flags("CXXFLAGS", bolt_cxx_flags)
295+
append_env_flags("NVCC_PREPEND_FLAGS", nvcc_host_flags)
296+
append_env_flags("LDFLAGS", nvcc_linker_flags)
297+
298+
build_signature = "\n".join(
299+
(
300+
f"gencode={gencode}",
301+
f"cxx={env['CXX']}",
302+
f"bolt_cxx_flags={' '.join(bolt_cxx_flags)}",
303+
f"bolt_linker_flags={' '.join(nvcc_linker_flags)}",
304+
)
305+
)
306+
previous_signature = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None
307+
if not nccl_ep_shared_lib.exists() or previous_signature != build_signature:
308+
if nccl_ep_shared_lib.exists() and previous_signature != build_signature:
309+
print("[NCCL EP] build configuration changed; rebuilding NCCL EP libraries")
285310
subprocess.check_call(
286311
["make", "-C", "nccl_ep", "clean"],
287312
cwd=str(nccl_root),
@@ -295,7 +320,7 @@ def build_nccl_ep_submodule() -> str:
295320
env=env,
296321
)
297322
gencode_stamp.parent.mkdir(parents=True, exist_ok=True)
298-
gencode_stamp.write_text(gencode)
323+
gencode_stamp.write_text(build_signature)
299324

300325
return nccl_home
301326

‎transformer_engine/common/CMakeLists.txt‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,71 @@ endif()
1515
# Transformer Engine library
1616
project(transformer_engine LANGUAGES CUDA CXX)
1717

18+
# Enable BOLT-compatible builds. ARM64 targets require additional compiler flags.
19+
string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _nvte_system_processor)
20+
set(NVTE_TARGET_IS_ARM64 OFF)
21+
if(_nvte_system_processor MATCHES "^(aarch64|arm64)$")
22+
set(NVTE_TARGET_IS_ARM64 ON)
23+
endif()
24+
unset(_nvte_system_processor)
25+
26+
include(CheckCXXCompilerFlag)
27+
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
28+
check_cxx_compiler_flag("-fno-reorder-blocks-and-partition"
29+
NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION)
30+
endif()
31+
check_cxx_compiler_flag("-fno-jump-tables"
32+
NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES)
33+
if(NVTE_TARGET_IS_ARM64)
34+
check_cxx_compiler_flag("-mno-fix-cortex-a53-835769"
35+
NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769)
36+
check_cxx_compiler_flag("-mno-fix-cortex-a53-843419"
37+
NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419)
38+
endif()
39+
if((CMAKE_CXX_COMPILER_ID STREQUAL "GNU"
40+
AND NOT NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION)
41+
OR NOT NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES
42+
OR (NVTE_TARGET_IS_ARM64
43+
AND (NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769
44+
OR NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419)))
45+
message(FATAL_ERROR
46+
"The host C++ compiler does not support the flags required for "
47+
"BOLT-compatible builds")
48+
endif()
49+
50+
# BOLT consumes the regular symbol table and emitted relocations. Shadow the
51+
# cached strip tool for this configure without changing it for later builds.
52+
set(CMAKE_STRIP "")
53+
54+
function(nvte_enable_bolt_compatible_compile_options TARGET_NAME)
55+
target_compile_options(
56+
${TARGET_NAME}
57+
PRIVATE
58+
$<$<COMPILE_LANGUAGE:CXX>:-fno-jump-tables>
59+
$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-jump-tables>
60+
)
61+
62+
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
63+
target_compile_options(
64+
${TARGET_NAME}
65+
PRIVATE
66+
$<$<COMPILE_LANGUAGE:CXX>:-fno-reorder-blocks-and-partition>
67+
$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-reorder-blocks-and-partition>
68+
)
69+
endif()
70+
71+
if(NVTE_TARGET_IS_ARM64)
72+
target_compile_options(
73+
${TARGET_NAME}
74+
PRIVATE
75+
$<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-835769>
76+
$<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-843419>
77+
$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-835769>
78+
$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-843419>
79+
)
80+
endif()
81+
endfunction()
82+
1883
# CUDA Toolkit
1984
find_package(CUDAToolkit REQUIRED)
2085
if (CUDAToolkit_VERSION VERSION_LESS 12.1)
@@ -363,13 +428,28 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources)
363428
endforeach()
364429

365430
add_library(transformer_engine SHARED ${transformer_engine_SOURCES})
431+
nvte_enable_bolt_compatible_compile_options(transformer_engine)
366432

367433
# This is TE-specific and should not apply to all targets
368434
target_link_options(
369435
transformer_engine
370436
PRIVATE
371437
"LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version"
372438
)
439+
target_link_options(
440+
transformer_engine
441+
PRIVATE
442+
"LINKER:--emit-relocs"
443+
"LINKER:-z,now"
444+
)
445+
if(NVTE_TARGET_IS_ARM64)
446+
# The Cortex-A53 843419 workaround is applied by the linker.
447+
target_link_options(
448+
transformer_engine
449+
PRIVATE
450+
"LINKER:--no-fix-cortex-a53-843419"
451+
)
452+
endif()
373453

374454
# Disable CMake's automatic architecture flag injection.
375455
# All architectures are handled explicitly via per-source COMPILE_OPTIONS
@@ -435,6 +515,7 @@ endif()
435515
option(NVTE_ENABLE_NVSHMEM "Compile with NVSHMEM library" OFF)
436516
if (NVTE_ENABLE_NVSHMEM)
437517
add_subdirectory(nvshmem_api)
518+
nvte_enable_bolt_compatible_compile_options(nvshmemapi)
438519
target_link_libraries(transformer_engine PUBLIC nvshmemapi)
439520
target_include_directories(transformer_engine PUBLIC ${NVSHMEMAPI_INCLUDE_DIR})
440521
endif()

0 commit comments

Comments
 (0)