diff --git a/build_tools/jax.py b/build_tools/jax.py index 031432e6f90..e54c668c5b4 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -15,6 +15,7 @@ all_files_in_dir, cudnn_frontend_include_path, debug_build_enabled, + get_bolt_build_flags, setup_mpi_flags, nccl_include_path, nccl_lib_path, @@ -113,6 +114,9 @@ def setup_jax_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + setup_mpi_flags(include_dirs, cxx_flags) if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): @@ -135,5 +139,6 @@ def setup_jax_extension( sources=[str(path) for path in sources], include_dirs=[str(path) for path in include_dirs], extra_compile_args=cxx_flags, + extra_link_args=linker_flags, **kwargs, ) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 98331ccbd86..18342dc14b4 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -15,6 +15,7 @@ cuda_version, get_cuda_include_dirs, debug_build_enabled, + get_bolt_build_flags, nccl_ep_enabled, setup_mpi_flags, ) @@ -76,6 +77,9 @@ def setup_pytorch_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + # Version-dependent CUDA options try: version = cuda_version() @@ -122,6 +126,7 @@ def setup_pytorch_extension( sources=[str(src) for src in sources], include_dirs=[str(inc) for inc in include_dirs], extra_compile_args={"cxx": cxx_flags}, + extra_link_args=linker_flags, libraries=[str(lib) for lib in libraries], library_dirs=[str(lib_dir) for lib_dir in library_dirs], ) diff --git a/build_tools/utils.py b/build_tools/utils.py index 3b02fcea2b7..6841a595802 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -9,6 +9,7 @@ import importlib import os import re +import shlex import shutil import subprocess import sys @@ -43,6 +44,72 @@ def debug_build_enabled() -> bool: return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0"))) +@functools.lru_cache(maxsize=None) +def _cxx_compiler_predefined_macros() -> Tuple[str, ...]: + """Predefined macros from the configured C++ compiler.""" + cxx = shlex.split(os.getenv("CXX", "c++")) + if not cxx: + return () + try: + result = subprocess.run( + [*cxx, "-dM", "-E", "-x", "c++", "-"], + input="", + capture_output=True, + text=True, + ) + except OSError: + return () + if result.returncode != 0: + return () + return tuple( + fields[1] + for line in result.stdout.splitlines() + if len(fields := line.split(maxsplit=2)) >= 2 and fields[0] == "#define" + ) + + +@functools.lru_cache(maxsize=None) +def cxx_compiler_supports_flag(flag: str) -> bool: + """Whether the configured C++ compiler accepts a command-line flag.""" + cxx = shlex.split(os.getenv("CXX", "c++")) + if not cxx: + return False + try: + result = subprocess.run( + [*cxx, flag, "-x", "c++", "-fsyntax-only", "-"], + input="", + capture_output=True, + text=True, + ) + except OSError: + return False + return result.returncode == 0 + + +def cxx_compiler_targets_arm64() -> bool: + """Whether the configured C++ compiler targets Arm64.""" + macros = _cxx_compiler_predefined_macros() + return any(macro in macros for macro in ("__aarch64__", "__arm64__", "_M_ARM64")) + + +def get_bolt_build_flags() -> Tuple[List[str], List[str]]: + """BOLT-compatible host compiler and linker flags.""" + compiler_flags = ["-fno-jump-tables"] + if cxx_compiler_supports_flag("-fno-reorder-blocks-and-partition"): + compiler_flags.append("-fno-reorder-blocks-and-partition") + linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"] + if cxx_compiler_targets_arm64(): + compiler_flags.extend( + [ + "-mno-fix-cortex-a53-835769", + "-mno-fix-cortex-a53-843419", + ] + ) + # The Cortex-A53 843419 workaround is applied by the linker. + linker_flags.append("-mno-fix-cortex-a53-843419") + return compiler_flags, linker_flags + + @functools.lru_cache(maxsize=None) def get_max_jobs_for_parallel_build() -> int: """Number of parallel jobs for Nina build""" diff --git a/setup.py b/setup.py index 944cba060e4..90379c25fe8 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ cuda_home_path, cuda_version, cudnn_frontend_include_path, + get_bolt_build_flags, get_frameworks, remove_dups, min_python_version_str, @@ -255,6 +256,9 @@ def build_nccl_ep_submodule() -> str: nproc = get_max_jobs_for_parallel_build() env = os.environ.copy() + # get_bolt_build_flags() defaults to `c++` when CXX is unset. Export the + # same default so Make does not independently select its `g++` default. + env.setdefault("CXX", "c++") if (cuda_home := cuda_home_path()) is not None: env.setdefault("CUDA_HOME", str(cuda_home)) if (nvcc_bin := nvcc_path()) is not None: @@ -266,13 +270,34 @@ def build_nccl_ep_submodule() -> str: env["NCCL_HOME"] = nccl_home env["NCCL_EP_BUILDDIR"] = str(build_dir) - prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None - if not nccl_ep_shared_lib.exists() or prev_gencode != gencode: - if nccl_ep_shared_lib.exists() and prev_gencode != gencode: - print( - f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); " - "rebuilding NCCL EP libraries" - ) + bolt_cxx_flags, bolt_linker_flags = get_bolt_build_flags() + nvcc_host_flags = [f"-Xcompiler={flag}" for flag in bolt_cxx_flags] + nvcc_linker_flags = [] + if bolt_linker_flags: + nvcc_linker_flags.extend(["-Xlinker=--emit-relocs", "-Xlinker=-z", "-Xlinker=now"]) + if "-mno-fix-cortex-a53-843419" in bolt_linker_flags: + nvcc_linker_flags.append("-Xlinker=--no-fix-cortex-a53-843419") + + def append_env_flags(name: str, flags: List[str]) -> None: + if flags: + env[name] = " ".join([env.get(name, ""), *flags]).strip() + + append_env_flags("CXXFLAGS", bolt_cxx_flags) + append_env_flags("NVCC_PREPEND_FLAGS", nvcc_host_flags) + append_env_flags("LDFLAGS", nvcc_linker_flags) + + build_signature = "\n".join( + ( + f"gencode={gencode}", + f"cxx={env['CXX']}", + f"bolt_cxx_flags={' '.join(bolt_cxx_flags)}", + f"bolt_linker_flags={' '.join(nvcc_linker_flags)}", + ) + ) + previous_signature = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None + if not nccl_ep_shared_lib.exists() or previous_signature != build_signature: + if nccl_ep_shared_lib.exists() and previous_signature != build_signature: + print("[NCCL EP] build configuration changed; rebuilding NCCL EP libraries") subprocess.check_call( ["make", "-C", "nccl_ep", "clean"], cwd=str(nccl_root), @@ -286,7 +311,7 @@ def build_nccl_ep_submodule() -> str: env=env, ) gencode_stamp.parent.mkdir(parents=True, exist_ok=True) - gencode_stamp.write_text(gencode) + gencode_stamp.write_text(build_signature) return nccl_home diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 3503941cf4f..43b0a64559e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -15,6 +15,71 @@ endif() # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) +# Enable BOLT-compatible builds. ARM64 targets require additional compiler flags. +string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _nvte_system_processor) +set(NVTE_TARGET_IS_ARM64 OFF) +if(_nvte_system_processor MATCHES "^(aarch64|arm64)$") + set(NVTE_TARGET_IS_ARM64 ON) +endif() +unset(_nvte_system_processor) + +include(CheckCXXCompilerFlag) +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + check_cxx_compiler_flag("-fno-reorder-blocks-and-partition" + NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION) +endif() +check_cxx_compiler_flag("-fno-jump-tables" + NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES) +if(NVTE_TARGET_IS_ARM64) + check_cxx_compiler_flag("-mno-fix-cortex-a53-835769" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769) + check_cxx_compiler_flag("-mno-fix-cortex-a53-843419" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419) +endif() +if((CMAKE_CXX_COMPILER_ID STREQUAL "GNU" + AND NOT NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION) + OR NOT NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES + OR (NVTE_TARGET_IS_ARM64 + AND (NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769 + OR NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419))) + message(FATAL_ERROR + "The host C++ compiler does not support the flags required for " + "BOLT-compatible builds") +endif() + +# BOLT consumes the regular symbol table and emitted relocations. Shadow the +# cached strip tool for this configure without changing it for later builds. +set(CMAKE_STRIP "") + +function(nvte_enable_bolt_compatible_compile_options TARGET_NAME) + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-fno-jump-tables> + $<$:-Xcompiler=-fno-jump-tables> + ) + + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-fno-reorder-blocks-and-partition> + $<$:-Xcompiler=-fno-reorder-blocks-and-partition> + ) + endif() + + if(NVTE_TARGET_IS_ARM64) + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-mno-fix-cortex-a53-835769> + $<$:-mno-fix-cortex-a53-843419> + $<$:-Xcompiler=-mno-fix-cortex-a53-835769> + $<$:-Xcompiler=-mno-fix-cortex-a53-843419> + ) + endif() +endfunction() + # CUDA Toolkit find_package(CUDAToolkit REQUIRED) if (CUDAToolkit_VERSION VERSION_LESS 12.1) @@ -333,6 +398,7 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) endforeach() add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) +nvte_enable_bolt_compatible_compile_options(transformer_engine) # This is TE-specific and should not apply to all targets target_link_options( @@ -340,6 +406,20 @@ target_link_options( PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version" ) +target_link_options( + transformer_engine + PRIVATE + "LINKER:--emit-relocs" + "LINKER:-z,now" +) +if(NVTE_TARGET_IS_ARM64) + # The Cortex-A53 843419 workaround is applied by the linker. + target_link_options( + transformer_engine + PRIVATE + "LINKER:--no-fix-cortex-a53-843419" + ) +endif() # Disable CMake's automatic architecture flag injection. # All architectures are handled explicitly via per-source COMPILE_OPTIONS @@ -396,6 +476,7 @@ endif() option(NVTE_ENABLE_NVSHMEM "Compile with NVSHMEM library" OFF) if (NVTE_ENABLE_NVSHMEM) add_subdirectory(nvshmem_api) + nvte_enable_bolt_compatible_compile_options(nvshmemapi) target_link_libraries(transformer_engine PUBLIC nvshmemapi) target_include_directories(transformer_engine PUBLIC ${NVSHMEMAPI_INCLUDE_DIR}) endif()