Skip to content

[CI] Minor CMake refactoring #3227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
fheinecke wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

[CI] Minor CMake refactoring #3227

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
28fe8ee
Simplify NCCL header/lib CMake logic
fheinecke Jul 21, 2026
b6303b8
Add additional NCCL header search paths
fheinecke Jul 21, 2026
30d92a7
Make TE linker script be provided when linking TE
fheinecke Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 22 additions & 23 deletions transformer_engine/common/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,6 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug")
set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -g -G")
endif()

# Hide non-necessary symbols in shared object.
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version")

# Transformer Engine library
project(transformer_engine LANGUAGES CUDA CXX)

Expand Down Expand Up @@ -103,17 +99,20 @@ set(CUTLASS_TOOLS_INCLUDE_DIR
# Python
find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR)
find_path(_nvte_nccl_include_dir
find_path(NCCL_INCLUDE_DIR
NAMES nccl.h
HINTS "${Python_SITEARCH}/nvidia/nccl"
"/opt/nvidia/nccl"
"/usr/local/nccl"
PATH_SUFFIXES include
REQUIRED)

file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_major_line
function(get_nccl_version OUT_VERSION INCLUDE_DIR)
file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_major_line
REGEX "^#define NCCL_MAJOR[ \t]+[0-9]+$")
file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_minor_line
file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_minor_line
REGEX "^#define NCCL_MINOR[ \t]+[0-9]+$")
file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_patch_line
file(STRINGS "${INCLUDE_DIR}/nccl.h" _nvte_nccl_patch_line
REGEX "^#define NCCL_PATCH[ \t]+[0-9]+$")

string(REGEX REPLACE "^#define NCCL_MAJOR[ \t]+([0-9]+)$" "\\1"
Expand All @@ -127,15 +126,16 @@ function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR)
OR "${_nvte_nccl_minor}" STREQUAL ""
OR "${_nvte_nccl_patch}" STREQUAL "")
message(FATAL_ERROR
"Failed to parse NCCL version from ${_nvte_nccl_include_dir}/nccl.h")
"Failed to parse NCCL version from ${INCLUDE_DIR}/nccl.h")
endif()

set(${OUT_VERSION}
"${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch}"
PARENT_SCOPE)
set(${OUT_INCLUDE_DIR} "${_nvte_nccl_include_dir}" PARENT_SCOPE)
endfunction()

get_nccl_version(NCCL_VERSION "${NCCL_INCLUDE_DIR}")

function(find_cublasmp_version OUT_VERSION OUT_INCLUDE_DIR SEARCH_DIR)
find_path(_nvte_cublasmp_include_dir
NAMES cublasmp.h
Expand Down Expand Up @@ -329,6 +329,14 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources)
endforeach()

add_library(transformer_engine SHARED ${transformer_engine_SOURCES})

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

# Disable CMake's automatic architecture flag injection.
# All architectures are handled explicitly via per-source COMPILE_OPTIONS
# using NVTE_STANDARD_ARCHS, NVTE_GENERIC_ARCHS, and NVTE_SPECIFIC_ARCHS above.
Expand Down Expand Up @@ -363,6 +371,7 @@ target_link_libraries(transformer_engine PUBLIC

target_include_directories(transformer_engine PRIVATE
${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
target_include_directories(transformer_engine PRIVATE ${NCCL_INCLUDE_DIR})
target_include_directories(transformer_engine SYSTEM PRIVATE
${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}/cccl)
target_include_directories(transformer_engine PRIVATE "${CUDNN_FRONTEND_INCLUDE_DIR}")
Expand Down Expand Up @@ -391,7 +400,6 @@ if (NVTE_WITH_CUBLASMP)

target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP)
target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include)
find_nccl_version(NCCL_VERSION NCCL_INCLUDE_DIR)
find_cublasmp_version(CUBLASMP_VERSION CUBLASMP_INCLUDE_DIR ${CUBLASMP_DIR})
find_library(CUBLASMP_LIB
NAMES cublasmp libcublasmp.so libcublasmp.so.0
Expand Down Expand Up @@ -470,15 +478,7 @@ find_file(NCCL_EP_LIB
NO_DEFAULT_PATH
REQUIRED)

# -- NCCL core: nccl.h + libnccl.so -----------------------------------------
# setup.py passes -DNCCL_INCLUDE_DIR; standalone CMake falls back to probing
# well-known NCCL install prefixes.
find_path(NCCL_INCLUDE_DIR nccl.h
HINTS /opt/nvidia/nccl/include /usr/local/nccl/include)
if(NOT NCCL_INCLUDE_DIR)
message(FATAL_ERROR
"nccl.h not found. Pass -DNCCL_INCLUDE_DIR=<prefix>/include.")
endif()
# -- NCCL core library -------------------------------------------------------
if(NOT NCCL_LIB)
find_library(NCCL_LIB
NAMES nccl libnccl
Expand All @@ -487,8 +487,7 @@ if(NOT NCCL_LIB)
endif()

target_include_directories(transformer_engine PRIVATE
${NCCL_EP_INCLUDE_DIR}
${NCCL_INCLUDE_DIR})
${NCCL_EP_INCLUDE_DIR})

# libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's
# DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow
Expand Down
Loading