# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

cmake_minimum_required(VERSION 3.21)

# Language options
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
if (CMAKE_BUILD_TYPE STREQUAL "Debug")
  set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -g -G")
endif()

# Transformer Engine library
project(transformer_engine LANGUAGES CUDA CXX)

# CUDA Toolkit
find_package(CUDAToolkit REQUIRED)
if (CUDAToolkit_VERSION VERSION_LESS 12.1)
  message(FATAL_ERROR "CUDA 12.1+ is required, but found CUDA ${CUDAToolkit_VERSION}")
endif()

# Process GPU architectures
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
  if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.0)
    set(CMAKE_CUDA_ARCHITECTURES 75 80 89 90 100 120)
  elseif (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8)
    set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90 100 120)
  else ()
    set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90)
  endif()
endif()

# Process CMAKE_CUDA_ARCHITECTURES to separate standard, generic, and specific architectures.
# - NVTE_STANDARD_ARCHS: pre-Blackwell archs (e.g. 75, 80, 89, 90). Applied to all CUDA sources.
# - NVTE_GENERIC_ARCHS: Blackwell family heads (e.g. 100, 120). Applied to non-arch-specific sources only.
# - NVTE_SPECIFIC_ARCHS: Blackwell specific targets (e.g. 100a, 120f). Applied to arch-specific sources only.
set(NVTE_STANDARD_ARCHS)
set(NVTE_GENERIC_ARCHS)
set(NVTE_SPECIFIC_ARCHS)

# Check for architecture 100
list(FIND CMAKE_CUDA_ARCHITECTURES "100" arch_100_index)
if(NOT arch_100_index EQUAL -1)
  list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "100")
  list(APPEND NVTE_GENERIC_ARCHS "100")
  list(APPEND NVTE_SPECIFIC_ARCHS "100a")
  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9)
    list(APPEND NVTE_SPECIFIC_ARCHS "103a")
  endif()
  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.4)
    list(APPEND NVTE_SPECIFIC_ARCHS "107a")
  endif()
endif()

# Check for architecture 101 (if we see this we are in toolkit <= 12.9)
list(FIND CMAKE_CUDA_ARCHITECTURES "101" arch_101_index)
if(NOT arch_101_index EQUAL -1)
  list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "101")
  list(APPEND NVTE_GENERIC_ARCHS "101")
  list(APPEND NVTE_SPECIFIC_ARCHS "101a")
endif()

# Check for architecture 110 (if we see this we are in toolkit >= 13.0)
list(FIND CMAKE_CUDA_ARCHITECTURES "110" arch_110_index)
if(NOT arch_110_index EQUAL -1)
  list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "110")
  list(APPEND NVTE_GENERIC_ARCHS "110")
  list(APPEND NVTE_SPECIFIC_ARCHS "110f")
endif()

# Check for architecture 120
list(FIND CMAKE_CUDA_ARCHITECTURES "120" arch_120_index)
if(NOT arch_120_index EQUAL -1)
  list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "120")
  list(APPEND NVTE_GENERIC_ARCHS "120")
  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9)
    list(APPEND NVTE_SPECIFIC_ARCHS "120f")
  else()
    list(APPEND NVTE_SPECIFIC_ARCHS "120a")
  endif()
endif()

# Move remaining standard (pre-Blackwell) architectures into NVTE_STANDARD_ARCHS.
# These are applied to all CUDA sources (both generic and arch-specific).
set(NVTE_STANDARD_ARCHS ${CMAKE_CUDA_ARCHITECTURES})

# cuDNN frontend API
if(NOT EXISTS "${CUDNN_FRONTEND_INCLUDE_DIR}/cudnn_frontend.h")
    message(FATAL_ERROR
            "Could not find cudnn_frontend.h in ${CUDNN_FRONTEND_INCLUDE_DIR}. "
            "Install nvidia-cudnn-frontend or set CUDNN_FRONTEND_INCLUDE_DIR.")
endif()
include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/cuDNN.cmake)

set(CUTLASS_INCLUDE_DIR
  "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cutlass/include")
