Produce BOLT-compatible libraries

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,
     )

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],
     )

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"""

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
 

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
+    $<$<COMPILE_LANGUAGE:CXX>:-fno-jump-tables>
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-jump-tables>
+  )
+
+  if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+    target_compile_options(
+      ${TARGET_NAME}
+      PRIVATE
+      $<$<COMPILE_LANGUAGE:CXX>:-fno-reorder-blocks-and-partition>
+      $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-reorder-blocks-and-partition>
+    )
+  endif()
+
+  if(NVTE_TARGET_IS_ARM64)
+    target_compile_options(
+      ${TARGET_NAME}
+      PRIVATE
+      $<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-835769>
+      $<$<COMPILE_LANGUAGE:CXX>:-mno-fix-cortex-a53-843419>
+      $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-835769>
+      $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-mno-fix-cortex-a53-843419>
+    )
+  endif()
+endfunction()
+
 # CUDA Toolkit
 find_package(CUDAToolkit REQUIRED)
 if (CUDAToolkit_VERSION VERSION_LESS 12.1)
@@ -333,13 +398,28 @@ 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(
   transformer_engine
   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()