set(CUTLASS_TOOLS_INCLUDE_DIR
  "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cutlass/tools/util/include")

# Python
find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

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

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 "${INCLUDE_DIR}/nccl.h" _nvte_nccl_minor_line
       REGEX "^#define NCCL_MINOR[ \t]+[0-9]+$")
  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"
         _nvte_nccl_major "${_nvte_nccl_major_line}")
  string(REGEX REPLACE "^#define NCCL_MINOR[ \t]+([0-9]+)$" "\\1"
         _nvte_nccl_minor "${_nvte_nccl_minor_line}")
  string(REGEX REPLACE "^#define NCCL_PATCH[ \t]+([0-9]+)$" "\\1"
         _nvte_nccl_patch "${_nvte_nccl_patch_line}")

  if ("${_nvte_nccl_major}" STREQUAL ""
      OR "${_nvte_nccl_minor}" STREQUAL ""
      OR "${_nvte_nccl_patch}" STREQUAL "")
    message(FATAL_ERROR
            "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)
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
            HINTS "${SEARCH_DIR}/include"
            PATH_SUFFIXES include
            REQUIRED)

  file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_major_line
       REGEX "^#define CUBLASMP_VER_MAJOR[ \t]+[0-9]+$")
  file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_minor_line
       REGEX "^#define CUBLASMP_VER_MINOR[ \t]+[0-9]+$")
  file(STRINGS "${_nvte_cublasmp_include_dir}/cublasmp.h" _nvte_cublasmp_patch_line
       REGEX "^#define CUBLASMP_VER_PATCH[ \t]+[0-9]+$")

  string(REGEX REPLACE "^#define CUBLASMP_VER_MAJOR[ \t]+([0-9]+)$" "\\1"
         _nvte_cublasmp_major "${_nvte_cublasmp_major_line}")
  string(REGEX REPLACE "^#define CUBLASMP_VER_MINOR[ \t]+([0-9]+)$" "\\1"
         _nvte_cublasmp_minor "${_nvte_cublasmp_minor_line}")
  string(REGEX REPLACE "^#define CUBLASMP_VER_PATCH[ \t]+([0-9]+)$" "\\1"
         _nvte_cublasmp_patch "${_nvte_cublasmp_patch_line}")

  if ("${_nvte_cublasmp_major}" STREQUAL ""
      OR "${_nvte_cublasmp_minor}" STREQUAL ""
      OR "${_nvte_cublasmp_patch}" STREQUAL "")
    message(FATAL_ERROR
            "Failed to parse cuBLASMp version from ${_nvte_cublasmp_include_dir}/cublasmp.h")
  endif()

  set(${OUT_VERSION}
      "${_nvte_cublasmp_major}.${_nvte_cublasmp_minor}.${_nvte_cublasmp_patch}"
      PARENT_SCOPE)
  set(${OUT_INCLUDE_DIR} "${_nvte_cublasmp_include_dir}" PARENT_SCOPE)
endfunction()

# Configure Transformer Engine library
include_directories(${PROJECT_SOURCE_DIR}/..)
set(transformer_engine_SOURCES)
set(transformer_engine_cpp_sources)
set(transformer_engine_cuda_sources)
set(transformer_engine_cuda_arch_specific_sources)

list(APPEND transformer_engine_cpp_sources
     cudnn_utils.cpp
     transformer_engine.cpp
     fused_attn/fused_attn.cpp
     gemm/config.cpp
     normalization/common.cpp
     normalization/rtc_dispatch.cpp
     normalization/layernorm/ln_api.cpp
     normalization/rmsnorm/rmsnorm_api.cpp
     util/cuda_driver.cpp
     util/cuda_nvml.cpp
     util/cuda_runtime.cpp
     util/multi_stream.cpp
     util/rtc.cpp
     comm_gemm/comm_gemm.cpp
     comm_gemm_overlap/userbuffers/ipcsocket.cc
     comm_gemm_overlap/userbuffers/userbuffers-host.cpp
     comm_gemm_overlap/comm_gemm_overlap.cpp
     newton_schulz/newton_schulz.cpp
     )

list(APPEND transformer_engine_cuda_sources
     common.cu
     multi_tensor/adam.cu
     multi_tensor/l2norm.cu
     multi_tensor/scale.cu
     multi_tensor/sgd.cu
     transpose/cast_transpose.cu
     transpose/transpose.cu
     transpose/cast_transpose_fusion.cu
     transpose/transpose_fusion.cu
     transpose/multi_cast_transpose.cu
     transpose/quantize_transpose_vector_blockwise.cu
     transpose/swap_first_dims.cu
     dropout/dropout.cu
     fused_attn/context_parallel.cu
     fused_attn/kv_cache.cu
     fused_attn/fused_attn_f16_arbitrary_seqlen.cu
     fused_attn/fused_attn_fp8.cu
     fused_attn/utils.cu
     gemm/cublaslt_gemm.cu
     gemm/cublaslt_grouped_gemm.cu
     normalization/layernorm/ln_bwd_semi_cuda_kernel.cu
     normalization/layernorm/ln_fwd_cuda_kernel.cu
     normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu
     normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu
     permutation/permutation.cu
     util/utils.cu
     util/padding.cu
     util/splits_to_offsets.cu
     util/topk.cu
     swizzle/swizzle.cu
     swizzle/swizzle_block_scaling.cu
     fused_softmax/scaled_masked_softmax.cu
     fused_softmax/scaled_upper_triang_masked_softmax.cu
     fused_softmax/scaled_aligned_causal_masked_softmax.cu
     fused_rope/fused_rope.cu
     fused_router/fused_moe_aux_loss.cu
     fused_router/fused_score_for_moe_aux_loss.cu
     fused_router/fused_topk_with_score_function.cu
     recipe/current_scaling.cu
     recipe/delayed_scaling.cu
     recipe/fp8_block_scaling.cu
     comm_gemm_overlap/userbuffers/userbuffers.cu)

list(APPEND transformer_engine_cuda_arch_specific_sources
     fused_attn/flash_attn.cu
     activation/gelu.cu
     activation/gelu_dbias.cu
     activation/gelu_grouped.cu
     activation/gelu_grouped_dbias.cu
     activation/glu.cu
     activation/relu.cu
     activation/relu_dbias.cu
     activation/relu_grouped.cu
     activation/relu_grouped_dbias.cu
     activation/scaled_activation.cu
     activation/scaled_srelu.cu
     activation/scaled_swiglu.cu
     activation/swiglu.cu
     activation/swiglu_dbias.cu
     activation/swiglu_grouped.cu
     activation/swiglu_grouped_dbias.cu
     cast/cast.cu
     cast/cast_dbias.cu
     cast/cast_grouped.cu
     cast/cast_grouped_dbias.cu
     gemm/cutlass_grouped_gemm.cu
     hadamard_transform/group_hadamard_transform.cu
     hadamard_transform/graph_safe_group_hadamard_transform.cu
     hadamard_transform/hadamard_transform.cu
     hadamard_transform/hadamard_transform_cast_fusion.cu
     hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu
     hadamard_transform/group_hadamard_transform_cast_fusion.cu
     hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
     hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu
     multi_tensor/compute_scale.cu
     recipe/mxfp8_scaling.cu
     recipe/nvfp4.cu
     transpose/quantize_transpose_square_blockwise.cu
     transpose/quantize_transpose_vector_blockwise_fp4.cu)

# Compiling the files with the worst compilation time first to hopefully overlap
# better with the faster-compiling cpp files
list(APPEND transformer_engine_SOURCES ${transformer_engine_cuda_arch_specific_sources}
                                       ${transformer_engine_cuda_sources}
                                       ${transformer_engine_cpp_sources})

# Set compile options for CUDA sources with generic architectures.
# These get standard archs (pre-Blackwell) + generic Blackwell family heads.
foreach(cuda_source IN LISTS transformer_engine_cuda_sources)
  set(arch_compile_options)
  foreach(arch IN LISTS NVTE_STANDARD_ARCHS)
    list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}")
  endforeach()
  foreach(arch IN LISTS NVTE_GENERIC_ARCHS)
    list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}")
  endforeach()

  if(arch_compile_options)
    set_property(
      SOURCE ${cuda_source}
      APPEND
      PROPERTY
      COMPILE_OPTIONS ${arch_compile_options}
    )
  endif()
endforeach()

# Set compile options for CUDA sources with arch-specific features.
# These get standard archs (pre-Blackwell) + Blackwell specific targets (a/f suffix).
# They must NOT get generic Blackwell archs, as they use family/arch-specific PTX features.
foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources)
  set(arch_compile_options)
  foreach(arch IN LISTS NVTE_STANDARD_ARCHS)
    list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}")
  endforeach()
  foreach(arch IN LISTS NVTE_SPECIFIC_ARCHS)
    list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}")
  endforeach()

  if(arch_compile_options)
    set_property(
      SOURCE ${cuda_source}
      APPEND
      PROPERTY
      COMPILE_OPTIONS ${arch_compile_options}
    )
  endif()
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.
set_target_properties(transformer_engine PROPERTIES CUDA_ARCHITECTURES OFF)
target_include_directories(transformer_engine PUBLIC
                           "${CMAKE_CURRENT_SOURCE_DIR}/include")

# Grouped GEMM kernels require SM90a
set_property(
  SOURCE gemm/cutlass_grouped_gemm.cu
  APPEND
  PROPERTY
  COMPILE_OPTIONS "--generate-code=arch=compute_90a,code=sm_90a")

# CUTLASS kernels could cause hang in debug build
set(CUTLASS_KERNEL_SOURCES
    gemm/cutlass_grouped_gemm.cu
    hadamard_transform/group_hadamard_transform_cast_fusion.cu
    hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
    hadamard_transform/hadamard_transform_cast_fusion.cu)
set_property(
  SOURCE ${CUTLASS_KERNEL_SOURCES}
  APPEND
  PROPERTY
  COMPILE_OPTIONS "-g0;-dopt=on")

# Configure dependencies
target_link_libraries(transformer_engine PUBLIC
                      CUDA::cublas
                      CUDA::cudart
                      CUDNN::cudnn_all)

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}")
target_include_directories(transformer_engine PRIVATE
                          ${CUTLASS_INCLUDE_DIR}
                          ${CUTLASS_TOOLS_INCLUDE_DIR})

# Compiling Userbuffers with native MPI bootstrapping requires linking against MPI
option(NVTE_UB_WITH_MPI "Bootstrap Userbuffers with MPI" OFF)
if (NVTE_UB_WITH_MPI)
    find_package(MPI REQUIRED)
    target_link_libraries(transformer_engine PUBLIC MPI::MPI_CXX)
    target_include_directories(transformer_engine PRIVATE ${MPI_CXX_INCLUDES})
    target_compile_definitions(transformer_engine PUBLIC NVTE_UB_WITH_MPI)
endif()

option(NVTE_ENABLE_NVSHMEM "Compile with NVSHMEM library" OFF)
if (NVTE_ENABLE_NVSHMEM)
    add_subdirectory(nvshmem_api)
    target_link_libraries(transformer_engine PUBLIC nvshmemapi)
    target_include_directories(transformer_engine PUBLIC ${NVSHMEMAPI_INCLUDE_DIR})
endif()

option(NVTE_WITH_CUBLASMP "Use cuBLASMp for tensor parallel GEMMs" OFF)
if (NVTE_WITH_CUBLASMP)

    target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP)
    target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include)
    find_cublasmp_version(CUBLASMP_VERSION CUBLASMP_INCLUDE_DIR ${CUBLASMP_DIR})
    find_library(CUBLASMP_LIB
                 NAMES cublasmp libcublasmp.so libcublasmp.so.0
                 PATHS ${CUBLASMP_DIR}
                 PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/sbsa-linux-gnu lib/x86_64-linux-gnu
                 REQUIRED)
    find_library(NCCL_LIB
                 NAMES nccl libnccl
                 PATH_SUFFIXES lib
                 REQUIRED)
    # cuBLASMp 0.8 is the first release with CUDA-graph-safe overlap algos,
    # and NCCL 2.30 is the first release with graph-safe one-sided RMA
    # primitives (ncclPutSignal/ncclWaitSignal) that those algos use.
    if (CUBLASMP_VERSION VERSION_LESS 0.8.0)
      message(FATAL_ERROR
              "NVTE_WITH_CUBLASMP requires cuBLASMp >= 0.8.0, but found cuBLASMp "
              "${CUBLASMP_VERSION} in ${CUBLASMP_INCLUDE_DIR}/cublasmp.h")
    endif()
    if (NCCL_VERSION VERSION_LESS 2.30.0)
      message(FATAL_ERROR
              "NVTE_WITH_CUBLASMP requires NCCL >= 2.30.0 (for graph-capture-safe "
              "one-sided RMA primitives used by cuBLASMp's overlap algorithms), but "
              "found NCCL ${NCCL_VERSION} in ${NCCL_INCLUDE_DIR}/nccl.h")
    endif()
    target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CUBLASMP_LIB})
    message(STATUS "Using cuBLASMp ${CUBLASMP_VERSION} at: ${CUBLASMP_DIR}")
    message(STATUS "Using NCCL ${NCCL_VERSION} at: ${NCCL_LIB}")
endif()

option(NVTE_WITH_CUSOLVERMP "Use cuSolverMp for distributed Newton-Schulz" OFF)
if (NVTE_WITH_CUSOLVERMP)
    target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUSOLVERMP)
    target_include_directories(transformer_engine PRIVATE ${CUSOLVERMP_DIR}/include)
    find_library(CUSOLVERMP_LIB
                 NAMES cusolverMp libcusolverMp
                 PATHS ${CUSOLVERMP_DIR}
                 PATH_SUFFIXES lib
                 REQUIRED)
    find_library(NCCL_LIB
                 NAMES nccl libnccl
                 PATH_SUFFIXES lib
                 REQUIRED)
    target_link_libraries(transformer_engine PRIVATE ${NCCL_LIB} ${CUSOLVERMP_LIB})
    message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}")
endif()

# -- NCCL EP (on by default, HT mode only) ---------------------------------
# Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to
# skip NCCL EP entirely - useful on older images whose system NCCL is below
# the 2.30.4 EP minimum.
option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON)
if(NVTE_WITH_NCCL_EP)
# SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize.
# -- NCCL EP headers --------------------------------------------------------
# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build
# (auto-built by setup.py via build_nccl_ep_submodule).
set(NCCL_EP_SUBMODULE_ROOT
    "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions")
set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include")
if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h")
  message(FATAL_ERROR
    "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. "
    "Run `git submodule update --init --recursive` and rebuild TE.")
endif()
message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}")

# -- libnccl_ep.a -----------------------------------------------------------
# Statically linked into libtransformer_engine.so. EPBackend::initialize checks
# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep
# imports stay unresolved (and harmless) under default ELF lazy binding when
# the gate trips. LD_BIND_NOW environments lose this property.
set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib")
find_file(NCCL_EP_LIB
    NAMES libnccl_ep.a
    HINTS ${NCCL_EP_LIB_DIR}
    NO_DEFAULT_PATH
    REQUIRED)

# -- NCCL core library -------------------------------------------------------
if(NOT NCCL_LIB)
  find_library(NCCL_LIB
      NAMES nccl libnccl
      PATH_SUFFIXES lib lib64
      REQUIRED)
endif()

target_include_directories(transformer_engine PRIVATE
    ${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
# the static archive on the link line so --as-needed records libcuda.so.1.
target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB})
target_link_libraries(transformer_engine PRIVATE
    -Wl,--whole-archive ${NCCL_EP_LIB} -Wl,--no-whole-archive
    CUDA::cuda_driver)

target_sources(transformer_engine PRIVATE
    ep/ep_backend.cpp
    ep/ep_api.cpp)
target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP)

message(STATUS "NCCL EP enabled (static link): ${NCCL_EP_LIB}")
message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}")
else()
  # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs.
  target_sources(transformer_engine PRIVATE ep/ep_api.cpp)
  message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) - using nvte_ep_* stubs")
endif()

# Number of philox4x32 rounds for stochastic rounding (build-time constant).
set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR $ENV{NVTE_BUILD_NUM_PHILOX_ROUNDS})
if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR)
  set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR "10")
endif()
if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR MATCHES "^[1-9][0-9]*$")
  message(FATAL_ERROR
          "Environment variable NVTE_BUILD_NUM_PHILOX_ROUNDS must be a positive integer, "
          "but got '${NVTE_BUILD_NUM_PHILOX_ROUNDS_STR}'.")
endif()
set(NVTE_BUILD_NUM_PHILOX_ROUNDS ${NVTE_BUILD_NUM_PHILOX_ROUNDS_STR})

target_compile_definitions(transformer_engine
                           PUBLIC NVTE_BUILD_NUM_PHILOX_ROUNDS=${NVTE_BUILD_NUM_PHILOX_ROUNDS})
message(STATUS "Philox rounds for stochastic rounding: ${NVTE_BUILD_NUM_PHILOX_ROUNDS}")

# Hack to enable dynamic loading in cuDNN frontend
target_compile_definitions(transformer_engine PUBLIC NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING)

# Helper functions to make header files with C++ strings
function(make_string_header STRING STRING_NAME)
    configure_file(util/string_header.h.in
                   "string_headers/${STRING_NAME}.h"
                   @ONLY)
endfunction()
function(make_string_header_from_file file_ STRING_NAME)
    file(READ "${file_}" STRING)
    configure_file(util/string_header.h.in
                   "string_headers/${STRING_NAME}.h"
                   @ONLY)
endfunction()

# Header files with C++ strings
list(GET CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES 0 cuda_include_path)
make_string_header("${cuda_include_path}"
                   string_path_cuda_include)
make_string_header_from_file(transpose/rtc/cast_transpose_fusion.cu
                             string_code_transpose_rtc_cast_transpose_fusion_cu)
make_string_header_from_file(transpose/rtc/cast_transpose.cu
                             string_code_transpose_rtc_cast_transpose_cu)
make_string_header_from_file(transpose/rtc/transpose.cu
                             string_code_transpose_rtc_transpose_cu)
make_string_header_from_file(transpose/rtc/swap_first_dims.cu
                             string_code_transpose_rtc_swap_first_dims_cu)
make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu
                             string_code_fused_softmax_scaled_masked_softmax_cu)
make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu
                             string_code_fused_softmax_scaled_upper_triang_masked_softmax_cu)
make_string_header_from_file(fused_softmax/scaled_aligned_causal_masked_softmax.cu
                             string_code_fused_softmax_scaled_aligned_causal_masked_softmax_cu)
make_string_header_from_file(utils.cuh
                             string_code_utils_cuh)
make_string_header_from_file(util/math.h
                             string_code_util_math_h)

# Norm NVRTC bundled headers + RTC source files
make_string_header_from_file(normalization/kernel_params.h
                             string_code_normalization_kernel_params_h)
make_string_header_from_file(normalization/kernel_traits.h
                             string_code_normalization_kernel_traits_h)
make_string_header_from_file(normalization/layernorm/ln_fwd_kernels.cuh
                             string_code_normalization_layernorm_ln_fwd_kernels_cuh)
make_string_header_from_file(normalization/layernorm/ln_bwd_kernels.cuh
                             string_code_normalization_layernorm_ln_bwd_kernels_cuh)
make_string_header_from_file(normalization/rmsnorm/rmsnorm_fwd_kernels.cuh
                             string_code_normalization_rmsnorm_rmsnorm_fwd_kernels_cuh)
make_string_header_from_file(normalization/rmsnorm/rmsnorm_bwd_kernels.cuh
                             string_code_normalization_rmsnorm_rmsnorm_bwd_kernels_cuh)
make_string_header_from_file(normalization/layernorm/rtc/ln_fwd_kernel.cu
                             string_code_normalization_layernorm_rtc_ln_fwd_kernel_cu)
make_string_header_from_file(normalization/layernorm/rtc/ln_bwd_kernel.cu
                             string_code_normalization_layernorm_rtc_ln_bwd_kernel_cu)
make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu
                             string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu)
make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu
                             string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu)
target_include_directories(transformer_engine PRIVATE
                           "${CMAKE_CURRENT_BINARY_DIR}/string_headers")

option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX
       "Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback"
       OFF)
target_compile_definitions(transformer_engine
                           PRIVATE
                           NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX=$<BOOL:${NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX}>)

# Default OFF: LayerNorm/RMSNorm require NVRTC. Set ON to additionally compile
# the legacy static template instantiations, which are selected only when
# NVTE_DISABLE_NVRTC=1 at runtime. This compiles all 556
# REGISTER_NORM_LAUNCHER variants up front.
option(NVTE_BUILD_LEGACY_STATIC_NORM
       "Also compile static norm kernels for NVTE_DISABLE_NVRTC fallback"
       OFF)
target_compile_definitions(transformer_engine
                           PRIVATE
                           NVTE_BUILD_LEGACY_STATIC_NORM=$<BOOL:${NVTE_BUILD_LEGACY_STATIC_NORM}>)

# Compiler options
set(nvte_sources_with_fast_math)
list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu
                                        fused_softmax/scaled_upper_triang_masked_softmax.cu
                                        fused_softmax/scaled_aligned_causal_masked_softmax.cu
                                        multi_tensor/adam.cu
                                        multi_tensor/compute_scale.cu
                                        multi_tensor/l2norm.cu
                                        multi_tensor/scale.cu
                                        multi_tensor/sgd.cu
                                        fused_attn/flash_attn.cu
                                        fused_attn/context_parallel.cu
                                        fused_attn/kv_cache.cu)

option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF)
if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH)
  list(APPEND nvte_sources_with_fast_math activation/gelu.cu
                                          activation/gelu_dbias.cu
                                          activation/gelu_grouped.cu
                                          activation/gelu_grouped_dbias.cu
                                          activation/glu.cu
                                          activation/relu.cu
                                          activation/relu_dbias.cu
                                          activation/relu_grouped.cu
                                          activation/relu_grouped_dbias.cu
                                          activation/scaled_activation.cu
                                          activation/scaled_srelu.cu
                                          activation/scaled_swiglu.cu
                                          activation/swiglu.cu
                                          activation/swiglu_dbias.cu
                                          activation/swiglu_grouped.cu
                                          activation/swiglu_grouped_dbias.cu)
endif()

foreach(cuda_source IN LISTS nvte_sources_with_fast_math)
  set_property(
    SOURCE ${cuda_source}
    APPEND
    PROPERTY
    COMPILE_OPTIONS "--use_fast_math")
endforeach()

set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -O3")

# Number of parallel build jobs
if($ENV{MAX_JOBS})
  set(BUILD_JOBS_STR $ENV{MAX_JOBS})
elseif($ENV{NVTE_BUILD_MAX_JOBS})
  set(BUILD_JOBS_STR $ENV{NVTE_BUILD_MAX_JOBS})
else()
  set(BUILD_JOBS_STR "max")
endif()
message(STATUS "Parallel build jobs: ${BUILD_JOBS_STR}")

# Number of threads per parallel build job
set(BUILD_THREADS_PER_JOB $ENV{NVTE_BUILD_THREADS_PER_JOB})
if (NOT BUILD_THREADS_PER_JOB)
  set(BUILD_THREADS_PER_JOB 1)
endif()
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --threads ${BUILD_THREADS_PER_JOB}")
message(STATUS "Threads per parallel build job: ${BUILD_THREADS_PER_JOB}")

# Install library
install(TARGETS transformer_engine DESTINATION .)
